fix(sdk): persist recovered partial to accumulator and snapshot on error

Commit the recovered partial to the canonical accumulator on the chat.agent
error path so the next turn and the reboot snapshot both carry it, matching
the success path. Fold queued response parts into the manual-loop error
partial too. Add a continuation regression test.

Author: matt-aitken

packages/trigger-sdk/src/v3/ai.ts

@@ -7944,6 +7944,23 @@ function chatAgent<
                 ? [...erroredUIMessages, partialResponse]
                 : erroredUIMessages;
 
+            // Commit the recovered partial to the canonical accumulator so the
+            // partial survives past this hook: the run stays alive after an
+            // error, so the next turn sees it, and the error-path snapshot below
+            // (the recovery source for non-hydrate apps) persists it for reboot.
+            // Matches the success path, which accumulates the response. Guard the
+            // model-message conversion so a secondary failure here can't crash
+            // the still-alive run.
+            if (partialResponse) {
+              accumulatedUIMessages = erroredUIMessagesWithPartial as TUIMessage[];
+              locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
+              try {
+                accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+              } catch {
+                // Keep the prior model accumulator if conversion fails.
+              }
+            }
+
             // 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` carries any partial
@@ -8007,7 +8024,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,
@@ -9788,9 +9805,19 @@ function createChatSession(
                   // accumulate it (mirroring the stop path) so `turn.uiMessages`
                   // reflects it and the caller can persist it after catching,
                   // then rethrow. Without this the partial pipeAndCapture
-                  // reconstructed is silently dropped on rethrow.
+                  // reconstructed is silently dropped on rethrow. Fold in any
+                  // data parts queued via chat.response / writer.write() this
+                  // turn, same as the success path below. cleanupAbortedParts is
+                  // intentionally skipped: it only runs on a user stop, not on a
+                  // hard transport error (which is not an abort).
                   if (captured.message) {
-                    await accumulator.addResponse(captured.message);
+                    const partial = 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;
                 }

packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts

@@ -3,7 +3,10 @@
 import { mockChatAgent } from "../src/v3/test/index.js";
 
 import { describe, expect, it } from "vitest";
-import type { UIMessage } from "ai";
+import type { ModelMessage, UIMessage } from "ai";
+import { simulateReadableStream, streamText } from "ai";
+import { MockLanguageModelV3 } from "ai/test";
+import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
 import { chat } from "../src/v3/ai.js";
 import type { TurnCompleteEvent } from "../src/v3/ai.js";
 
@@ -96,6 +99,67 @@ describe("chat.agent managed loop — source-stream failure", () => {
       await harness.close();
     }
   });
+
+  it("carries the recovered partial into the next turn's accumulated messages", async () => {
+    let turn = 0;
+    let turn2Messages: ModelMessage[] | undefined;
+
+    const okStream = () =>
+      simulateReadableStream({
+        chunks: [
+          { type: "text-start", id: "t2" },
+          { type: "text-delta", id: "t2", delta: "second answer" },
+          { type: "text-end", id: "t2" },
+          {
+            type: "finish",
+            finishReason: { unified: "stop", raw: "stop" },
+            usage: {
+              inputTokens: { total: 5, noCache: 5, cacheRead: undefined, cacheWrite: undefined },
+              outputTokens: { total: 5, text: 5, reasoning: undefined },
+            },
+          },
+        ] as LanguageModelV3StreamPart[],
+      });
+
+    const agent = chat.agent({
+      id: "chatAgent.source-stream-error-continuation",
+      run: async ({ messages }) => {
+        turn++;
+        if (turn === 1) {
+          return erroringSource("UND_ERR_BODY_TIMEOUT") as never;
+        }
+        // Second turn: the failed turn's partial assistant output must be in
+        // the accumulated history the model now sees.
+        turn2Messages = messages;
+        return streamText({
+          model: new MockLanguageModelV3({ doStream: async () => ({ stream: okStream() }) }),
+          messages,
+        });
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-source-error-cont" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await harness.sendMessage(userMessage("still there?", "u-2"));
+      await waitFor(() => turn2Messages !== undefined);
+
+      const assistantText = turn2Messages!
+        .filter((m) => m.role === "assistant")
+        .map((m) =>
+          typeof m.content === "string"
+            ? m.content
+            : (m.content as Array<{ type: string; text?: string }>)
+                .filter((p) => p.type === "text")
+                .map((p) => p.text ?? "")
+                .join("")
+        )
+        .join("");
+      expect(assistantText).toContain("partial answer");
+    } finally {
+      await harness.close();
+    }
+  });
 });
 
 describe("chat.createSession turn.complete() — source-stream failure", () => {