triggerdotdev · matt-aitken · Jul 23, 2026 · Jul 23, 2026 · Jul 23, 2026 · Jul 23, 2026
diff --git a/.changeset/chat-agent-preserve-partial-on-stream-failure.md b/.changeset/chat-agent-preserve-partial-on-stream-failure.md
@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Preserve the partial assistant message when a chat turn's model stream fails mid-response (e.g. a transport timeout). Previously the streamed-so-far output was dropped: `chat.agent`'s `onTurnComplete` fired with `responseMessage: undefined`, and `chat.createSession`'s `turn.complete()` rethrew without keeping the partial. Now the recovered partial is passed to `onTurnComplete` (on `responseMessage`, `uiMessages`, and `newUIMessages`) and accumulated before `turn.complete()` rethrows, so it survives for persistence even when `hydrateMessages` disables boot-time replay recovery. The turn is still reported as errored.
diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts
@@ -6389,6 +6389,15 @@ function chatAgent<
           // 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: { off: () => void } | undefined;
+          // Declared at turn scope (not inside the span callback) so the error
+          // handler below can recover the partial assistant output when a
+          // source-stream failure abandons the turn before onTurnComplete's
+          // success-path capture runs. `capturedPartialResponse` holds the
+          // onFinish message if it fired; otherwise the buffered chunks are
+          // reconstructed as the fallback (mirrors chat.pipeAndCapture).
+          let capturedPartialResponse: TUIMessage | undefined;
+          let responseCommitted = false;
+          const turnBufferedChunks: UIMessageChunk[] = [];
           try {
             // Extract turn-level context before entering the span. Slim
             // wire: at most one delta message per record. `headStartMessages`
@@ -7175,11 +7184,17 @@ function chatAgent<
                           finishReason?: FinishReason;
                         }) => {
                           capturedResponseMessage = responseMessage as TUIMessage;
+                          capturedPartialResponse = responseMessage as TUIMessage;
                           capturedFinishReason = finishReason;
                           resolveOnFinish!();
                         },
                       });
-                      await pipeChat(uiStream, {
+                      // Buffer chunks as they flow to the pipe so a source-stream
+                      // transport failure — which abandons the UI stream before
+                      // onFinish fires — can still reconstruct the partial in the
+                      // error handler instead of dropping it. The happy path pays
+                      // nothing extra; reassembly only runs in the error fallback.
+                      await pipeChat(tapUIMessageChunks(uiStream, turnBufferedChunks), {
                         signal: combinedSignal,
                         spanName: "stream response",
                       });
@@ -7390,6 +7405,12 @@ function chatAgent<
                     }
                   }
 
+                  if (capturedResponseMessage) {
+                    responseCommitted = true;
+                    capturedPartialResponse = capturedResponseMessage;
+                    turnBufferedChunks.length = 0;
+                  }
+
                   if (runSignal.aborted) return "exit";
 
                   // Await deferred background work (e.g. DB writes from onTurnStart)
@@ -7611,6 +7632,7 @@ function chatAgent<
                         parts: [...(msg.parts ?? []), ...lateParts],
                       } as TUIMessage;
                       capturedResponseMessage = accumulatedUIMessages[idx] as TUIMessage;
+                      capturedPartialResponse = capturedResponseMessage;
                       turnCompleteEvent.responseMessage = capturedResponseMessage;
                       turnCompleteEvent.uiMessages = accumulatedUIMessages;
                     }
@@ -7904,10 +7926,91 @@ function chatAgent<
                 ? [...accumulatedUIMessages, erroredWireMessage]
                 : accumulatedUIMessages;
 
