test: add streaming chat early disconnect e2e#31
Merged
Conversation
📝 WalkthroughWalkthroughAdds test support for upstream disconnects during SSE streaming, a helper to assert premature stream termination, refactors proxy streaming to handle empty or prematurely-closed upstream streams without treating them as hard errors, and makes a rate-limit hook tolerant when token usage is absent. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client as Client
participant Proxy as Proxy
participant Upstream as Upstream (mock)
participant Hook as HookManager
Client->>Proxy: POST /v1/chat/completions (stream: true)
Proxy->>Upstream: open streaming request
activate Upstream
Upstream-->>Proxy: SSE event 1 (data)
Proxy->>Client: SSE event 1 (data)
Note right of Proxy: saw_chunk = true
Upstream-->>Proxy: SSE event 2 (data)
Proxy->>Client: SSE event 2 (data)
Upstream--x Proxy: socket destroyed (disconnectAfterEvents reached)
deactivate Upstream
Note right of Proxy: finalize stream (no [DONE] emitted if saw_chunk=false)
Proxy->>Hook: post_call_streaming (includes total_tokens/defaults)
Proxy->>Client: stream ends
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (2 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![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
🧹 Nitpick comments (1)
tests/fixtures/mock-upstream.ts (1)
19-20: NormalizedisconnectAfterEventsand make zero-event disconnect timing deterministic.
disconnectAfterEventscurrently accepts anynumber, so negative/NaN/non-integer values can produce surprising behavior. Also, the=== 0path destroys the socket immediately, unlike the mid-stream path that yields one tick; aligning both reduces timing flakiness.♻️ Proposed refactor
+ const disconnectAfterEvents = + Number.isInteger(current.disconnectAfterEvents) && + (current.disconnectAfterEvents ?? -1) >= 0 + ? current.disconnectAfterEvents + : undefined; + if (requestStream(bodyJson)) { res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', Connection: 'keep-alive', }); - if (current.disconnectAfterEvents === 0) { + if (disconnectAfterEvents === 0) { res.flushHeaders(); + await new Promise((resolve) => setImmediate(resolve)); res.socket?.destroy(); return; } let sentEvents = 0; for (const event of current.streamEvents ?? defaultStreamEvents(model)) { @@ - if ( - current.disconnectAfterEvents !== undefined && - sentEvents >= current.disconnectAfterEvents - ) { + if ( + disconnectAfterEvents !== undefined && + sentEvents >= disconnectAfterEvents + ) { await new Promise((resolve) => setImmediate(resolve)); res.socket?.destroy(); return; }Also applies to: 469-492
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/fixtures/mock-upstream.ts` around lines 19 - 20, Normalize the disconnectAfterEvents option to a non-negative integer when it is read/assigned (coerce NaN to 0, floor fractional values, clamp negatives to 0) so invalid values won't produce weird behavior; then unify the zero-case disconnect timing by scheduling the socket destroy on the next tick (e.g., via process.nextTick or setTimeout(...,0)) instead of destroying immediately so the === 0 path matches the mid-stream delayed disconnect path. Apply this normalization and timing change wherever disconnectAfterEvents is used (including the other occurrence mentioned).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/proxy/chat-completions.test.ts`:
- Around line 393-396: The test expects a 200 with an empty SSE body but the
chat_completions handler currently converts an upstream empty stream into
ProviderError::ServiceError (BAD_GATEWAY) which becomes HTTP 502; update the
handler in src/proxy/handlers/chat_completions/mod.rs so that when the upstream
stream yields zero events (empty stream) it is treated as a successful empty SSE
response instead of an error—adjust the logic around the
ProviderError::ServiceError branch (lines ~111–119) or the IntoResponse
conversion so empty-stream cases return HTTP 200 with an empty text/event-stream
body.
---
Nitpick comments:
In `@tests/fixtures/mock-upstream.ts`:
- Around line 19-20: Normalize the disconnectAfterEvents option to a
non-negative integer when it is read/assigned (coerce NaN to 0, floor fractional
values, clamp negatives to 0) so invalid values won't produce weird behavior;
then unify the zero-case disconnect timing by scheduling the socket destroy on
the next tick (e.g., via process.nextTick or setTimeout(...,0)) instead of
destroying immediately so the === 0 path matches the mid-stream delayed
disconnect path. Apply this normalization and timing change wherever
disconnectAfterEvents is used (including the other occurrence mentioned).
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d5ea14a9-ab70-454c-95e1-4a0447f1d9e8
📒 Files selected for processing (4)
tests/fixtures/mock-upstream.tstests/proxy/chat-completions.test.tstests/proxy/timeout.test.tstests/utils/stream-assert.ts
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/proxy/handlers/chat_completions/mod.rs (1)
159-231:⚠️ Potential issue | 🟠 MajorFinalize hooks on stream errors too.
The new EOF paths call
HOOK_MANAGER.post_call_streaming, but theSome(Err(err))branch still returnsNoneimmediately. For the early-disconnect scenario this PR adds, that skips post-stream bookkeeping entirely, so hooks likesrc/proxy/hooks/rate_limit/mod.rsnever run their final accounting on truncated streams.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/handlers/chat_completions/mod.rs` around lines 159 - 231, The Some(Err(err)) branch currently returns early and never runs the final hook bookkeeping; update that branch to call HOOK_MANAGER.post_call_streaming(&mut hook_ctx, HOOK_FILTER_ALL).await (logging any error) before dropping span and returning None so hooks like post_call_streaming and rate_limit final accounting always run on stream errors; keep the existing error log for the stream itself and preserve dropping span and returning None after awaiting the hook call.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/proxy/handlers/chat_completions/mod.rs`:
- Around line 159-231: The Some(Err(err)) branch currently returns early and
never runs the final hook bookkeeping; update that branch to call
HOOK_MANAGER.post_call_streaming(&mut hook_ctx, HOOK_FILTER_ALL).await (logging
any error) before dropping span and returning None so hooks like
post_call_streaming and rate_limit final accounting always run on stream errors;
keep the existing error log for the stream itself and preserve dropping span and
returning None after awaiting the hook call.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d53b5b26-1b6c-44f4-a7d2-6bae0046bd29
📒 Files selected for processing (3)
src/proxy/handlers/chat_completions/mod.rssrc/proxy/hooks/rate_limit/mod.rstests/proxy/chat-completions.test.ts
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/proxy/handlers/chat_completions/mod.rs (1)
171-187: Refactor suggestion: centralize stream-finalization logic.The no-chunk branch duplicates the
post_call_streaming + drop(span)sequence already used in the done path, which can drift over time. Consider extracting a small finalization helper used by both branches.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/proxy/handlers/chat_completions/mod.rs` around lines 171 - 187, The code duplicates finalization steps in the None branch (sending a trailing SseEvent::default().data("[DONE]") vs executing HOOK_MANAGER.post_call_streaming(&mut hook_ctx, HOOK_FILTER_ALL).await and drop(span)); extract a small helper (e.g., finalize_stream or finish_streaming) that performs the post_call_streaming call, logs any error, drops the span, and returns the appropriate Option/tuple or signaling value, then replace the duplicated sequences in both the done path and the no-chunk path to call that helper from the iterator state handling logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@src/proxy/handlers/chat_completions/mod.rs`:
- Around line 171-187: The code duplicates finalization steps in the None branch
(sending a trailing SseEvent::default().data("[DONE]") vs executing
HOOK_MANAGER.post_call_streaming(&mut hook_ctx, HOOK_FILTER_ALL).await and
drop(span)); extract a small helper (e.g., finalize_stream or finish_streaming)
that performs the post_call_streaming call, logs any error, drops the span, and
returns the appropriate Option/tuple or signaling value, then replace the
duplicated sequences in both the done path and the no-chunk path to call that
helper from the iterator state handling logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b3ba5749-8a8e-4560-b116-13d412e7ba23
📒 Files selected for processing (2)
src/proxy/handlers/chat_completions/mod.rstests/proxy/timeout.test.ts
✅ Files skipped from review due to trivial changes (1)
- tests/proxy/timeout.test.ts
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Bug Fixes
Tests