@trigger.dev/sdk 4.8.1 → 4.8.2

package/dist/esm/v3/ai.js

@@ -1,4 +1,4 @@
-import { accessoryAttributes, apiClientManager, controlSubtype, generateJWT, getSchemaParseFn, headerValue, InputStreamOncePromise, isAdditionalApiKey, isSchemaZodEsque, logger, ManualWaitpointPromise, OutOfMemoryError, resourceCatalog, SemanticInternalAttributes, SESSION_IN_CONSUMED_ID_HEADER, SESSION_IN_EVENT_ID_HEADER, sessionStreams, SessionChannelRouter, InputStreamTimeoutError, taskContext, TRIGGER_CONTROL_SUBTYPE, SESSION_CLOSED_HEADER, SESSION_CLOSED_REASON_HEADER, tryCatch, } from "@trigger.dev/core/v3";
+import { accessoryAttributes, apiClientManager, controlSubtype, generateJWT, getSchemaParseFn, headerValue, InputStreamOncePromise, isAdditionalApiKey, isRunToken, isSchemaZodEsque, logger, ManualWaitpointPromise, OutOfMemoryError, resourceCatalog, SemanticInternalAttributes, SESSION_IN_CONSUMED_ID_HEADER, SESSION_IN_EVENT_ID_HEADER, sessionStreams, SessionChannelRouter, InputStreamTimeoutError, taskContext, TRIGGER_CONTROL_SUBTYPE, SESSION_CLOSED_HEADER, SESSION_CLOSED_REASON_HEADER, tryCatch, } from "@trigger.dev/core/v3";
 import { chatEvent, normalizeKeyString } from "./webhooks.js";
 import { fetchPublicFile } from "./modelFileDownload.js";
 // Runtime VALUES go through the ESM/CJS shim so the CJS build can `require`
@@ -12,7 +12,8 @@ import { responseAfterCompaction } from "./compactionResponse.js";
 import { withToolResultsInCallOrder } from "./toolResultOrder.js";
 import { ManagedChatResponse, createOrderedChatWriter } from "./managedChatResponse.js";
 import { reduceUIMessageChunks } from "./uiMessageChunks.js";
