fix(streaming): flush answer tail and sources on any stream exit - #749
Conversation
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe SSE source-filtering stream now performs one terminal flush for clean completion, upstream closure, and mid-stream errors. It preserves buffered content and filtered sources, marks unclean termination as truncated, closes upstream iterators, and always emits a downstream ChangesSource filtering stream finalization
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Upstream as Upstream SSE stream
participant Filtering as stream_with_source_filtering
participant Client as Downstream client
Upstream->>Filtering: SSE content and optional data: [DONE]
Filtering->>Filtering: Capture termination and flush buffered tail
Filtering->>Client: Tail content and filtered extra.sources
Filtering->>Client: Finish metadata with truncated status
Filtering->>Client: data: [DONE]
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
There was a problem hiding this comment.
-
The truncated flag can silently disappear in a narrow edge case. I reproduced this directly: if the upstream connection ends before any chunk with content or a finish_reason ever arrives (e.g. only a bare role-preamble chunk, then disconnect), chunk_template stays None. Both the tail-chunk and finish-chunk blocks are gated on if chunk_template:, so neither ever fires the client gets data: [DONE] with no extra.truncated anywhere, even though the server logged the warning internally. It's a narrow case (needs a disconnect before the very first content token), and the client at least doesn't hang, but the "distinguishable from a clean completion" guarantee the PR advertises doesn't hold here.
-
More significant — real timeouts/connection drops likely bypass this fix entirely. I checked how llm_stream actually terminates on a network problem: both #vllm_client.py:161 and ollama_client.py:112 wrap the streaming loop in try/except httpx.ConnectError / except httpx.TimeoutException, re-raising as InferenceConnectionError/InferenceTimeoutError. I confirmed with a direct repro that when the upstream generator raises mid-iteration (rather than ending cleanly), the exception propagates straight out of stream_with_source_filtering the new post-loop flush code never runs, because Python doesn't execute the code after a for loop when the loop raises. That exception is instead caught by chat.py's existing except OpenRAGError / except Exception handlers, which yield a generic _make_sse_error(...) chunk the buffered answer tail and sources are lost there too, just as before this PR, only now replaced by an error message instead of a clean tail.
Since the PR explicitly names "timeout" as one of the three root causes of #715, and a mid-stream read timeout is exactly what gets converted to InferenceTimeoutError in the existing client code, this fix likely does not address that particular named cause only the flavor of disconnect where the transport ends with a bare EOF and no distinguishable HTTP-level error (plausible for "proxy drop" / some "worker restart" cases, where the connection just closes without httpx raising). Worth confirming against whatever telemetry/repro originally motivated #715 — if it was specifically an httpx timeout, this PR may not close the loop on it.
…nly close Address review on #749: - A real read timeout / connection drop surfaces as the upstream generator *raising* mid-iteration (InferenceTimeoutError / InferenceConnectionError), not a clean close, so the post-loop flush was skipped entirely and the buffered answer tail + sources were lost — the exact 'timeout' cause named in #715. Wrap the loop in try/except: CancelledError/GeneratorExit still propagate (client-disconnect semantics); any other error falls through to the flush so the partial answer is delivered with extra.truncated=true. If nothing was ever streamed the error is re-raised so the router surfaces the real failure instead of a silent empty [DONE]. - When the stream died before any content/finish chunk (e.g. only a role preamble), chunk_template stayed None and neither flush block fired, so the truncated flag never reached the client. Track a last_chunk fallback and template the flush off (chunk_template or last_chunk) so a preamble-only stream is still flagged truncated. Adds regression tests for all three paths.
|
@aditykris thanks — both are spot-on. Fixed in 5869529. 2 (the significant one — mid-stream raise bypasses the flush): confirmed. A read timeout / connection drop surfaces as the upstream generator raising mid-iteration (
1 (truncated flag disappears when Regression tests added for all three paths ( |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 `@openrag/core/utils/source_filtering.py`:
- Around line 143-148: Update the finish_reason/content handling in the
streaming filtering logic so content is appended to pending even when
finish_reason is also present, preserving the terminal chunk’s final token.
Ensure chunk_template and last_finish_reason behavior remains unchanged.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4cfeb61a-8be2-44fd-b37b-fd31c28476da
📒 Files selected for processing (2)
openrag/core/utils/source_filtering.pytests/unit/core/utils/test_source_filtering.py
…nly close Address review on #749: - A real read timeout / connection drop surfaces as the upstream generator *raising* mid-iteration (InferenceTimeoutError / InferenceConnectionError), not a clean close, so the post-loop flush was skipped entirely and the buffered answer tail + sources were lost — the exact 'timeout' cause named in #715. Wrap the loop in try/except: CancelledError/GeneratorExit still propagate (client-disconnect semantics); any other error falls through to the flush so the partial answer is delivered with extra.truncated=true. If nothing was ever streamed the error is re-raised so the router surfaces the real failure instead of a silent empty [DONE]. - When the stream died before any content/finish chunk (e.g. only a role preamble), chunk_template stayed None and neither flush block fired, so the truncated flag never reached the client. Track a last_chunk fallback and template the flush off (chunk_template or last_chunk) so a preamble-only stream is still flagged truncated. Adds regression tests for all three paths.
5869529 to
c9e6b27
Compare
… a chunk Some OpenAI-compatible providers (and litellm passthroughs) pack the last token and the terminal finish_reason into the same chunk. The loop gated content on `elif content` after `if finish_reason`, so that final token was silently dropped (e.g. `Hello ` + a final `world` chunk yielded only `Hello`). Handle content and finish_reason independently: always append content to the buffer, and record finish_reason separately. A finish-only chunk is still deferred to the terminal flush so it can carry extra.sources. Reported by @hedhoud and CodeRabbit on #749. Adds a regression test.
… a chunk Some OpenAI-compatible providers (and litellm passthroughs) pack the last token and the terminal finish_reason into the same chunk. The loop gated content on `elif content` after `if finish_reason`, so that final token was silently dropped (e.g. `Hello ` + a final `world` chunk yielded only `Hello`). Handle content and finish_reason independently: always append content to the buffer, and record finish_reason separately. A finish-only chunk is still deferred to the terminal flush so it can carry extra.sources. Reported by @hedhoud and CodeRabbit on #749. Adds a regression test.
0dd35e9 to
e937f62
Compare
…connection Breaking out of the `async for` on `data: [DONE]` leaves the vLLM/Ollama `stream_chat` generator suspended inside its `async with response`, so the pooled httpx connection stays checked out until GC (an `async for` does not close its iterator on break — unlike `with`). Under concurrent chat traffic that holds connections longer than necessary and can exhaust the pool. Close `llm_stream` in a `finally` (awaiting is safe there — only yielding during teardown is forbidden) so the connection is released promptly on every exit path. The getattr guard tolerates non-generator iterables (e.g. tests). Reported by @hedhoud on #749. Regression test added: it holds a reference to the upstream generator so a pass can only mean an explicit aclose() ran.
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
LGTM. The missing-[DONE], mid-stream error, final-token, and upstream cleanup paths are now covered. The focused tests and CI are green.
…nly close Address review on #749: - A real read timeout / connection drop surfaces as the upstream generator *raising* mid-iteration (InferenceTimeoutError / InferenceConnectionError), not a clean close, so the post-loop flush was skipped entirely and the buffered answer tail + sources were lost — the exact 'timeout' cause named in #715. Wrap the loop in try/except: CancelledError/GeneratorExit still propagate (client-disconnect semantics); any other error falls through to the flush so the partial answer is delivered with extra.truncated=true. If nothing was ever streamed the error is re-raised so the router surfaces the real failure instead of a silent empty [DONE]. - When the stream died before any content/finish chunk (e.g. only a role preamble), chunk_template stayed None and neither flush block fired, so the truncated flag never reached the client. Track a last_chunk fallback and template the flush off (chunk_template or last_chunk) so a preamble-only stream is still flagged truncated. Adds regression tests for all three paths.
… a chunk Some OpenAI-compatible providers (and litellm passthroughs) pack the last token and the terminal finish_reason into the same chunk. The loop gated content on `elif content` after `if finish_reason`, so that final token was silently dropped (e.g. `Hello ` + a final `world` chunk yielded only `Hello`). Handle content and finish_reason independently: always append content to the buffer, and record finish_reason separately. A finish-only chunk is still deferred to the terminal flush so it can carry extra.sources. Reported by @hedhoud and CodeRabbit on #749. Adds a regression test.
…connection Breaking out of the `async for` on `data: [DONE]` leaves the vLLM/Ollama `stream_chat` generator suspended inside its `async with response`, so the pooled httpx connection stays checked out until GC (an `async for` does not close its iterator on break — unlike `with`). Under concurrent chat traffic that holds connections longer than necessary and can exhaust the pool. Close `llm_stream` in a `finally` (awaiting is safe there — only yielding during teardown is forbidden) so the connection is released promptly on every exit path. The getattr guard tolerates non-generator iterables (e.g. tests). Reported by @hedhoud on #749. Regression test added: it holds a reference to the upstream generator so a pass can only mean an explicit aclose() ran.
90d132b to
86c3f46
Compare
There was a problem hiding this comment.
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 `@openrag/core/utils/source_filtering.py`:
- Around line 158-175: Clear finish_reason from the choice emitted by the
intermediate content path in the source-filtering loop, matching the existing
protection used for tail chunks. Update the choice construction around
_strip_sources_tags and emitted_len so intermediate content chunks cannot appear
terminal, while preserving the emitted delta content and subsequent finish
handling.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 14fb8c84-c387-4b30-a948-7c5b2151cb53
📒 Files selected for processing (2)
openrag/core/utils/source_filtering.pytests/unit/core/utils/test_source_filtering.py
The mid-stream emit path built its chunk with {**choice, ...}, inheriting
finish_reason from the source chunk. Once content is handled independently
of finish_reason, a provider that packs the last token + finish_reason into
one chunk can hit this path (when pending has overflowed buffer_size) and
emit a *mid-stream* content chunk carrying finish_reason='stop'. A spec
client treats that as terminal, drops the delta, and ignores the tail/finish
chunks — truncating a long answer. Clear finish_reason on the emitted chunk,
the same guard the terminal tail chunk already applies.
Reported by CodeRabbit on #749. Regression test added (fails without the fix).
- Move the terminal flush (tail content + extra.sources) out of the data: [DONE] branch into a post-loop block that runs whether the loop exits via [DONE] or the upstream connection simply closes. - Always emit a downstream data: [DONE] so clients don't hang when upstream never sent one. - Flag extra.truncated=true when the stream ended without [DONE], so a silently truncated answer is distinguishable from a normal one. Fixes #715 Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…nly close Address review on #749: - A real read timeout / connection drop surfaces as the upstream generator *raising* mid-iteration (InferenceTimeoutError / InferenceConnectionError), not a clean close, so the post-loop flush was skipped entirely and the buffered answer tail + sources were lost — the exact 'timeout' cause named in #715. Wrap the loop in try/except: CancelledError/GeneratorExit still propagate (client-disconnect semantics); any other error falls through to the flush so the partial answer is delivered with extra.truncated=true. If nothing was ever streamed the error is re-raised so the router surfaces the real failure instead of a silent empty [DONE]. - When the stream died before any content/finish chunk (e.g. only a role preamble), chunk_template stayed None and neither flush block fired, so the truncated flag never reached the client. Track a last_chunk fallback and template the flush off (chunk_template or last_chunk) so a preamble-only stream is still flagged truncated. Adds regression tests for all three paths.
Emit a single, visible `Answer truncated` warning at the point the `truncated` flag is set — with the cause (upstream error vs. connection closed), model, delivered length and source count folded into the message so the context actually renders (the previous 'flushing buffered tail' warnings passed kwargs with no placeholders, so nothing showed). The no-content error path keeps its own distinct log before re-raising.
… a chunk Some OpenAI-compatible providers (and litellm passthroughs) pack the last token and the terminal finish_reason into the same chunk. The loop gated content on `elif content` after `if finish_reason`, so that final token was silently dropped (e.g. `Hello ` + a final `world` chunk yielded only `Hello`). Handle content and finish_reason independently: always append content to the buffer, and record finish_reason separately. A finish-only chunk is still deferred to the terminal flush so it can carry extra.sources. Reported by @hedhoud and CodeRabbit on #749. Adds a regression test.
…connection Breaking out of the `async for` on `data: [DONE]` leaves the vLLM/Ollama `stream_chat` generator suspended inside its `async with response`, so the pooled httpx connection stays checked out until GC (an `async for` does not close its iterator on break — unlike `with`). Under concurrent chat traffic that holds connections longer than necessary and can exhaust the pool. Close `llm_stream` in a `finally` (awaiting is safe there — only yielding during teardown is forbidden) so the connection is released promptly on every exit path. The getattr guard tolerates non-generator iterables (e.g. tests). Reported by @hedhoud on #749. Regression test added: it holds a reference to the upstream generator so a pass can only mean an explicit aclose() ran.
The mid-stream emit path built its chunk with {**choice, ...}, inheriting
finish_reason from the source chunk. Once content is handled independently
of finish_reason, a provider that packs the last token + finish_reason into
one chunk can hit this path (when pending has overflowed buffer_size) and
emit a *mid-stream* content chunk carrying finish_reason='stop'. A spec
client treats that as terminal, drops the delta, and ignores the tail/finish
chunks — truncating a long answer. Clear finish_reason on the emitted chunk,
the same guard the terminal tail chunk already applies.
Reported by CodeRabbit on #749. Regression test added (fails without the fix).
55a27d9 to
6ce8bac
Compare
4 participants
Summary
data: [DONE](timeout, proxy drop, worker restart) silently dropped the buffered answer tail and all ofextra.sources, with no error surfaced.stream_with_source_filteringnow performs the terminal flush (tail content +extra.sources) in a single post-loop block that runs whether the loop exits via[DONE]or the upstream connection simply closing, instead of only inside thedata: [DONE]branch.data: [DONE], even if upstream never sent one, so clients (e.g. the OpenAI SDK) don't hang waiting for stream termination.[DONE], the final chunk'sextranow includes"truncated": trueso a silently-truncated answer is distinguishable from a normal completion instead of looking indistinguishable from a clean one.Test plan
tests/unit/core/utils/test_source_filtering.py::TestStreamClosedWithoutDonecovering: tail/sources still flushed without[DONE],truncatedflag set when[DONE]is missing, and notruncatedflag on a normal[DONE]-terminated stream.uv run pytest tests/unit/core/utils/test_source_filtering.py— 42 passed.uv run pytest tests/unit/— 1888 passed.uv run ruff check/uv run ruff format --checkon changed files.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
[DONE], upstream close without[DONE], and mid-stream exceptions).extra.truncatedand warning behavior based on whether termination was clean.Tests
finish_reason, role-only preamble closes, missing[DONE], upstream exceptions, and disconnect behavior.