feat(chat): add composable primitives, raw task example, and task mode switcher

Author: ericallam

packages/trigger-sdk/src/v3/ai.ts

@@ -1567,6 +1567,156 @@ function cleanupAbortedParts(message: UIMessage): UIMessage {
   };
 }
 
+// ---------------------------------------------------------------------------
+// Composable primitives for raw task chat
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a managed stop signal wired to the chat stop input stream.
+ *
+ * Call once at the start of your run. Use `signal` as the abort signal for
+ * `streamText`. Call `reset()` at the start of each turn to get a fresh
+ * per-turn signal. Call `cleanup()` when the run ends.
+ *
+ * @example
+ * ```ts
+ * const stop = chat.createStopSignal();
+ * for (let turn = 0; turn < 100; turn++) {
+ *   stop.reset();
+ *   const result = streamText({ model, messages, abortSignal: stop.signal });
+ *   await chat.pipe(result);
+ *   // ...
+ * }
+ * stop.cleanup();
+ * ```
+ */
+function createStopSignal(): { readonly signal: AbortSignal; reset: () => void; cleanup: () => void } {
+  let controller = new AbortController();
+  const sub = stopInput.on((data) => {
+    controller.abort(data?.message || "stopped");
+  });
+  return {
+    get signal() { return controller.signal; },
+    reset() { controller = new AbortController(); },
+    cleanup() { sub.off(); },
+  };
+}
+
+/**
+ * Signal the frontend that the current turn is complete.
+ *
+ * The `TriggerChatTransport` intercepts this to close the ReadableStream
+ * for the current turn. Call after piping the response stream.
+ *
+ * @example
+ * ```ts
+ * await chat.pipe(result);
+ * await chat.writeTurnComplete();
+ * ```
+ */
+async function chatWriteTurnComplete(options?: { publicAccessToken?: string }): Promise<void> {
+  await writeTurnCompleteChunk(undefined, options?.publicAccessToken);
+}
+
+/**
+ * Pipe a `StreamTextResult` (or similar) to the chat stream and capture
+ * the assistant's response message via `onFinish`.
+ *
+ * Combines `toUIMessageStream()` + `onFinish` callback + `chat.pipe()`.
+ * Returns the captured `UIMessage`, or `undefined` if capture failed.
+ *
+ * @example
+ * ```ts
+ * const result = streamText({ model, messages, abortSignal: signal });
+ * const response = await chat.pipeAndCapture(result, { signal });
+ * if (response) conversation.addResponse(response);
+ * ```
+ */
+async function pipeChatAndCapture(
+  source: UIMessageStreamable,
+  options?: { signal?: AbortSignal; spanName?: string }
+): Promise<UIMessage | undefined> {
+  let captured: UIMessage | undefined;
+  let resolveOnFinish: () => void;
+  const onFinishPromise = new Promise<void>((r) => { resolveOnFinish = r; });
+
+  const uiStream = source.toUIMessageStream({
+    onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => {
+      captured = responseMessage;
+      resolveOnFinish!();
+    },
+  });
+
+  await pipeChat(uiStream, { signal: options?.signal, spanName: options?.spanName ?? "stream response" });
+  await onFinishPromise;
+
+  return captured;
+}
+
+/**
+ * Accumulates conversation messages across turns.
+ *
+ * Handles the transport protocol: turn 0 sends full history (replace),
+ * subsequent turns send only new messages (append), regenerate sends
+ * full history minus last assistant message (replace).
+ *
+ * @example
+ * ```ts
+ * const conversation = new chat.MessageAccumulator();
+ * for (let turn = 0; turn < 100; turn++) {
+ *   const messages = await conversation.addIncoming(payload.messages, payload.trigger, turn);
+ *   const result = streamText({ model, messages });
+ *   const response = await chat.pipeAndCapture(result);
+ *   if (response) await conversation.addResponse(response);
+ * }
+ * ```
+ */
+class ChatMessageAccumulator {
+  modelMessages: ModelMessage[] = [];
+  uiMessages: UIMessage[] = [];
+
+  /**
+   * Add incoming messages from the transport payload.
+   * Returns the full accumulated model messages for `streamText`.
+   */
+  async addIncoming(
+    messages: UIMessage[],
+    trigger: string,
+    turn: number
+  ): Promise<ModelMessage[]> {
+    const cleaned = messages.map((m) =>
+      m.role === "assistant" ? cleanupAbortedParts(m) : m
+    );
+    const model = await convertToModelMessages(cleaned);
+
+    if (turn === 0 || trigger === "regenerate-message") {
+      this.modelMessages = model;
+      this.uiMessages = [...cleaned];
+    } else {
+      this.modelMessages.push(...model);
+      this.uiMessages.push(...cleaned);
+    }
+    return this.modelMessages;
+  }
+
+  /**
+   * Add the assistant's response to the accumulator.
+   * Call after `pipeAndCapture` with the captured response.
+   */
+  async addResponse(response: UIMessage): Promise<void> {
+    if (!response.id) {
+      response = { ...response, id: generateMessageId() };
+    }
+    this.uiMessages.push(response);
+    try {
+      const msgs = await convertToModelMessages([stripProviderMetadata(response)]);
+      this.modelMessages.push(...msgs);
+    } catch {
+      // Conversion failed — skip model message accumulation for this response
+    }
+  }
+}
+
 // ---------------------------------------------------------------------------
 // chat.local — per-run typed data with Proxy access
 // ---------------------------------------------------------------------------
