Skip to content

fix: add context cancellation to prevent resource leaks in streaming requests - #732

Merged
akshaydeo merged 1 commit into
mainfrom
11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports
Nov 3, 2025
Merged

fix: add context cancellation to prevent resource leaks in streaming requests#732
akshaydeo merged 1 commit into
mainfrom
11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports

Conversation

@Pratham-Mishra04

Copy link
Copy Markdown
Collaborator

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

  • Bug fix
  • Feature
  • Refactor
  • Documentation
  • Chore/CI

Affected areas

  • Core (Go)
  • Transports (HTTP)
  • Providers/Integrations
  • Plugins
  • UI (Next.js)
  • Docs

How to test

Verify that streaming requests are properly cancelled when clients disconnect:

# 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
  • 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

  • I added/updated tests where appropriate
  • I verified builds succeed (Go and UI)
  • I verified the CI pipeline passes locally if applicable

@coderabbitai

coderabbitai Bot commented Nov 3, 2025

Copy link
Copy Markdown
Contributor

Caution

Review failed

The pull request is closed.

📝 Walkthrough

Summary by CodeRabbit

  • New Features

    • Added utility functions for origin validation and model string parsing.
  • Bug Fixes

    • Improved context cancellation handling to ensure proper cleanup when clients disconnect or requests fail, reducing potential resource leaks in streaming operations.
  • Refactor

    • Streamlined internal error handling by standardizing logger usage across HTTP handlers.

Walkthrough

ConvertToBifrostContext 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

Cohort / File(s) Summary
Context factory
transports/bifrost-http/lib/ctx.go
ConvertToBifrostContext now returns (*context.Context, context.CancelFunc) and creates a cancellable base context with context.WithCancel.
Router: streaming & request wiring
transports/bifrost-http/integrations/router.go
createHandler and streaming flows unpack (bifrostCtx, cancel). Non-streaming branches derive local requestCtx and defer cancel(). handleStreamingRequest/handleStreaming signatures accept cancel and invoke it on write errors/disconnects.
Inference & streaming handlers
transports/bifrost-http/handlers/inference.go
Removed per-instance logger from CompletionHandler and constructor. Streaming handlers accept a cancel func, derive a local streamCtx from bifrostCtx, propagate cancel upstream, and call cancel() on write/errors; non-streaming paths defer cancel(). Logger-based calls removed.
MCP handler
transports/bifrost-http/handlers/mcp.go
MCPHandler logger removed; constructor signature updated. Uses (bifrostCtx, cancel) and defers cancel(); SendError/SendJSON calls no longer take a logger.
Handler constructors (removed per-instance logger)
transports/bifrost-http/handlers/cache.go, .../config.go, .../governance.go, .../health.go, .../integrations.go, .../logging.go, .../plugins.go, .../providers.go, .../websocket.go, .../server.go
Removed logger fields from many Handler structs and dropped logger parameters from New*Handler constructors. Call sites updated to construct handlers without passing per-instance loggers; package-level logger used where applicable.
Response helpers & utilities
transports/bifrost-http/handlers/utils.go
SendJSON, SendJSONWithStatus, SendError, SendBifrostError, SendSSEError signatures changed to remove logger parameters. Added helpers: IsOriginAllowed, isLocalhostOrigin, matchesWildcardPattern, and exported ParseModel.
Middleware error handling
transports/bifrost-http/handlers/middlewares.go
Error path updated to call SendError without passing a logger.
WebSocket & logging handlers
transports/bifrost-http/handlers/websocket.go, transports/bifrost-http/handlers/logging.go
Removed per-instance logger fields and constructor logger params; replaced h.logger uses with package-level logger or removed them.

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
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

  • Check consistent use of defer cancel() on all non-streaming early returns (router.go, inference.go, mcp.go).
  • Verify every streaming callsite receives the new cancel param and that no old-signature calls remain (especially server.go and integration points).
  • Ensure removal of per-instance logger left no unresolved h.logger references and package-level logger usage is appropriate.

Poem

🐰 I hopped through code with a tiny drum,
Cancel in pocket to hush the hum,
Streams now quiet when clients roam,
Goroutines rest and fields come home,
Carrots saved — the heap is calm.

Pre-merge checks and finishing touches

