Fix goroutine leaks on context cancellation

[pjcdawkins]

Summary

Fixes #828: Workers and streaming goroutines do not exit and leak resources on context cancellation.

This PR addresses two critical goroutine leak issues:

1. Worker loops that ignore context cancellation
2. Streaming handlers that block indefinitely on I/O

Changes

Worker Loop Fix

Refactored requestWorker in core/bifrost.go to monitor bifrost.ctx.Done() using a select statement. Workers now exit immediately when the context is cancelled, instead of only exiting when the queue channel closes.

Before:

for req := range queue {
    // Process request
}

After:

for {
    select {
    case req, ok := <-queue:
        if !ok {
            return
        }
        bifrost.processRequest(provider, config, req)
    case <-bifrost.ctx.Done():
        return
    }
}

Streaming Handler Fixes

Added context monitoring goroutines to all streaming handlers that force-close the response body stream when context is cancelled. This unblocks scanner.Scan() calls that were preventing goroutines from exiting.

Pattern applied to all streaming methods:

// Monitor context cancellation and force-close response body to unblock scanner.
done := make(chan struct{})
defer close(done)
go func() {
    select {
    case <-ctx.Done():
        resp.CloseBodyStream()
    case <-done:
    }
}()

Fixed in 11 streaming methods across 4 providers:

• OpenAI: 5 methods (TextCompletionStream, ChatCompletionStream, ResponsesStream, SpeechStream, TranscriptionStream)
• Anthropic: 2 methods (ChatCompletionStream, ResponsesStream)
• Cohere: 2 methods (ChatCompletionStream, ResponsesStream)
• Gemini: 2 methods (SpeechStream, TranscriptionStream)

Testing

Added TestBifrostContextCancellationCleanup to verify goroutines clean up properly after context timeout.

Test Results:

• Before fixes: 5 goroutines leaked
• After fixes: 3 goroutines leaked (40% improvement)

Type of change

• Bug fix

Affected areas

• Core (Go)

How to test

# Run the goroutine leak test
go test -v -run TestBifrostContextCancellationCleanup ./core

# Run all tests
go test ./...

Breaking changes

• Yes
• No

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I verified builds succeed
• I verified the CI pipeline passes locally if applicable

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (5)

core/providers/gemini/gemini.go (1)

749-758: TranscriptionStream uses the same robust cancellation pattern

This block mirrors the SpeechStream watcher and provides the same cancellation and cleanup guarantees; no additional issues here.

core/providers/openai/openai.go (4)

798-812: Same cancellation watcher pattern applied to chat streams

Chat SSE streaming now uses the same ctx.Done() → CloseBodyStream watcher and done channel to avoid stuck scanners; consistent with the text completion path.

---

1196-1210: Responses streaming cancellation handling matches other handlers

The responses stream handler uses the same watcher pattern, which should prevent leaks when callers cancel the context mid‑stream.

---

1593-1607: Speech streaming gains proper cancellation unblocking

For speech SSE, the new watcher ensures the underlying body stream is closed on ctx.Done(), so the scanner loop can exit promptly instead of hanging on I/O.

---

1878-1892: Transcription streaming cancellation is now aligned with other streams

Transcription SSE streaming now also force‑closes the body stream on context cancellation, bringing it in line with the other OpenAI streaming handlers and avoiding goroutine leaks.

🧹 Nitpick comments (3)

core/bifrost.go (1)

2058-2071: Align key-selection error send with context-aware response sending

In the key-selection failure path you still use a bare send:

req.Err <- schemas.BifrostError{ ... }
return

Whereas later you’ve switched to select-based, context-aware sends with a timeout for both responses and streams. While this error path is expected to be consumed by the same caller, using the same non-blocking pattern here would make behavior more consistent and avoid a potential hard block if the caller has already stopped listening.

You could mirror the later pattern:

-        if err != nil {
-            bifrost.logger.Warn("error selecting key for model %s: %v", model, err)
-            req.Err <- schemas.BifrostError{
-                IsBifrostError: false,
-                Error: &schemas.ErrorField{
-                    Message: err.Error(),
-                    Error:   err,
-                },
-            }
-            return
-        }
+        if err != nil {
+            bifrost.logger.Warn("error selecting key for model %s: %v", model, err)
+            bifrostErr := schemas.BifrostError{
+                IsBifrostError: false,
+                Error: &schemas.ErrorField{
+                    Message: err.Error(),
+                    Error:   err,
+                },
+            }
+            select {
+            case req.Err <- bifrostErr:
+            case <-req.Context.Done():
+                bifrost.logger.Debug("Client context cancelled while sending key-selection error")
+            case <-time.After(5 * time.Second):
+                bifrost.logger.Warn("Timeout while sending key-selection error, client may have disconnected")
+            }
+            return
+        }

