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

.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;