-import { convertSteeredMessages, retainStepMessages, steeringMarkers, } from "./steeringContext.js";
+import { forwardSubAgentChunks } from "./subAgentStream.js";
+import { convertSteeredMessages, retainStepMessages, steeringMarkers, withoutInjectionMarkers, } from "./steeringContext.js";
 let transcriptStorageOverride;
 /**
  * Test-only override for the storage `chat.agent` persists through, so a
@@ -22,7 +23,7 @@ let transcriptStorageOverride;
 export function __setTranscriptStorageForTests(storage) {
     transcriptStorageOverride = storage;
 }
-import { PENDING_MESSAGE_INJECTED_TYPE, upsertIncomingMessage, chatRunTags, } from "./ai-shared.js";
+import { BACKGROUND_INJECTION_TYPE, PENDING_MESSAGE_INJECTED_TYPE, upsertIncomingMessage, chatRunTags, } from "./ai-shared.js";
 import { auth } from "./auth.js";
 import { locals } from "./locals.js";
 import { metadata } from "./metadata.js";
@@ -924,6 +925,64 @@ const chatResponse = {
         managedResponse().writeData(part);
     },
 };
+/**
+ * Stream a sub-agent's reply to the frontend from inside a parent tool call,
+ * and resolve with the sub-agent's final message once it finishes.
+ *
+ * Each chunk the sub-agent streams is written once to the parent chat as a
+ * transient `data-subagent` part, so the bytes written grow with the length of
+ * the reply. Render it while it runs with `useSubAgentMessages` from
+ * `@trigger.dev/sdk/chat/react`. The tool's output is the final message, which
+ * persists with the parent's history the way a tool result always does.
+ *
+ * Prefer this over yielding `stream.messages()` snapshots as preliminary tool
+ * results: every snapshot repeats the whole message so far, so that pattern
+ * writes bytes in proportion to the square of the reply length.
+ *
+ * Call it inside a `chat.agent` (or `chat.customAgent`) run.
+ *
+ * @example
+ * ```ts
+ * const research = tool({
+ *   description: "Delegate research to the research agent.",
+ *   inputSchema: z.object({ topic: z.string() }),
+ *   execute: async ({ topic }, { toolCallId, abortSignal }) => {
+ *     const agent = new AgentChat<typeof researchAgent>({ agent: "research-agent" });
+ *     const stream = await agent.sendMessage(topic, { abortSignal });
+ *     return await chat.forwardSubAgent(stream, { toolCallId });
+ *   },
+ *   toModelOutput: ({ output }) => ({ type: "text", value: lastText(output) }),
+ * });
+ * ```
+ */
+async function forwardSubAgent(stream, options) {
+    const source = isReadableStream(stream)
+        ? stream
+        : stream.stream;
+    let outcome = { message: undefined };
+    const { waitUntilComplete } = chatStream.writer({
+        spanName: "forward sub-agent",
+        collapsed: true,
+        execute: async ({ write }) => {
+            try {
+                outcome = {
+                    message: await forwardSubAgentChunks(source, {
+                        toolCallId: options.toolCallId,
+                        flushIntervalMs: options.flushIntervalMs ?? 50,
+                        write,
+                    }),
+                };
+            }
+            catch (error) {
+                outcome = { error };
+            }
+        },
+    });
+    await waitUntilComplete();
+    if ("error" in outcome)
+        throw outcome.error;
+    return outcome.message;
+}
 /**
  * Creates a lazy ChatWriter that only opens a realtime stream on first use.
  * Call `flush()` after the callback returns to await stream completion.
@@ -1918,12 +1977,6 @@ function spliceHandoverPartial(modelMessages, uiMessages, signal) {
  * @internal
  */
 const chatBackgroundQueueKey = locals.create("chat.backgroundQueue");
