Skip to content

Commit 6aeb2b9

Browse files
committed
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.
1 parent ec33b97 commit 6aeb2b9

2 files changed

Lines changed: 53 additions & 12 deletions

File tree

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

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7936,6 +7936,9 @@ function chatAgent<
79367936
let partialResponse: TUIMessage | undefined =
79377937
capturedPartialResponse ??
79387938
((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined);
7939+
if (partialResponse) {
7940+
partialResponse = cleanupAbortedParts(partialResponse);
7941+
}
79397942

79407943
// Build the complete error UI state. A HITL/tool continuation partial
79417944
// reuses an existing assistant id, so replace it in place (like the
@@ -7999,6 +8002,7 @@ function chatAgent<
79998002
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
80008003
} catch {
80018004
// Keep the prior model accumulator if conversion fails.
8005+
erroredNewModelMessages = [];
80028006
}
80038007
}
80048008

@@ -9840,17 +9844,8 @@ function createChatSession(
98409844
// Surface a genuine stream failure to the caller. A user stop
98419845
// (status "aborted") falls through so the partial is accumulated.
98429846
if (captured.status === "error") {
9843-
// Preserve the partial the model streamed before the failure:
9844-
// accumulate it (mirroring the stop path) so `turn.uiMessages`
9845-
// reflects it and the caller can persist it after catching,
9846-
// then rethrow. Without this the partial pipeAndCapture
9847-
// reconstructed is silently dropped on rethrow. Fold in any
9848-
// data parts queued via chat.response / writer.write() this
9849-
// turn, same as the success path below. cleanupAbortedParts is
9850-
// intentionally skipped: it only runs on a user stop, not on a
9851-
// hard transport error (which is not an abort).
98529847
if (captured.message) {
9853-
const partial = captured.message;
9848+
const partial = cleanupAbortedParts(captured.message);
98549849
const queuedParts = locals.get(chatResponsePartsKey);
98559850
if (queuedParts && queuedParts.length > 0) {
98569851
(partial as any).parts = [...(partial.parts ?? []), ...queuedParts];

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

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,13 +48,17 @@ function erroringSource(errorMessage: string) {
4848
{ type: "text-start", id: "t1" },
4949
{ type: "text-delta", id: "t1", delta: "partial answer" },
5050
];
51+
return sourceFromChunks(partialChunks, errorMessage);
52+
}
53+
54+
function sourceFromChunks(chunks: unknown[], errorMessage: string) {
5155
return {
5256
toUIMessageStream() {
5357
let i = 0;
5458
return new ReadableStream({
5559
pull(controller) {
56-
if (i < partialChunks.length) {
57-
controller.enqueue(partialChunks[i++]);
60+
if (i < chunks.length) {
61+
controller.enqueue(chunks[i++]);
5862
} else {
5963
controller.error(new Error(errorMessage));
6064
}
@@ -319,6 +323,48 @@ describe("chat.agent managed loop — source-stream failure", () => {
319323
await harness.close();
320324
}
321325
});
326+
327+
it("cleans dangling tool parts from the recovered partial while keeping its text", async () => {
328+
const turnCompletes: TurnCompleteEvent<unknown, UIMessage>[] = [];
329+
330+
const agent = chat.agent({
331+
id: "chatAgent.error-partial-cleanup",
332+
run: async () =>
333+
sourceFromChunks(
334+
[
335+
{ type: "start", messageId: "a-tool" },
336+
{ type: "text-start", id: "t1" },
337+
{ type: "text-delta", id: "t1", delta: "thinking" },
338+
{ type: "text-end", id: "t1" },
339+
{ type: "tool-input-start", toolCallId: "tc1", toolName: "search" },
340+
{
341+
type: "tool-input-available",
342+
toolCallId: "tc1",
343+
toolName: "search",
344+
input: { q: "x" },
345+
},
346+
],
347+
"UND_ERR_BODY_TIMEOUT"
348+
) as never,
349+
onTurnComplete: async (event) => {
350+
turnCompletes.push(event);
351+
},
352+
});
353+
354+
const harness = mockChatAgent(agent, { chatId: "cae-partial-cleanup" });
355+
try {
356+
await harness.sendMessage(userMessage("hi", "u-1"));
357+
await waitFor(() => turnCompletes.length >= 1);
358+
359+
const evt = turnCompletes[0]!;
360+
expect(evt.responseMessage).toBeDefined();
361+
const parts = evt.responseMessage!.parts as Array<{ type: string }>;
362+
expect(extractText(evt.responseMessage)).toBe("thinking");
363+
expect(parts.some((p) => p.type.startsWith("tool-"))).toBe(false);
364+
} finally {
365+
await harness.close();
366+
}
367+
});
322368
});
323369

324370
describe("chat.createSession turn.complete() — source-stream failure", () => {

0 commit comments

Comments
 (0)