feat: add context propagation to provider requests#52
Conversation
WalkthroughThis update introduces context propagation across the request lifecycle, allowing for improved cancellation and timeout handling in provider operations. Provider interfaces and implementations are refactored to accept Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Bifrost
participant Provider
participant HTTPAPI
Client->>Bifrost: Send request (with context)
Bifrost->>Provider: Call TextCompletion/ChatCompletion(ctx, ...)
Provider->>HTTPAPI: Make HTTP request with context
HTTPAPI-->>Provider: Response or error (respects context cancellation)
Provider-->>Bifrost: Return result or cancellation error
Bifrost-->>Client: Return response or error
sequenceDiagram
participant TestRunner
participant SetupAllRequests
participant Goroutine
participant WaitGroup
TestRunner->>SetupAllRequests: Start test setup
SetupAllRequests->>WaitGroup: Add(1) per goroutine
SetupAllRequests->>Goroutine: Launch test goroutine (with context)
Goroutine-->>WaitGroup: Done() on completion
SetupAllRequests->>TestRunner: Wait for all or context cancel
Note over SetupAllRequests,TestRunner: On interrupt, context is cancelled and waits for goroutines to finish
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
core/bifrost.go (1)
206-215: 🛠️ Refactor suggestionPool object retains previous
Context– potential memory leak
releaseChannelMessageclearsResponseandErrbut leavesContextpopulated.
The next borrower may unexpectedly inherit an unrelated context and, worse, the old context (with all its values) is kept alive in the pool.Clear it before returning to the pool.
func (bifrost *Bifrost) releaseChannelMessage(msg *ChannelMessage) { // Put channels back in pools bifrost.responseChannelPool.Put(msg.Response) bifrost.errorChannelPool.Put(msg.Err) // Clear references and return to pool msg.Response = nil msg.Err = nil + msg.Context = nil bifrost.channelMessagePool.Put(msg) }
🧹 Nitpick comments (2)
core/providers/azure.go (1)
170-178: Typo in error message reduces clarity
"deployment if not found for model %s"should read"deployment not found for model %s".
Tiny issue, but logs/searches are easier when messages are consistent and grammatically correct.- Message: fmt.Sprintf("deployment if not found for model %s", model), + Message: fmt.Sprintf("deployment not found for model %s", model),core/providers/bedrock.go (1)
208-220:json.Marshalwill never returncontext.Canceled/DeadlineExceeded
json.Marshalis CPU-bound and has no access to the passedctx; the extraerrors.Ischecks therefore cannot be hit and add noise.Consider removing this block or moving cancel/timeout handling to the network call where it is relevant.
- if err != nil { - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - ... - } - ... - } + if err != nil { + return nil, &schemas.BifrostError{ + IsBifrostError: true, + Error: schemas.ErrorField{ + Message: schemas.ErrProviderJSONMarshaling, + Error: err, + }, + } + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
core/bifrost.go(13 hunks)core/providers/anthropic.go(7 hunks)core/providers/azure.go(7 hunks)core/providers/bedrock.go(7 hunks)core/providers/cohere.go(4 hunks)core/providers/openai.go(5 hunks)core/providers/utils.go(2 hunks)core/providers/vertex.go(5 hunks)core/schemas/bifrost.go(2 hunks)core/schemas/provider.go(2 hunks)core/tests/anthropic_test.go(0 hunks)core/tests/azure_test.go(0 hunks)core/tests/bedrock_test.go(0 hunks)core/tests/cohere_test.go(0 hunks)core/tests/openai_test.go(1 hunks)core/tests/tests.go(11 hunks)core/tests/vertex_test.go(0 hunks)
💤 Files with no reviewable changes (5)
- core/tests/anthropic_test.go
- core/tests/azure_test.go
- core/tests/bedrock_test.go
- core/tests/vertex_test.go
- core/tests/cohere_test.go
🧰 Additional context used
🧬 Code Graph Analysis (7)
core/providers/utils.go (2)
core/schemas/bifrost.go (4)
BifrostError(299-305)ErrorField(308-315)RequestCancelled(295-295)Message(143-149)core/schemas/provider.go (1)
ErrProviderRequest(20-20)
core/schemas/provider.go (1)
core/schemas/bifrost.go (4)
ModelParameters(75-89)BifrostResponse(162-172)BifrostError(299-305)Message(143-149)
core/providers/cohere.go (1)
core/schemas/bifrost.go (4)
ModelParameters(75-89)BifrostResponse(162-172)BifrostError(299-305)Message(143-149)
core/providers/vertex.go (2)
core/schemas/bifrost.go (6)
ModelParameters(75-89)BifrostResponse(162-172)BifrostError(299-305)Message(143-149)ErrorField(308-315)RequestCancelled(295-295)core/providers/utils.go (1)
StrPtr(277-279)
core/providers/openai.go (1)
core/schemas/bifrost.go (4)
ModelParameters(75-89)BifrostResponse(162-172)BifrostError(299-305)Message(143-149)
core/providers/azure.go (1)
core/schemas/bifrost.go (4)
BifrostError(299-305)ModelParameters(75-89)BifrostResponse(162-172)Message(143-149)
core/providers/anthropic.go (1)
core/schemas/bifrost.go (4)
BifrostError(299-305)ModelParameters(75-89)BifrostResponse(162-172)Message(143-149)
🔇 Additional comments (30)
core/providers/utils.go (2)
6-6: LGTM! Context import added for new functionality.The addition of the
contextpackage import is necessary for the newmakeRequestWithContextfunction.
108-149: Well-implemented context-aware HTTP request handling.The
makeRequestWithContextfunction provides excellent context cancellation support with clear documentation about its limitations. The implementation correctly:
- Uses a goroutine to handle the blocking
client.Docall- Implements proper channel-based communication between goroutines
- Returns appropriate
BifrostErrortypes for different scenarios- Documents that it doesn't truly cancel the underlying fasthttp request
The function follows good practices by using a buffered channel to prevent goroutine leaks and providing clear error messages.
core/tests/openai_test.go (1)
30-30:Details
✅ Verification successful
Verify the Anthropic model version exists.
The model version "claude-3-7-sonnet-20250219" appears unusual. Anthropic typically uses "claude-3-5-sonnet" series models. Please verify this model exists and is available.
🌐 Web query:
Does Anthropic have a model called "claude-3-7-sonnet-20250219"? What are the latest Claude 3 model versions available?💡 Result:
Yes, Anthropic has a model named "claude-3-7-sonnet-20250219." This model, released on February 24, 2025, is part of the Claude 3 family and is known for its hybrid reasoning capabilities, allowing users to choose between rapid responses and more detailed, step-by-step reasoning. (docs.anthropic.com)
As of June 2025, the latest versions in the Claude 3 series are:
Claude 3.7 Sonnet: Released on February 24, 2025, this model offers enhanced reasoning capabilities and is available on all Claude plans, including Free, Pro, Team, and Enterprise. (reuters.com)
Claude 3.5 Sonnet v2: An upgraded version of the previous Sonnet model, offering improved performance and capabilities. (docs.anthropic.com)
Claude 3.5 Haiku: The fastest model in the Claude 3 lineup, optimized for quick responses and cost-effectiveness. (docs.anthropic.com)
Claude 3 Opus: Designed for complex reasoning tasks, this model excels in deep analysis and understanding. (docs.anthropic.com)
These models are available through Anthropic's API, Amazon Bedrock, and Google Cloud's Vertex AI. (docs.anthropic.com)
Citations:
- 1: https://docs.anthropic.com/en/docs/about-claude/models/all-models?utm_source=openai
- 2: https://www.reuters.com/technology/artificial-intelligence/anthropic-launches-advanced-ai-hybrid-reasoning-model-2025-02-24/?utm_source=openai
- 3: https://docs.anthropic.com/en/docs/about-claude/models/all-models?utm_source=openai
- 4: https://docs.anthropic.com/en/docs/about-claude/models/all-models?utm_source=openai
- 5: https://docs.anthropic.com/en/docs/about-claude/models/all-models?utm_source=openai
- 6: https://docs.anthropic.com/en/docs/about-claude/models/all-models?utm_source=openai
Verified Anthropic model version
The model
claude-3-7-sonnet-20250219is a valid Claude 3 release (launched Feb 24 2025) and is available on all plans. No further changes are needed.core/schemas/bifrost.go (2)
146-146: LGTM! ToolCallID field enhances message metadata.The addition of the optional
ToolCallIDfield to the Message struct properly supports tool call identification with consistent pointer typing for optional fields.
294-296: LGTM! Standardized cancellation error constant.The
RequestCancelledconstant provides a consistent way to identify request cancellation errors across the system, which aligns perfectly with the context propagation and cancellation features being introduced.core/schemas/provider.go (2)
4-7: LGTM! Context import added for interface updates.The addition of the
contextpackage import is necessary for the updated Provider interface method signatures.
146-148:Details
✅ Verification successful
LGTM! Interface updated for context propagation.
The Provider interface method signatures are correctly updated to include
context.Contextas the first parameter, following Go conventions. This breaking change is necessary to enable request cancellation and timeout handling throughout the provider system.Note that this change requires updating all provider implementations to match the new interface.
🏁 Script executed:
#!/bin/bash # Description: Verify all provider implementations have been updated to match the new interface. # Search for provider implementations that might need updating ast-grep --pattern $'func ($$$) TextCompletion($$$) ($$$, $$$) { $$$ }' ast-grep --pattern $'func ($$$) ChatCompletion($$$) ($$$, $$$) { $$$ }'Length of output: 47895
All provider implementations now include context propagation
All
TextCompletionandChatCompletionmethods across all providers have been updated to acceptctx context.Contextas their first parameter, enabling request cancellation and timeout handling as intended. No further changes are required.core/providers/cohere.go (4)
6-6: LGTM: Context import added for context propagation support.The context package import is correctly added to support the new context-aware method signatures.
132-132: LGTM: Method signature updated to accept context parameter.The
TextCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, maintaining consistency with the updated Provider interface.
144-144: LGTM: Method signature updated to accept context parameter.The
ChatCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, enabling context-aware request handling.
226-229: LGTM: HTTP request execution updated to use context-aware helper.The HTTP request execution is correctly updated to use
makeRequestWithContexthelper function, which should handle context cancellation and timeouts properly. The error handling is also updated to propagate errors from the helper function.core/providers/vertex.go (5)
8-8: LGTM: Errors import added for context error checking.The errors package import is correctly added to support context error detection (
context.Canceled,context.DeadlineExceeded).
80-80: LGTM: Method signature updated to accept context parameter.The
TextCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, maintaining consistency with the updated Provider interface.
92-92: LGTM: Method signature updated to accept context parameter.The
ChatCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, enabling context-aware request handling.
156-156: LGTM: HTTP request creation updated to use context.The HTTP request creation is correctly updated to use
http.NewRequestWithContext, which associates the request with the provided context for cancellation and timeout handling.
171-180: LGTM: Enhanced error handling for context cancellation.The error handling correctly detects context cancellation (
context.Canceled) and timeout (context.DeadlineExceeded) scenarios, returning appropriateBifrostErrorwith theRequestCancelledtype. The error message includes the context error for debugging purposes.core/providers/openai.go (5)
6-6: LGTM: Context import added for context propagation support.The context package import is correctly added to support the new context-aware method signatures.
105-105: LGTM: Method signature updated to accept context parameter.The
TextCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, maintaining consistency with the updated Provider interface.
117-117: LGTM: Method signature updated to accept context parameter.The
ChatCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, enabling context-aware request handling.
149-152: LGTM: HTTP request execution updated to use context-aware helper.The HTTP request execution is correctly updated to use
makeRequestWithContexthelper function, which should handle context cancellation and timeouts properly. The error handling is also updated to propagate errors from the helper function.
241-250: LGTM: Message preparation enhanced with tool_call_id support.The message preparation logic is correctly enhanced to include the
tool_call_idfield whenmsg.ToolCallIDis present. This aligns with theMessagestruct incore/schemas/bifrost.go(lines 143-149) that includes theToolCallIDfield.core/providers/anthropic.go (7)
6-6: LGTM: Context import added for context propagation support.The context package import is correctly added to support the new context-aware method signatures.
172-172: LGTM: Method signature updated to accept context parameter.The
completeRequestmethod signature is correctly updated to acceptcontext.Contextas the first parameter, enabling context-aware HTTP request handling.
199-202: LGTM: HTTP request execution updated to use context-aware helper.The HTTP request execution is correctly updated to use
makeRequestWithContexthelper function, which should handle context cancellation and timeouts properly. The error handling is also updated to propagate errors from the helper function.
226-226: LGTM: Method signature updated to accept context parameter.The
TextCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, maintaining consistency with the updated Provider interface.
235-235: LGTM: Context parameter correctly passed to completeRequest.The context parameter is correctly passed through to the
completeRequestmethod, enabling end-to-end context propagation for request cancellation and timeout handling.
280-280: LGTM: Method signature updated to accept context parameter.The
ChatCompletionmethod signature is correctly updated to acceptcontext.Contextas the first parameter, enabling context-aware request handling.
289-289: LGTM: Context parameter correctly passed to completeRequest.The context parameter is correctly passed through to the
completeRequestmethod, ensuring consistent context propagation for chat completion requests.core/providers/azure.go (1)
206-211: ConfirmmakeRequestWithContextreleases resources on cancellation
fasthttpdoes not natively respectContext; cancelling the context while a request is in-flight will leak the connection unlessmakeRequestWithContextforce-closes it (e.g. by callingreq.SetTimeout...orconn.Close()).Please double-check the helper to ensure:
req/respare released even whenctx.Done()fires.- the underlying TCP connection is closed to unblock
client.Do.- a
BifrostErroris returned withType == schemas.RequestCancelled.If any of these are missing, cancellation will not work as intended.
core/bifrost.go (1)
378-381:RequestCancelledbypass correctly implemented – kudos!Breaking the retry loop when the provider signals
RequestCancelledprevents unnecessary retries and honours client intent.
| allDoneChan := make(chan struct{}) | ||
| go func() { | ||
| wg.Wait() | ||
| close(allDoneChan) | ||
| }() | ||
|
|
||
| select { | ||
| case <-ctx.Done(): | ||
| log.Println("Context cancelled, test setup winding down.") | ||
| time.Sleep(1 * time.Second) | ||
| case <-allDoneChan: | ||
| log.Println("All test goroutines completed.") | ||
| } | ||
| log.Println("Test setup finished.") | ||
| bifrostClient.Cleanup() |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Race between cleanup and goroutine completion
When ctx.Done() is selected the code waits one second and proceeds to bifrostClient.Cleanup() without a wg.Wait().
Slow goroutines that have not yet noticed the cancellation may still write to closed request queues, causing panics.
Synchronise both branches:
- case <-ctx.Done():
- log.Println("Context cancelled, test setup winding down.")
- time.Sleep(1 * time.Second)
- case <-allDoneChan:
- log.Println("All test goroutines completed.")
+ case <-ctx.Done():
+ log.Println("Context cancelled, waiting for goroutines to exit…")
+ wg.Wait()
+ case <-allDoneChan:
+ log.Println("All test goroutines completed.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| allDoneChan := make(chan struct{}) | |
| go func() { | |
| wg.Wait() | |
| close(allDoneChan) | |
| }() | |
| select { | |
| case <-ctx.Done(): | |
| log.Println("Context cancelled, test setup winding down.") | |
| time.Sleep(1 * time.Second) | |
| case <-allDoneChan: | |
| log.Println("All test goroutines completed.") | |
| } | |
| log.Println("Test setup finished.") | |
| bifrostClient.Cleanup() | |
| allDoneChan := make(chan struct{}) | |
| go func() { | |
| wg.Wait() | |
| close(allDoneChan) | |
| }() | |
| select { | |
| case <-ctx.Done(): | |
| log.Println("Context cancelled, waiting for goroutines to exit…") | |
| wg.Wait() | |
| case <-allDoneChan: | |
| log.Println("All test goroutines completed.") | |
| } | |
| log.Println("Test setup finished.") | |
| bifrostClient.Cleanup() |
🤖 Prompt for AI Agents
In core/tests/tests.go around lines 356 to 370, the code waits one second after
ctx.Done() but does not wait for all goroutines to finish before calling
bifrostClient.Cleanup(), risking race conditions. Modify the select statement so
that in both cases, after ctx.Done() or allDoneChan, the code waits for
wg.Wait() to ensure all goroutines have completed before proceeding to cleanup.
This synchronizes cleanup with goroutine completion and prevents panics from
concurrent writes.
# Add context propagation and request cancellation support This PR adds support for context propagation and request cancellation throughout the Bifrost system. It allows clients to cancel in-flight requests when their context is cancelled or times out, improving resource management and responsiveness. Key changes: - Added context propagation from client requests to provider API calls - Implemented request cancellation handling in all providers - Added a new `RequestCancelled` error type to identify cancelled requests - Modified the `ChannelMessage` struct to include the request context - Improved the `Cleanup` method to gracefully shut down without signal handling - Added proper context handling in test utilities to ensure clean test termination - Refactored provider interfaces to accept context in method signatures The implementation ensures that when a client cancels a request context, Bifrost will stop processing that request and return an appropriate cancellation error. This is particularly important for long-running requests or when clients need to implement timeouts.
2 participants
Add context propagation and request cancellation support
This PR adds support for context propagation and request cancellation throughout the Bifrost system. It allows clients to cancel in-flight requests when their context is cancelled or times out, improving resource management and responsiveness.
Key changes:
RequestCancellederror type to identify cancelled requestsChannelMessagestruct to include the request contextCleanupmethod to gracefully shut down without signal handlingThe implementation ensures that when a client cancels a request context, Bifrost will stop processing that request and return an appropriate cancellation error. This is particularly important for long-running requests or when clients need to implement timeouts.