fix(sdk): preserve partial assistant message on chat stream failure (#4348)

## Summary

When a `chat.agent` (or `chat.createSession`) turn's model stream fails
mid-response (e.g. a transport timeout like `UND_ERR_BODY_TIMEOUT`), the
assistant output that already streamed was dropped: `onTurnComplete`
fired with `responseMessage: undefined`, and the manual loop's
`turn.complete()` rethrew without keeping the partial. Apps that
register `hydrateMessages` are hit hardest, since boot-time tail-replay
recovery is off by design.

This preserves the streamed-so-far assistant output while still
reporting the turn as errored, so persistence and recovery keep it.

## Scope of behavior change

Only the **error path** changes. Successful turns are unaffected: the
same chunks stream to the client in the same order, and
backpressure/cancel behave as before. Everything here is a correctness
improvement on a turn that hit a source-stream failure.

## What it does

Follow-up to #4304 (`chat.pipeAndCapture`), extending the same
partial-recovery to the two loops that lacked it:

- **`chat.agent`**: taps the response stream (via a `TransformStream`,
so pass-through backpressure and cancel are preserved) to buffer chunks,
and on a source-stream failure reconstructs the partial (preferring the
`onFinish` message). It's surfaced on the error-path `onTurnComplete`
(`responseMessage`, `rawResponseMessage`, `uiMessages`, `newUIMessages`,
`newMessages`) and committed to the accumulator so the next turn and the
reboot snapshot keep it.
- **`chat.createSession` / `turn.complete()`**: the reconstructed
partial is accumulated (so `turn.uiMessages` reflects it and the caller
can persist after catching) before `turn.complete()` rethrows.

`onBeforeTurnComplete` stays skipped on the error path (it hands out a
writer for a stream that has already broken).

## Correctness properties (each covered by a regression test)

Each test below was confirmed to fail without its fix:

- The recovered partial reaches `onTurnComplete` and the next turn's
accumulated messages.
- An already-committed (possibly enriched) response is not overwritten
if a post-response hook then throws.
- Incomplete tool parts are cleaned from the recovered partial (text
kept), so the UI and model views agree and the next turn isn't poisoned.
- A prior turn's model-only compaction survives an errored turn (append
only the new tail, don't reconvert the full history).
- A reconstructed fragment that reuses an existing message id does not
clobber the complete message.
- Queued `chat.response` data parts are folded into the recovered
partial, matching the success path.
- `newMessages` (model delta) stays symmetric with `newUIMessages`.

## Tests

New `chat-agent-source-stream-error.test.ts` covers the cases above. The
full `@trigger.dev/sdk` unit suite passes and the package build is green
across all supported runtimes (Node 20 to 26, Bun, Deno, Cloudflare
Workers).

.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. `chat.agent` now passes the recovered partial to `onTurnComplete`, and `chat.createSession`'s `turn.complete()` keeps it before rethrowing, instead of dropping the streamed-so-far output.

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

@@ -6389,6 +6389,9 @@ 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;
+          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 +7178,12 @@ function chatAgent<
                           finishReason?: FinishReason;
                         }) => {
                           capturedResponseMessage = responseMessage as TUIMessage;
+                          capturedPartialResponse = responseMessage as TUIMessage;
                           capturedFinishReason = finishReason;
                           resolveOnFinish!();
                         },
                       });
-                      await pipeChat(uiStream, {
+                      await pipeChat(tapUIMessageChunks(uiStream, turnBufferedChunks), {
                         signal: combinedSignal,
                         spanName: "stream response",
                       });
@@ -7390,6 +7394,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 +7621,7 @@ function chatAgent<
                         parts: [...(msg.parts ?? []), ...lateParts],
                       } as TUIMessage;
                       capturedResponseMessage = accumulatedUIMessages[idx] as TUIMessage;
+                      capturedPartialResponse = capturedResponseMessage;
                       turnCompleteEvent.responseMessage = capturedResponseMessage;
                       turnCompleteEvent.uiMessages = accumulatedUIMessages;
                     }
