@trigger.dev/sdk 4.5.7 → 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,
                                 });
@@ -5640,28 +5705,21 @@ async function chatWriteTurnComplete(options) {
  * can return, and propagates a source error to the consumer after buffering
  * whatever streamed first. See {@link pipeChatAndCapture} for why.
  */
-async function* tapUIMessageChunks(source, buffer) {
+function tapUIMessageChunks(source, buffer) {
     if (isReadableStream(source)) {
-        const reader = source.getReader();
-        try {
-            while (true) {
-                const { done, value } = await reader.read();
-                if (done)
-                    break;
-                buffer.push(value);
-                yield value;
-            }
-        }
-        finally {
-            reader.releaseLock();
-        }
+        return source.pipeThrough(new TransformStream({
+            transform(chunk, controller) {
+                buffer.push(chunk);
+                controller.enqueue(chunk);
+            },
+        }));
     }
-    else {
+    return (async function* () {
         for await (const chunk of source) {
             buffer.push(chunk);
             yield chunk;
         }
-    }
+    })();
 }
 /**
  * Reconstruct a partial assistant `UIMessage` from the raw chunks that
@@ -6312,6 +6370,15 @@ function createChatSession(payload, options) {
                                 // 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;
@@ -6856,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,
     });