Not strictly required for correctness, but it would make this path consistent with the new response/stream handling.

core/bifrost_test.go (1)

965-1048: Goroutine leak test is well‑targeted; consider flakiness and runtime cost

This test does a good job exercising a streaming request under context timeout and asserting goroutines return close to baseline. Two minor considerations:

• Using GC + time.Sleep(100 * time.Millisecond) twice plus a 2s timeout and 2s post‑Shutdown sleep can make the test relatively slow; you might be able to reduce the sleeps/timeouts once you’re confident in the fix.
• The ≤ 2 goroutine tolerance is quite strict; if CI ever shows sporadic failures due to unrelated background goroutines, slightly widening the tolerance (or aggregating over a few runs) would make it more robust while still catching real leaks.

Functionally this looks correct and gives valuable regression coverage for #828.

core/providers/openai/openai.go (1)

398-411: Context cancellation watcher correctly unblocks text‑completion stream

The added done/ctx.Done() watcher that calls resp.CloseBodyStream() gives the scanner a reliable unblock path on cancellation while ensuring the watcher goroutine exits when the stream routine returns. This directly addresses stuck streaming goroutines.

If you find yourself touching this pattern again, a small helper (e.g., withStreamCancellationWatcher(ctx, resp, func(...) { ... })) shared across providers could reduce duplication and keep behavior consistent.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 18a9284 and e4c277f.

📒 Files selected for processing (15)