@@ -7904,10 +7915,77 @@ function chatAgent<
                 ? [...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: TUIMessage | undefined =
+              capturedPartialResponse ??
+              ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined);
+            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() } as TUIMessage;
+            }
+            if (partialResponse && !responseCommitted) {
+              const queuedParts = locals.get(chatResponsePartsKey);
+              if (queuedParts && queuedParts.length > 0) {
+                partialResponse = {
+                  ...partialResponse,
+                  parts: [...partialResponse.parts, ...(queuedParts as UIMessage["parts"])],
+                } as TUIMessage;
+                locals.set(chatResponsePartsKey, []);
+              }
+            }
+            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[] = [];
+
+            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(
@@ -7917,11 +7995,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 +8045,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 +8965,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 +9820,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;

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

@@ -0,0 +1,435 @@
+// Import the test harness FIRST — this installs the resource catalog so
+// `chat.agent()` calls below register their task functions correctly.
+import { mockChatAgent } from "../src/v3/test/index.js";
+
+import { describe, expect, it } from "vitest";
+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";
+
+// ── Helpers ────────────────────────────────────────────────────────────
+
+function userMessage(text: string, id: string): UIMessage {
+  return { id, role: "user", parts: [{ type: "text", text }] };
+}
+
+function extractText(message: UIMessage | undefined): string {
+  if (!message) return "";
+  return (message.parts as Array<{ type: string; text?: string }>)
+    .filter((p) => p.type === "text")
+    .map((p) => p.text ?? "")
+    .join("");
+}
+
+async function waitFor(check: () => boolean, timeoutMs = 5_000) {
+  const start = Date.now();
+  while (Date.now() - start < timeoutMs) {
+    if (check()) return;
+    await new Promise((r) => setTimeout(r, 20));
+  }
+  throw new Error("waitFor timed out");
+}
+
+/**
+ * A `run()` return value that looks like a `StreamTextResult` (has
+ * `toUIMessageStream()`) but whose UI stream emits a partial assistant
+ * message and then errors — reproducing a source-stream transport failure
+ * (e.g. `UND_ERR_BODY_TIMEOUT`) mid-turn. `onFinish` is never invoked, which
+ * is exactly what happens on a hard transport error. Chunks are delivered
+ * one-per-pull before the error so they aren't discarded (calling
+ * `controller.error()` in the same tick as `enqueue()` resets the queue).
+ */
+function erroringSource(errorMessage: string) {
+  const partialChunks = [
+    { type: "start", messageId: "a-err" },
+    { type: "text-start", id: "t1" },
+    { type: "text-delta", id: "t1", delta: "partial answer" },
+  ];
+  return sourceFromChunks(partialChunks, errorMessage);
+}
+
+function sourceFromChunks(chunks: unknown[], errorMessage: string) {
+  return {
+    toUIMessageStream() {
+      let i = 0;
+      return new ReadableStream({
+        pull(controller) {
+          if (i < chunks.length) {
+            controller.enqueue(chunks[i++]);
+          } else {
+            controller.error(new Error(errorMessage));
+          }
+        },
+      });
+    },
+  };
+}
+
+// ── Tests ──────────────────────────────────────────────────────────────
+
+describe("chat.agent managed loop — source-stream failure", () => {
+  it("preserves the partial assistant message on onTurnComplete when the source stream fails", async () => {
+    const turnCompletes: TurnCompleteEvent<unknown, UIMessage>[] = [];
+
+    const agent = chat.agent({
+      id: "chatAgent.source-stream-error",
+      run: async () => erroringSource("UND_ERR_BODY_TIMEOUT") as never,
+      onTurnComplete: async (event) => {
+        turnCompletes.push(event);
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-source-error" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => turnCompletes.length >= 1);
+
+      const evt = turnCompletes[0]!;
+
+      expect(evt.finishReason).toBe("error");
+      expect(evt.error).toBeInstanceOf(Error);
+      expect((evt.error as Error).message).toBe("UND_ERR_BODY_TIMEOUT");
+
+      expect(evt.responseMessage).toBeDefined();
+      expect(extractText(evt.responseMessage)).toBe("partial answer");
+
+      const newAssistantText = (evt.newMessages as ModelMessage[])
+        .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(newAssistantText).toBe("partial answer");
+
+      expect((evt.newMessages as ModelMessage[]).some((m) => m.role === "user")).toBe(true);
+      expect((evt.newUIMessages as UIMessage[]).some((m) => m.role === "user")).toBe(true);
+    } finally {
+      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;
+        }
+        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();
+    }
+  });
+
+  it("does not overwrite an already-committed enriched response when a post-response hook throws", async () => {
+    const events: TurnCompleteEvent<unknown, UIMessage>[] = [];
+
+    const okModel = (text: string) =>
+      new MockLanguageModelV3({
+        doStream: async () => ({
+          stream: simulateReadableStream({
+            chunks: [
+              { type: "text-start", id: "t1" },
+              { type: "text-delta", id: "t1", delta: text },
+              { type: "text-end", id: "t1" },
+              {
+                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.post-commit-hook-throw",
+      run: async ({ messages }) => {
+        chat.response.write({ type: "data-marker", data: { kept: true } } as never);
+        return streamText({ model: okModel("full response"), messages });
+      },
+      onTurnComplete: async (event) => {
+        events.push(event);
+        if (event.error == null) {
+          throw new Error("hook boom after commit");
+        }
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-post-commit-throw" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => events.some((e) => e.error != null));
+
+      const errorEvent = events.find((e) => e.error != null)!;
+      const assistant = (errorEvent.uiMessages as UIMessage[]).find((m) => m.role === "assistant");
+      expect(assistant).toBeDefined();
+      expect(
+        (assistant!.parts as Array<{ type: string }>).some((p) => p.type === "data-marker")
+      ).toBe(true);
+      expect(extractText(assistant)).toBe("full response");
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("does not clobber an existing message when a reconstructed fragment reuses its id", async () => {
+    let turn = 0;
+    let firstAssistantId: string | undefined;
+    const events: TurnCompleteEvent<unknown, UIMessage>[] = [];
+
+    const okModel = () =>
+      new MockLanguageModelV3({
+        doStream: async () => ({
+          stream: simulateReadableStream({
+            chunks: [
+              { type: "text-start", id: "t1" },
+              { type: "text-delta", id: "t1", delta: "first answer" },
+              { type: "text-end", id: "t1" },
+              {
+                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 collidingErroringSource = (id: string) => ({
+      toUIMessageStream() {
+        const chunks = [
+          { type: "start", messageId: id },
+          { type: "text-start", id: "t2" },
+          { type: "text-delta", id: "t2", delta: "clobber" },
+        ];
+        let i = 0;
+        return new ReadableStream({
+          pull(controller) {
+            if (i < chunks.length) controller.enqueue(chunks[i++]);
+            else controller.error(new Error("UND_ERR_BODY_TIMEOUT"));
+          },
+        });
+      },
+    });
+
+    const agent = chat.agent({
+      id: "chatAgent.fragment-id-collision",
+      run: async ({ messages }) => {
+        turn++;
+        if (turn === 1) {
+          return streamText({ model: okModel(), messages });
+        }
+        return collidingErroringSource(firstAssistantId!) as never;
+      },
+      onTurnComplete: async (event) => {
+        events.push(event);
+        if (event.error == null && event.responseMessage) {
+          firstAssistantId = event.responseMessage.id;
+        }
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-fragment-collision" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => firstAssistantId !== undefined);
+      await harness.sendMessage(userMessage("again", "u-2"));
+      await waitFor(() => events.some((e) => e.error != null));
+
+      const errorEvent = events.find((e) => e.error != null)!;
+      const preserved = (errorEvent.uiMessages as UIMessage[]).find(
+        (m) => m.id === firstAssistantId
+      );
+      expect(preserved).toBeDefined();
+      expect(extractText(preserved)).toBe("first answer");
+      expect(
+        (errorEvent.uiMessages as UIMessage[]).some((m) => extractText(m).includes("clobber"))
+      ).toBe(false);
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("cleans dangling tool parts from the recovered partial while keeping its text", async () => {
+    const turnCompletes: TurnCompleteEvent<unknown, UIMessage>[] = [];
+
+    const agent = chat.agent({
+      id: "chatAgent.error-partial-cleanup",
+      run: async () =>
+        sourceFromChunks(
+          [
+            { type: "start", messageId: "a-tool" },
+            { type: "text-start", id: "t1" },
+            { type: "text-delta", id: "t1", delta: "thinking" },
+            { type: "text-end", id: "t1" },
+            { type: "tool-input-start", toolCallId: "tc1", toolName: "search" },
+            {
+              type: "tool-input-available",
+              toolCallId: "tc1",
+              toolName: "search",
+              input: { q: "x" },
+            },
+          ],
+          "UND_ERR_BODY_TIMEOUT"
+        ) as never,
+      onTurnComplete: async (event) => {
+        turnCompletes.push(event);
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-partial-cleanup" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => turnCompletes.length >= 1);
+
+      const evt = turnCompletes[0]!;
+      expect(evt.responseMessage).toBeDefined();
+      const parts = evt.responseMessage!.parts as Array<{ type: string }>;
+      expect(extractText(evt.responseMessage)).toBe("thinking");
+      expect(parts.some((p) => p.type.startsWith("tool-"))).toBe(false);
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("folds queued response data parts into the recovered partial", async () => {
+    const turnCompletes: TurnCompleteEvent<unknown, UIMessage>[] = [];
+
+    const agent = chat.agent({
+      id: "chatAgent.error-queued-parts",
+      run: async () => {
+        chat.response.write({ type: "data-marker", data: { kept: true } } as never);
+        return erroringSource("UND_ERR_BODY_TIMEOUT") as never;
+      },
+      onTurnComplete: async (event) => {
+        turnCompletes.push(event);
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cae-error-queued-parts" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => turnCompletes.length >= 1);
+
+      const evt = turnCompletes[0]!;
+      expect(evt.responseMessage).toBeDefined();
+      const parts = evt.responseMessage!.parts as Array<{ type: string }>;
+      expect(extractText(evt.responseMessage)).toBe("partial answer");
+      expect(parts.some((p) => p.type === "data-marker")).toBe(true);
+    } finally {
+      await harness.close();
+    }
+  });
+});
+
+describe("chat.createSession turn.complete() — source-stream failure", () => {
+  it("accumulates the partial before rethrowing so the caller can persist it", async () => {
+    let caughtError: unknown;
+    let uiMessagesAfterError: UIMessage[] = [];
+
+    const agent = chat.customAgent({
+      id: "createSession.source-stream-error",
+      run: async (payload) => {
+        const session = chat.createSession(payload, {
+          signal: new AbortController().signal,
+          idleTimeoutInSeconds: 2,
+        });
+        for await (const turn of session) {
+          try {
+            await turn.complete(erroringSource("UND_ERR_BODY_TIMEOUT") as never);
+          } catch (err) {
+            caughtError = err;
+            uiMessagesAfterError = [...turn.uiMessages];
+            await turn.done();
+          }
+        }
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "cs-source-error" });
+    try {
+      await harness.sendMessage(userMessage("hi", "u-1"));
+      await waitFor(() => caughtError !== undefined);
+
+      expect(caughtError).toBeInstanceOf(Error);
+      expect((caughtError as Error).message).toBe("UND_ERR_BODY_TIMEOUT");
+
+      const partial = uiMessagesAfterError.find((m) => m.role === "assistant");
+      expect(partial).toBeDefined();
+      expect(extractText(partial)).toBe("partial answer");
+    } finally {
+      await harness.close();
+    }
+  });
+});