handle ctx cancel before handling read errors in streaming - #3522
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (18)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR systematically updates streaming resource cleanup and cancellation handling across providers: ReleaseStreamingResponse now accepts a Bifrost context and stream defers mark/skip draining when connection-closed; stream read loops check ctx.Err() first and return immediately on cancellation before non-EOF error processing. ChangesStream Cancellation & Context-Aware Resource Release
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/utils/utils.go (1)
2051-2077:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winOnly mark
connection_closedwhen a close path actually ran.Line 2077 sets the flag even when
bodyStreamimplements neitherio.ClosernorstreamCloserWithError. In that caseReleaseStreamingResponse()will hit the early return on Line 2423 and skip the old drain/close fallback entirely, which can leave the body unread and hurt connection reuse for wrapped readers. Mirror thectx.Done()branch and gate the flag on whether a close branch was actually taken.🛠️ Suggested fix
case <-done: // If context was also cancelled (race between done and ctx.Done), // still close the body stream to unblock the drain in ReleaseStreamingResponse. if ctx.Err() != nil { + closedStream := false if closer, ok := bodyStream.(io.Closer); ok { if err := closer.Close(); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) } + closedStream = true } else if wce, ok := bodyStream.(streamCloserWithError); ok { if err := wce.CloseWithError(ctx.Err()); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) } + closedStream = true } - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) + if closedStream { + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) + } } }🤖 Prompt for 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. In `@core/providers/utils/utils.go` around lines 2051 - 2077, The code sets schemas.BifrostContextKeyConnectionClosed unconditionally in the done branch even when no close occurred; change the done-case to mirror the ctx.Done() branch by only setting ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) when a close path actually ran (i.e., when bodyStream asserts to io.Closer or streamCloserWithError and the Close/CloseWithError call is executed), so gate the flag behind the same type-assert checks for bodyStream (io.Closer and streamCloserWithError) used above and avoid setting the flag when neither branch matched; this will ensure ReleaseStreamingResponse behavior (and connection reuse) remains correct for wrapped readers.
🤖 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.
Outside diff comments:
In `@core/providers/utils/utils.go`:
- Around line 2051-2077: The code sets schemas.BifrostContextKeyConnectionClosed
unconditionally in the done branch even when no close occurred; change the
done-case to mirror the ctx.Done() branch by only setting
ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) when a close path
actually ran (i.e., when bodyStream asserts to io.Closer or
streamCloserWithError and the Close/CloseWithError call is executed), so gate
the flag behind the same type-assert checks for bodyStream (io.Closer and
streamCloserWithError) used above and avoid setting the flag when neither branch
matched; this will ensure ReleaseStreamingResponse behavior (and connection
reuse) remains correct for wrapped readers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 066c8260-f22d-4522-bc97-7e4867bc8f17
📒 Files selected for processing (18)
core/providers/anthropic/anthropic.gocore/providers/azure/azure.gocore/providers/cohere/cohere.gocore/providers/elevenlabs/elevenlabs.gocore/providers/gemini/gemini.gocore/providers/huggingface/huggingface.gocore/providers/mistral/mistral.gocore/providers/openai/openai.gocore/providers/replicate/replicate.gocore/providers/replicate/utils.gocore/providers/utils/utils.gocore/providers/vertex/vertex.gocore/providers/vllm/vllm.gocore/schemas/bifrost.goframework/go.modplugins/otel/go.modtransports/bifrost-http/handlers/inference.gotransports/bifrost-http/lib/validator.go
Confidence Score: 4/5Safe to merge with one fix: the ctx.Done() branch in SetupStreamCancellation sets the ConnectionClosed flag even when close returns an error, which can cause ReleaseStreamingResponse to skip draining on a stream that was never actually closed. The ctx.Done() branch in SetupStreamCancellation unconditionally sets BifrostContextKeyConnectionClosed after calling closer.Close(), regardless of whether that call succeeded. The done branch was corrected in this same PR to only set the flag on successful close, but the same fix was not applied to ctx.Done(). When close fails in that branch, skipping the drain in ReleaseStreamingResponse can leave unread bytes on the connection, reintroducing the connection-corruption bug the drain is explicitly guarding against. core/providers/utils/utils.go — the ctx.Done() case in SetupStreamCancellation Important Files Changed
|
6fe17ce to
8e9a74c
Compare
|
|
Merge activity
|
## Summary This PR fixes a resource leak and incorrect error handling during streaming response cleanup when a client context is cancelled or times out. When a stream is forcibly closed due to context cancellation, the body stream is already closed, so attempting to drain it before releasing the response is unnecessary and can cause errors. A new context key (`BifrostContextKeyConnectionClosed`) is introduced to signal that the connection has already been closed, allowing `ReleaseStreamingResponse` to skip the drain step safely. Additionally, context cancellation is now checked *before* the `io.EOF` check in all SSE/stream read loops, ensuring that a cancelled context causes an immediate clean exit rather than potentially logging spurious errors or sending error events downstream. ## Changes - `ReleaseStreamingResponse` now accepts a `*schemas.BifrostContext` and skips draining the body stream if `BifrostContextKeyConnectionClosed` is set to `true`. - `SetupStreamCancellation` now accepts a `*schemas.BifrostContext` (instead of `context.Context`) and sets `BifrostContextKeyConnectionClosed` on the context when the body stream is closed due to cancellation or timeout. - All call sites across every provider (Anthropic, Azure, Cohere, ElevenLabs, Gemini, HuggingFace, Mistral, OpenAI, Replicate, Vertex, vLLM) are updated to pass the `BifrostContext` to `ReleaseStreamingResponse`. - In all SSE read loops, the `ctx.Err() != nil` early-return check is moved to execute *before* the `io.EOF` guard, so context cancellation is handled immediately regardless of the read error type. - In passthrough stream handlers (OpenAI, Gemini, Vertex), `io.EOF` read errors no longer incorrectly trigger `ProcessAndSendError`; the error path is now guarded with `readErr != io.EOF`. ## Type of change - [x] Bug fix - [ ] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test Initiate a streaming request and cancel the client context mid-stream (e.g., by closing the HTTP connection early). Verify that: - No "whitespace in header" or drain-related panics appear in logs. - No spurious stream error events are sent to the response channel after cancellation. - Response objects are properly released without goroutine leaks. ```sh go test ./... ``` ## Breaking changes - [ ] Yes - [x] No ## Related issues ## Security considerations None. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
## Summary This PR fixes a resource leak and incorrect error handling during streaming response cleanup when a client context is cancelled or times out. When a stream is forcibly closed due to context cancellation, the body stream is already closed, so attempting to drain it before releasing the response is unnecessary and can cause errors. A new context key (`BifrostContextKeyConnectionClosed`) is introduced to signal that the connection has already been closed, allowing `ReleaseStreamingResponse` to skip the drain step safely. Additionally, context cancellation is now checked *before* the `io.EOF` check in all SSE/stream read loops, ensuring that a cancelled context causes an immediate clean exit rather than potentially logging spurious errors or sending error events downstream. ## Changes - `ReleaseStreamingResponse` now accepts a `*schemas.BifrostContext` and skips draining the body stream if `BifrostContextKeyConnectionClosed` is set to `true`. - `SetupStreamCancellation` now accepts a `*schemas.BifrostContext` (instead of `context.Context`) and sets `BifrostContextKeyConnectionClosed` on the context when the body stream is closed due to cancellation or timeout. - All call sites across every provider (Anthropic, Azure, Cohere, ElevenLabs, Gemini, HuggingFace, Mistral, OpenAI, Replicate, Vertex, vLLM) are updated to pass the `BifrostContext` to `ReleaseStreamingResponse`. - In all SSE read loops, the `ctx.Err() != nil` early-return check is moved to execute *before* the `io.EOF` guard, so context cancellation is handled immediately regardless of the read error type. - In passthrough stream handlers (OpenAI, Gemini, Vertex), `io.EOF` read errors no longer incorrectly trigger `ProcessAndSendError`; the error path is now guarded with `readErr != io.EOF`. ## Type of change - [x] Bug fix - [ ] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test Initiate a streaming request and cancel the client context mid-stream (e.g., by closing the HTTP connection early). Verify that: - No "whitespace in header" or drain-related panics appear in logs. - No spurious stream error events are sent to the response channel after cancellation. - Response objects are properly released without goroutine leaks. ```sh go test ./... ``` ## Breaking changes - [ ] Yes - [x] No ## Related issues ## Security considerations None. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable

Summary
This PR fixes a resource leak and incorrect error handling during streaming response cleanup when a client context is cancelled or times out. When a stream is forcibly closed due to context cancellation, the body stream is already closed, so attempting to drain it before releasing the response is unnecessary and can cause errors. A new context key (
BifrostContextKeyConnectionClosed) is introduced to signal that the connection has already been closed, allowingReleaseStreamingResponseto skip the drain step safely.Additionally, context cancellation is now checked before the
io.EOFcheck in all SSE/stream read loops, ensuring that a cancelled context causes an immediate clean exit rather than potentially logging spurious errors or sending error events downstream.Changes
ReleaseStreamingResponsenow accepts a*schemas.BifrostContextand skips draining the body stream ifBifrostContextKeyConnectionClosedis set totrue.SetupStreamCancellationnow accepts a*schemas.BifrostContext(instead ofcontext.Context) and setsBifrostContextKeyConnectionClosedon the context when the body stream is closed due to cancellation or timeout.BifrostContexttoReleaseStreamingResponse.ctx.Err() != nilearly-return check is moved to execute before theio.EOFguard, so context cancellation is handled immediately regardless of the read error type.io.EOFread errors no longer incorrectly triggerProcessAndSendError; the error path is now guarded withreadErr != io.EOF.Type of change
Affected areas
How to test
Initiate a streaming request and cancel the client context mid-stream (e.g., by closing the HTTP connection early). Verify that:
go test ./...Breaking changes
Related issues
Security considerations
None.
Checklist
docs/contributing/README.mdand followed the guidelines