fix: persist final step when saveStreamDeltas.returnImmediately is true (#265)

[sethconvex]

Summary

Fixes #265: when streamText is called with saveStreamDeltas: { returnImmediately: true }, the final assistant message was never persisted and the stream stayed stuck in "streaming" status forever. Only intermediate (tool-call) steps were saved.

Root cause

The deferred-save path added for atomic stream finish (#181) doesn't work when streamText returns without awaiting stream consumption:

1. onStepFinish runs only when the caller later drains the stream (e.g. via result.toTextStreamResponse()).
2. With returnImmediately: true, streamText skips its own await stream block and exits.
3. The post-await if (pendingFinalStep && streamer) block runs synchronously after exit — but onStepFinish hasn't fired yet, so pendingFinalStep is still undefined.
4. The caller drains the stream later, onStepFinish fires, sets pendingFinalStep and calls markFinishedExternally() — but no one is reading pendingFinalStep anymore, and consumeStream skips finish() due to the external-finish flag.
5. Net effect: final message dropped, stream stuck.

Fix

Branch on a new willAwaitStream flag in onStepFinish:

• If streamText will await consumption (saveStreamDeltas === true or returnImmediately: false): defer the save via pendingFinalStep as before (preserves the atomic-finish behavior from Sub-messages flow problem #181).
• If not (returnImmediately: true): save the final step inline using streamer.getOrCreateStreamId() so the message persistence and stream-finish land atomically — before streamText returns.

Also guards DeltaStreamer.addParts against late deltas after markFinishedExternally. The stream record is already in "finished" state and streams.addDelta would silently drop those writes anyway; skipping them is just an optimization, but it avoids racing the inline save.

Provenance

These two pieces (the willAwaitStream branch and the addParts guard) are lifted from the open PR #259 ("Add httpStreamText and useHttpStream"), which bundles them with a larger HTTP streaming surface (~2,800 LOC across 21 files). This PR cherry-picks just the bug fix so #265 can ship independently of that feature work — which is also currently blocked on merge conflicts with main.

Notes for reviewers

• Minor behavioral nuance worth knowing: with the new addParts guard, delta chunks emitted by the model after onStepFinish fires but before consumeStream has finished pumping them are silently dropped on the returnImmediately path. The persisted message is unaffected — it carries the complete final-step text via call.save({step}). A UI subscribed to live deltas could see slightly-truncated streaming content until it reconciles against the saved message, but no canonical data is lost.
• The fix preserves the Sub-messages flow problem #181 atomic-finish guarantee for the original path (when we do await stream consumption).
• The full HTTP streaming PR Add httpStreamText and useHttpStream #259 should still land for its other improvements, but is intentionally out of scope here.

Test plan

• New regression test src/client/streamText.test.ts reproduces the issue: calls streamText with returnImmediately: true, drains the result, then asserts (a) a persisted assistant message exists with the model's text and (b) no stream is left in "streaming" status.
• Test fails on main with expected 0 to be greater than 0 (no assistant message persisted).
• Test passes with the fix applied.
• npm run lint clean.
• npx tsc --noEmit clean.
• Full vitest suite: 257/257 unit tests pass. (The 3 example/convex/*.test.ts failures are pre-existing on main — they need npm run build to populate dist/ first; unrelated to this change.)

[pkg-pr-new]

Open in StackBlitz

npm i https://pkg.pr.new/@convex-dev/agent@266

commit: b602a46

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: get-convex/coderabbit/.coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 267a30eb-0f66-4d76-a489-4783a346b45f

📥 Commits

Reviewing files that changed from the base of the PR and between 10cd3cf and b602a46.

📒 Files selected for processing (2)

• example/convex/setup.test.ts
• src/client/streaming.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• src/client/streaming.ts

---

📝 Walkthrough

Walkthrough

This PR fixes issue #265 where saveStreamDeltas with returnImmediately: true never persisted the final step to the messages table. The fix introduces willAwaitStream logic in streamText.ts to determine whether stream consumption will be awaited, then conditionally defers or immediately saves the final step in onStepFinish. When returnImmediately is true, the final step is saved immediately using streamer.getOrCreateStreamId() rather than deferred. In streaming.ts, DeltaStreamer.addParts now guards against queuing parts after the stream is marked externally finished. A new test validates that the final assistant message is persisted and the stream is properly finished in the returnImmediately scenario.

Sequence Diagram(s)

sequenceDiagram
  participant Agent
  participant Streamer as DeltaStreamer
  participant StreamsAPI as streams.addDelta/finish
  participant Messages as MessagesTable

  Agent->>Streamer: streamText start (saveStreamDeltas:{returnImmediately})
  Streamer->>Streamer: compute willAwaitStream
  alt willAwaitStream == false (returnImmediately)
    Streamer->>Messages: save final step inline (getOrCreateStreamId + save)
    Streamer->>Agent: return immediately (stream consumed in background)
    Agent->>StreamsAPI: consumeStream -> finish
    StreamsAPI->>Messages: finalization (if needed)
  else willAwaitStream == true
    Streamer->>Agent: return a streaming result
    Agent->>Streamer: consumer drains stream
    Streamer->>StreamsAPI: addDelta/... finish
    StreamsAPI->>Messages: persist deferred final step and mark finished
  end

Loading

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main fix: persisting the final step when returnImmediately is true, and references the related issue #265.
Description check	✅ Passed	The description is comprehensive and directly addresses the changeset, including root cause analysis, the fix approach, test plan, and behavioral notes.
Linked Issues check	✅ Passed	All code changes directly address #265: the willAwaitStream branch ensures final steps are saved inline when returnImmediately is true, and the addParts guard prevents late deltas after external finish.
Out of Scope Changes check	✅ Passed	The example/convex/setup.test.ts typecheck fix is a pre-existing CI issue unrelated to #265 but documented as necessary to unblock this PR's CI; the streaming.ts race-condition fix is also documented and necessary.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/issue-265-returnImmediately-final-step-save

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/client/streaming.ts`:
- Around line 290-296: The in-flight delta path can reach `#sendDelta` after
this.#finishedExternally is set and currently treats addDelta returning false as
an async failure that aborts the stream; instead, update the logic in the
delta-send flow (the caller of addDelta and the `#sendDelta` handler) to treat a
false return from addDelta as a benign "delta dropped because stream already
finished" condition: when addDelta(...) === false, simply return (optionally log
at debug) and do not call abort/fail the stream or set error state; keep the
existing behavior of aborting only for real errors/exceptions thrown by addDelta
or network failures. Ensure references to `#sendDelta`, addDelta, and
`#finishedExternally` are used so you change the correct branch.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository: get-convex/coderabbit/.coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a0996dde-57a0-4641-9956-0487da9974a0

📥 Commits

Reviewing files that changed from the base of the PR and between 84534d3 and 10cd3cf.

📒 Files selected for processing (3)

• src/client/streamText.test.ts
• src/client/streamText.ts
• src/client/streaming.ts

[sethconvex]

Pushed b602a46 addressing two things:

CodeRabbit's #sendDelta race (valid, fixed) — confirmed in code at src/client/streaming.ts:351. An in-flight #sendDelta started before markFinishedExternally() would call addDelta, get success === false because the row is already finished, and the current path would onAsyncAbort + abort() — turning a successfully externally-finished stream into the failure path. Fixed by skipping the abort branch when #finishedExternally is set; the late-write miss is benign.

CI failure (unrelated to this PR, but fixed here) — the Test and lint failure on commit 10cd3cf was in the tsc -p example/convex step, complaining about @convex-dev/workflow/test and @convex-dev/rate-limiter/test typing register against the pre-generic SchemaDefinition<GenericSchema, boolean>. Last 3 CI runs on main itself failed with the same error (runs 25619373541, 24437925179, 24047631068) — it's a pre-existing break, not introduced here. Included the drive-by typecheck cast from PR #259 to get this PR's CI green.

Local verification: npm run lint clean, npm run typecheck clean, npx vitest run 257/257 passing (streaming tests included).