fix(sdk): preserve partial assistant message on chat stream failure#4348
fix(sdk): preserve partial assistant message on chat stream failure#4348matt-aitken wants to merge 10 commits into
Conversation
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.
🦋 Changeset detectedLatest commit: 1916507 The changes in this PR will be included in the next version bump. This PR includes changesets to release 26 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe SDK buffers streamed chunks and captures partial assistant responses during chat turns. When a source stream fails, 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
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.
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).
@trigger.dev/build
trigger.dev
@trigger.dev/core
@trigger.dev/python
@trigger.dev/react-hooks
@trigger.dev/redis-worker
@trigger.dev/rsc
@trigger.dev/schema-to-json
@trigger.dev/sdk
commit: |
…tep 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.
|
Additional hardening from an adversarial self-review of the error path: Fixed (real): the Considered, no change needed:
319 SDK tests pass, 0 typecheck errors. |
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.
… 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.
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.
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.
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.
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.
Summary
When a
chat.agent(orchat.createSession) turn's model stream fails mid-response (e.g. a transport timeout likeUND_ERR_BODY_TIMEOUT), the assistant output that already streamed was dropped.onTurnCompletefired withresponseMessage: undefined, and the manual loop'sturn.complete()rethrew without keeping the partial. Apps that registerhydrateMessagesare hit hardest: boot-time tail-replay recovery is off by design, so the partial couldn't be reclaimed later either.This preserves the partial while still reporting the turn as errored, so persistence keeps the streamed-so-far output.
Fix
Follow-up to #4304, which introduced
chat.pipeAndCapture(buffer chunks, classify instead of throw, reconstruct the partial from buffered chunks). That fix landed for custom agents but not the two other loops:chat.agent): buffer chunks as they flow to the pipe, and on the error path reconstruct the partial (falling back to theonFinishmessage if it fired). It's now passed toonTurnCompleteviaresponseMessage,rawResponseMessage,uiMessages, andnewUIMessages. Empty for non-stream failures, preserving prior behavior.chat.createSession/turn.complete()):pipeAndCapturealready reconstructs the partial, but it was discarded on the error branch. Now it's accumulated (soturn.uiMessagesreflects it and the caller can persist after catching) beforeturn.complete()rethrows.onBeforeTurnCompleteis intentionally still skipped on the error path: it's the "stream still open" hook that hands out a writer, which is meaningless once the source stream has broken.Tests
New
chat-agent-source-stream-error.test.tsdrives a source that streams a partial then errors, for both loops. Both cases fail against the unpatched code and pass with the fix. Full@trigger.dev/sdksuite passes (317 tests).