@@ -1825,6 +1975,16 @@ export const chat = {
   defer: chatDefer,
   /** Typed chat output stream for writing custom chunks or piping from subtasks. */
   stream: chatStream,
+  /** Pre-built input stream for receiving messages from the transport. */
+  messages: messagesInput,
+  /** Create a managed stop signal wired to the stop input stream. See {@link createStopSignal}. */
+  createStopSignal,
+  /** Signal the frontend that the current turn is complete. See {@link chatWriteTurnComplete}. */
+  writeTurnComplete: chatWriteTurnComplete,
+  /** Pipe a stream and capture the response message. See {@link pipeChatAndCapture}. */
+  pipeAndCapture: pipeChatAndCapture,
+  /** Message accumulator class for raw task chat. See {@link ChatMessageAccumulator}. */
+  MessageAccumulator: ChatMessageAccumulator,
 };
 
 /**

references/ai-chat/src/app/actions.ts

@@ -4,7 +4,8 @@ import { chat } from "@trigger.dev/sdk/ai";
 import type { aiChat } from "@/trigger/chat";
 import { prisma } from "@/lib/prisma";
 
-export const getChatToken = async () => chat.createAccessToken<typeof aiChat>("ai-chat");
+export const getChatToken = async (taskId?: string) =>
+  chat.createAccessToken<typeof aiChat>((taskId ?? "ai-chat") as any);
 
 export async function getChatList() {
   const chats = await prisma.chat.findMany({

references/ai-chat/src/app/page.tsx

@@ -25,6 +25,7 @@ export default function Home() {
     Record<string, { runId: string; publicAccessToken: string; lastEventId?: string }>
   >({});
   const [loaded, setLoaded] = useState(false);
+  const [taskMode, setTaskMode] = useState<string>("ai-chat");
 
   useEffect(() => {
     async function load() {
@@ -50,6 +51,9 @@ export default function Home() {
 
   return (
     <ChatApp
+      key={taskMode}
+      taskMode={taskMode}
+      onTaskModeChange={setTaskMode}
       initialChatList={chatList}
       initialActiveChatId={activeChatId}
       initialMessages={initialMessages}

references/ai-chat/src/components/chat-app.tsx

@@ -32,13 +32,17 @@ type SessionInfo = {
 };
 
 type ChatAppProps = {
+  taskMode: string;
+  onTaskModeChange: (mode: string) => void;
   initialChatList: ChatMeta[];
   initialActiveChatId: string | null;
   initialMessages: UIMessage[];
   initialSessions: Record<string, SessionInfo>;
 };
 
 export function ChatApp({
+  taskMode,
+  onTaskModeChange,
   initialChatList,
   initialActiveChatId,
   initialMessages,
@@ -70,9 +74,9 @@ export function ChatApp({
     []
   );
 
-  const transport = useTriggerChatTransport<typeof aiChat>({
-    task: "ai-chat",
-    accessToken: getChatToken,
+  const transport = useTriggerChatTransport({
+    task: taskMode,
+    accessToken: () => getChatToken(taskMode),
     baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
     sessions: initialSessions,
     onSessionChange: handleSessionChange,
@@ -157,6 +161,8 @@ export function ChatApp({
         onPreloadChange={setPreloadEnabled}
         warmTimeoutInSeconds={warmTimeoutInSeconds}
         onWarmTimeoutChange={setWarmTimeoutInSeconds}
+        taskMode={taskMode}
+        onTaskModeChange={onTaskModeChange}
       />
       <div className="flex-1">
         {activeChatId ? (

references/ai-chat/src/components/chat-sidebar.tsx

@@ -28,6 +28,8 @@ type ChatSidebarProps = {
   onPreloadChange: (enabled: boolean) => void;
   warmTimeoutInSeconds: number;
   onWarmTimeoutChange: (seconds: number) => void;
+  taskMode: string;
+  onTaskModeChange: (mode: string) => void;
 };
 
 export function ChatSidebar({
@@ -40,6 +42,8 @@ export function ChatSidebar({
   onPreloadChange,
   warmTimeoutInSeconds,
   onWarmTimeoutChange,
+  taskMode,
+  onTaskModeChange,
 }: ChatSidebarProps) {
   const sorted = [...chats].sort((a, b) => b.updatedAt - a.updatedAt);
 
@@ -108,6 +112,17 @@ export function ChatSidebar({
           />
           <span>s</span>
         </div>
+        <div className="flex items-center gap-2 text-xs text-gray-500">
+          <span className="shrink-0">Task</span>
+          <select
+            value={taskMode}
+            onChange={(e) => onTaskModeChange(e.target.value)}
+            className="flex-1 rounded border border-gray-300 px-1.5 py-0.5 text-xs text-gray-600 outline-none focus:border-blue-500"
+          >
+            <option value="ai-chat">ai-chat (chat.task)</option>
+            <option value="ai-chat-raw">ai-chat-raw (raw task)</option>
+          </select>
+        </div>
       </div>
     </div>
   );