-/**
- * Background injections a step-boundary drain handed to the model this turn,
- * with the transcript message they followed. Reconciled into the model lane
- * and the persisted injections once the turn's response is in.
- */
-const chatPendingBackgroundKey = locals.create("chat.pendingBackground");
 /**
  * System-role context injected mid-conversation, held for the instructions lane.
  *
@@ -2886,6 +2939,25 @@ const EMPTY_DRAIN = { injected: [], claimed: [] };
 function modelFormOf(m, batch, injected) {
     return batch[0] === m ? injected : [];
 }
+/**
+ * Record a prepared injection and write its marker into the response, so the
+ * model lane rebuilds it at the step boundary where the model received it.
+ */
+function markInjection(steps, messageIds, messages, markerMessages, kind) {
+    const id = generateMessageId();
+    const injections = locals.get(chatSteeringInjectionsKey) ?? new Map();
+    injections.set(id, { id, messageIds, messages, ...(kind ? { kind } : {}) });
+    locals.set(chatSteeringInjectionsKey, injections);
+    const response = managedResponse();
+    // prepareStep can run ahead of the UI consumer. Admit the marker only
+    // once the actual preceding finish-step chunk has entered this stream.
+    response.afterToolStep(steps.at(-1)?.content.flatMap((part) => (part.toolCallId ? [part.toolCallId] : [])) ?? []);
+    response.writeData({
+        type: kind === "background" ? BACKGROUND_INJECTION_TYPE : PENDING_MESSAGE_INJECTED_TYPE,
+        id,
+        data: { messageIds, messages: markerMessages },
+    });
+}
 /**
  * Drain the steering queue as a batch. Calls `shouldInject` once with all
  * pending messages. If it returns true, calls `prepareMessages` once to
@@ -2899,8 +2971,6 @@ function modelFormOf(m, batch, injected) {
  * @internal
  */
 async function drainSteeringQueue(config, messages, steps, queueOverride) {
-    if (steps.length === 0)
-        managedResponse().beginGeneration();
     const queue = queueOverride ?? locals.get(chatSteeringQueueKey);
     if (!queue || queue.length === 0)
         return EMPTY_DRAIN;
@@ -3033,23 +3103,7 @@ async function drainSteeringQueue(config, messages, steps, queueOverride) {
             locals.set(chatPendingSteerKey, pendingSteer);
         }
         if (injected.length > 0) {
-            const id = generateMessageId();
-            const messageIds = claimedUIMessages.map((m) => m.id);
-            const injections = locals.get(chatSteeringInjectionsKey) ?? new Map();
-            injections.set(id, { id, messageIds, messages: injected });
-            locals.set(chatSteeringInjectionsKey, injections);
-            const response = managedResponse();
-            // prepareStep can run ahead of the UI consumer. Admit the marker only
-            // once the actual preceding finish-step chunk has entered this stream.
-            response.afterStep(steps.length);
-            response.writeData({
-                type: PENDING_MESSAGE_INJECTED_TYPE,
-                id,
-                data: {
-                    messageIds,
-                    messages: claimedUIMessages.map((m) => ({ id: m.id, text: textOfUIMessage(m) })),
-                },
-            });
+            markInjection(steps, claimedUIMessages.map((m) => m.id), injected, claimedUIMessages.map((m) => ({ id: m.id, text: textOfUIMessage(m) })));
         }
         // Fire onInjected callback
         if (config.onInjected && injected.length > 0) {
@@ -3614,12 +3668,7 @@ function toStreamTextOptions(options, retain = true) {
             if (bgQueue && bgQueue.length > 0) {
                 const injected = bgQueue.splice(0); // drain
                 resultMessages = [...(resultMessages ?? messages), ...injected];
-                const pendingBackground = locals.get(chatPendingBackgroundKey) ?? [];
-                pendingBackground.push({
-                    afterId: locals.get(chatCurrentUIMessagesKey)?.at(-1)?.id ?? "",
-                    messages: injected,
-                });
-                locals.set(chatPendingBackgroundKey, pendingBackground);
+                markInjection(steps, [], injected, [], "background");
             }
             // 4. Parallel tool results in call order, the order the next turn's history
             // will list them in. See `withToolResultsInCallOrder`.
@@ -4969,24 +5018,6 @@ function chatAgent(options) {
             // durable snapshot + `session.out` replay (or `hydrateMessages` if
             // registered) — the wire is delta-only now, no longer a seed.
             let accumulatedMessages = [];
-            /**
-             * Give the model accumulator the background injections a step-boundary
-             * drain handed to the model this turn, and record them for persistence.
-             * Returns how many model messages were appended.
-             */
-            const reconcilePendingBackground = () => {
-                const pending = locals.get(chatPendingBackgroundKey);
-                if (!pending || pending.length === 0)
-                    return 0;
-                locals.set(chatPendingBackgroundKey, []);
-                let appended = 0;
-                for (const entry of pending) {
-                    accumulatedMessages.push(...entry.messages);
-                    laneInjections.push(entry);
-                    appended += entry.messages.length;
-                }
-                return appended;
-            };
             /**
              * Give the model accumulator the steering messages a drain consumed,
              * in the form the model actually received. Appended, never reconverted
@@ -5010,6 +5041,15 @@ function chatAgent(options) {
             // Accumulated UI messages for persistence. Mirrors the model accumulator
             // but in frontend-friendly UIMessage format (with parts, id, etc.).
             let accumulatedUIMessages = [];
+            const cleanIncomingMessage = (msg) => {
+                if (msg.role !== "assistant")
+                    return withoutInjectionMarkers(msg);
+                const cleaned = cleanupAbortedParts(msg);
+                // Markers are the agent's own; only a known assistant message keeps them.
+                return accumulatedUIMessages.some((m) => m.id === msg.id)
+                    ? cleaned
+                    : withoutInjectionMarkers(cleaned);
+            };
             // ── Snapshot + replay (gated on prior-state signals) ─────────────
             //
             // With `hydrateMessages` registered the customer owns history — the
@@ -6526,7 +6566,7 @@ function chatAgent(options) {
                         // Cleaning happens once here so `extractLastUserMessageText` and
                         // every downstream consumer see the same message shape — and
                         // `cleanupAbortedParts` no longer has to be re-applied below.
-                        const cleanedIncomingMessages = incomingMessages.map((msg) => msg.role === "assistant" ? cleanupAbortedParts(msg) : msg);
+                        const cleanedIncomingMessages = incomingMessages.map(cleanIncomingMessage);
                         const lastUserMessage = extractLastUserMessageText(cleanedIncomingMessages);
                         channelTurnRan =
                             !!channelConn &&
@@ -6587,7 +6627,6 @@ function chatAgent(options) {
                             locals.set(chatOutGateKey, createChatOutGate());
                             locals.set(chatCompactionStateKey, undefined);
                             locals.set(chatSteeringQueueKey, []);
-                            locals.set(chatPendingBackgroundKey, []);
                             await locals.get(chatManagedResponseKey)?.close();
                             locals.set(chatManagedResponseKey, undefined);
                             locals.set(chatManagedResponseActiveKey, true);
@@ -7686,7 +7725,7 @@ function chatAgent(options) {
                                 const steerTailThisTurn = reconcilePendingSteer({
                                     turnNew: turnNewModelMessages,
                                     response: capturedResponseMessage,
-                                }).reduce((n, e) => n + e.model.length, 0) + reconcilePendingBackground();
+                                }).reduce((n, e) => n + e.model.length, 0);
                                 // 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.
@@ -8569,7 +8608,8 @@ function chatAgent(options) {
                         // yet (a pre-run hook threw), so fold the wire message in for the
                         // error event + snapshot — the cursor has already advanced past it,
                         // so otherwise it survives in neither the snapshot nor the `.in` tail.
-                        const erroredWireMessage = currentWirePayload.message;
+                        const rawWireMessage = currentWirePayload.message;
+                        const erroredWireMessage = rawWireMessage && cleanIncomingMessage(rawWireMessage);
                         const erroredUIMessages = erroredWireMessage &&
                             !accumulatedUIMessages.some((m) => m.id === erroredWireMessage.id)
                             ? [...accumulatedUIMessages, erroredWireMessage]
@@ -8626,7 +8666,6 @@ function chatAgent(options) {
                         let erroredNewUIMessages = buildErroredNew();
                         let erroredNewModelMessages = [];
                         const reconciledSteer = reconcilePendingSteer({ response: partialResponse });
-                        const backgroundTailThisTurn = reconcilePendingBackground();
                         if (!responseCommitted) {
                             try {
                                 if (erroredNewUIMessages.length > 0) {
@@ -8644,8 +8683,7 @@ function chatAgent(options) {
                                         accumulatedMessages.push(...(await toModelMessages(appended.map((m) => stripProviderMetadata(m)))));
                                     }
                                     else {
-                                        const ok = await replaceModelRun(accumulatedMessages, erroredUIMessages[partialIdx], partialResponse, reconciledSteer.reduce((n, e) => n + e.model.length, 0) +
-                                            backgroundTailThisTurn);
+                                        const ok = await replaceModelRun(accumulatedMessages, erroredUIMessages[partialIdx], partialResponse, reconciledSteer.reduce((n, e) => n + e.model.length, 0));
                                         if (!ok) {
                                             logger.warn("chat.agent: replaced partial not found at the model lane tail; reconverting the lane");
                                             accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial, erroredUIMessagesWithPartial);
@@ -11045,10 +11083,10 @@ function resolveChatStartBaseURL(endpoint, chatId, option) {
     const raw = typeof option === "function" ? option({ endpoint, chatId }) : (option ?? fallback);
     return raw.replace(/\/$/, "");
 }
-function overrideRequestHeaders(accessToken) {
+function overrideRequestHeaders(authHeaders) {
     const headers = {
         "Content-Type": "application/json",
-        Authorization: `Bearer ${accessToken}`,
+        ...authHeaders,
         "x-trigger-source": "sdk",
     };
     // Forward the preview-branch hint so override-mode requests land on the
@@ -11061,15 +11099,16 @@ function overrideRequestHeaders(accessToken) {
     return headers;
 }
 async function callSessionsCreateWithOverride(args) {
-    const accessToken = apiClientManager.accessToken;
-    if (!accessToken) {
+    const apiClient = apiClientManager.client;
+    if (!apiClient) {
         throw new Error("chat.createStartSessionAction: no API access token configured. Set TRIGGER_SECRET_KEY or call apiClientManager.setGlobalAPIClientConfiguration before invoking the action.");
     }
+    const authHeaders = apiClient.getAuthenticationHeaders("createSession");
     const ctx = { endpoint: "sessions", chatId: args.chatId };
     const url = `${resolveChatStartBaseURL("sessions", args.chatId, args.baseURLOption)}/api/v1/sessions`;
     const init = {
         method: "POST",
-        headers: overrideRequestHeaders(accessToken),
+        headers: overrideRequestHeaders(authHeaders),
         // This path bypasses `sessions.start`, so it resolves the pin itself.
         body: JSON.stringify(withResolvedExternalDeploymentId(args.body)),
     };
@@ -11084,18 +11123,20 @@ async function callSessionsCreateWithOverride(args) {
     return json;
 }
 async function mintPublicTokenWithOverride(args) {
-    const accessToken = apiClientManager.accessToken;
-    if (!accessToken) {
+    const apiClient = apiClientManager.client;
+    if (!apiClient) {
         throw new Error("chat.createStartSessionAction: no API access token configured for JWT mint.");
     }
+    const authHeaders = apiClient.getAuthenticationHeaders("createPublicToken");
+    const accessToken = authHeaders.Authorization.slice("Bearer ".length);
     const ctx = { endpoint: "auth", chatId: args.chatId };
     const scopes = [`read:sessions:${args.chatId}`, `write:sessions:${args.chatId}`];
-    const serverMint = isAdditionalApiKey(accessToken);
+    const serverMint = isAdditionalApiKey(accessToken) || isRunToken(accessToken);
     const endpoint = serverMint ? "/api/v1/auth/public-tokens" : "/api/v1/auth/jwt/claims";
     const url = `${resolveChatStartBaseURL("auth", args.chatId, args.baseURLOption)}${endpoint}`;
     const init = {
         method: "POST",
-        headers: overrideRequestHeaders(accessToken),
+        headers: overrideRequestHeaders(authHeaders),
         ...(serverMint
             ? {
                 body: JSON.stringify({
@@ -11111,11 +11152,9 @@ async function mintPublicTokenWithOverride(args) {
         ? await args.fetchOverride(url, init, ctx)
         : await fetch(url, init);
     if (!response.ok) {
-        // An additional API key cannot self-sign, so it must use the server mint
-        // endpoint. On a server too old to expose it, explain the recovery path
-        // instead of surfacing a bare 404 (mirrors auth.createServerPublicToken).
         if (serverMint && response.status === 404) {
-            throw new Error("This additional API key cannot self-sign public tokens, and the server does not support public-token minting. Upgrade the server or use the root API key.");
+            const credentialType = isRunToken(accessToken) ? "run token" : "additional API key";
+            throw new Error(`This ${credentialType} cannot self-sign public tokens, and the server does not support public-token minting. Upgrade the server or use the root API key.`);
         }
         const text = await response.text().catch(() => "");
         throw new Error(`auth.createPublicToken failed: ${response.status} ${text}`);
@@ -11229,6 +11268,8 @@ export const chat = {
     stream: chatStream,
     /** Write data parts that persist to the response message. See {@link chatResponse}. */
     response: chatResponse,
+    /** Stream a sub-agent's reply through a parent tool call. See {@link forwardSubAgent}. */
+    forwardSubAgent,
     /** Pre-built input stream for receiving messages from the transport. */
     messages: messagesInput,
     /** The current chat.agent run's Session handle. See {@link SessionHandle}. */