+            // Recover the partial assistant output the model streamed before the
+            // failure. A source-stream transport error (e.g. UND_ERR_BODY_TIMEOUT)
+            // abandons the UI stream before onFinish fires, so fall back to
+            // reconstructing the partial from the buffered chunks. Surfacing it on
+            // the error-path event lets persistence keep the partial instead of
+            // losing it — critical when hydrateMessages disables boot-time tail
+            // replay, so the recovery path can't reclaim it later. Empty for
+            // non-stream failures (nothing buffered), preserving prior behavior.
+            let partialResponse: TUIMessage | undefined =
+              capturedPartialResponse ??
+              ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined);
+            if (partialResponse) {
+              partialResponse = cleanupAbortedParts(partialResponse);
+            }
+
+            // Build the complete error UI state. A HITL/tool continuation partial
+            // reuses an existing assistant id, so replace it in place (like the
+            // success path) instead of dropping it as a dup; otherwise append.
+            // `erroredWireMessage` was already folded into `erroredUIMessages`
+            // above when the pre-run merge hadn't happened.
+            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() } as TUIMessage;
+            }
+            const includePartial = partialResponse != null && !responseCommitted;
+            let erroredUIMessagesWithPartial: TUIMessage[] = !includePartial
+              ? erroredUIMessages
+              : partialIdx === -1
+                ? [...erroredUIMessages, partialResponse]
+                : (erroredUIMessages.map((m, i) =>
+                    i === partialIdx ? partialResponse : m
+                  ) as TUIMessage[]);
+
+            let erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : [];
+            if (includePartial) {
+              erroredNewUIMessages.push(partialResponse!);
+            }
+
+            let erroredNewModelMessages: ModelMessage[] = [];
+
+            // Commit the complete error state to the canonical accumulator so the
+            // errored user message and any recovered partial survive past this
+            // hook: the run stays alive after an error, so the next turn sees
+            // them, and the error-path snapshot below (the recovery source for
+            // non-hydrate apps) persists them for reboot. Matches the success
+            // path. Skip when nothing changed so an errored turn never needlessly
+            // reconverts. When the only change is an appended partial, push just
+            // its model messages to preserve a prior turn's compaction (mirrors
+            // the success path's append branch); otherwise reconvert from UI.
+            // Guard the conversion so a secondary failure can't crash the run.
+            if (!responseCommitted && erroredUIMessagesWithPartial !== accumulatedUIMessages) {
+              try {
+                if (includePartial) {
+                  erroredNewModelMessages = await toModelMessages([
+                    stripProviderMetadata(partialResponse!),
+                  ]);
+                }
+                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 {
+                // Keep the prior model accumulator if conversion fails.
+                erroredNewModelMessages = [];
+                erroredUIMessagesWithPartial = erroredUIMessages;
+                erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : [];
+              }
+            }
+
             // 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.
+            // mark the turn failed. `responseMessage` carries any partial
+            // recovered above and `error` carries the thrown value.
             if (onTurnComplete) {
               try {
                 await tracer.startActiveSpan(
@@ -7917,11 +8020,11 @@ function chatAgent<
                       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: "",
@@ -7967,7 +8070,7 @@ function chatAgent<
                 await writeChatSnapshot<TUIMessage>(sessionIdForSnapshot, {
                   version: 1,
                   savedAt: Date.now(),
-                  messages: erroredUIMessages,
+                  messages: erroredUIMessagesWithPartial,
                   lastOutEventId: errorTurnCompleteResult?.lastEventId,
                   lastInEventId:
                     errorSnapshotInCursor !== undefined ? String(errorSnapshotInCursor) : undefined,
@@ -8887,28 +8990,26 @@ export type PipeAndCaptureResult = {
  * can return, and propagates a source error to the consumer after buffering
  * whatever streamed first. See {@link pipeChatAndCapture} for why.
  */
-async function* tapUIMessageChunks(
+function tapUIMessageChunks(
   source: AsyncIterable<unknown> | ReadableStream<unknown>,
   buffer: UIMessageChunk[]
-): AsyncGenerator<unknown> {
+): ReadableStream<unknown> | AsyncGenerator<unknown> {
   if (isReadableStream(source)) {
-    const reader = source.getReader();
-    try {
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) break;
-        buffer.push(value as UIMessageChunk);
-        yield value;
-      }
-    } finally {
-      reader.releaseLock();
-    }
-  } else {
+    return source.pipeThrough(
+      new TransformStream<unknown, unknown>({
+        transform(chunk, controller) {
+          buffer.push(chunk as UIMessageChunk);
+          controller.enqueue(chunk);
+        },
+      })
+    );
+  }
+  return (async function* () {
     for await (const chunk of source) {
       buffer.push(chunk as UIMessageChunk);
       yield chunk;
     }
-  }
+  })();
 }
 
 /**
@@ -9744,6 +9845,15 @@ function createChatSession(
                 // 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 as any).parts = [...(partial.parts ?? []), ...queuedParts];
+                      locals.set(chatResponsePartsKey, []);
+                    }
+                    await accumulator.addResponse(partial);
+                  }
                   throw captured.error;
                 }
                 response = captured.message;