@trigger.dev/sdk 4.6.0 → 4.6.2

package/dist/esm/v3/ai.js

@@ -2,6 +2,7 @@ import { accessoryAttributes, apiClientManager, controlSubtype, generateJWT, get
 // 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";
+import { traceSessionIdle } from "./sessionTracing.js";
 import { tool as aiTool, convertToModelMessages, dynamicTool, generateId as generateMessageId, getToolName, isToolUIPart, jsonSchema, readUIMessageStream, streamText as aiStreamText, zodSchema, } from "../imports/ai-runtime.js";
 import { createTranscriptShadow, defaultStorage, diffTranscript, fingerprintMessage, parseTranscriptRuntimeState, restoreModelLane, } from "./transcriptStorage.js";
 let transcriptStorageOverride;
@@ -1087,7 +1088,7 @@ async function waitOnChatRoute(route, options) {
     return tracer.startActiveSpan(options.spanName ?? `chat.${route}.wait()`, async (span) => {
         const idleMs = (options.idleTimeoutInSeconds ?? 0) * 1000;
         if (idleMs > 0) {
-            const warm = await router.next(route, { timeoutMs: idleMs });
+            const warm = await traceSessionIdle(session.id, idleMs / 1000, () => router.next(route, { timeoutMs: idleMs }));
             if (warm) {
                 span.setAttribute("wait.resolved", "idle");
                 return { ok: true, output: warm.data, record: warm };
@@ -1141,10 +1142,6 @@ async function waitOnChatRoute(route, options) {
             session: session.id,
             io: "in",
             route,
-            ...accessoryAttributes({
-                items: [{ text: `${session.id}.in:${route}`, variant: "normal" }],
-                style: "codepath",
-            }),
         },
     });
 }
@@ -1820,6 +1817,14 @@ const chatHandoverPartialKey = locals.create("chat.handoverPartial");
  * @internal
  */
 const chatHandoverMessageIdKey = locals.create("chat.handoverMessageId");
+/**
+ * The model messages a turn-0 head-start splice contributed to the model lane,
+ * keyed by the UI message id it was synthesized under. When the agent's response
+ * completes that message under the same id, this is the run to replace: the UI
+ * form alone converts to something else (its pending tool calls drop out, the
+ * approval round was never on it), so it cannot locate the run itself.
+ */
+const chatHandoverSplicedRunKey = locals.create("chat.handoverSplicedRun");
 /**
  * Run-scoped slot indicating that the customer's step-1 head-start
  * response is the FINAL turn response. When true, turn 0 runs through
@@ -1903,16 +1908,19 @@ function synthesizeHandoverUIMessage(partial, messageId) {
  */
 function spliceHandoverPartial(modelMessages, uiMessages, signal) {
     if (!signal.partialAssistantMessage || signal.partialAssistantMessage.length === 0) {
-        return;
+        return undefined;
     }
     // Skip if the hydrated chain already persisted the partial under this id.
     const alreadyInChain = signal.messageId !== undefined && uiMessages.some((m) => m.id === signal.messageId);
     if (alreadyInChain)
-        return;
-    modelMessages.push(...signal.partialAssistantMessage);
+        return undefined;
+    const run = [...signal.partialAssistantMessage];
+    modelMessages.push(...run);
     const partialUI = synthesizeHandoverUIMessage(signal.partialAssistantMessage, signal.messageId);
-    if (partialUI)
-        uiMessages.push(partialUI);
+    if (!partialUI)
+        return undefined;
+    uiMessages.push(partialUI);
+    return { id: partialUI.id, run };
 }
 /**
  * Per-turn background context queue. Messages added via `chat.backgroundWork.inject()`
@@ -3652,8 +3660,8 @@ function isActionTurn(value) {
  * tail does not match the old message's conversion, nothing is changed and
  * `false` is returned so the caller can fall back to a full reconversion.
  */
-async function replaceModelRun(lane, oldUi, newUi, tailAfter) {
-    const oldRun = await toModelMessages([stripProviderMetadata(oldUi)]);
+async function replaceModelRun(lane, oldUi, newUi, tailAfter, knownOldRun) {
+    const oldRun = knownOldRun ?? (await toModelMessages([stripProviderMetadata(oldUi)]));
     const newRun = await toModelMessages([stripProviderMetadata(newUi)]);
     // A message that converts to nothing (a pending tool call with no output yet,
     // which `ignoreIncompleteToolCalls` drops) locates no run in the lane. Matching
@@ -4785,7 +4793,7 @@ function chatAgent(options) {
                     const preloadResult = await messagesInput.waitWithIdleTimeout({
                         idleTimeoutInSeconds: effectivePreloadIdleTimeout,
                         timeout: effectivePreloadTimeout,
-                        spanName: "waiting for first message",
+                        spanName: "first message",
                         skipSuspend: exitAfterPreloadIdle,
                         onSuspend: onChatSuspend
                             ? async () => {
@@ -4937,7 +4945,7 @@ function chatAgent(options) {
                         const continuationResult = await messagesInput.waitWithIdleTimeout({
                             idleTimeoutInSeconds: effectiveIdleTimeout,
                             timeout: effectiveTurnTimeout,
-                            spanName: "waiting for first message (continuation)",
+                            spanName: "first message (continuation)",
                             onSuspend: onChatSuspend
                                 ? async () => {
                                     await tracer.startActiveSpan("onChatSuspend()", async () => {
@@ -5466,10 +5474,12 @@ function chatAgent(options) {
                                     // `UIMessageStreamError: No tool invocation found`.
                                     const pendingHandoverPartial = locals.get(chatHandoverPartialKey);
                                     if (pendingHandoverPartial && pendingHandoverPartial.length > 0) {
-                                        spliceHandoverPartial(accumulatedMessages, accumulatedUIMessages, {
+                                        const spliced = spliceHandoverPartial(accumulatedMessages, accumulatedUIMessages, {
                                             partialAssistantMessage: pendingHandoverPartial,
                                             messageId: locals.get(chatHandoverMessageIdKey),
                                         });
+                                        if (spliced)
+                                            locals.set(chatHandoverSplicedRunKey, spliced);
                                         locals.set(chatHandoverPartialKey, []); // consume once
                                         splicedHandoverPartial = true;
                                     }
@@ -5814,6 +5824,7 @@ function chatAgent(options) {
                                 // never reports final usage), which would block the turn loop
                                 // from ever firing onTurnComplete / writeTurnComplete.
                                 let turnUsage;
+                                let lastStepUsage;
                                 if (runResult != null &&
                                     typeof runResult.totalUsage?.then === "function") {
                                     try {
@@ -5826,6 +5837,18 @@ function chatAgent(options) {
                                         /* non-fatal — usage capture failed */
                                     }
                                 }
+                                const lastStepUsagePromise = runResult != null ? runResult.usage : undefined;
+                                if (typeof lastStepUsagePromise?.then === "function") {
+                                    try {
+                                        lastStepUsage = (await Promise.race([
+                                            lastStepUsagePromise,
+                                            new Promise((r) => setTimeout(() => r(undefined), 2_000)),
+                                        ]));
+                                    }
+                                    catch {
+                                        /* non-fatal — usage capture failed */
+                                    }
+                                }
                                 if (turnUsage) {
                                     cumulativeUsage = addUsage(cumulativeUsage, turnUsage);
                                     previousTurnUsage = turnUsage;
@@ -5971,8 +5994,14 @@ function chatAgent(options) {
                                                 stripProviderMetadata(capturedResponseMessage),
                                             ]);
                                             if (existingIdx !== -1) {
+                                                const spliced = locals.get(chatHandoverSplicedRunKey);
+                                                const splicedRun = spliced && previousAtIdx && spliced.id === previousAtIdx.id
+                                                    ? spliced.run
+                                                    : undefined;
                                                 const ok = previousAtIdx !== undefined &&
-                                                    (await replaceModelRun(accumulatedMessages, previousAtIdx, capturedResponseMessage, steerTailThisTurn));
+                                                    (await replaceModelRun(accumulatedMessages, previousAtIdx, capturedResponseMessage, steerTailThisTurn, splicedRun));
+                                                if (splicedRun)
+                                                    locals.set(chatHandoverSplicedRunKey, undefined);
                                                 if (!ok) {
                                                     logger.warn("chat.agent: replaced response not found at the model lane tail; reconverting the lane");
                                                     accumulatedMessages = await toModelMessages(accumulatedUIMessages);
@@ -6032,12 +6061,14 @@ function chatAgent(options) {
                                 const outerCompaction = locals.get(chatAgentCompactionKey);
                                 const innerCompactionState = locals.get(chatCompactionStateKey);
                                 if (outerCompaction && !innerCompactionState && turnUsage && !wasStopped) {
+                                    const contextUsage = lastStepUsage ?? turnUsage;
                                     const shouldTrigger = await outerCompaction.shouldCompact({
                                         messages: accumulatedMessages,
-                                        totalTokens: turnUsage.totalTokens,
-                                        inputTokens: turnUsage.inputTokens,
-                                        outputTokens: turnUsage.outputTokens,
-                                        usage: turnUsage,
+                                        totalTokens: contextUsage.totalTokens,
+                                        inputTokens: contextUsage.inputTokens,
+                                        outputTokens: contextUsage.outputTokens,
+                                        usage: contextUsage,
+                                        turnUsage,
                                         totalUsage: cumulativeUsage,
                                         chatId: currentWirePayload.chatId,
                                         turn,
@@ -6362,7 +6393,7 @@ function chatAgent(options) {
                             const next = await messagesInput.waitWithIdleTimeout({
                                 idleTimeoutInSeconds: effectiveIdleTimeout,
                                 timeout: effectiveTurnTimeout,
-                                spanName: "waiting for next message",
+                                spanName: "next message",
                                 onSuspend: onChatSuspend
                                     ? async () => {
                                         await tracer.startActiveSpan("onChatSuspend()", async () => {
@@ -6699,7 +6730,7 @@ function chatAgent(options) {
                         const next = await messagesInput.waitWithIdleTimeout({
                             idleTimeoutInSeconds: effectiveIdleTimeout,
                             timeout: effectiveTurnTimeout,
-                            spanName: "waiting for next message (after error)",
+                            spanName: "next message (after error)",
                         });
                         if (!next.ok) {
                             return; // Timed out — end run gracefully
@@ -7739,6 +7770,8 @@ class ChatMessageAccumulator {
     modelMessages = [];
     uiMessages = [];
     _compaction;
+    /** The run a spliced head-start partial contributed, until its response replaces it. */
+    _handoverRun;
     _pendingMessages;
     _steeringQueue = [];
     constructor(options) {
@@ -7782,7 +7815,7 @@ class ChatMessageAccumulator {
      * `consumeHandover` for the wait+seed+apply convenience.
      */
     applyHandover(signal) {
-        spliceHandoverPartial(this.modelMessages, this.uiMessages, signal);
+        this._handoverRun = spliceHandoverPartial(this.modelMessages, this.uiMessages, signal);
     }
     /**
      * One-call `chat.headStart` handover for a custom-agent loop: waits for the
@@ -7825,8 +7858,13 @@ class ChatMessageAccumulator {
         if (existingIdx !== -1) {
             const previous = this.uiMessages[existingIdx];
             this.uiMessages[existingIdx] = response;
+            const handoverRun = this._handoverRun && this._handoverRun.id === previous.id
+                ? this._handoverRun.run
+                : undefined;
+            if (handoverRun)
+                this._handoverRun = undefined;
             try {
-                if (!(await replaceModelRun(this.modelMessages, previous, response, 0))) {
+                if (!(await replaceModelRun(this.modelMessages, previous, response, 0, handoverRun))) {
                     this.modelMessages = await toModelMessages(this.uiMessages.map((m) => stripProviderMetadata(m)));
                 }
             }
@@ -7929,8 +7967,10 @@ class ChatMessageAccumulator {
     }
     /**
      * Run outer-loop compaction if needed. Call after adding the response
-     * and capturing usage. Applies `compactModelMessages` and `compactUIMessages`
-     * callbacks if configured.
+     * and capturing usage. Pass the LAST step's usage (`result.usage`), which is
+     * the context the model held on its final call; `result.totalUsage` sums every
+     * step of a tool-using turn and belongs in `context.turnUsage`. Applies
+     * `compactModelMessages` and `compactUIMessages` callbacks if configured.
      *
      * @returns `true` if compaction was performed, `false` otherwise.
      */
@@ -7943,6 +7983,7 @@ class ChatMessageAccumulator {
             inputTokens: usage.inputTokens,
             outputTokens: usage.outputTokens,
             usage,
+            turnUsage: context?.turnUsage,
             totalUsage: context?.totalUsage,
             chatId: context?.chatId,
             turn: context?.turn,
@@ -8175,8 +8216,8 @@ function createChatSession(payload, options) {
                             idleTimeoutInSeconds: sessionIdleTimeoutOpt ?? currentPayload.idleTimeoutInSeconds ?? 30,
                             timeout,
                             spanName: currentPayload.trigger === "preload"
-                                ? "waiting for first message"
-                                : "waiting for first message (continuation)",
+                                ? "first message"
+                                : "first message (continuation)",
                         });
                         if (!result.ok || runSignal.aborted) {
                             stop.cleanup();
@@ -8213,7 +8254,7 @@ function createChatSession(payload, options) {
                         const next = await messagesInput.waitWithIdleTimeout({
                             idleTimeoutInSeconds,
                             timeout,
-                            spanName: "waiting for next message",
+                            spanName: "next message",
                         });
                         if (!next.ok || runSignal.aborted) {
                             stop.cleanup();
@@ -8424,6 +8465,7 @@ function createChatSession(payload, options) {
                             // indefinitely, which would wedge the turn loop (same guard as
                             // chat.agent's turn loop).
                             let turnUsage;
+                            let lastStepUsage;
                             if (typeof source.totalUsage?.then === "function") {
                                 try {
                                     const usage = (await Promise.race([
@@ -8440,14 +8482,28 @@ function createChatSession(payload, options) {
                                     /* non-fatal */
                                 }
                             }
+                            const lastStepUsagePromise = source.usage;
+                            if (typeof lastStepUsagePromise?.then === "function") {
+                                try {
+                                    lastStepUsage = (await Promise.race([
+                                        lastStepUsagePromise,
+                                        new Promise((r) => setTimeout(() => r(undefined), 2_000)),
+                                    ]));
+                                }
+                                catch {
+                                    /* non-fatal */
+                                }
+                            }
                             // Outer-loop compaction (same logic as chat.agent)
                             if (sessionCompaction && turnUsage && !turnObj.stopped) {
+                                const contextUsage = lastStepUsage ?? turnUsage;
                                 const shouldTrigger = await sessionCompaction.shouldCompact({
                                     messages: accumulator.modelMessages,
-                                    totalTokens: turnUsage.totalTokens,
-                                    inputTokens: turnUsage.inputTokens,
-                                    outputTokens: turnUsage.outputTokens,
-                                    usage: turnUsage,
+                                    totalTokens: contextUsage.totalTokens,
+                                    inputTokens: contextUsage.inputTokens,
+                                    outputTokens: contextUsage.outputTokens,
+                                    usage: contextUsage,
+                                    turnUsage,
                                     totalUsage: cumulativeUsage,
                                     chatId: currentPayload.chatId,
                                     turn,