• core/bifrost.go (1 hunks)
• core/bifrost_test.go (2 hunks)
• core/providers/anthropic/anthropic.go (2 hunks)
• core/providers/anthropic/chat.go (1 hunks)
• core/providers/azure/types.go (1 hunks)
• core/providers/bedrock/chat.go (2 hunks)
• core/providers/cohere/chat.go (1 hunks)
• core/providers/cohere/cohere.go (2 hunks)
• core/providers/cohere/models.go (1 hunks)
• core/providers/gemini/chat.go (1 hunks)
• core/providers/gemini/gemini.go (2 hunks)
• core/providers/gemini/models.go (2 hunks)
• core/providers/openai/openai.go (5 hunks)
• core/providers/openai/types.go (1 hunks)
• core/providers/vertex/models.go (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (7)

core/providers/gemini/chat.go (2)

core/providers/gemini/types.go (1)

• Type (769-769)

core/schemas/chatcompletions.go (1)

• ChatToolChoiceTypeFunction (250-250)

core/providers/cohere/models.go (1)

core/schemas/models.go (1)

• Model (47-66)

core/bifrost.go (4)

core/schemas/provider.go (4)

• Provider (207-234)
• ProviderConfig (161-169)
• CustomProviderConfig (142-148)
• PostHookRunner (204-204)

core/schemas/bifrost.go (9)

• BifrostRequest (140-150)
• BifrostResponse (243-253)
• BifrostStream (315-322)
• BifrostError (350-359)
• RequestType (81-81)
• ErrorField (368-375)
• BifrostContextKeySelectedKeyID (107-107)
• BifrostContextKeySelectedKeyName (108-108)
• BifrostErrorExtraFields (417-421)

core/schemas/account.go (1)

• Key (8-17)

core/utils.go (1)

• IsStreamRequestType (190-192)

core/bifrost_test.go (7)

core/schemas/bifrost.go (2)

• OpenAI (35-35)
• BifrostConfig (21-29)

core/bifrost.go (1)

• Init (86-214)

core/schemas/account.go (1)

• Account (53-69)

core/schemas/logger.go (2)

• Logger (28-55)
• LogLevelError (13-13)

core/logger.go (1)

• NewDefaultLogger (40-49)

core/schemas/chatcompletions.go (3)

• BifrostChatRequest (11-18)
• ChatMessage (346-355)
• ChatMessageContent (358-361)

core/schemas/provider.go (1)

• Provider (207-234)

core/providers/gemini/models.go (2)

core/schemas/models.go (1)

• Model (47-66)

core/providers/gemini/types.go (1)

• GeminiModel (1255-1269)

core/providers/vertex/models.go (1)

core/schemas/models.go (1)

• Model (47-66)

core/providers/bedrock/chat.go (1)

core/schemas/chatcompletions.go (1)

• ChatAssistantMessageToolCall (483-488)

🔇 Additional comments (16)

core/providers/bedrock/chat.go (2)

90-91: LGTM! Proper string pointer initialization for tool calls.

The formatting changes correctly populate the Type and ID fields as string pointers, matching the schema requirements. The use of &toolUseID is safe since toolUseID is a local copy created on line 85 to avoid range loop variable capture issues.

---

188-188: LGTM! Standard variable declaration.

The variable declaration formatting is standard and appropriate for Go.

core/providers/azure/types.go (1)

20-30: No-op struct formatting change

Only field alignment on AzureModel.CreatedAt; tags and types are unchanged, so there’s no behavioral impact. Looks good as-is.

core/providers/anthropic/chat.go (1)

704-732: Comment wording/whitespace tweak is accurate

The updated comment above the AnthropicStreamEventTypeContentBlockStart branch still correctly describes emitting tool-call metadata; no code behavior change and the comment remains consistent with the logic.

core/providers/cohere/models.go (1)

10-12: Slice init formatting only

Data: make([]schemas.Model, 0, len(response.Models)) is unchanged semantically from before; preallocation is appropriate and behavior is unaffected.

core/providers/gemini/chat.go (1)

70-80: Tool call struct literal spacing only

The Gemini tool-call construction (ID and Type fields) is unchanged in value and behavior; this is purely a formatting tweak and is consistent with the rest of the file.

core/providers/cohere/chat.go (1)

271-278: Formatting-only change in tool call mapping

The bifrostToolCall literal still maps ID and function fields exactly as before; this is a whitespace-only adjustment with no behavioral impact.

core/providers/gemini/gemini.go (1)

490-499: Context cancellation watcher for SpeechStream looks correct

The done-channel + watcher goroutine cleanly forces resp.CloseBodyStream() on ctx.Done() and guarantees the watcher exits when the main streaming goroutine ends. This should unblock the scanner without introducing leaks.

core/providers/cohere/cohere.go (2)

411-420: ChatCompletionStream cancellation watcher is well-scoped

The added watcher cleanly unblocks the streaming scanner on context cancellation and terminates via the done channel, so workers shouldn’t leak even if callers drop the context mid-stream.

---

634-643: ResponsesStream cancellation watcher mirrors ChatCompletionStream

Same pattern and guarantees as in ChatCompletionStream; it looks sound and consistent with the rest of the codebase.

core/providers/gemini/models.go (1)

15-16: Formatting-only struct literal changes

The adjustments to slice literal alignment are cosmetic; capacities and types are unchanged, so runtime behavior is identical.

Also applies to: 42-43

core/providers/anthropic/anthropic.go (2)

469-478: Anthropic chat streaming now correctly unblocks on context cancellation

The watcher goroutine plus done channel ensure the BodyStream is closed when ctx.Done() fires and that the watcher terminates with the main streaming goroutine. This aligns Anthropic’s behavior with other providers and should prevent stuck scanners.

---

735-744: Anthropic ResponsesStream adopts the same safe cancellation pattern

Same watcher pattern as in HandleAnthropicChatCompletionStreaming, providing consistent cancellation behavior across Anthropic streaming APIs.

core/providers/vertex/models.go (1)

10-12: Slice preallocation formatting only

Data still preallocates to len(response.Models); behavior unchanged and correct.

core/providers/openai/types.go (1)

128-134: Struct field whitespace cleanup is safe

Created keeps the same pointer type and JSON tag; no impact on encoding/decoding.

core/bifrost_test.go (1)

3-13: New test imports are appropriate and scoped

runtime and testify/assert are used only for the goroutine-leak test; import block remains tidy.

[Pratham-Mishra04]

Hey @pjcdawkins, thanks for the PR! It’s looking great. There are a few solid Code Rabbit comments to address though. And while you’re at it, please add your update to core/changelog.md and double-check that all tests in tests/core-providers are still passing.

[codecov]

Codecov Report

❌ Patch coverage is 22.28916% with 129 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
core/providers/openai/openai.go	0.00%	38 Missing ⚠️
core/bifrost.go	51.38%	32 Missing and 3 partials ⚠️
core/providers/gemini/gemini.go	0.00%	16 Missing ⚠️
core/providers/anthropic/anthropic.go	0.00%	15 Missing ⚠️
core/providers/cohere/cohere.go	0.00%	15 Missing ⚠️
core/providers/bedrock/chat.go	0.00%	3 Missing ⚠️
core/providers/gemini/chat.go	0.00%	2 Missing ⚠️
core/providers/gemini/models.go	0.00%	2 Missing ⚠️
core/providers/cohere/chat.go	0.00%	1 Missing ⚠️
core/providers/cohere/models.go	0.00%	1 Missing ⚠️
... and 1 more

📢 Thoughts on this report? Let us know!

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

core/changelog.md (1)

1-1: Changelog entry accurately captures the fix.

The entry correctly documents the goroutine leak fix across both the worker loop and streaming handlers, addressing the PR objectives. The concise format is appropriate for a changelog.

Minor enhancement: Consider mentioning "context cancellation" for clarity, since that's the mechanism driving the fix—e.g., "fix: Goroutine leaks during context cancellation in worker loop and streaming request handlers". This helps users understand the specific scenario that was fixed.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8028533 and 6d849f0.

📒 Files selected for processing (1)

• core/changelog.md (1 hunks)

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

core/bifrost.go (1)

2564-2569: Document expected Shutdown usage to avoid race condition.

The shutdown sequence (cancel context → close queues → wait for workers) is intentional to signal streaming goroutines to unblock. However, this creates a brief window where new requests could be enqueued after context cancellation but before queues are closed.

Recommend documenting that Shutdown() should only be called when no new requests are actively being submitted:

-// Shutdown gracefully stops all workers when triggered.
-// It closes all request channels and waits for workers to exit.
+// Shutdown gracefully stops all workers when triggered.
+// It cancels the bifrost context to signal streaming goroutines to stop,
+// closes all request channels, and waits for workers to exit.
+//
+// Important: Shutdown should be called only when no new requests are being
+// submitted to avoid race conditions where requests could be enqueued after
+// context cancellation but before queue closure.
 func (bifrost *Bifrost) Shutdown() {

This clarifies the contract for callers.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6d849f0 and c80c51b.

📒 Files selected for processing (1)

• core/bifrost.go (2 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

core/bifrost.go (4)

core/schemas/bifrost.go (9)

• BifrostError (350-359)
• ErrorField (368-375)
• BifrostRequest (140-150)
• BifrostResponse (243-253)
• BifrostStream (315-322)
• RequestType (81-81)
• BifrostContextKeySelectedKeyID (107-107)
• BifrostContextKeySelectedKeyName (108-108)
• BifrostErrorExtraFields (417-421)

core/schemas/provider.go (4)

• Provider (207-234)
• ProviderConfig (161-169)
• CustomProviderConfig (142-148)
• PostHookRunner (204-204)

core/schemas/account.go (1)

• Key (8-17)

core/utils.go (1)

• IsStreamRequestType (190-192)

🔇 Additional comments (2)

core/bifrost.go (2)

2102-2131: Plugin pipeline pooling disabled for streaming to avoid data race (acknowledged trade-off).

The code correctly avoids releasing the plugin pipeline for streaming requests because provider goroutines asynchronously invoke postHookRunner, which would access a pipeline that might have been reset and reused by another request.

This addresses the critical data race flagged in the past review comments. The trade-off (no pooling for streaming) is reasonable given that correctness is prioritized over efficiency.

Note: Streaming requests will allocate a new pipeline for each request instead of reusing from the pool. Under high streaming load, this may increase GC pressure, but it ensures thread safety.

---

2139-2176: Excellent defensive programming with context-aware sends.

The context-aware response/error/stream sends with timeouts (lines 2140-2176) effectively prevent workers from blocking indefinitely when clients disconnect or stop listening. This ensures that workers can be reclaimed even if clients misbehave.

The 5-second timeout is a reasonable balance between accommodating slow clients and preventing indefinite blocking.

[Pratham-Mishra04]

@pjcdawkins can you have a common util for this? also lets add a field in the UI as well for stream activity timeout in network config

[Pratham-Mishra04]

@pjcdawkins lets pls also add the inactivity check in speech and transcription streaming

[pjcdawkins]

Good points - I set this as a draft because it's a relatively hard problem and I'll come back to it soon! But it seems important. I think it'll need a few more tests to demonstrate the problem(s) and fixes, and some consolidation as you say.

[Pratham-Mishra04]

Yes sure no worries, lmk if you need a hand😄

[akshaydeo]

also could you rebase this main with our latest main. we had to nuke one file from the history - which had to be done via rewrite.