@trigger.dev/sdk 4.5.6 → 4.5.8

package/dist/esm/v3/ai.js

@@ -1,4 +1,4 @@
-import { accessoryAttributes, apiClientManager, controlSubtype, generateJWT, getSchemaParseFn, headerValue, InputStreamOncePromise, isSchemaZodEsque, logger, ManualWaitpointPromise, OutOfMemoryError, resourceCatalog, SemanticInternalAttributes, SESSION_IN_EVENT_ID_HEADER, sessionStreams, taskContext, TRIGGER_CONTROL_SUBTYPE, } from "@trigger.dev/core/v3";
+import { accessoryAttributes, apiClientManager, controlSubtype, generateJWT, getSchemaParseFn, headerValue, InputStreamOncePromise, isAdditionalApiKey, isSchemaZodEsque, logger, ManualWaitpointPromise, OutOfMemoryError, resourceCatalog, SemanticInternalAttributes, SESSION_IN_EVENT_ID_HEADER, sessionStreams, taskContext, TRIGGER_CONTROL_SUBTYPE, } from "@trigger.dev/core/v3";
 // Runtime VALUES go through the ESM/CJS shim so the CJS build can `require`
 // ESM-only `ai@7` (see ../imports/ai-runtime.ts).
 import { trace } from "@opentelemetry/api";
@@ -3630,6 +3630,9 @@ function chatAgent(options) {
                     // Declared here so the finally can detach it — a handler leaked past
                     // its turn duplicates every mid-stream message into the shared buffer.
                     let turnMsgSub;
+                    let capturedPartialResponse;
+                    let responseCommitted = false;
+                    const turnBufferedChunks = [];
                     try {
                         // Extract turn-level context before entering the span. Slim
                         // wire: at most one delta message per record. `headStartMessages`
@@ -4307,11 +4310,12 @@ function chatAgent(options) {
                                             generateMessageId: resolvedOptions.generateMessageId ?? generateMessageId,
                                             onFinish: ({ responseMessage, finishReason, }) => {
                                                 capturedResponseMessage = responseMessage;
+                                                capturedPartialResponse = responseMessage;
                                                 capturedFinishReason = finishReason;
                                                 resolveOnFinish();
                                             },
                                         });
-                                        await pipeChat(uiStream, {
+                                        await pipeChat(tapUIMessageChunks(uiStream, turnBufferedChunks), {
                                             signal: combinedSignal,
                                             spanName: "stream response",
                                         });
@@ -4505,6 +4509,11 @@ function chatAgent(options) {
                                         locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                                     }
                                 }
+                                if (capturedResponseMessage) {
+                                    responseCommitted = true;
+                                    capturedPartialResponse = capturedResponseMessage;
+                                    turnBufferedChunks.length = 0;
+                                }
                                 if (runSignal.aborted)
                                     return "exit";
                                 // Await deferred background work (e.g. DB writes from onTurnStart)
@@ -4695,6 +4704,7 @@ function chatAgent(options) {
                                             parts: [...(msg.parts ?? []), ...lateParts],
                                         };
                                         capturedResponseMessage = accumulatedUIMessages[idx];
+                                        capturedPartialResponse = capturedResponseMessage;
                                         turnCompleteEvent.responseMessage = capturedResponseMessage;
                                         turnCompleteEvent.uiMessages = accumulatedUIMessages;
                                     }
@@ -4937,10 +4947,65 @@ function chatAgent(options) {
                             !accumulatedUIMessages.some((m) => m.id === erroredWireMessage.id)
                             ? [...accumulatedUIMessages, erroredWireMessage]
                             : accumulatedUIMessages;
-                        // Fire onTurnComplete on the error path too — the docs promise it
-                        // runs "after every turn, successful or errored" so customers can
-                        // mark the turn failed. `responseMessage` is undefined/partial and
-                        // `error` carries the thrown value.
+                        let partialResponse = capturedPartialResponse ??
+                            (await assemblePartialFromChunks(turnBufferedChunks));
+                        if (partialResponse) {
+                            partialResponse = cleanupAbortedParts(partialResponse);
+                        }
+                        let partialIdx = partialResponse?.id
+                            ? erroredUIMessages.findIndex((m) => m.id === partialResponse.id)
+                            : -1;
+                        if (partialResponse && capturedPartialResponse === undefined && partialIdx !== -1) {
+                            partialResponse = undefined;
+                            partialIdx = -1;
+                        }
+                        if (partialResponse && !partialResponse.id) {
+                            partialResponse = { ...partialResponse, id: generateMessageId() };
+                        }
+                        if (partialResponse && !responseCommitted) {
+                            const queuedParts = locals.get(chatResponsePartsKey);
+                            if (queuedParts && queuedParts.length > 0) {
+                                partialResponse = {
+                                    ...partialResponse,
+                                    parts: [...partialResponse.parts, ...queuedParts],
+                                };
+                                locals.set(chatResponsePartsKey, []);
+                            }
+                        }
+                        const includePartial = partialResponse != null && !responseCommitted;
+                        let erroredUIMessagesWithPartial = !includePartial
+                            ? erroredUIMessages
+                            : partialIdx === -1
+                                ? [...erroredUIMessages, partialResponse]
+                                : erroredUIMessages.map((m, i) => i === partialIdx ? partialResponse : m);
+                        let erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : [];
+                        if (includePartial) {
+                            erroredNewUIMessages.push(partialResponse);
+                        }
+                        let erroredNewModelMessages = [];
+                        if (!responseCommitted) {
+                            try {
+                                if (erroredNewUIMessages.length > 0) {
+                                    erroredNewModelMessages = await toModelMessages(erroredNewUIMessages.map((m) => stripProviderMetadata(m)));
+                                }
+                                if (erroredUIMessagesWithPartial !== accumulatedUIMessages) {
+                                    if (partialIdx === -1) {
+                                        const appended = erroredUIMessagesWithPartial.slice(accumulatedUIMessages.length);
+                                        accumulatedMessages.push(...(await toModelMessages(appended.map((m) => stripProviderMetadata(m)))));
+                                    }
+                                    else {
+                                        accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial);
+                                    }
+                                    accumulatedUIMessages = erroredUIMessagesWithPartial;
+                                    locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
+                                }
+                            }
+                            catch {
+                                erroredNewModelMessages = [];
+                                erroredUIMessagesWithPartial = erroredUIMessages;
+                                erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : [];
+                            }
+                        }
                         if (onTurnComplete) {
                             try {
                                 await tracer.startActiveSpan("onTurnComplete()", async () => {
@@ -4948,11 +5013,11 @@ function chatAgent(options) {
                                         ctx,
                                         chatId: currentWirePayload.chatId,
                                         messages: accumulatedMessages,
-                                        uiMessages: erroredUIMessages,
-                                        newMessages: [],
-                                        newUIMessages: erroredWireMessage ? [erroredWireMessage] : [],
-                                        responseMessage: undefined,
-                                        rawResponseMessage: undefined,
+                                        uiMessages: erroredUIMessagesWithPartial,
+                                        newMessages: erroredNewModelMessages,
+                                        newUIMessages: erroredNewUIMessages,
+                                        responseMessage: partialResponse,
+                                        rawResponseMessage: partialResponse,
                                         turn,
                                         runId: ctx.run.id,
                                         chatAccessToken: "",
@@ -4996,7 +5061,7 @@ function chatAgent(options) {
                                 await writeChatSnapshot(sessionIdForSnapshot, {
                                     version: 1,
                                     savedAt: Date.now(),
-                                    messages: erroredUIMessages,
+                                    messages: erroredUIMessagesWithPartial,
                                     lastOutEventId: errorTurnCompleteResult?.lastEventId,
                                     lastInEventId: errorSnapshotInCursor !== undefined ? String(errorSnapshotInCursor) : undefined,
                                 });
@@ -5606,31 +5671,110 @@ function createStopSignal() {
  * The `TriggerChatTransport` intercepts this to close the ReadableStream
  * for the current turn. Call after piping the response stream.
  *
+ * Returns two resume cursors for the turn boundary, both saveable from the
+ * task instead of round-tripping them back from the client:
+ * - `lastEventId` — the turn-complete control record's seq_num on
+ *   `session.out`; where the next turn's output stream resumes.
+ * - `sessionInEventId` — the committed-consume cursor on `session.in` as of
+ *   this turn-complete, letting a raw loop correlate the boundary with the
+ *   exact input record it acknowledged. Trigger owns input-cursor recovery,
+ *   so this is for correlation / out-of-sync detection, not required.
+ *
+ * Either is `undefined` when the corresponding cursor isn't available.
+ *
  * @example
  * ```ts
  * await chat.pipe(result);
- * await chat.writeTurnComplete();
+ * const { lastEventId, sessionInEventId } = await chat.writeTurnComplete();
+ * await db.chats.update(chatId, { lastEventId, sessionInEventId });
  * ```
  */
 async function chatWriteTurnComplete(options) {
-    await writeTurnCompleteChunk(undefined, options?.publicAccessToken);
+    const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken);
+    // Same cursor written to the `session-in-event-id` header inside
+    // `writeTurnCompleteChunk`; surfaced here so the caller can persist it.
+    const inCursor = getChatSession().in.lastDispatchedSeqNum();
+    return {
+        lastEventId: result?.lastEventId,
+        ...(inCursor !== undefined ? { sessionInEventId: String(inCursor) } : {}),
+    };
+}
+/**
+ * Pass every chunk through untouched while recording it in `buffer`. Handles
+ * both the `AsyncIterable` and `ReadableStream` shapes `toUIMessageStream()`
+ * can return, and propagates a source error to the consumer after buffering
+ * whatever streamed first. See {@link pipeChatAndCapture} for why.
+ */
+function tapUIMessageChunks(source, buffer) {
+    if (isReadableStream(source)) {
+        return source.pipeThrough(new TransformStream({
+            transform(chunk, controller) {
+                buffer.push(chunk);
+                controller.enqueue(chunk);
+            },
+        }));
+    }
+    return (async function* () {
+        for await (const chunk of source) {
+            buffer.push(chunk);
+            yield chunk;
+        }
+    })();
+}
+/**
+ * Reconstruct a partial assistant `UIMessage` from the raw chunks that
+ * streamed before a failure — the fallback for {@link pipeChatAndCapture}
+ * when a transport error abandons the stream before `onFinish` runs. Uses the
+ * same `readUIMessageStream` reducer as the boot-time replay path. Returns
+ * `undefined` if there's nothing to assemble or the reducer throws.
+ */
+async function assemblePartialFromChunks(chunks) {
+    const relevant = chunks.filter((c) => {
+        const type = c.type;
+        return typeof type === "string" && !type.startsWith("trigger:");
+    });
+    if (relevant.length === 0)
+        return undefined;
+    try {
+        const stream = new ReadableStream({
+            start(controller) {
+                for (const chunk of relevant)
+                    controller.enqueue(chunk);
+                controller.close();
+            },
+        });
+        let last;
+        for await (const message of readUIMessageStream({ stream })) {
+            last = message;
+        }
+        return last;
+    }
+    catch {
+        return undefined;
+    }
 }
 /**
  * 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.
+ * Never throws on a stopped or failed stream: it returns a
+ * {@link PipeAndCaptureResult} whose `message` holds any partial output
+ * captured before the stop/failure, alongside a typed `status` (and `error`
+ * on failure). Save the partial after a stop or error without separate
+ * capture logic.
  *
  * @example
  * ```ts
  * const result = streamText({ model, messages, abortSignal: signal });
- * const response = await chat.pipeAndCapture(result, { signal });
- * if (response) conversation.addResponse(response);
+ * const { message, status, error } = await chat.pipeAndCapture(result, { signal });
+ * if (message) conversation.addResponse(message);
+ * if (status === "error") logger.error("turn failed", { error });
  * ```
  */
 async function pipeChatAndCapture(source, options) {
     let captured;
+    let capturedFinishReason;
     let resolveOnFinish;
     const onFinishPromise = new Promise((r) => {
         resolveOnFinish = r;
@@ -5648,17 +5792,57 @@ async function pipeChatAndCapture(source, options) {
         // the frontend replaces the partial message — wiping the
         // pre-injection text from the UI and the captured response.
         generateMessageId: resolvedOptions.generateMessageId ?? generateMessageId,
-        onFinish: ({ responseMessage }) => {
+        onFinish: ({ responseMessage, finishReason, }) => {
             captured = responseMessage;
+            capturedFinishReason = finishReason;
             resolveOnFinish();
         },
     });
-    await pipeChat(uiStream, {
-        signal: options?.signal,
-        spanName: options?.spanName ?? "stream response",
-    });
-    await onFinishPromise;
-    return captured;
+    // Buffer chunks as they flow to the pipe so a transport failure that
+    // abandons the UI stream before `onFinish` fires — the one termination path
+    // that skips `onFinish` — can still reconstruct the partial. This retains
+    // chunk references only; the reassembly runs solely in the failure fallback
+    // below, so the happy path pays nothing extra.
+    const bufferedChunks = [];
+    const tappedStream = tapUIMessageChunks(uiStream, bufferedChunks);
+    let status = "complete";
+    let error;
+    try {
+        await pipeChat(tappedStream, {
+            signal: options?.signal,
+            spanName: options?.spanName ?? "stream response",
+        });
+        // The pipe can drain cleanly on a stop — the source stream just ends
+        // early — so classify by the signal rather than relying on a throw.
+        if (options?.signal?.aborted) {
+            status = "aborted";
+        }
+    }
+    catch (err) {
+        if ((err instanceof Error && err.name === "AbortError") || options?.signal?.aborted) {
+            status = "aborted";
+        }
+        else {
+            status = "error";
+            error = err;
+        }
+    }
+    // `onFinish` fires even on abort, carrying the partial — but a hard stop can
+    // prevent it from firing at all, so race it against a timeout to avoid
+    // hanging the caller. Mirrors chat.agent's capture path.
+    await Promise.race([onFinishPromise, new Promise((r) => setTimeout(r, 2_000))]);
+    // A transport failure can abandon the UI stream before `onFinish` fires, so
+    // reconstruct the partial from the buffered chunks rather than losing output
+    // that already streamed. Only runs when `onFinish` produced nothing.
+    if (!captured && bufferedChunks.length > 0) {
+        captured = await assemblePartialFromChunks(bufferedChunks);
+    }
+    return {
+        message: captured,
+        status,
+        finishReason: capturedFinishReason,
+        ...(error !== undefined ? { error } : {}),
+    };
 }
 /**
  * Accumulates conversation messages across turns.
@@ -5673,8 +5857,8 @@ async function pipeChatAndCapture(source, options) {
  * 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);
+ *   const { message } = await chat.pipeAndCapture(result);
+ *   if (message) await conversation.addResponse(message);
  * }
  * ```
  */
@@ -6168,7 +6352,7 @@ function createChatSession(payload, options) {
                             }
                             let response;
                             try {
-                                response = await pipeChatAndCapture(source, {
+                                const captured = await pipeChatAndCapture(source, {
                                     signal: combinedSignal,
                                     // On a non-final handover turn, thread the spliced partial so a
                                     // resumed tool round's tool-output chunks merge into the
@@ -6177,20 +6361,27 @@ function createChatSession(payload, options) {
                                     // fresh response into the prior assistant message).
                                     ...(handoverThisTurn ? { originalMessages: accumulator.uiMessages } : {}),
                                 });
-                            }
-                            catch (error) {
-                                if (error instanceof Error && error.name === "AbortError") {
-                                    if (runSignal.aborted) {
-                                        // Full cancel — don't accumulate
-                                        sessionMsgSub.off();
-                                        await chatWriteTurnComplete();
-                                        return undefined;
-                                    }
-                                    // Stop — fall through to accumulate partial response
+                                if (runSignal.aborted) {
+                                    // Full cancel — don't accumulate
+                                    sessionMsgSub.off();
+                                    await chatWriteTurnComplete();
+                                    return undefined;
                                 }
-                                else {
-                                    throw error;
+                                // Surface a genuine stream failure to the caller. A user stop
+                                // (status "aborted") falls through so the partial is accumulated.
+                                if (captured.status === "error") {
+                                    if (captured.message) {
+                                        const partial = cleanupAbortedParts(captured.message);
+                                        const queuedParts = locals.get(chatResponsePartsKey);
+                                        if (queuedParts && queuedParts.length > 0) {
+                                            partial.parts = [...(partial.parts ?? []), ...queuedParts];
+                                            locals.set(chatResponsePartsKey, []);
+                                        }
+                                        await accumulator.addResponse(partial);
+                                    }
+                                    throw captured.error;
                                 }
+                                response = captured.message;
                             }
                             finally {
                                 // Detach at stream end (like the agent loop): the steering queue
@@ -6732,24 +6923,49 @@ async function mintPublicTokenWithOverride(args) {
         throw new Error("chat.createStartSessionAction: no API access token configured for JWT mint.");
     }
     const ctx = { endpoint: "auth", chatId: args.chatId };
-    const url = `${resolveChatStartBaseURL("auth", args.chatId, args.baseURLOption)}/api/v1/auth/jwt/claims`;
+    const scopes = [`read:sessions:${args.chatId}`, `write:sessions:${args.chatId}`];
+    const serverMint = isAdditionalApiKey(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),
+        ...(serverMint
+            ? {
+                body: JSON.stringify({
+                    scopes,
+                    expirationTime: args.expirationTime instanceof Date
+                        ? Math.floor(args.expirationTime.getTime() / 1000)
+                        : args.expirationTime,
+                }),
+            }
+            : {}),
     };
     const response = args.fetchOverride
         ? 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 text = await response.text().catch(() => "");
         throw new Error(`auth.createPublicToken failed: ${response.status} ${text}`);
     }
-    const claims = (await response.json());
+    const responseBody = (await response.json());
+    if (serverMint) {
+        if (typeof responseBody.token !== "string") {
+            throw new Error("auth.createPublicToken failed: server response did not include a token");
+        }
+        return responseBody.token;
+    }
     return generateJWT({
         secretKey: accessToken,
         payload: {
-            ...claims,
-            scopes: [`read:sessions:${args.chatId}`, `write:sessions:${args.chatId}`],
+            ...responseBody,
+            scopes,
         },
         expirationTime: args.expirationTime,
     });