✅ Passed checks (2 passed)
Check name Status Explanation
Title Check ✅ Passed The PR title "fix: add context cancellation to prevent resource leaks in streaming requests" clearly identifies the primary objective of this changeset. The title accurately reflects the main changes visible in the raw summary, which center on implementing context cancellation through the updated ConvertToBifrostContext() function and threading cancel functions through HTTP handlers and streaming paths. While the PR also includes refactoring to remove logger fields from multiple handlers, this appears to be a supporting change related to the context cancellation implementation rather than the primary focus. The title is concise, specific, and avoids generic or misleading language.
Description Check ✅ Passed The pull request description comprehensively follows the provided template structure. It includes a clear summary explaining the purpose of implementing context cancellation to prevent resource leaks, a detailed changes section outlining the key modifications, properly marked type of change (Bug fix and Refactor), identified affected areas (Core Go and Transports HTTP), and concrete testing steps with example commands. The description also addresses breaking changes (marked as None), includes security considerations explaining the resource management improvements, and completes the checklist items. All critical template sections are present and substantively filled out rather than left blank or incomplete.

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c7002b9 and 835b846.

📒 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)

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

Pratham-Mishra04 commented Nov 3, 2025

Copy link
Copy Markdown
Collaborator Author

@Pratham-Mishra04
Pratham-Mishra04 changed the base branch from 10-31-refactor_integration_router_refactored to graphite-base/732 November 3, 2025 07:08
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from 87b0515 to b018959 Compare November 3, 2025 07:08
@Pratham-Mishra04
Pratham-Mishra04 changed the base branch from graphite-base/732 to 10-31-feat_reqest_level_key_skipping_headers_raw_request_passing_and_response_gzip_decompression_added November 3, 2025 07:08
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from b018959 to f9cd27a Compare November 3, 2025 07:28

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 := false

Note: 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.WithCancel must 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 notified

The 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 ConvertToBifrostContext still 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.Context by value instead of *context.Context to 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

📥 Commits

Reviewing files that changed from the base of the PR and between ed452fc and f9cd27a.

📒 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.

@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from f9cd27a to c7002b9 Compare November 3, 2025 09:06

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ConvertToBifrostContext we 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 invoking cancel(), so the context and any upstream work tied to it linger—exactly the leak this PR is trying to eliminate. Please ensure every exit between ConvertToBifrostContext and 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:

  1. The context pointer is already available as bifrostCtx
  2. The dereferencing and copying adds cognitive overhead
  3. The copy is immediately used in the closure that captures it

Consider simplifying by passing *bifrostCtx directly 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

📥 Commits

Reviewing files that changed from the base of the PR and between f9cd27a and c7002b9.

📒 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: Ensure logger is still defined before these calls

We dropped the logger parameter, yet logger.Warn is 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 with undefined: logger. Please confirm where the shared logger now lives or wire it back in.

transports/bifrost-http/handlers/integrations.go (1)

18-26: Undefined logger after signature change

After removing the logger parameter, this function still forwards logger to each router. Unless a package-level logger has 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 does logger now come from?

LoggingHandler dropped its logger field, yet logger.Error is 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 ConvertToBifrostContext and cleanup is ensured via defer 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.Warn calls 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:

  1. Cancels the upstream context on getStream() errors (line 927)
  2. Cancels on write errors when the client disconnects (lines 969, 975, 981, 988, 997)
  3. Includes detailed comments explaining the cancellation strategy (lines 913-915, 1002)
  4. 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: cancel is passed to handleStreamingResponse, which calls it when write errors indicate client disconnection (lines detecting fmt.Fprintf and w.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."

@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from c7002b9 to 8e7d3a0 Compare November 3, 2025 10:23
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 10-31-feat_reqest_level_key_skipping_headers_raw_request_passing_and_response_gzip_decompression_added branch 2 times, most recently from 79dca19 to d576419 Compare November 3, 2025 10:38
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from 8e7d3a0 to 2fa41dd Compare November 3, 2025 10:38
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 10-31-feat_reqest_level_key_skipping_headers_raw_request_passing_and_response_gzip_decompression_added branch from d576419 to ea67394 Compare November 3, 2025 10:48
@Pratham-Mishra04
Pratham-Mishra04 force-pushed the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch from 2fa41dd to 835b846 Compare November 3, 2025 10:48

akshaydeo commented Nov 3, 2025

Copy link
Copy Markdown
Contributor

Merge activity

  • Nov 3, 10:49 AM UTC: A user started a stack merge that includes this pull request via Graphite.
  • Nov 3, 10:51 AM UTC: @akshaydeo merged this pull request with Graphite.

@akshaydeo
akshaydeo changed the base branch from 10-31-feat_reqest_level_key_skipping_headers_raw_request_passing_and_response_gzip_decompression_added to graphite-base/732 November 3, 2025 10:50
@akshaydeo
akshaydeo changed the base branch from graphite-base/732 to main November 3, 2025 10:50
@akshaydeo
akshaydeo merged commit 408a175 into main Nov 3, 2025
3 checks passed
@akshaydeo
akshaydeo deleted the 11-03-feat_propgate_ctx_cancellation_from_connection_drop_in_transports branch November 3, 2025 10:51
akshaydeo added a commit that referenced this pull request Nov 17, 2025
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants