@trigger.dev/sdk 0.0.0-prerelease-20260306173130 → 0.0.0-prerelease-20260310150337

package/dist/commonjs/v3/ai.js

@@ -14,6 +14,7 @@ const chat_constants_js_1 = require("./chat-constants.js");
 Object.defineProperty(exports, "CHAT_MESSAGES_STREAM_ID", { enumerable: true, get: function () { return chat_constants_js_1.CHAT_MESSAGES_STREAM_ID; } });
 Object.defineProperty(exports, "CHAT_STOP_STREAM_ID", { enumerable: true, get: function () { return chat_constants_js_1.CHAT_STOP_STREAM_ID; } });
 const METADATA_KEY = "tool.execute.options";
+const chatTurnContextKey = locals_js_1.locals.create("chat.turnContext");
 function toolFromTask(task, options) {
     if (("schema" in task && !task.schema) || ("jsonSchema" in task && !task.jsonSchema)) {
         throw new Error("Cannot convert this task to to a tool because the task has no schema. Make sure to either use schemaTask or a task with an input jsonSchema.");
@@ -22,12 +23,43 @@ function toolFromTask(task, options) {
         description: task.description,
         inputSchema: convertTaskSchemaToToolParameters(task),
         execute: async (input, options) => {
-            const serializedOptions = options ? JSON.parse(JSON.stringify(options)) : undefined;
+            // Build tool metadata — skip messages (can be large) and abortSignal (non-serializable)
+            const toolMeta = {
+                toolCallId: options?.toolCallId ?? "",
+            };
+            if (options?.experimental_context !== undefined) {
+                try {
+                    toolMeta.experimental_context = JSON.parse(JSON.stringify(options.experimental_context));
+                }
+                catch {
+                    // Non-serializable context — skip
+                }
+            }
+            // Auto-detect chat context from the parent turn
+            const chatCtx = locals_js_1.locals.get(chatTurnContextKey);
+            if (chatCtx) {
+                toolMeta.chatId = chatCtx.chatId;
+                toolMeta.turn = chatCtx.turn;
+                toolMeta.continuation = chatCtx.continuation;
+                toolMeta.clientData = chatCtx.clientData;
+            }
+            // Serialize initialized chat.local values for subtask hydration
+            const chatLocals = {};
+            for (const entry of chatLocalRegistry) {
+                const value = locals_js_1.locals.get(entry.key);
+                if (value !== undefined) {
+                    chatLocals[entry.id] = value;
+                }
+            }
+            if (Object.keys(chatLocals).length > 0) {
+                toolMeta.chatLocals = chatLocals;
+            }
             return await task
                 .triggerAndWait(input, {
                 metadata: {
-                    [METADATA_KEY]: serializedOptions,
+                    [METADATA_KEY]: toolMeta,
                 },
+                tags: options?.toolCallId ? [`toolCallId:${options.toolCallId}`] : undefined,
             })
                 .unwrap();
         },
@@ -42,6 +74,53 @@ function getToolOptionsFromMetadata() {
     }
     return tool;
 }
+/**
+ * Get the current tool call ID from inside a subtask invoked via `ai.tool()`.
+ * Returns `undefined` if not running as a tool subtask.
+ */
+function getToolCallId() {
+    return getToolOptionsFromMetadata()?.toolCallId;
+}
+/**
+ * Get the chat context from inside a subtask invoked via `ai.tool()` within a `chat.task`.
+ * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`.
+ * Returns `undefined` if the parent is not a chat task.
+ *
+ * @example
+ * ```ts
+ * const ctx = ai.chatContext<typeof myChat>();
+ * // ctx?.clientData is typed based on myChat's clientDataSchema
+ * ```
+ */
+function getToolChatContext() {
+    const opts = getToolOptionsFromMetadata();
+    if (!opts?.chatId)
+        return undefined;
+    return {
+        chatId: opts.chatId,
+        turn: opts.turn ?? 0,
+        continuation: opts.continuation ?? false,
+        clientData: opts.clientData,
+    };
+}
+/**
+ * Get the chat context from inside a subtask, throwing if not in a chat context.
+ * Pass `typeof yourChatTask` as the type parameter to get typed `clientData`.
+ *
+ * @example
+ * ```ts
+ * const ctx = ai.chatContextOrThrow<typeof myChat>();
+ * // ctx.chatId, ctx.clientData are guaranteed non-null
+ * ```
+ */
+function getToolChatContextOrThrow() {
+    const ctx = getToolChatContext();
+    if (!ctx) {
+        throw new Error("ai.chatContextOrThrow() called outside of a chat.task context. " +
+            "This helper can only be used inside a subtask invoked via ai.tool() from a chat.task.");
+    }
+    return ctx;
+}
 function convertTaskSchemaToToolParameters(task) {
     if ("schema" in task) {
         // If TaskSchema is ArkTypeEsque, use ai.jsonSchema to convert it to a Schema
@@ -61,6 +140,12 @@ function convertTaskSchemaToToolParameters(task) {
 exports.ai = {
     tool: toolFromTask,
     currentToolOptions: getToolOptionsFromMetadata,
+    /** Get the tool call ID from inside a subtask invoked via `ai.tool()`. */
+    toolCallId: getToolCallId,
+    /** Get chat context (chatId, turn, clientData, etc.) from inside a subtask of a `chat.task`. Returns undefined if not in a chat context. */
+    chatContext: getToolChatContext,
+    /** Get chat context or throw if not in a chat context. Pass `typeof yourChatTask` for typed clientData. */
+    chatContextOrThrow: getToolChatContextOrThrow,
 };
 /**
  * Creates a public access token for a chat task.
@@ -91,15 +176,48 @@ function createChatAccessToken(taskId) {
  * use this key by default.
  */
 exports.CHAT_STREAM_KEY = chat_constants_js_1.CHAT_STREAM_KEY;
+/**
+ * Typed chat output stream. Provides `.writer()`, `.pipe()`, `.append()`,
+ * and `.read()` methods pre-bound to the chat stream key and typed to `UIMessageChunk`.
+ *
+ * Use from within a `chat.task` run to write custom chunks:
+ * ```ts
+ * const { waitUntilComplete } = chat.stream.writer({
+ *   execute: ({ write }) => {
+ *     write({ type: "text-start", id: "status-1" });
+ *     write({ type: "text-delta", id: "status-1", delta: "Processing..." });
+ *     write({ type: "text-end", id: "status-1" });
+ *   },
+ * });
+ * await waitUntilComplete();
+ * ```
+ *
+ * Use from a subtask to stream back to the parent chat:
+ * ```ts
+ * chat.stream.pipe(myStream, { target: "root" });
+ * ```
+ */
+const chatStream = streams_js_1.streams.define({ id: chat_constants_js_1.CHAT_STREAM_KEY });
 // Input streams for bidirectional chat communication
 const messagesInput = streams_js_1.streams.input({ id: chat_constants_js_1.CHAT_MESSAGES_STREAM_ID });
 const stopInput = streams_js_1.streams.input({ id: chat_constants_js_1.CHAT_STOP_STREAM_ID });
+/**
+ * Per-turn deferred promises. Registered via `chat.defer()`, awaited
+ * before `onTurnComplete` fires. Reset each turn.
+ * @internal
+ */
+const chatDeferKey = locals_js_1.locals.create("chat.defer");
 /**
  * Run-scoped pipe counter. Stored in locals so concurrent runs in the
  * same worker don't share state.
  * @internal
  */
 const chatPipeCountKey = locals_js_1.locals.create("chat.pipeCount");
+const chatStopControllerKey = locals_js_1.locals.create("chat.stopController");
+/** Static (task-level) UIMessageStream options, set once during chatTask setup. @internal */
+const chatUIStreamStaticKey = locals_js_1.locals.create("chat.uiMessageStreamOptions.static");
+/** Per-turn UIMessageStream options, set via chat.setUIMessageStreamOptions(). @internal */
+const chatUIStreamPerTurnKey = locals_js_1.locals.create("chat.uiMessageStreamOptions.perTurn");
 function isUIMessageStreamable(value) {
     return (typeof value === "object" &&
         value !== null &&
@@ -206,7 +324,7 @@ async function pipeChat(source, options) {
  * ```
  */
 function chatTask(options) {
-    const { run: userRun, clientDataSchema, onChatStart, onTurnStart, onTurnComplete, maxTurns = 100, turnTimeout = "1h", warmTimeoutInSeconds = 30, chatAccessTokenTTL = "1h", ...restOptions } = options;
+    const { run: userRun, clientDataSchema, onPreload, onChatStart, onTurnStart, onTurnComplete, maxTurns = 100, turnTimeout = "1h", warmTimeoutInSeconds = 30, chatAccessTokenTTL = "1h", preloadWarmTimeoutInSeconds, preloadTimeout, uiMessageStreamOptions, ...restOptions } = options;
     const parseClientData = clientDataSchema
         ? (0, v3_1.getSchemaParseFn)(clientDataSchema)
         : undefined;
@@ -218,7 +336,14 @@ function chatTask(options) {
             if (activeSpan) {
                 activeSpan.setAttribute("gen_ai.conversation.id", payload.chatId);
             }
+            // Store static UIMessageStream options in locals so resolveUIMessageStreamOptions() can read them
+            if (uiMessageStreamOptions) {
+                locals_js_1.locals.set(chatUIStreamStaticKey, uiMessageStreamOptions);
+            }
             let currentWirePayload = payload;
+            const continuation = payload.continuation ?? false;
+            const previousRunId = payload.previousRunId;
+            const preloaded = payload.trigger === "preload";
             // Accumulated model messages across turns. Turn 1 initialises from the
             // full history the frontend sends; subsequent turns append only the new
             // user message(s) and the captured assistant response.
@@ -234,6 +359,67 @@ function chatTask(options) {
                 currentStopController?.abort(data?.message || "stopped");
             });
             try {
+                // Handle preloaded runs — fire onPreload, then wait for the first real message
+                if (preloaded) {
+                    if (activeSpan) {
+                        activeSpan.setAttribute("chat.preloaded", true);
+                    }
+                    const currentRunId = v3_1.taskContext.ctx?.run.id ?? "";
+                    let preloadAccessToken = "";
+                    if (currentRunId) {
+                        try {
+                            preloadAccessToken = await auth_js_1.auth.createPublicToken({
+                                scopes: {
+                                    read: { runs: currentRunId },
+                                    write: { inputStreams: currentRunId },
+                                },
+                                expirationTime: chatAccessTokenTTL,
+                            });
+                        }
+                        catch {
+                            // Token creation failed
+                        }
+                    }
+                    // Parse client data for the preload hook
+                    const preloadClientData = (parseClientData
+                        ? await parseClientData(payload.metadata)
+                        : payload.metadata);
+                    // Fire onPreload hook
+                    if (onPreload) {
+                        await tracer_js_1.tracer.startActiveSpan("onPreload()", async () => {
+                            await onPreload({
+                                chatId: payload.chatId,
+                                runId: currentRunId,
+                                chatAccessToken: preloadAccessToken,
+                                clientData: preloadClientData,
+                            });
+                        }, {
+                            attributes: {
+                                [v3_1.SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
+                                [v3_1.SemanticInternalAttributes.COLLAPSED]: true,
+                                "chat.id": payload.chatId,
+                                "chat.preloaded": true,
+                            },
+                        });
+                    }
+                    // Wait for the first real message — use preload-specific timeouts if configured
+                    const effectivePreloadWarmTimeout = payload.warmTimeoutInSeconds
+                        ?? preloadWarmTimeoutInSeconds
+                        ?? warmTimeoutInSeconds;
+                    const effectivePreloadTimeout = metadata_js_1.metadata.get(TURN_TIMEOUT_METADATA_KEY)
+                        ?? preloadTimeout
+                        ?? turnTimeout;
+                    const preloadResult = await messagesInput.waitWithWarmup({
+                        warmTimeoutInSeconds: effectivePreloadWarmTimeout,
+                        timeout: effectivePreloadTimeout,
+                        spanName: "waiting for first message",
+                    });
+                    if (!preloadResult.ok) {
+                        return; // Timed out waiting for first message — end run
+                    }
+                    let firstMessage = preloadResult.output;
+                    currentWirePayload = firstMessage;
+                }
                 for (let turn = 0; turn < maxTurns; turn++) {
                     // Extract turn-level context before entering the span
                     const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload;
@@ -266,9 +452,18 @@ function chatTask(options) {
                     }
                     const turnResult = await tracer_js_1.tracer.startActiveSpan(`chat turn ${turn + 1}`, async () => {
                         locals_js_1.locals.set(chatPipeCountKey, 0);
+                        locals_js_1.locals.set(chatDeferKey, new Set());
+                        // Store chat context for auto-detection by ai.tool subtasks
+                        locals_js_1.locals.set(chatTurnContextKey, {
+                            chatId: currentWirePayload.chatId,
+                            turn,
+                            continuation,
+                            clientData,
+                        });
                         // Per-turn stop controller (reset each turn)
                         const stopController = new AbortController();
                         currentStopController = stopController;
+                        locals_js_1.locals.set(chatStopControllerKey, stopController);
                         // Three signals for the user's run function
                         const stopSignal = stopController.signal;
                         const cancelSignal = runSignal;
@@ -278,20 +473,26 @@ function chatTask(options) {
                         const msgSub = messagesInput.on((msg) => {
                             pendingMessages.push(msg);
                         });
+                        // Clean up any incomplete tool parts in the incoming history.
+                        // When a previous run was stopped mid-tool-call, the frontend's
+                        // useChat state may still contain assistant messages with tool parts
+                        // in partial/input-available state. These cause API errors (e.g.
+                        // Anthropic requires every tool_use to have a matching tool_result).
+                        const cleanedUIMessages = uiMessages.map((msg) => msg.role === "assistant" ? cleanupAbortedParts(msg) : msg);
                         // Convert the incoming UIMessages to model messages and update the accumulator.
                         // Turn 1: full history from the frontend → replaces the accumulator.
                         // Turn 2+: only the new message(s) → appended to the accumulator.
-                        const incomingModelMessages = await (0, ai_1.convertToModelMessages)(uiMessages);
+                        const incomingModelMessages = await (0, ai_1.convertToModelMessages)(cleanedUIMessages);
                         // Track new messages for this turn (user input + assistant response).
                         const turnNewModelMessages = [];
                         const turnNewUIMessages = [];
                         if (turn === 0) {
                             accumulatedMessages = incomingModelMessages;
-                            accumulatedUIMessages = [...uiMessages];
+                            accumulatedUIMessages = [...cleanedUIMessages];
                             // On first turn, the "new" messages are just the last user message
                             // (the rest is history). We'll add the response after streaming.
-                            if (uiMessages.length > 0) {
-                                turnNewUIMessages.push(uiMessages[uiMessages.length - 1]);
+                            if (cleanedUIMessages.length > 0) {
+                                turnNewUIMessages.push(cleanedUIMessages[cleanedUIMessages.length - 1]);
                                 const lastModel = incomingModelMessages[incomingModelMessages.length - 1];
                                 if (lastModel)
                                     turnNewModelMessages.push(lastModel);
@@ -301,15 +502,15 @@ function chatTask(options) {
                             // Regenerate: frontend sent full history with last assistant message
                             // removed. Reset the accumulator to match.
                             accumulatedMessages = incomingModelMessages;
-                            accumulatedUIMessages = [...uiMessages];
+                            accumulatedUIMessages = [...cleanedUIMessages];
                             // No new user messages for regenerate — just the response (added below)
                         }
                         else {
                             // Submit: frontend sent only the new user message(s). Append to accumulator.
                             accumulatedMessages.push(...incomingModelMessages);
-                            accumulatedUIMessages.push(...uiMessages);
+                            accumulatedUIMessages.push(...cleanedUIMessages);
                             turnNewModelMessages.push(...incomingModelMessages);
-                            turnNewUIMessages.push(...uiMessages);
+                            turnNewUIMessages.push(...cleanedUIMessages);
                         }
                         // Mint a scoped public access token once per turn, reused for
                         // onChatStart, onTurnStart, onTurnComplete, and the turn-complete chunk.
@@ -338,11 +539,19 @@ function chatTask(options) {
                                     clientData,
                                     runId: currentRunId,
                                     chatAccessToken: turnAccessToken,
+                                    continuation,
+                                    previousRunId,
+                                    preloaded,
                                 });
                             }, {
                                 attributes: {
                                     [v3_1.SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
                                     [v3_1.SemanticInternalAttributes.COLLAPSED]: true,
+                                    "chat.id": currentWirePayload.chatId,
+                                    "chat.messages.count": accumulatedMessages.length,
+                                    "chat.continuation": continuation,
+                                    "chat.preloaded": preloaded,
+                                    ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}),
                                 },
                             });
                         }
@@ -358,21 +567,42 @@ function chatTask(options) {
                                     runId: currentRunId,
                                     chatAccessToken: turnAccessToken,
                                     clientData,
+                                    continuation,
+                                    previousRunId,
+                                    preloaded,
                                 });
                             }, {
                                 attributes: {
                                     [v3_1.SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
                                     [v3_1.SemanticInternalAttributes.COLLAPSED]: true,
+                                    "chat.id": currentWirePayload.chatId,
+                                    "chat.turn": turn + 1,
+                                    "chat.messages.count": accumulatedMessages.length,
+                                    "chat.trigger": currentWirePayload.trigger,
+                                    "chat.continuation": continuation,
+                                    "chat.preloaded": preloaded,
+                                    ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}),
                                 },
                             });
                         }
                         // Captured by the onFinish callback below — works even on abort/stop.
                         let capturedResponseMessage;
+                        // Promise that resolves when the AI SDK's onFinish fires.
+                        // On abort, the stream's cancel() handler calls onFinish
+                        // asynchronously AFTER pipeChat resolves, so we must await
+                        // this to avoid a race where we check capturedResponseMessage
+                        // before it's been set.
+                        let resolveOnFinish;
+                        const onFinishPromise = new Promise((r) => { resolveOnFinish = r; });
+                        let onFinishAttached = false;
                         try {
                             const result = await userRun({
                                 ...restWire,
                                 messages: accumulatedMessages,
                                 clientData,
+                                continuation,
+                                previousRunId,
+                                preloaded,
                                 signal: combinedSignal,
                                 cancelSignal,
                                 stopSignal,
@@ -381,9 +611,12 @@ function chatTask(options) {
                             // but only if pipeChat() wasn't already called manually during this turn.
                             // We call toUIMessageStream ourselves to attach onFinish for response capture.
                             if ((locals_js_1.locals.get(chatPipeCountKey) ?? 0) === 0 && isUIMessageStreamable(result)) {
+                                onFinishAttached = true;
                                 const uiStream = result.toUIMessageStream({
+                                    ...resolveUIMessageStreamOptions(),
                                     onFinish: ({ responseMessage }) => {
                                         capturedResponseMessage = responseMessage;
+                                        resolveOnFinish();
                                     },
                                 });
                                 await pipeChat(uiStream, { signal: combinedSignal, spanName: "stream response" });
@@ -404,10 +637,24 @@ function chatTask(options) {
                         finally {
                             msgSub.off();
                         }
+                        // Wait for onFinish to fire — on abort this may resolve slightly
+                        // after pipeChat, since the stream's cancel() handler is async.
+                        if (onFinishAttached) {
+                            await onFinishPromise;
+                        }
+                        // Determine if the user stopped generation this turn (not a full run cancel).
+                        const wasStopped = stopController.signal.aborted && !runSignal.aborted;
                         // Append the assistant's response (partial or complete) to the accumulator.
                         // The onFinish callback fires even on abort/stop, so partial responses
                         // from stopped generation are captured correctly.
+                        let rawResponseMessage;
                         if (capturedResponseMessage) {
+                            // Keep the raw message before cleanup for users who want custom handling
+                            rawResponseMessage = capturedResponseMessage;
+                            // Clean up aborted parts (streaming tool calls, reasoning) when stopped
+                            if (wasStopped) {
+                                capturedResponseMessage = cleanupAbortedParts(capturedResponseMessage);
+                            }
                             // Ensure the response message has an ID (the stream's onFinish
                             // may produce a message with an empty ID since IDs are normally
                             // assigned by the frontend's useChat).
@@ -436,6 +683,15 @@ function chatTask(options) {
                         // Write turn-complete control chunk so frontend closes its stream.
                         // Capture the lastEventId from the stream writer for resume support.
                         const turnCompleteResult = await writeTurnCompleteChunk(currentWirePayload.chatId, turnAccessToken);
+                        // Await deferred background work (e.g. DB writes from onTurnStart)
+                        // before firing onTurnComplete so hooks can rely on the work being done.
+                        const deferredWork = locals_js_1.locals.get(chatDeferKey);
+                        if (deferredWork && deferredWork.size > 0) {
+                            await Promise.race([
+                                Promise.allSettled(deferredWork),
+                                new Promise((r) => setTimeout(r, 5_000)),
+                            ]);
+                        }
                         // Fire onTurnComplete after response capture
                         if (onTurnComplete) {
                             await tracer_js_1.tracer.startActiveSpan("onTurnComplete()", async () => {
@@ -446,16 +702,30 @@ function chatTask(options) {
                                     newMessages: turnNewModelMessages,
                                     newUIMessages: turnNewUIMessages,
                                     responseMessage: capturedResponseMessage,
+                                    rawResponseMessage,
                                     turn,
                                     runId: currentRunId,
                                     chatAccessToken: turnAccessToken,
                                     lastEventId: turnCompleteResult.lastEventId,
                                     clientData,
+                                    stopped: wasStopped,
+                                    continuation,
+                                    previousRunId,
+                                    preloaded,
                                 });
                             }, {
                                 attributes: {
                                     [v3_1.SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete",
                                     [v3_1.SemanticInternalAttributes.COLLAPSED]: true,
+                                    "chat.id": currentWirePayload.chatId,
+                                    "chat.turn": turn + 1,
+                                    "chat.stopped": wasStopped,
+                                    "chat.continuation": continuation,
+                                    "chat.preloaded": preloaded,
+                                    ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}),
+                                    "chat.messages.count": accumulatedMessages.length,
+                                    "chat.response.parts.count": capturedResponseMessage?.parts?.length ?? 0,
+                                    "chat.new_messages.count": turnNewUIMessages.length,
                                 },
                             });
                         }
@@ -464,28 +734,15 @@ function chatTask(options) {
                             currentWirePayload = pendingMessages[0];
                             return "continue";
                         }
-                        // Phase 1: Keep the run warm for quick response to the next message.
-                        // The run stays active (using compute) during this window.
+                        // Wait for the next message — stay warm briefly, then suspend
                         const effectiveWarmTimeout = metadata_js_1.metadata.get(WARM_TIMEOUT_METADATA_KEY) ?? warmTimeoutInSeconds;
-                        if (effectiveWarmTimeout > 0) {
-                            const warm = await messagesInput.once({
-                                timeoutMs: effectiveWarmTimeout * 1000,
-                                spanName: "waiting (warm)",
-                            });
-                            if (warm.ok) {
-                                // Message arrived while warm — respond instantly
-                                currentWirePayload = warm.output;
-                                return "continue";
-                            }
-                        }
-                        // Phase 2: Suspend the task (frees compute) until the next message arrives
                         const effectiveTurnTimeout = metadata_js_1.metadata.get(TURN_TIMEOUT_METADATA_KEY) ?? turnTimeout;
-                        const next = await messagesInput.wait({
+                        const next = await messagesInput.waitWithWarmup({
+                            warmTimeoutInSeconds: effectiveWarmTimeout,
                             timeout: effectiveTurnTimeout,
-                            spanName: "waiting (suspended)",
+                            spanName: "waiting for next message",
                         });
                         if (!next.ok) {
-                            // Timed out waiting for the next message — end the conversation
                             return "exit";
                         }
                         currentWirePayload = next.output;
@@ -589,6 +846,424 @@ function setTurnTimeoutInSeconds(seconds) {
 function setWarmTimeoutInSeconds(seconds) {
     metadata_js_1.metadata.set(WARM_TIMEOUT_METADATA_KEY, seconds);
 }
+/**
+ * Override the `toUIMessageStream()` options for the current turn.
+ *
+ * These options control how the `StreamTextResult` is converted to a
+ * `UIMessageChunk` stream — error handling, reasoning/source visibility,
+ * message metadata, etc.
+ *
+ * Per-turn options are merged on top of the static `uiMessageStreamOptions`
+ * set on `chat.task()`. Per-turn values win on conflicts.
+ *
+ * @example
+ * ```ts
+ * run: async ({ messages, signal }) => {
+ *   chat.setUIMessageStreamOptions({
+ *     sendReasoning: true,
+ *     onError: (error) => error instanceof Error ? error.message : "An error occurred.",
+ *   });
+ *   return streamText({ model, messages, abortSignal: signal });
+ * }
+ * ```
+ */
+function setUIMessageStreamOptions(options) {
+    locals_js_1.locals.set(chatUIStreamPerTurnKey, options);
+}
+/**
+ * Resolve the effective UIMessageStream options by merging:
+ * 1. Static task-level options (from `chat.task({ uiMessageStreamOptions })`)
+ * 2. Per-turn overrides (from `chat.setUIMessageStreamOptions()`)
+ *
+ * Per-turn values win on conflicts. Clears the per-turn override after reading
+ * so it doesn't leak into subsequent turns.
+ * @internal
+ */
+function resolveUIMessageStreamOptions() {
+    const staticOptions = locals_js_1.locals.get(chatUIStreamStaticKey) ?? {};
+    const perTurnOptions = locals_js_1.locals.get(chatUIStreamPerTurnKey) ?? {};
+    // Clear per-turn override so it doesn't leak into subsequent turns
+    locals_js_1.locals.set(chatUIStreamPerTurnKey, undefined);
+    return { ...staticOptions, ...perTurnOptions };
+}
+// ---------------------------------------------------------------------------
+// Stop detection
+// ---------------------------------------------------------------------------
+/**
+ * Check whether the user stopped generation during the current turn.
+ *
+ * Works from **anywhere** inside a `chat.task` run — including inside
+ * `streamText`'s `onFinish` callback — without needing to thread the
+ * `stopSignal` through closures.
+ *
+ * This is especially useful when the AI SDK's `isAborted` flag is unreliable
+ * (e.g. when using `createUIMessageStream` + `writer.merge()`).
+ *
+ * @example
+ * ```ts
+ * onFinish: ({ isAborted }) => {
+ *   const wasStopped = isAborted || chat.isStopped();
+ *   if (wasStopped) {
+ *     // handle stop
+ *   }
+ * }
+ * ```
+ */
+function isStopped() {
+    const controller = locals_js_1.locals.get(chatStopControllerKey);
+    return controller?.signal.aborted ?? false;
+}
+// ---------------------------------------------------------------------------
+// Per-turn deferred work
+// ---------------------------------------------------------------------------
+/**
+ * Register a promise that runs in the background during the current turn.
+ *
+ * Use this to move non-blocking work (DB writes, analytics, etc.) out of
+ * the critical path. The promise runs in parallel with streaming and is
+ * awaited (with a 5 s timeout) before `onTurnComplete` fires.
+ *
+ * @example
+ * ```ts
+ * onTurnStart: async ({ chatId, uiMessages }) => {
+ *   // Persist messages without blocking the LLM call
+ *   chat.defer(db.chat.update({ where: { id: chatId }, data: { messages: uiMessages } }));
+ * },
+ * ```
+ */
+function chatDefer(promise) {
+    const promises = locals_js_1.locals.get(chatDeferKey);
+    if (promises) {
+        promises.add(promise);
+    }
+}
+// ---------------------------------------------------------------------------
+// Aborted message cleanup
+// ---------------------------------------------------------------------------
+/**
+ * Clean up a UIMessage that was captured during an aborted/stopped turn.
+ *
+ * When generation is stopped mid-stream, the captured message may contain:
+ * - Tool parts stuck in incomplete states (`partial-call`, `input-available`,
+ *   `input-streaming`) that cause permanent UI spinners
+ * - Reasoning parts with `state: "streaming"` instead of `"done"`
+ * - Text parts with `state: "streaming"` instead of `"done"`
+ *
+ * This function returns a cleaned copy with:
+ * - Incomplete tool parts removed entirely
+ * - Reasoning and text parts marked as `"done"`
+ *
+ * `chat.task` calls this automatically when stop is detected before passing
+ * the response to `onTurnComplete`. Use this manually when calling `pipeChat`
+ * directly and capturing response messages yourself.
+ *
+ * @example
+ * ```ts
+ * onTurnComplete: async ({ responseMessage, stopped }) => {
+ *   // Already cleaned automatically by chat.task — but if you captured
+ *   // your own message via pipeChat, clean it manually:
+ *   const cleaned = chat.cleanupAbortedParts(myMessage);
+ *   await db.messages.save(cleaned);
+ * }
+ * ```
+ */
+function cleanupAbortedParts(message) {
+    if (!message.parts)
+        return message;
+    const isToolPart = (part) => part.type === "tool-invocation" ||
+        part.type?.startsWith("tool-") ||
+        part.type === "dynamic-tool";
+    return {
+        ...message,
+        parts: message.parts
+            .filter((part) => {
+            if (!isToolPart(part))
+                return true;
+            // Remove tool parts that never completed execution.
+            // partial-call: input was still streaming when aborted.
+            // input-available: input was complete but tool never ran.
+            // input-streaming: input was mid-stream.
+            const state = part.toolInvocation?.state ?? part.state;
+            return state !== "partial-call" && state !== "input-available" && state !== "input-streaming";
+        })
+            .map((part) => {
+            // Mark streaming reasoning as done
+            if (part.type === "reasoning" && part.state === "streaming") {
+                return { ...part, state: "done" };
+            }
+            // Mark streaming text as done
+            if (part.type === "text" && part.state === "streaming") {
+                return { ...part, state: "done" };
+            }
+            return part;
+        }),
+    };
+}
+// ---------------------------------------------------------------------------
+// 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() {
+    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) {
+    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, options) {
+    let captured;
+    let resolveOnFinish;
+    const onFinishPromise = new Promise((r) => { resolveOnFinish = r; });
+    const uiStream = source.toUIMessageStream({
+        ...resolveUIMessageStreamOptions(),
+        onFinish: ({ responseMessage }) => {
+            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 = [];
+    uiMessages = [];
+    /**
+     * Add incoming messages from the transport payload.
+     * Returns the full accumulated model messages for `streamText`.
+     */
+    async addIncoming(messages, trigger, turn) {
+        const cleaned = messages.map((m) => m.role === "assistant" ? cleanupAbortedParts(m) : m);
+        const model = await (0, ai_1.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) {
+        if (!response.id) {
+            response = { ...response, id: (0, ai_1.generateId)() };
+        }
+        this.uiMessages.push(response);
+        try {
+            const msgs = await (0, ai_1.convertToModelMessages)([stripProviderMetadata(response)]);
+            this.modelMessages.push(...msgs);
+        }
+        catch {
+            // Conversion failed — skip model message accumulation for this response
+        }
+    }
+}
+/**
+ * Create a chat session that yields turns as an async iterator.
+ *
+ * Handles: preload wait, stop signals, message accumulation, turn-complete
+ * signaling, and warm/suspend between turns. You control: initialization,
+ * model/tool selection, persistence, and any custom per-turn logic.
+ *
+ * @example
+ * ```ts
+ * import { task } from "@trigger.dev/sdk";
+ * import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
+ * import { streamText } from "ai";
+ * import { openai } from "@ai-sdk/openai";
+ *
+ * export const myChat = task({
+ *   id: "my-chat",
+ *   run: async (payload: ChatTaskWirePayload, { signal }) => {
+ *     const session = chat.createSession(payload, { signal });
+ *
+ *     for await (const turn of session) {
+ *       const result = streamText({
+ *         model: openai("gpt-4o"),
+ *         messages: turn.messages,
+ *         abortSignal: turn.signal,
+ *       });
+ *       await turn.complete(result);
+ *     }
+ *   },
+ * });
+ * ```
+ */
+function createChatSession(payload, options) {
+    const { signal: runSignal, warmTimeoutInSeconds = 30, timeout = "1h", maxTurns = 100, } = options;
+    return {
+        [Symbol.asyncIterator]() {
+            let currentPayload = payload;
+            let turn = -1;
+            const stop = createStopSignal();
+            const accumulator = new ChatMessageAccumulator();
+            return {
+                async next() {
+                    turn++;
+                    // First turn: handle preload — wait for the first real message
+                    if (turn === 0 && currentPayload.trigger === "preload") {
+                        const result = await messagesInput.waitWithWarmup({
+                            warmTimeoutInSeconds: currentPayload.warmTimeoutInSeconds ?? warmTimeoutInSeconds,
+                            timeout,
+                            spanName: "waiting for first message",
+                        });
+                        if (!result.ok || runSignal.aborted) {
+                            stop.cleanup();
+                            return { done: true, value: undefined };
+                        }
+                        currentPayload = result.output;
+                    }
+                    // Subsequent turns: wait for the next message
+                    if (turn > 0) {
+                        const next = await messagesInput.waitWithWarmup({
+                            warmTimeoutInSeconds,
+                            timeout,
+                            spanName: "waiting for next message",
+                        });
+                        if (!next.ok || runSignal.aborted) {
+                            stop.cleanup();
+                            return { done: true, value: undefined };
+                        }
+                        currentPayload = next.output;
+                    }
+                    // Check limits
+                    if (turn >= maxTurns || runSignal.aborted) {
+                        stop.cleanup();
+                        return { done: true, value: undefined };
+                    }
+                    // Reset stop signal for this turn
+                    stop.reset();
+                    // Accumulate messages
+                    const messages = await accumulator.addIncoming(currentPayload.messages, currentPayload.trigger, turn);
+                    const combinedSignal = AbortSignal.any([runSignal, stop.signal]);
+                    const turnObj = {
+                        number: turn,
+                        chatId: currentPayload.chatId,
+                        trigger: currentPayload.trigger,
+                        clientData: currentPayload.metadata,
+                        messages,
+                        uiMessages: accumulator.uiMessages,
+                        signal: combinedSignal,
+                        get stopped() { return stop.signal.aborted && !runSignal.aborted; },
+                        continuation: currentPayload.continuation ?? false,
+                        async complete(source) {
+                            let response;
+                            try {
+                                response = await pipeChatAndCapture(source, { signal: combinedSignal });
+                            }
+                            catch (error) {
+                                if (error instanceof Error && error.name === "AbortError") {
+                                    if (runSignal.aborted) {
+                                        // Full cancel — don't accumulate
+                                        await chatWriteTurnComplete();
+                                        return undefined;
+                                    }
+                                    // Stop — fall through to accumulate partial response
+                                }
+                                else {
+                                    throw error;
+                                }
+                            }
+                            if (response) {
+                                const cleaned = (stop.signal.aborted && !runSignal.aborted)
+                                    ? cleanupAbortedParts(response)
+                                    : response;
+                                await accumulator.addResponse(cleaned);
+                            }
+                            await chatWriteTurnComplete();
+                            return response;
+                        },
+                        async addResponse(response) {
+                            await accumulator.addResponse(response);
+                        },
+                        async done() {
+                            await chatWriteTurnComplete();
+                        },
+                    };
+                    return { done: false, value: turnObj };
+                },
+                async return() {
+                    stop.cleanup();
+                    return { done: true, value: undefined };
+                },
+            };
+        },
+    };
+}
 // ---------------------------------------------------------------------------
 // chat.local — per-run typed data with Proxy access
 // ---------------------------------------------------------------------------
@@ -596,8 +1271,25 @@ function setWarmTimeoutInSeconds(seconds) {
 const CHAT_LOCAL_KEY = Symbol("chatLocalKey");
 /** @internal Symbol for storing the dirty-tracking locals key. */
 const CHAT_LOCAL_DIRTY_KEY = Symbol("chatLocalDirtyKey");
-/** @internal Counter for generating unique locals IDs. */
-let chatLocalCounter = 0;
+const chatLocalRegistry = new Set();
+/** @internal Run-scoped flag to ensure hydration happens at most once per run. */
+const chatLocalsHydratedKey = locals_js_1.locals.create("chat.locals.hydrated");
+/**
+ * Hydrate chat.local values from subtask metadata (set by toolFromTask).
+ * Runs once per run — subsequent calls are no-ops.
+ * @internal
+ */
+function hydrateLocalsFromMetadata() {
+    if (locals_js_1.locals.get(chatLocalsHydratedKey))
+        return;
+    locals_js_1.locals.set(chatLocalsHydratedKey, true);
+    const opts = metadata_js_1.metadata.get(METADATA_KEY);
+    if (!opts?.chatLocals)
+        return;
+    for (const [id, value] of Object.entries(opts.chatLocals)) {
+        locals_js_1.locals.set(locals_js_1.locals.create(id), value);
+    }
+}
 /**
  * Creates a per-run typed data object accessible from anywhere during task execution.
  *
@@ -606,12 +1298,16 @@ let chatLocalCounter = 0;
  *
  * Multiple locals can coexist — each gets its own isolated run-scoped storage.
  *
+ * The `id` is required and must be unique across all `chat.local()` calls in
+ * your project. It's used to serialize values into subtask metadata so that
+ * `ai.tool()` subtasks can auto-hydrate parent locals (read-only).
+ *
  * @example
  * ```ts
  * import { chat } from "@trigger.dev/sdk/ai";
  *
- * const userPrefs = chat.local<{ theme: string; language: string }>();
- * const gameState = chat.local<{ score: number; streak: number }>();
+ * const userPrefs = chat.local<{ theme: string; language: string }>({ id: "userPrefs" });
+ * const gameState = chat.local<{ score: number; streak: number }>({ id: "gameState" });
  *
  * export const myChat = chat.task({
  *   id: "my-chat",
@@ -635,9 +1331,11 @@ let chatLocalCounter = 0;
  * });
  * ```
  */
-function chatLocal() {
-    const localKey = locals_js_1.locals.create(`chat.local.${chatLocalCounter++}`);
-    const dirtyKey = locals_js_1.locals.create(`chat.local.${chatLocalCounter++}.dirty`);
+function chatLocal(options) {
+    const id = `chat.local.${options.id}`;
+    const localKey = locals_js_1.locals.create(id);
+    const dirtyKey = locals_js_1.locals.create(`${id}.dirty`);
+    chatLocalRegistry.add({ key: localKey, id });
     const target = {};
     target[CHAT_LOCAL_KEY] = localKey;
     target[CHAT_LOCAL_DIRTY_KEY] = dirtyKey;
@@ -664,7 +1362,11 @@ function chatLocal() {
             }
             if (prop === "get") {
                 return () => {
-                    const current = locals_js_1.locals.get(localKey);
+                    let current = locals_js_1.locals.get(localKey);
+                    if (current === undefined) {
+                        hydrateLocalsFromMetadata();
+                        current = locals_js_1.locals.get(localKey);
+                    }
                     if (current === undefined) {
                         throw new Error("local.get() called before initialization. Call local.init() first.");
                     }
@@ -674,11 +1376,20 @@ function chatLocal() {
             // toJSON for serialization (JSON.stringify(local))
             if (prop === "toJSON") {
                 return () => {
-                    const current = locals_js_1.locals.get(localKey);
+                    let current = locals_js_1.locals.get(localKey);
+                    if (current === undefined) {
+                        hydrateLocalsFromMetadata();
+                        current = locals_js_1.locals.get(localKey);
+                    }
                     return current ? { ...current } : undefined;
                 };
             }
-            const current = locals_js_1.locals.get(localKey);
+            let current = locals_js_1.locals.get(localKey);
+            if (current === undefined) {
+                // Auto-hydrate from parent metadata in subtask context
+                hydrateLocalsFromMetadata();
+                current = locals_js_1.locals.get(localKey);
+            }
             if (current === undefined)
                 return undefined;
             return current[prop];
@@ -699,17 +1410,29 @@ function chatLocal() {
         has(_target, prop) {
             if (typeof prop === "symbol")
                 return prop in _target;
-            const current = locals_js_1.locals.get(localKey);
+            let current = locals_js_1.locals.get(localKey);
+            if (current === undefined) {
+                hydrateLocalsFromMetadata();
+                current = locals_js_1.locals.get(localKey);
+            }
             return current !== undefined && prop in current;
         },
         ownKeys() {
-            const current = locals_js_1.locals.get(localKey);
+            let current = locals_js_1.locals.get(localKey);
+            if (current === undefined) {
+                hydrateLocalsFromMetadata();
+                current = locals_js_1.locals.get(localKey);
+            }
             return current ? Reflect.ownKeys(current) : [];
         },
         getOwnPropertyDescriptor(_target, prop) {
             if (typeof prop === "symbol")
                 return undefined;
-            const current = locals_js_1.locals.get(localKey);
+            let current = locals_js_1.locals.get(localKey);
+            if (current === undefined) {
+                hydrateLocalsFromMetadata();
+                current = locals_js_1.locals.get(localKey);
+            }
             if (current === undefined || !(prop in current))
                 return undefined;
             return {
@@ -736,6 +1459,28 @@ exports.chat = {
     setTurnTimeoutInSeconds,
     /** Override the warm timeout at runtime. See {@link setWarmTimeoutInSeconds}. */
     setWarmTimeoutInSeconds,
+    /** Override toUIMessageStream() options for the current turn. See {@link setUIMessageStreamOptions}. */
+    setUIMessageStreamOptions,
+    /** Check if the current turn was stopped by the user. See {@link isStopped}. */
+    isStopped,
+    /** Clean up aborted parts from a UIMessage. See {@link cleanupAbortedParts}. */
+    cleanupAbortedParts,
+    /** Register background work that runs in parallel with streaming. See {@link chatDefer}. */
+    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,
+    /** Create a chat session (async iterator). See {@link createChatSession}. */
+    createSession: createChatSession,
 };
 /**
  * Writes a turn-complete control chunk to the chat output stream.