fix: add context cancellation to prevent resource leaks in streaming requests - #732
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughSummary by CodeRabbit
WalkthroughConvertToBifrostContext now returns a cancellable context and CancelFunc; cancel is propagated and invoked across router and handler streaming/non‑streaming flows. Many handlers had per-instance logger fields removed and constructors updated; Send* helper signatures dropped logger parameters. Streaming paths now explicitly cancel upstream on write errors or client disconnects. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router
participant CtxFactory as ConvertToBifrostContext
participant Handler
participant Bifrost as Bifrost API
participant Upstream
Client->>Router: HTTP/SSE/WebSocket request
Router->>CtxFactory: ConvertToBifrostContext(req)
CtxFactory-->>Router: (bifrostCtx, cancel)
Router->>Handler: invoke with bifrostCtx (and cancel)
Handler->>Bifrost: call (uses requestCtx derived from bifrostCtx)
Bifrost->>Upstream: open/proxy stream (inherits requestCtx)
rect rgb(255,245,230)
Note over Client,Handler: Client disconnect or SSE write failure
Client--xHandler: connection lost / write error
Handler->>Handler: invoke cancel()
Handler->>Upstream: cancellation via context
end
Upstream-->>Bifrost: stream closed
Bifrost-->>Handler: stream end/error
Handler->>Handler: defer cancel() cleanup on exit
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Poem
Pre-merge checks and finishing touches✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (16)
Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
87b0515 to
b018959
Compare
b018959 to
f9cd27a
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
transports/bifrost-http/integrations/router.go (1)
645-837: Fix: Missing cancel() on normal stream completion.The cancel function is called on write errors throughout the streaming logic but never called when the stream completes normally. This leaks context resources.
When using
context.WithCancel(), the cancel function MUST be called in all paths to release associated resources (goroutines, timers, etc.). The comments at lines 733 and 836 claiming "Bifrost handles cleanup internally" are incorrect for Go context lifecycle management.Apply this diff to ensure cancel is called in all paths:
func (g *GenericRouter) handleStreaming(ctx *fasthttp.RequestCtx, config RouteConfig, streamChan chan *schemas.BifrostStream, cancel context.CancelFunc) { // Use streaming response writer ctx.Response.SetBodyStreamWriter(func(w *bufio.Writer) { + defer cancel() // Ensure cleanup in all paths: errors, client disconnect, or normal completion defer w.Flush() includeEventType := falseNote: With this defer, the individual
cancel()calls scattered throughout (lines 699, 716, 723, 730, 768, 783, 797, 813, 821, 831) become redundant but harmless. Consider removing them for cleaner code, or leave them for explicit documentation of the cancellation points.transports/bifrost-http/handlers/inference.go (1)
915-1002: Critical: Missing defer cancel() causes context resource leak in streaming requests.The cancel function is only called on error paths (lines 926, 968, 974, 980, 987, 996) but never called when the stream completes normally. This violates Go's context management contract: every cancel function returned by
context.WithCancelmust be called exactly once to release resources.Impact:
- Every streaming request that completes successfully leaks context resources
- Memory accumulates over time as contexts are never released
- Any goroutines waiting on
ctx.Done()may not be properly notifiedThe comment on lines 912-914 suggests this is intentional ("Bifrost handles cleanup internally"), but this doesn't absolve us from calling cancel(). Even if Bifrost cleans up its own state, the context created in
ConvertToBifrostContextstill needs explicit cancellation.Apply this fix:
func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, getStream func() (chan *schemas.BifrostStream, *schemas.BifrostError), cancel context.CancelFunc) { + // Ensure context cleanup on all paths (normal completion, errors, client disconnect) + defer cancel() + // Set SSE headers ctx.SetContentType("text/event-stream")Then remove the explicit cancel() calls in error paths since defer will handle all cases:
stream, bifrostErr := getStream() if bifrostErr != nil { - // Cancel stream context since we're not proceeding - cancel() SendBifrostError(ctx, bifrostErr, h.logger) return }And update error handling in the write paths:
if _, err := fmt.Fprintf(w, "event: %s\n", eventType); err != nil { - cancel() // Client disconnected (write error), cancel upstream stream return }(Apply similar changes to all other explicit cancel() calls within the function—they become redundant with defer.)
🧹 Nitpick comments (1)
transports/bifrost-http/handlers/inference.go (1)
847-857: Note: Context dereferencing pattern is safe.All streaming handlers dereference the context pointer (
streamCtx := *bifrostCtx) before passing to Bifrost client methods. This is safe in Go—contexts are designed to be copied by value.However, consider whether the upstream signature could accept
context.Contextby value instead of*context.Contextto align with Go idioms. This would be a future refactor, not required for this PR.Also applies to: 860-870, 873-883, 886-896, 899-909
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
transports/bifrost-http/handlers/inference.go(11 hunks)transports/bifrost-http/handlers/mcp.go(1 hunks)transports/bifrost-http/integrations/router.go(18 hunks)transports/bifrost-http/lib/ctx.go(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
transports/bifrost-http/handlers/inference.go (3)
transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)transports/bifrost-http/handlers/utils.go (1)
SendError(35-44)core/schemas/bifrost.go (2)
BifrostStream(291-298)BifrostError(326-335)
transports/bifrost-http/handlers/mcp.go (1)
transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)
transports/bifrost-http/integrations/router.go (2)
transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)core/schemas/bifrost.go (4)
BifrostRequest(134-144)ListModelsRequest(82-82)TextCompletionRequest(83-83)BifrostStream(291-298)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (14)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (8)
transports/bifrost-http/handlers/mcp.go (1)
60-61: LGTM! Context cancellation properly implemented.The handler correctly captures the cancel function and defers cleanup. This ensures resources are released when the request completes, whether normally or via error path.
transports/bifrost-http/lib/ctx.go (3)
73-77: LGTM! Cancellable context correctly implemented.The function now properly creates a cancellable context that enables cleanup when clients disconnect. The base context and cancel function are correctly set up with
context.WithCancel.
53-70: Good documentation updates.The doc comments clearly explain the new cancellable context behavior and show proper usage with
defer cancel(). The example demonstrates the correct pattern.
253-253: LGTM! Return signature correctly updated.The function now returns both the context pointer and the cancel function, enabling callers to clean up resources.
transports/bifrost-http/integrations/router.go (3)
318-365: LGTM! Ownership transfer pattern correctly implemented.The createHandler properly manages cancel ownership:
- Captures cancel from ConvertToBifrostContext
- Non-streaming path: defers cancel for cleanup on exit
- Streaming path: passes cancel to handleStreamingRequest which takes ownership
371-377: Good explanation of non-streaming context usage.The comments clearly explain why a cancellable context is useful even for non-streaming requests, noting that it allows provider-level cancellation even though client disconnects can't be detected until write time.
550-597: LGTM! Streaming initialization correctly handles errors.The function properly:
- Creates streamCtx from the cancellable bifrostCtx
- Calls cancel() when streaming setup fails (lines 577, 585)
- Drains the stream channel to prevent goroutine leaks (lines 587-590)
The defensive cleanup on error paths is well implemented.
transports/bifrost-http/handlers/inference.go (1)
296-297: LGTM: Proper context cleanup pattern.The defer cancel() immediately after obtaining the cancellable context ensures cleanup in all code paths.
f9cd27a to
c7002b9
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
transports/bifrost-http/integrations/router.go (1)
317-366: Always cancel bifrostCtx on early exits in createHandler.After
ConvertToBifrostContextwe now hold a cancellable context, but several error returns (pre-callback failure, conversion errors, fallback parsing, etc.) happen before we know whether streaming is requested. In those paths we bail out without ever invokingcancel(), so the context and any upstream work tied to it linger—exactly the leak this PR is trying to eliminate. Please ensure every exit betweenConvertToBifrostContextand the streaming check triggers the cancel function, while still keeping the streaming path in control so we don't tear down live streams prematurely.A pattern like this keeps non-streaming/exiting paths safe while letting the streaming path own cancellation:
- bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, g.handlerStore.ShouldAllowDirectKeys()) + bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, g.handlerStore.ShouldAllowDirectKeys()) + cleanup := cancel + defer func() { + if cleanup != nil { + cleanup() + } + }() … - if isStreaming { - g.handleStreamingRequest(ctx, config, bifrostReq, bifrostCtx, cancel) + if isStreaming { + cleanup = nil // streaming path owns cancel() + g.handleStreamingRequest(ctx, config, bifrostReq, bifrostCtx, cancel) return } else { - defer cancel() // Ensure cleanup on function exit g.handleNonStreamingRequest(ctx, config, req, bifrostReq, bifrostCtx) }That way any failure before the streaming branch still calls
cancel(), while streaming continues to manage disconnect-triggered cancellation.
🧹 Nitpick comments (1)
transports/bifrost-http/handlers/inference.go (1)
848-858: Consider simplifying context copying pattern.Lines 851, 864, 877, 890, and 903 create a local copy of the context via
streamCtx := *bifrostCtx. This pattern seems unnecessary since:
- The context pointer is already available as
bifrostCtx- The dereferencing and copying adds cognitive overhead
- The copy is immediately used in the closure that captures it
Consider simplifying by passing
*bifrostCtxdirectly to the client methods instead of creating a copy:-func (h *CompletionHandler) handleStreamingTextCompletion(ctx *fasthttp.RequestCtx, req *schemas.BifrostTextCompletionRequest, bifrostCtx *context.Context, cancel context.CancelFunc) { - // Use the cancellable context from ConvertToBifrostContext - // See router.go for detailed explanation of why we need a cancellable context - streamCtx := *bifrostCtx - getStream := func() (chan *schemas.BifrostStream, *schemas.BifrostError) { - return h.client.TextCompletionStreamRequest(streamCtx, req) + return h.client.TextCompletionStreamRequest(*bifrostCtx, req) } h.handleStreamingResponse(ctx, getStream, cancel) }Apply similar changes to the other streaming handlers (lines 861-871, 874-884, 887-897, 900-910).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
transports/bifrost-http/handlers/cache.go(2 hunks)transports/bifrost-http/handlers/config.go(11 hunks)transports/bifrost-http/handlers/governance.go(29 hunks)transports/bifrost-http/handlers/health.go(2 hunks)transports/bifrost-http/handlers/inference.go(18 hunks)transports/bifrost-http/handlers/integrations.go(1 hunks)transports/bifrost-http/handlers/logging.go(4 hunks)transports/bifrost-http/handlers/mcp.go(3 hunks)transports/bifrost-http/handlers/middlewares.go(1 hunks)transports/bifrost-http/handlers/plugins.go(5 hunks)transports/bifrost-http/handlers/providers.go(14 hunks)transports/bifrost-http/handlers/server.go(3 hunks)transports/bifrost-http/handlers/utils.go(2 hunks)transports/bifrost-http/handlers/websocket.go(7 hunks)transports/bifrost-http/integrations/router.go(18 hunks)transports/bifrost-http/lib/ctx.go(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (14)
transports/bifrost-http/handlers/integrations.go (2)
core/bifrost.go (1)
Bifrost(34-52)transports/bifrost-http/lib/config.go (1)
HandlerStore(34-37)
transports/bifrost-http/handlers/logging.go (2)
plugins/logging/utils.go (1)
LogManager(16-25)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)
transports/bifrost-http/handlers/governance.go (4)
plugins/governance/main.go (1)
GovernancePlugin(39-58)framework/configstore/store.go (1)
ConfigStore(17-134)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)framework/configstore/tables/utils.go (1)
ParseDuration(9-43)
transports/bifrost-http/handlers/cache.go (2)
plugins/semanticcache/main.go (1)
Plugin(137-144)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)
transports/bifrost-http/handlers/config.go (3)
transports/bifrost-http/lib/config.go (1)
Config(135-165)transports/bifrost-http/handlers/utils.go (2)
SendJSON(16-22)SendError(35-44)framework/configstore/store.go (1)
ConfigStore(17-134)
transports/bifrost-http/handlers/providers.go (3)
transports/bifrost-http/lib/config.go (2)
Config(135-165)ValidateCustomProviderUpdate(2270-2299)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)core/utils.go (2)
IsStandardProvider(177-180)IsSupportedBaseProvider(163-166)
transports/bifrost-http/handlers/plugins.go (3)
framework/configstore/store.go (1)
ConfigStore(17-134)transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)framework/configstore/tables/plugin.go (2)
TablePlugin(12-22)TablePlugin(25-25)
transports/bifrost-http/handlers/mcp.go (4)
transports/bifrost-http/lib/config.go (1)
Config(135-165)transports/bifrost-http/handlers/utils.go (3)
SendError(35-44)SendBifrostError(47-62)SendJSON(16-22)transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)core/schemas/mcp.go (3)
MCPConfig(11-13)MCPClient(58-63)MCPClientConfig(16-28)
transports/bifrost-http/handlers/middlewares.go (1)
transports/bifrost-http/handlers/utils.go (1)
SendError(35-44)
transports/bifrost-http/handlers/health.go (1)
transports/bifrost-http/handlers/utils.go (2)
SendError(35-44)SendJSON(16-22)
transports/bifrost-http/handlers/inference.go (2)
transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)transports/bifrost-http/handlers/utils.go (3)
SendError(35-44)SendBifrostError(47-62)SendJSON(16-22)
transports/bifrost-http/handlers/server.go (12)
transports/bifrost-http/handlers/logging.go (1)
NewLoggingHandler(24-28)transports/bifrost-http/handlers/governance.go (2)
GovernanceHandler(22-26)NewGovernanceHandler(29-39)transports/bifrost-http/handlers/cache.go (2)
CacheHandler(11-13)NewCacheHandler(15-24)transports/bifrost-http/handlers/websocket.go (2)
WebSocketHandler(27-35)NewWebSocketHandler(38-47)transports/bifrost-http/handlers/health.go (1)
NewHealthHandler(19-23)transports/bifrost-http/handlers/providers.go (1)
NewProviderHandler(29-34)transports/bifrost-http/handlers/inference.go (1)
NewInferenceHandler(33-39)transports/bifrost-http/handlers/mcp.go (1)
NewMCPHandler(25-30)transports/bifrost-http/handlers/integrations.go (1)
NewIntegrationHandler(18-31)transports/bifrost-http/handlers/config.go (1)
NewConfigHandler(38-43)transports/bifrost-http/handlers/plugins.go (1)
NewPluginsHandler(29-34)transports/bifrost-http/handlers/utils.go (1)
SendError(35-44)
transports/bifrost-http/integrations/router.go (2)
transports/bifrost-http/lib/ctx.go (1)
ConvertToBifrostContext(73-254)core/schemas/bifrost.go (4)
BifrostRequest(134-144)ListModelsRequest(82-82)TextCompletionRequest(83-83)BifrostStream(291-298)
transports/bifrost-http/handlers/utils.go (1)
core/schemas/bifrost.go (2)
BifrostError(326-335)ErrorField(344-351)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (14)
transports/bifrost-http/handlers/utils.go (1)
17-31: Ensureloggeris still defined before these callsWe dropped the logger parameter, yet
logger.Warnis still used here (and in SendBifrostError/SendSSEError). I don't see a package-level logger introduced in this package within this PR; without one the build will fail withundefined: logger. Please confirm where the shared logger now lives or wire it back in.transports/bifrost-http/handlers/integrations.go (1)
18-26: Undefinedloggerafter signature changeAfter removing the logger parameter, this function still forwards
loggerto each router. Unless a package-levelloggerhas been introduced elsewhere, this will not compile. Please point me to the shared logger or adjust the constructor wiring accordingly.transports/bifrost-http/handlers/logging.go (1)
143-147: Where doesloggernow come from?
LoggingHandlerdropped its logger field, yetlogger.Erroris still invoked here. Without a new package-level logger, this is an undefined identifier and will break the build. Please either restore the field/parameter or hook these calls to the new shared logger.transports/bifrost-http/handlers/cache.go (1)
15-24: LGTM! Clean logger removal.The constructor signature simplification and logger field removal are consistent with the project-wide refactor to use a package-level logger instead of per-instance loggers.
transports/bifrost-http/handlers/mcp.go (1)
58-59: LGTM! Proper context cancellation implementation.The cancellable context is correctly obtained from
ConvertToBifrostContextand cleanup is ensured viadefer cancel(). This prevents resource leaks when clients disconnect during MCP tool execution.transports/bifrost-http/handlers/config.go (1)
38-43: LGTM! Constructor simplified correctly.The logger parameter removal aligns with the shift to package-level logging. The remaining
logger.Warncalls throughout the file use the package-level logger, which is the intended pattern.transports/bifrost-http/handlers/websocket.go (1)
38-46: LGTM! WebSocket handler updated correctly.The logger removal is consistent with the project-wide refactor. The WebSocket connection lifecycle is managed independently of the HTTP request context, so the existing cleanup mechanisms remain appropriate.
transports/bifrost-http/handlers/health.go (1)
19-23: LGTM! Health handler simplified appropriately.The logger removal is clean and the health check logic with timeout context for store pings remains correct.
transports/bifrost-http/handlers/providers.go (1)
29-34: LGTM! Provider handler refactored correctly.The logger field removal is clean and the remaining logger calls use the package-level logger appropriately. Provider management endpoints correctly omit context cancellation as they're not long-running inference operations.
transports/bifrost-http/handlers/plugins.go (1)
29-34: LGTM! Plugins handler updated consistently.The logger parameter removal is correct and the plugin CRUD operations use the package-level logger appropriately. Plugin management endpoints correctly don't implement context cancellation.
transports/bifrost-http/handlers/inference.go (4)
294-299: LGTM! Proper cancellation in listModels.The context is correctly obtained with cancellation support and cleanup is ensured via
defer cancel().
916-1003: Excellent streaming cancellation implementation!The streaming handler correctly:
- Cancels the upstream context on
getStream()errors (line 927)- Cancels on write errors when the client disconnects (lines 969, 975, 981, 988, 997)
- Includes detailed comments explaining the cancellation strategy (lines 913-915, 1002)
- Properly handles both event-type responses and standard SSE format
The comment on line 1002 correctly notes that normal completion doesn't need explicit cancellation since Bifrost handles cleanup internally. This is a well-designed implementation that prevents resource leaks.
620-625: LGTM! Embeddings handler updated correctly.The context cancellation with
defer cancel()is properly placed before the request, ensuring cleanup for both success and error paths.
392-405: Context cancellation is properly implemented—no changes needed.Verification confirms the code handles context cleanup correctly for both streaming and non-streaming paths:
- Streaming path:
cancelis passed tohandleStreamingResponse, which calls it when write errors indicate client disconnection (lines detectingfmt.Fprintfandw.Flush()failures).- Non-streaming path:
defer cancel()ensures cleanup on normal function exit.The behavior is intentional and already documented in the function header comment: "The cancel function is called ONLY when client disconnects are detected via write errors."
c7002b9 to
8e7d3a0
Compare
79dca19 to
d576419
Compare
8e7d3a0 to
2fa41dd
Compare
d576419 to
ea67394
Compare
2fa41dd to
835b846
Compare
Merge activity
|
…requests (#732) ## Summary Implement proper context cancellation for HTTP requests to prevent resource leaks and wasted provider quota when clients disconnect, especially during streaming operations. ## Changes - Added cancellable context support in `ConvertToBifrostContext()` which now returns both a context and a cancel function - Ensured all HTTP handlers properly defer the cancel function to clean up resources - Modified streaming handlers to immediately cancel upstream provider requests when clients disconnect - Added detailed comments explaining the context cancellation pattern and its importance - Fixed potential resource leaks in streaming handlers by properly handling write errors ## Type of change - [x] Bug fix - [ ] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [ ] Plugins - [ ] UI (Next.js) - [ ] Docs ## How to test Verify that streaming requests are properly cancelled when clients disconnect: ```sh # Start Bifrost server go run cmd/bifrost/main.go # In another terminal, start a streaming request and then cancel it (Ctrl+C) curl -N http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Write a very long story"}], "stream": true}' # Verify in server logs that the upstream request was cancelled ``` ## Breaking changes - [ ] Yes - [x] No ## Related issues Addresses resource leaks and wasted provider quota when clients disconnect during streaming requests. ## Security considerations Improves resource management by ensuring provider requests are cancelled when no longer needed, preventing potential resource exhaustion. ## Checklist - [x] I added/updated tests where appropriate - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable

Summary
Implement proper context cancellation for HTTP requests to prevent resource leaks and wasted provider quota when clients disconnect, especially during streaming operations.
Changes
ConvertToBifrostContext()which now returns both a context and a cancel functionType of change
Affected areas
How to test
Verify that streaming requests are properly cancelled when clients disconnect:
Breaking changes
Related issues
Addresses resource leaks and wasted provider quota when clients disconnect during streaming requests.
Security considerations
Improves resource management by ensuring provider requests are cancelled when no longer needed, preventing potential resource exhaustion.
Checklist