From 8152444cfd4264c0bf1259485d4f88f404b58ae0 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 14:58:02 +0100 Subject: [PATCH 01/13] fix(sdk): preserve partial assistant message on chat stream failure When a chat turn's model stream fails mid-response (e.g. a transport timeout), the streamed-so-far output is no longer dropped. chat.agent passes the recovered partial to onTurnComplete, and chat.createSession accumulates it before turn.complete() rethrows, so it survives for persistence even when hydrateMessages disables boot-time replay recovery. The turn is still reported as errored. --- ...gent-preserve-partial-on-stream-failure.md | 5 + packages/trigger-sdk/src/v3/ai.ts | 62 +++++++- .../chat-agent-source-stream-error.test.ts | 142 ++++++++++++++++++ 3 files changed, 202 insertions(+), 7 deletions(-) create mode 100644 .changeset/chat-agent-preserve-partial-on-stream-failure.md create mode 100644 packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts diff --git a/.changeset/chat-agent-preserve-partial-on-stream-failure.md b/.changeset/chat-agent-preserve-partial-on-stream-failure.md new file mode 100644 index 00000000000..2c2236e6033 --- /dev/null +++ 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 index 4612302c4ae..2c470a7ee4b 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -6389,6 +6389,14 @@ 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; + const turnBufferedChunks: UIMessageChunk[] = []; try { // Extract turn-level context before entering the span. Slim // wire: at most one delta message per record. `headStartMessages` @@ -7175,11 +7183,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", }); @@ -7904,10 +7918,36 @@ 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. + const partialResponse: TUIMessage | undefined = + capturedPartialResponse ?? + ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined); + + // Include the partial in the UI-message views too, so customers who + // persist from uiMessages / newUIMessages (not responseMessage) keep + // it as well. Dedup by id against the accumulator to be safe. + const erroredNewUIMessages: TUIMessage[] = erroredWireMessage + ? [erroredWireMessage] + : []; + if (partialResponse) { + erroredNewUIMessages.push(partialResponse); + } + const erroredUIMessagesWithPartial = + partialResponse && !erroredUIMessages.some((m) => m.id === partialResponse.id) + ? [...erroredUIMessages, partialResponse] + : erroredUIMessages; + // 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 +7957,11 @@ function chatAgent< ctx, chatId: currentWirePayload.chatId, messages: accumulatedMessages, - uiMessages: erroredUIMessages, + uiMessages: erroredUIMessagesWithPartial, newMessages: [], - newUIMessages: erroredWireMessage ? [erroredWireMessage] : [], - responseMessage: undefined, - rawResponseMessage: undefined, + newUIMessages: erroredNewUIMessages, + responseMessage: partialResponse, + rawResponseMessage: partialResponse, turn, runId: ctx.run.id, chatAccessToken: "", @@ -9744,6 +9784,14 @@ 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") { + // Preserve the partial the model streamed before the failure: + // 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. + if (captured.message) { + await accumulator.addResponse(captured.message); + } throw captured.error; } response = captured.message; diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts new file mode 100644 index 00000000000..5b6d8d807cf --- /dev/null +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -0,0 +1,142 @@ +// 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 { UIMessage } from "ai"; +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 { + toUIMessageStream() { + let i = 0; + return new ReadableStream({ + pull(controller) { + if (i < partialChunks.length) { + controller.enqueue(partialChunks[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[] = []; + + 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]!; + + // The turn is reported as errored, carrying the thrown transport error. + expect(evt.finishReason).toBe("error"); + expect(evt.error).toBeInstanceOf(Error); + expect((evt.error as Error).message).toBe("UND_ERR_BODY_TIMEOUT"); + + // The partial assistant output that streamed before the failure must be + // preserved so persistence / recovery can keep it, instead of being + // dropped (responseMessage: undefined). + expect(evt.responseMessage).toBeDefined(); + expect(extractText(evt.responseMessage)).toBe("partial answer"); + } 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; + // The partial must be accumulated so persistence from the session + // state keeps it, rather than being lost on the rethrow. + 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(); + } + }); +}); From dc63a6e6266e5e0bd1542d0b1e58de9199a3ef4c Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 15:07:30 +0100 Subject: [PATCH 02/13] 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. --- packages/trigger-sdk/src/v3/ai.ts | 33 +++++++++- .../chat-agent-source-stream-error.test.ts | 66 ++++++++++++++++++- 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 2c470a7ee4b..2795e72a975 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/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(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; } diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index 5b6d8d807cf..eccc607c048 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/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", () => { From 1fe685cf3145677ff2d7ea47761e84962dd1c177 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 15:18:42 +0100 Subject: [PATCH 03/13] fix(sdk): mirror success-path accumulation on chat.agent error path Replace a same-id continuation partial in place instead of dropping it as a dup, commit the errored user message unconditionally so it reaches the next live turn, and push only the appended partial's model messages to preserve a prior turn's compaction (reconvert only when replacing or folding). --- packages/trigger-sdk/src/v3/ai.ts | 59 +++++++++++++++++++++---------- 1 file changed, 41 insertions(+), 18 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 2795e72a975..32844b5467d 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7930,32 +7930,55 @@ function chatAgent< capturedPartialResponse ?? ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined); - // Include the partial in the UI-message views too, so customers who - // persist from uiMessages / newUIMessages (not responseMessage) keep - // it as well. Dedup by id against the accumulator to be safe. + // 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. + const partialIdx = + partialResponse?.id != null + ? erroredUIMessages.findIndex((m) => m.id === partialResponse!.id) + : -1; + const erroredUIMessagesWithPartial: TUIMessage[] = !partialResponse + ? erroredUIMessages + : partialIdx === -1 + ? [...erroredUIMessages, partialResponse] + : (erroredUIMessages.map((m, i) => + i === partialIdx ? partialResponse : m + ) as TUIMessage[]); + const erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : []; if (partialResponse) { erroredNewUIMessages.push(partialResponse); } - const erroredUIMessagesWithPartial = - partialResponse && !erroredUIMessages.some((m) => m.id === partialResponse.id) - ? [...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[]; + + // 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 (erroredUIMessagesWithPartial !== accumulatedUIMessages) { + const onlyAppendedPartial = + partialResponse != null && + partialIdx === -1 && + erroredUIMessages === accumulatedUIMessages; + accumulatedUIMessages = erroredUIMessagesWithPartial; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); try { - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + if (onlyAppendedPartial) { + accumulatedMessages.push( + ...(await toModelMessages([stripProviderMetadata(partialResponse!)])) + ); + } else { + accumulatedMessages = await toModelMessages(accumulatedUIMessages); + } } catch { // Keep the prior model accumulator if conversion fails. } From b6ea02fd6c69095af84882ce2d482d9035e28b40 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 17:27:34 +0100 Subject: [PATCH 04/13] fix(sdk): don't overwrite a committed response when a post-response step throws The chat.agent try block spans the whole turn, so a throw from a post-response step (a customer onBeforeTurnComplete/onTurnComplete hook, a late conversion) lands in the error handler after the response was already committed. Track a per-turn responseCommitted flag and keep capturedPartialResponse pointed at the enriched committed message, so the error path reports it without re-recovering a raw partial and clobbering the committed message, its queued data parts, or a prior turn's compaction. Also use truthy id checks to match the success path. --- packages/trigger-sdk/src/v3/ai.ts | 18 ++++-- .../chat-agent-source-stream-error.test.ts | 60 +++++++++++++++++++ 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 32844b5467d..ffa62dc1d48 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -6396,6 +6396,7 @@ function chatAgent< // 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 @@ -7404,6 +7405,11 @@ function chatAgent< } } + if (capturedResponseMessage) { + responseCommitted = true; + capturedPartialResponse = capturedResponseMessage; + } + if (runSignal.aborted) return "exit"; // Await deferred background work (e.g. DB writes from onTurnStart) @@ -7625,6 +7631,7 @@ function chatAgent< parts: [...(msg.parts ?? []), ...lateParts], } as TUIMessage; capturedResponseMessage = accumulatedUIMessages[idx] as TUIMessage; + capturedPartialResponse = capturedResponseMessage; turnCompleteEvent.responseMessage = capturedResponseMessage; turnCompleteEvent.uiMessages = accumulatedUIMessages; } @@ -7935,10 +7942,9 @@ function chatAgent< // 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. - const partialIdx = - partialResponse?.id != null - ? erroredUIMessages.findIndex((m) => m.id === partialResponse!.id) - : -1; + const partialIdx = partialResponse?.id + ? erroredUIMessages.findIndex((m) => m.id === partialResponse!.id) + : -1; const erroredUIMessagesWithPartial: TUIMessage[] = !partialResponse ? erroredUIMessages : partialIdx === -1 @@ -7950,7 +7956,7 @@ function chatAgent< const erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : []; - if (partialResponse) { + if (partialResponse && !responseCommitted) { erroredNewUIMessages.push(partialResponse); } @@ -7964,7 +7970,7 @@ function chatAgent< // 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 (erroredUIMessagesWithPartial !== accumulatedUIMessages) { + if (!responseCommitted && erroredUIMessagesWithPartial !== accumulatedUIMessages) { const onlyAppendedPartial = partialResponse != null && partialIdx === -1 && diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index eccc607c048..ffb57630d00 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -160,6 +160,66 @@ describe("chat.agent managed loop — source-stream failure", () => { await harness.close(); } }); + + it("does not overwrite an already-committed enriched response when a post-response hook throws", async () => { + const events: TurnCompleteEvent[] = []; + + 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(); + } + }); }); describe("chat.createSession turn.complete() — source-stream failure", () => { From a225c369ea8ed3ca67249dd8c4b25238c9eadd73 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 19:09:23 +0100 Subject: [PATCH 05/13] fix(sdk): include the recovered partial in the error event's newMessages The error-path onTurnComplete added the partial to newUIMessages/uiMessages but left newMessages (the model-message delta) empty, so apps persisting the model delta missed the partial. Populate it from the same conversion. --- packages/trigger-sdk/src/v3/ai.ts | 13 +++++++++---- .../test/chat-agent-source-stream-error.test.ts | 13 +++++++++++++ 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index ffa62dc1d48..f328f8305d6 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7960,6 +7960,8 @@ function chatAgent< 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 @@ -7978,10 +7980,13 @@ function chatAgent< accumulatedUIMessages = erroredUIMessagesWithPartial; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); try { + if (partialResponse) { + erroredNewModelMessages = await toModelMessages([ + stripProviderMetadata(partialResponse), + ]); + } if (onlyAppendedPartial) { - accumulatedMessages.push( - ...(await toModelMessages([stripProviderMetadata(partialResponse!)])) - ); + accumulatedMessages.push(...erroredNewModelMessages); } else { accumulatedMessages = await toModelMessages(accumulatedUIMessages); } @@ -8004,7 +8009,7 @@ function chatAgent< chatId: currentWirePayload.chatId, messages: accumulatedMessages, uiMessages: erroredUIMessagesWithPartial, - newMessages: [], + newMessages: erroredNewModelMessages, newUIMessages: erroredNewUIMessages, responseMessage: partialResponse, rawResponseMessage: partialResponse, diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index ffb57630d00..a001ec5c21c 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -95,6 +95,19 @@ describe("chat.agent managed loop — source-stream failure", () => { // dropped (responseMessage: undefined). 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"); } finally { await harness.close(); } From 882364a4d8ab1a5b1063278b884a5facb2d9d0e1 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 20:11:47 +0100 Subject: [PATCH 06/13] fix(sdk): harden error-path partial recovery against fragment clobber and desync Drop a reconstructed-from-chunks partial that reuses an existing message id (a fragment can't be safely merged into a complete message the way an onFinish message can), and make the error-path accumulator update atomic so a failed model-message conversion can't leave the UI and model accumulators out of sync. --- packages/trigger-sdk/src/v3/ai.ts | 14 +-- .../chat-agent-source-stream-error.test.ts | 86 +++++++++++++++++++ 2 files changed, 95 insertions(+), 5 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index f328f8305d6..42bedf160c2 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7933,7 +7933,7 @@ function chatAgent< // 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. - const partialResponse: TUIMessage | undefined = + let partialResponse: TUIMessage | undefined = capturedPartialResponse ?? ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined); @@ -7942,9 +7942,13 @@ function chatAgent< // 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. - const partialIdx = partialResponse?.id + let partialIdx = partialResponse?.id ? erroredUIMessages.findIndex((m) => m.id === partialResponse!.id) : -1; + if (partialResponse && capturedPartialResponse === undefined && partialIdx !== -1) { + partialResponse = undefined; + partialIdx = -1; + } const erroredUIMessagesWithPartial: TUIMessage[] = !partialResponse ? erroredUIMessages : partialIdx === -1 @@ -7977,8 +7981,6 @@ function chatAgent< partialResponse != null && partialIdx === -1 && erroredUIMessages === accumulatedUIMessages; - accumulatedUIMessages = erroredUIMessagesWithPartial; - locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); try { if (partialResponse) { erroredNewModelMessages = await toModelMessages([ @@ -7988,8 +7990,10 @@ function chatAgent< if (onlyAppendedPartial) { accumulatedMessages.push(...erroredNewModelMessages); } else { - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); } + accumulatedUIMessages = erroredUIMessagesWithPartial; + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } catch { // Keep the prior model accumulator if conversion fails. } diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index a001ec5c21c..12f0ac3fbef 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -233,6 +233,92 @@ describe("chat.agent managed loop — source-stream failure", () => { 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[] = []; + + 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(); + } + }); }); describe("chat.createSession turn.complete() — source-stream failure", () => { From ec33b975d77c6b3231d5addbb786f77ee9a0c329 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 21:05:10 +0100 Subject: [PATCH 07/13] fix(sdk): preserve stream backpressure when tapping chat chunks Tap the response stream with a TransformStream instead of an async generator. The generator forced ensureReadableStream down its eager-drain wrapper, which enqueues without honoring desiredSize (dropping backpressure to the model stream) and never propagates consumer cancel to the source. piping through a TransformStream keeps the ReadableStream pass-through, so native backpressure and stop/cancel propagation are retained. Also stamp an id on a recovered partial that lacks one, matching the success path. --- packages/trigger-sdk/src/v3/ai.ts | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 42bedf160c2..afe0c6d5101 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7949,6 +7949,9 @@ function chatAgent< partialResponse = undefined; partialIdx = -1; } + if (partialResponse && !partialResponse.id) { + partialResponse = { ...partialResponse, id: generateMessageId() } as TUIMessage; + } const erroredUIMessagesWithPartial: TUIMessage[] = !partialResponse ? erroredUIMessages : partialIdx === -1 @@ -8982,28 +8985,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 | ReadableStream, buffer: UIMessageChunk[] -): AsyncGenerator { +): ReadableStream | AsyncGenerator { 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({ + 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; } - } + })(); } /** From 6aeb2b95a656b25e577ce3c9c28a5ef9e1c2b933 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 21:11:39 +0100 Subject: [PATCH 08/13] fix(sdk): clean incomplete parts from the recovered error partial Apply cleanupAbortedParts to the partial recovered on a source-stream failure, in both the managed and manual loops. A transport error leaves the same incomplete-part state as a user stop (streaming text not finalized, dangling tool calls), so the partial is cleaned the same way: text is kept and unfinished tool parts are dropped, keeping the persisted UI history and the model view consistent. Also reset the error event's newMessages delta when the model-message conversion fails, so it can't advertise messages that were rolled back. --- packages/trigger-sdk/src/v3/ai.ts | 15 ++---- .../chat-agent-source-stream-error.test.ts | 50 ++++++++++++++++++- 2 files changed, 53 insertions(+), 12 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index afe0c6d5101..6a73df3cf0a 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7936,6 +7936,9 @@ function chatAgent< 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 @@ -7999,6 +8002,7 @@ function chatAgent< locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } catch { // Keep the prior model accumulator if conversion fails. + erroredNewModelMessages = []; } } @@ -9840,17 +9844,8 @@ 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") { - // Preserve the partial the model streamed before the failure: - // 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. 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) { - const partial = captured.message; + const partial = cleanupAbortedParts(captured.message); const queuedParts = locals.get(chatResponsePartsKey); if (queuedParts && queuedParts.length > 0) { (partial as any).parts = [...(partial.parts ?? []), ...queuedParts]; diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index 12f0ac3fbef..61a14159bf1 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -48,13 +48,17 @@ function erroringSource(errorMessage: string) { { 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 < partialChunks.length) { - controller.enqueue(partialChunks[i++]); + if (i < chunks.length) { + controller.enqueue(chunks[i++]); } else { controller.error(new Error(errorMessage)); } @@ -319,6 +323,48 @@ describe("chat.agent managed loop — source-stream failure", () => { await harness.close(); } }); + + it("cleans dangling tool parts from the recovered partial while keeping its text", async () => { + const turnCompletes: TurnCompleteEvent[] = []; + + 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(); + } + }); }); describe("chat.createSession turn.complete() — source-stream failure", () => { From ecfede634a535d9513ac637254e427ddf980f564 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 21:12:43 +0100 Subject: [PATCH 09/13] perf(sdk): release the chat chunk buffer once the response is committed On a successful turn the buffered chunks are dead weight (the recovered-partial fallback prefers the onFinish message), yet they were retained through the end-of-turn idle wait. Clear the buffer once the response is committed so a completed turn no longer holds a full copy of its streamed output while idle. --- packages/trigger-sdk/src/v3/ai.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 6a73df3cf0a..cbc777d19e9 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7408,6 +7408,7 @@ function chatAgent< if (capturedResponseMessage) { responseCommitted = true; capturedPartialResponse = capturedResponseMessage; + turnBufferedChunks.length = 0; } if (runSignal.aborted) return "exit"; From 19165076b01a46baaeff377c50d642cc1f897571 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 21:18:17 +0100 Subject: [PATCH 10/13] fix(sdk): preserve compaction and avoid double-persist on the error path When committing the errored turn's state to the accumulator, append only the new tail's model messages (errored user message and/or recovered partial) instead of reconverting the whole UI history, so a prior turn's model-only compaction survives an errored turn. Skip re-adding the partial once the response was already committed (a later compaction may have folded it into a summary), and roll the event/snapshot views back if the model conversion fails so they never advertise a partial the accumulator didn't take. --- packages/trigger-sdk/src/v3/ai.ts | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index cbc777d19e9..f4c5959a6c0 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7956,7 +7956,8 @@ function chatAgent< if (partialResponse && !partialResponse.id) { partialResponse = { ...partialResponse, id: generateMessageId() } as TUIMessage; } - const erroredUIMessagesWithPartial: TUIMessage[] = !partialResponse + const includePartial = partialResponse != null && !responseCommitted; + let erroredUIMessagesWithPartial: TUIMessage[] = !includePartial ? erroredUIMessages : partialIdx === -1 ? [...erroredUIMessages, partialResponse] @@ -7964,11 +7965,9 @@ function chatAgent< i === partialIdx ? partialResponse : m ) as TUIMessage[]); - const erroredNewUIMessages: TUIMessage[] = erroredWireMessage - ? [erroredWireMessage] - : []; - if (partialResponse && !responseCommitted) { - erroredNewUIMessages.push(partialResponse); + let erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : []; + if (includePartial) { + erroredNewUIMessages.push(partialResponse!); } let erroredNewModelMessages: ModelMessage[] = []; @@ -7984,18 +7983,17 @@ function chatAgent< // 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) { - const onlyAppendedPartial = - partialResponse != null && - partialIdx === -1 && - erroredUIMessages === accumulatedUIMessages; try { - if (partialResponse) { + if (includePartial) { erroredNewModelMessages = await toModelMessages([ - stripProviderMetadata(partialResponse), + stripProviderMetadata(partialResponse!), ]); } - if (onlyAppendedPartial) { - accumulatedMessages.push(...erroredNewModelMessages); + if (partialIdx === -1) { + const appended = erroredUIMessagesWithPartial.slice(accumulatedUIMessages.length); + accumulatedMessages.push( + ...(await toModelMessages(appended.map((m) => stripProviderMetadata(m)))) + ); } else { accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); } @@ -8004,6 +8002,8 @@ function chatAgent< } catch { // Keep the prior model accumulator if conversion fails. erroredNewModelMessages = []; + erroredUIMessagesWithPartial = erroredUIMessages; + erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : []; } } From b5c37ee72749b9b66ba2b3cea4733dc96a05e45b Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 22:11:01 +0100 Subject: [PATCH 11/13] fix(sdk): include the errored user message in the error event's newMessages The error-path onTurnComplete reported the errored user message in newUIMessages but not in newMessages (the model-message delta), so a customer persisting the model delta lost it. Derive newMessages from the same set as newUIMessages so the two stay symmetric. --- packages/trigger-sdk/src/v3/ai.ts | 30 ++++++++++--------- .../chat-agent-source-stream-error.test.ts | 3 ++ 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index f4c5959a6c0..c790f390a8e 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7982,23 +7982,25 @@ function chatAgent< // 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) { + if (!responseCommitted) { 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)))) + if (erroredNewUIMessages.length > 0) { + erroredNewModelMessages = await toModelMessages( + erroredNewUIMessages.map((m) => stripProviderMetadata(m)) ); - } else { - accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); } - accumulatedUIMessages = erroredUIMessagesWithPartial; - locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); + 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 { // Keep the prior model accumulator if conversion fails. erroredNewModelMessages = []; diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index 61a14159bf1..b435b5621d2 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -112,6 +112,9 @@ describe("chat.agent managed loop — source-stream failure", () => { ) .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(); } From aaa51a37e9a011c8b1909a1ddcee82cfd5bb084d Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 22:16:47 +0100 Subject: [PATCH 12/13] fix(sdk): fold queued response data parts into the managed error partial The managed error path now folds data parts queued via chat.response / writer.write() into the recovered partial before persisting, matching the success path and the manual loop, so a data part queued just before a source-stream failure isn't dropped. --- packages/trigger-sdk/src/v3/ai.ts | 14 ++++++++- .../chat-agent-source-stream-error.test.ts | 29 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c790f390a8e..745c8160718 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7956,6 +7956,16 @@ function chatAgent< 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 @@ -7991,7 +8001,9 @@ function chatAgent< } if (erroredUIMessagesWithPartial !== accumulatedUIMessages) { if (partialIdx === -1) { - const appended = erroredUIMessagesWithPartial.slice(accumulatedUIMessages.length); + const appended = erroredUIMessagesWithPartial.slice( + accumulatedUIMessages.length + ); accumulatedMessages.push( ...(await toModelMessages(appended.map((m) => stripProviderMetadata(m)))) ); diff --git a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts index b435b5621d2..8ec6998ad59 100644 --- a/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts +++ b/packages/trigger-sdk/test/chat-agent-source-stream-error.test.ts @@ -368,6 +368,35 @@ describe("chat.agent managed loop — source-stream failure", () => { await harness.close(); } }); + + it("folds queued response data parts into the recovered partial", async () => { + const turnCompletes: TurnCompleteEvent[] = []; + + 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", () => { From 0316b05fe3c5aa57e7edeaf8e0215b103c09145b Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Thu, 23 Jul 2026 22:20:55 +0100 Subject: [PATCH 13/13] fix(sdk): satisfy the package build's stricter null-narrowing on the error partial The error-path UI-message array construction tripped the tshy build's stricter check (partialResponse not narrowed by the includePartial flag). Assert non-null where includePartial already guarantees it. --- packages/trigger-sdk/src/v3/ai.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 745c8160718..ed8aa0dfb79 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -7970,9 +7970,9 @@ function chatAgent< let erroredUIMessagesWithPartial: TUIMessage[] = !includePartial ? erroredUIMessages : partialIdx === -1 - ? [...erroredUIMessages, partialResponse] + ? [...erroredUIMessages, partialResponse!] : (erroredUIMessages.map((m, i) => - i === partialIdx ? partialResponse : m + i === partialIdx ? partialResponse! : m ) as TUIMessage[]); let erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : [];