move to schemas.BifrostContext#1262
Conversation
|
Warning Rate limit exceeded@akshaydeo has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 2 minutes and 41 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (132)
📝 WalkthroughWalkthroughReplaces stdlib context.Context with *schemas.BifrostContext across core APIs, providers, MCP/agent/tooling, HTTP router/converters, streaming utilities, plugins, tests, and examples; adds BifrostContext helpers (WithValue/SetValue/BlockRestrictedWrites) and threads the new context type through tracing, hooks, streaming, and cancellation. Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
transports/bifrost-http/integrations/router.go (4)
507-519: Critical: Undefined variablerequestCtxcauses compilation error.The function signature was updated to use
bifrostCtx *schemas.BifrostContext, but the code on line 519 (and subsequent lines 539, 562, 585, 608, 631, 665, 688) referencesrequestCtxwhich is never defined. This will fail to compile.Based on the pattern in
handleBatchRequest(lines 738, 760, etc.) wherebifrostCtxis passed directly to client methods, the references torequestCtxshould be replaced withbifrostCtx.🔎 Proposed fix
func (g *GenericRouter) handleNonStreamingRequest(ctx *fasthttp.RequestCtx, config RouteConfig, req interface{}, bifrostReq *schemas.BifrostRequest, bifrostCtx *schemas.BifrostContext) { - // Use the cancellable context from ConvertToBifrostContext - // While we can't detect client disconnects until we try to write, having a cancellable context - // allows providers that check ctx.Done() to cancel early if needed. This is less critical than - // streaming requests (where we actively detect write errors), but still provides a mechanism - // for providers to respect cancellation. var response interface{} - var err error switch { case bifrostReq.ListModelsRequest != nil: - listModelsResponse, bifrostErr := g.client.ListModelsRequest(requestCtx, bifrostReq.ListModelsRequest) + listModelsResponse, bifrostErr := g.client.ListModelsRequest(bifrostCtx, bifrostReq.ListModelsRequest)Apply similar changes for all occurrences of
requestCtx→bifrostCtxin this function.
857-858: Inconsistent context type:handleFileRequeststill uses*context.Context.The function signature on line 857 uses
bifrostCtx *context.Contextinstead of*schemas.BifrostContext, which is inconsistent with the rest of the refactoring in this file (e.g.,handleNonStreamingRequest,handleBatchRequest). Line 858 then dereferences it to getrequestCtx.This appears to be an incomplete migration. The signature should be updated to match the pattern used elsewhere in this PR.
🔎 Proposed fix
-func (g *GenericRouter) handleFileRequest(ctx *fasthttp.RequestCtx, config RouteConfig, req interface{}, fileReq *FileRequest, bifrostCtx *context.Context) { - requestCtx := *bifrostCtx +func (g *GenericRouter) handleFileRequest(ctx *fasthttp.RequestCtx, config RouteConfig, req interface{}, fileReq *FileRequest, bifrostCtx *schemas.BifrostContext) {Then update all internal calls to use
bifrostCtxdirectly instead ofrequestCtx.
1019-1040: Inconsistent context type:handleStreamingRequeststill uses*context.Context.The function signature on line 1019 uses
bifrostCtx *context.Contextinstead of*schemas.BifrostContext, and line 1040 dereferences it asstreamCtx := *bifrostCtx. This is inconsistent with the broader migration to*schemas.BifrostContext.🔎 Proposed fix
-func (g *GenericRouter) handleStreamingRequest(ctx *fasthttp.RequestCtx, config RouteConfig, bifrostReq *schemas.BifrostRequest, bifrostCtx *context.Context, cancel context.CancelFunc) { +func (g *GenericRouter) handleStreamingRequest(ctx *fasthttp.RequestCtx, config RouteConfig, bifrostReq *schemas.BifrostRequest, bifrostCtx *schemas.BifrostContext, cancel context.CancelFunc) {Then update all internal calls to use
bifrostCtxdirectly.
1134-1134: Inconsistent context type:handleStreamingstill uses*context.Context.The function signature uses
bifrostCtx *context.Context, which is inconsistent with the rest of the refactoring.🔎 Proposed fix
-func (g *GenericRouter) handleStreaming(ctx *fasthttp.RequestCtx, bifrostCtx *context.Context, config RouteConfig, streamChan chan *schemas.BifrostStream, cancel context.CancelFunc) { +func (g *GenericRouter) handleStreaming(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, config RouteConfig, streamChan chan *schemas.BifrostStream, cancel context.CancelFunc) {
🤖 Fix all issues with AI Agents
In @core/schemas/context.go:
- Around line 74-80: Update the function doc comment for
NewBifrostContextWithCancel to replace the incorrect "PluginContext" reference
with "BifrostContext" so the comment accurately describes the returned type;
keep the rest of the comment (purpose, convenience wrapper around
NewBifrostContext, and note about calling the cancel function) unchanged and
ensure the first line still starts with the function name.
🧹 Nitpick comments (2)
transports/bifrost-http/lib/account.go (1)
38-74: Consider simplifying pointer dereference syntax.Line 51 uses
(*ctx).Value(...)to call the method. In Go, methods can be called directly on pointer receivers, so this can be simplified toctx.Value(...). While both forms are valid, the explicit dereference is less idiomatic and unnecessarily verbose.🔎 Proposed simplification
- if v := (*ctx).Value(schemas.BifrostContextKey("bf-governance-include-only-keys")); v != nil { + if v := ctx.Value(schemas.BifrostContextKey("bf-governance-include-only-keys")); v != nil {core/mcp/toolmanager.go (1)
385-387: Consider using BifrostContext-native timeout method.Line 386 creates a standard
context.Contextwith timeout from the*schemas.BifrostContext, which may lose custom context features (like the custom value storage viaSetValue/Value). Consider whetherBifrostContextprovides its own timeout method to preserve the custom context chain.If
BifrostContextprovides a timeout method, consider using it:// Instead of stdlib context.WithTimeout, use BifrostContext method if available toolCtx, cancel := ctx.WithTimeout(toolExecutionTimeout) // OR toolCtx, cancel := schemas.NewBifrostContextWithTimeout(ctx, toolExecutionTimeout)
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (15)
core/bifrost.gocore/mcp/agent.gocore/mcp/agentadaptors.gocore/mcp/mcp.gocore/mcp/toolmanager.gocore/schemas/account.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/provider.gotransports/bifrost-http/integrations/router.gotransports/bifrost-http/integrations/utils.gotransports/bifrost-http/lib/account.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/schemas/context.gocore/schemas/account.gotransports/bifrost-http/lib/account.gotransports/bifrost-http/integrations/router.gocore/schemas/mcp.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/integrations/utils.gocore/schemas/bifrost.gocore/mcp/agentadaptors.gocore/mcp/agent.gocore/mcp/mcp.gocore/mcp/toolmanager.gotransports/bifrost-http/server/server.gocore/schemas/provider.gocore/bifrost.go
🧠 Learnings (4)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/schemas/context.gocore/schemas/account.gotransports/bifrost-http/lib/account.gotransports/bifrost-http/integrations/router.gocore/schemas/mcp.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/integrations/utils.gocore/schemas/bifrost.gocore/mcp/agentadaptors.gocore/mcp/agent.gocore/mcp/mcp.gocore/mcp/toolmanager.gotransports/bifrost-http/server/server.gocore/schemas/provider.gocore/bifrost.go
📚 Learning: 2025-12-29T11:54:55.836Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 1153
File: framework/configstore/rdb.go:2221-2246
Timestamp: 2025-12-29T11:54:55.836Z
Learning: In Go reviews, do not flag range-over-int patterns like for i := range n as compile-time errors, assuming Go 1.22+ semantics. Only flag actual range-capable values (slices, arrays, maps, channels, strings) and other compile-time issues. This applies to all Go files across the repository.
Applied to files:
core/schemas/context.gocore/schemas/account.gotransports/bifrost-http/lib/account.gotransports/bifrost-http/integrations/router.gocore/schemas/mcp.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/integrations/utils.gocore/schemas/bifrost.gocore/mcp/agentadaptors.gocore/mcp/agent.gocore/mcp/mcp.gocore/mcp/toolmanager.gotransports/bifrost-http/server/server.gocore/schemas/provider.gocore/bifrost.go
📚 Learning: 2025-12-12T08:25:02.629Z
Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.
Applied to files:
transports/bifrost-http/lib/account.gotransports/bifrost-http/integrations/router.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/integrations/utils.gotransports/bifrost-http/server/server.go
📚 Learning: 2025-12-29T09:14:16.633Z
Learnt from: akshaydeo
Repo: maximhq/bifrost PR: 888
File: transports/bifrost-http/handlers/middlewares.go:246-256
Timestamp: 2025-12-29T09:14:16.633Z
Learning: In the bifrost HTTP transport, fasthttp.RequestCtx is the primary context carrier and should be passed directly to functions that expect a context.Context. Do not convert to context.Context unless explicitly required. Ensure tracer implementations and related components are designed to accept fasthttp.RequestCtx directly, and document this architectural decision for maintainers.
Applied to files:
transports/bifrost-http/lib/account.gotransports/bifrost-http/integrations/router.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/integrations/utils.gotransports/bifrost-http/server/server.go
🧬 Code graph analysis (14)
core/schemas/account.go (2)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(30-30)
transports/bifrost-http/lib/account.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
ModelProvider(30-30)ui/lib/types/config.ts (1)
ModelProvider(206-209)
transports/bifrost-http/integrations/router.go (2)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (3)
BifrostContextKeySendBackRawResponse(136-136)BifrostContextKeyIntegrationType(137-137)BifrostContextKeyDirectKey(125-125)
core/schemas/mcp.go (1)
core/schemas/context.go (1)
BifrostContext(31-41)
transports/bifrost-http/lib/ctx.go (2)
core/schemas/context.go (2)
BifrostContext(31-41)NewBifrostContextWithCancel(77-80)core/schemas/bifrost.go (7)
BifrostContextKeyRequestID(123-123)BifrostContextKey(117-117)BifrostContextKeyVirtualKey(121-121)BifrostContextKeyAPIKeyName(122-122)BifrostContextKeySendBackRawResponse(136-136)BifrostContextKeyExtraHeaders(132-132)BifrostContextKeyDirectKey(125-125)
transports/bifrost-http/integrations/utils.go (2)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
BifrostError(490-499)
core/schemas/bifrost.go (1)
core/schemas/context.go (1)
BifrostContext(31-41)
core/mcp/agentadaptors.go (4)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (1)
BifrostError(490-499)core/schemas/chatcompletions.go (1)
BifrostChatRequest(10-17)core/schemas/responses.go (1)
BifrostResponsesRequest(30-37)
core/mcp/agent.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/bifrost.go (2)
BifrostContextKeyRequestID(123-123)BifrostMCPAgentOriginalRequestID(139-139)core/mcp/toolmanager.go (1)
ClientManager(16-20)
core/mcp/mcp.go (3)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/chatcompletions.go (2)
BifrostChatRequest(10-17)BifrostChatResponse(25-40)core/schemas/responses.go (2)
BifrostResponsesRequest(30-37)BifrostResponsesResponse(43-83)
core/mcp/toolmanager.go (2)
core/schemas/context.go (1)
BifrostContext(31-41)core/schemas/mcp.go (1)
MCPToolManagerConfig(27-31)
transports/bifrost-http/server/server.go (3)
core/schemas/context.go (4)
BifrostContext(31-41)NewBifrostContext(46-64)NoDeadline(10-10)NewBifrostContextWithCancel(77-80)core/schemas/bifrost.go (1)
ListModelsRequest(90-90)core/schemas/models.go (1)
BifrostListModelsRequest(21-32)
core/schemas/provider.go (1)
core/schemas/context.go (1)
BifrostContext(31-41)
core/bifrost.go (3)
core/schemas/context.go (2)
BifrostContext(31-41)NewBifrostContextWithCancel(77-80)core/schemas/bifrost.go (17)
BifrostContextKeyTracer(146-146)TextCompletionStreamRequest(92-92)ChatCompletionStreamRequest(94-94)ResponsesStreamRequest(96-96)SpeechStreamRequest(99-99)TranscriptionStreamRequest(101-101)BifrostContextKeyFallbackIndex(129-129)BifrostContextKeyRequestID(123-123)BifrostContextKeyFallbackRequestID(124-124)BifrostContextKeySpanID(143-143)BifrostContextKeyStreamEndIndicator(130-130)BifrostContextKeyStreamStartTime(145-145)RequestType(87-87)BifrostContextKeySelectedKeyID(126-126)BifrostContextKeySelectedKeyName(127-127)BifrostContextKeyPostHookSpanFinalizer(149-149)BifrostContextKeyDirectKey(125-125)core/utils.go (1)
IsStreamRequestType(197-199)
⏰ 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 (29)
core/schemas/mcp.go (1)
24-24: LGTM! Public API migration to BifrostContext is consistent.The function signature update aligns with the broader context migration across the codebase. This is a breaking change but appears coordinated across the stack.
core/schemas/bifrost.go (1)
13-13: LGTM! KeySelector migration to BifrostContext is correct.The signature update is consistent with the broader migration. Note that the original signature used
*context.Context(pointer to interface), while the new signature uses*BifrostContext(pointer to struct), which is the correct pattern for the custom context type.core/schemas/account.go (1)
89-89: LGTM! Account interface updated consistently with BifrostContext migration.The method signature change is part of the coordinated context migration. All implementations must be updated to match this interface change.
transports/bifrost-http/integrations/router.go (3)
101-211: LGTM on type signature updates.The converter function type signatures have been consistently updated from
context.Contextto*schemas.BifrostContext, aligning with the broader migration across the codebase.
396-399: LGTM on context value setting pattern.The migration from
context.WithValuetobifrostCtx.SetValue(...)is correctly applied for setting context values likeBifrostContextKeySendBackRawResponse,BifrostContextKeyIntegrationType, andBifrostContextKeyDirectKey.Also applies to: 435-436
728-854: LGTM onhandleBatchRequestimplementation.The function correctly uses
bifrostCtxdirectly in all client method calls (lines 738, 760, 782, 804, 826), properly leveraging the new*schemas.BifrostContexttype.transports/bifrost-http/lib/ctx.go (2)
80-94: LGTM onConvertToBifrostContextsignature and initialization.The function signature correctly returns
(*schemas.BifrostContext, context.CancelFunc), and the initialization properly usesschemas.NewBifrostContextWithCancel(ctx)to create the cancellable Bifrost context. TheSetValuepattern is consistently applied for setting context values.
157-407: LGTM on context value population.All
context.WithValuecalls have been uniformly replaced withbifrostCtx.SetValue(...)for various context keys including:
- Prometheus labels
- Maxim tracing headers
- MCP control headers
- Virtual key and API key headers
- Cache-related values
- Extra headers
- Direct keys
- Fallback context
The function correctly returns the bifrost context directly at line 407.
transports/bifrost-http/server/server.go (3)
86-86: LGTM on server context field type change.The
ctxfield inBifrostHTTPServeris correctly updated to*schemas.BifrostContext, aligning with the new context type used throughout the codebase.
930-934: LGTM onRefetchModelsForProvidercontext handling.The method correctly creates a new
BifrostContextusingschemas.NewBifrostContext(ctx, schemas.NoDeadline), defersbfCtx.Cancel()for proper cleanup, and passes the context toListModelsRequest.
1185-1185: LGTM on Bootstrap context initialization and usage.
- Line 1185: Correctly initializes server context using
schemas.NewBifrostContextWithCancel(ctx).- Lines 1242-1244: MCP config's
FetchNewRequestIDFuncsignature is properly updated to accept*schemas.BifrostContext.- Line 1264:
ListAllModelscorrectly usess.ctx(now a*schemas.BifrostContext).Also applies to: 1242-1244, 1264-1264
transports/bifrost-http/integrations/utils.go (1)
138-198: LGTM on utility function signature updates.The
sendStreamError,sendError, andsendSuccessfunctions have been correctly updated to acceptbifrostCtx *schemas.BifrostContext, aligning with the broader migration. Thecontextpackage import has been appropriately removed since it's no longer needed.core/mcp/agentadaptors.go (1)
54-54: LGTM on agent adaptor context type updates.The
agentAPIAdapterinterface and its implementations (chatAPIAdapter,responsesAPIAdapter) have been correctly updated:
- Interface method
makeLLMCallnow accepts*schemas.BifrostContextmakeReqfield types updated to use*schemas.BifrostContextin the callback signatures- Concrete implementations on lines 159-162 and 372-375 correctly use the new context type
The
contextpackage import has been appropriately removed.Also applies to: 69-70, 89-90, 159-162, 372-375
core/mcp/mcp.go (2)
133-147: LGTM onExecuteChatToolandExecuteResponsesToolsignature updates.Both methods now correctly accept
ctx *schemas.BifrostContext, aligning with the broader context type migration.
180-201: LGTM on agent execution entry point signatures.
CheckAndExecuteAgentForChatRequestandCheckAndExecuteAgentForResponsesRequesthave been correctly updated:
- The
ctxparameter now accepts*schemas.BifrostContext- The
makeReqcallback signatures are updated to accept*schemas.BifrostContextThis aligns with the changes in
agent.goandagentadaptors.go.Also applies to: 231-252
core/mcp/agent.go (5)
29-63: LGTM onExecuteAgentForChatRequestcontext type updates.The function signature and callback types are correctly updated to use
*schemas.BifrostContext, enabling the new context semantics throughout the agent execution flow.
82-116: LGTM onExecuteAgentForResponsesRequestcontext type updates.Consistent with the Chat API counterpart, this function correctly uses
*schemas.BifrostContextthroughout.
134-160: LGTM onexecuteAgenthelper context handling.The internal helper correctly:
- Accepts
*schemas.BifrostContextas the context parameter- Uses
ctx.Value()to retrieve the original request ID (line 156)- Uses
ctx.SetValue()to store the agent's original request ID (line 158)
336-343: LGTM on request ID propagation.The agent loop correctly uses
ctx.SetValue()to update the request ID before making subsequent LLM calls, maintaining proper request tracking across agent iterations.
428-430: No issues found.*schemas.BifrostContextfully implements thecontext.Contextinterface with all four required methods (Deadline(),Done(),Err(), andValue()), so passing it toGetToolPerClient()is valid and compatible.core/mcp/toolmanager.go (2)
322-322: LGTM: Systematic context type migration.The method signatures have been consistently updated to use
*schemas.BifrostContext. The changes maintain the same logical flow while adopting the new context type.Also applies to: 436-437, 481-485, 512-516
19-19: No action needed—BifrostContext properly implements context.Context interface.
*schemas.BifrostContextfully implements thecontext.Contextinterface with all required methods (Deadline,Done,Err,Value), so passing it toClientManager.GetToolPerClient(ctx context.Context)is valid and correct. The implementation is sound.core/schemas/provider.go (1)
314-365: LGTM: Systematic Provider interface migration to BifrostContext.The Provider interface and PostHookRunner have been consistently updated to use
*schemas.BifrostContextinstead ofcontext.Context. This is a breaking change that requires all provider implementations to be updated, but given this is part of a broader stack migration (PR #1262 building on #1242), the systematic approach is appropriate.Note: As per the coding guidelines to review in light of the whole stack, this change aligns with the broader BifrostContext migration across the codebase.
core/bifrost.go (6)
136-139: LGTM: Proper BifrostContext initialization.The
Initfunction correctly wraps the incomingcontext.Contextin aBifrostContextusingNewBifrostContextWithCancel. This establishes the context lifecycle for the entire Bifrost instance.
302-304: LGTM: Consistent nil context handling.The code defensively handles nil contexts by falling back to
bifrost.ctxthroughout public API methods. This prevents nil pointer panics and provides a sensible default context.Also applies to: 632-634, 728-730
2884-2926: LGTM: executeRequestWithRetries correctly uses BifrostContext.The generic
executeRequestWithRetriesfunction properly:
- Accepts
*schemas.BifrostContextas parameter (line 2885)- Uses
SetValueto track retry attempts (line 2898)- Retrieves tracer from context with type assertion (lines 2921-2925)
- Updates span context appropriately (line 2949)
3364-3380: Note: Plugin interface breaking change.The
PluginPipeline.RunPreHooksandRunPostHooksmethods now pass*schemas.BifrostContextto plugin hooks (lines 3380, 3428, 3446). This is a breaking change that requires all plugin implementations to update their hook signatures.This aligns with the broader migration strategy across the stack.
Also applies to: 3408-3446
3917-3939: LGTM: KeySelector signature updated for consistency.
WeightedRandomKeySelectorsignature updated to accept*schemas.BifrostContext(line 3917), though the context is not used in the implementation. This is appropriate for interface consistency with theKeySelectortype definition (line 75 in the broader code).
367-367: SetValue is already properly synchronized.The
BifrostContext.SetValue()method incore/schemas/context.gocorrectly implements thread-safe synchronization usingsync.RWMutex. The implementation properly acquiresbc.valuesMulock for the entire duration of the write operation and usesdeferto guarantee unlock on exit. The code even includes a comment confirming: "This is thread-safe and can be called concurrently." No concurrency issues exist with the current implementation.Likely an incorrect or invalid review comment.
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
f3775fb to
bd2418a
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
transports/bifrost-http/lib/account.go (1)
51-51: Remove explicit dereferencing in BifrostContext.Value() call.The
BifrostContext.Value()method is defined with a pointer receiver. In Go, you can call pointer receiver methods directly on the pointer variable without explicit dereferencing—usectx.Value(...)instead of(*ctx).Value(...)for idiomatic code. Go automatically handles the dereferencing.core/bifrost_test.go (1)
703-710: Update line 656 to remove unnecessary context wrapping.The
Initfunction signature isfunc Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error)and internally wraps the context viaschemas.NewBifrostContextWithCancel(ctx)at line 136 of core/bifrost.go. Line 656 unnecessarily pre-wrapscontext.Background()inschemas.NewBifrostContextbefore passing it toInit, causing double-wrapping. Lines 703, 729, 771, 831, 884, and 937 correctly passcontext.Background()directly. Update line 656 to match this pattern:ctx := context.Background() bifrost, err := Init(ctx, schemas.BifrostConfig{core/bifrost.go (1)
2353-2450: PluginPipeline reuse in streaming short-circuit can race after being returned to the poolIn
tryStreamRequest, aPluginPipelineis acquired and scheduled forreleasePluginPipelineviadefer, but in theshortCircuit.Stream != nilpath you start a goroutine that continues to callpipeline.RunPostHooksviapipelinePostHookRunneraftertryStreamRequestreturns. At that point the deferredreleasePluginPipelinehas already reset and put the pipeline back into the pool, so the goroutine can mutate a reused instance, causing data races and corrupted hook state.You’re already handling streaming lifetime correctly in
requestWorkerby tying pipeline release to a finalizer callback; here you need analogous ownership for the short-circuit streaming path.Proposed fix: release PluginPipeline only after streaming goroutine finishes
One way to fix this is to avoid deferring the release for the streaming short‑circuit branch and instead release the pipeline in the goroutine when the stream ends:
func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { - pipeline := bifrost.getPluginPipeline() - defer bifrost.releasePluginPipeline(pipeline) + pipeline := bifrost.getPluginPipeline() + pipelineReleased := false + defer func() { + if !pipelineReleased { + bifrost.releasePluginPipeline(pipeline) + } + }() @@ - if shortCircuit.Stream != nil { + if shortCircuit.Stream != nil { outputStream := make(chan *schemas.BifrostStream) @@ - pipelinePostHookRunner := func(ctx *schemas.BifrostContext, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) { + pipelinePostHookRunner := func(ctx *schemas.BifrostContext, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) { return pipeline.RunPostHooks(ctx, result, err, preCount) } - go func() { - defer close(outputStream) + go func() { + defer func() { + close(outputStream) + bifrost.releasePluginPipeline(pipeline) + }() @@ }() + // Ownership of the pipeline has moved to the goroutine. + pipelineReleased = true return outputStream, nil }This keeps the existing defer for all non-streaming and non-short‑circuit paths, but ensures that in the streaming short‑circuit case the pipeline is only reset and returned to the pool after the goroutine finishes processing the stream.
Also applies to: 2467-2564, 2727-2793
core/providers/groq/groq.go (1)
176-198: ChatCompletionStream should honor context-derived path overrides
ChatCompletionusesproviderUtils.GetPathFromContext(ctx, "/v1/chat/completions"), butChatCompletionStreamstill hardcodes"/v1/chat/completions"without going throughGetPathFromContext. That means any path overrides injected viaBifrostContext(e.g., integration-specific routing) will apply to non-streaming requests but not to streaming ones.To keep routing consistent and avoid subtle bugs when paths are customized, consider updating the streaming path to mirror the non-streaming variant:
Proposed fix
func (provider *GroqProvider) ChatCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostChatRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { @@ - return openai.HandleOpenAIChatCompletionStreaming( - ctx, - provider.client, - provider.networkConfig.BaseURL+"/v1/chat/completions", + return openai.HandleOpenAIChatCompletionStreaming( + ctx, + provider.client, + provider.networkConfig.BaseURL+providerUtils.GetPathFromContext(ctx, "/v1/chat/completions"), request,core/providers/cerebras/cerebras.go (1)
79-92: Align Cerebras streaming paths with context-aware routingFor Cerebras, non-streaming endpoints already use
providerUtils.GetPathFromContext, but the streaming variants still hardcode paths:
TextCompletionStreamusesBaseURL+"/v1/completions".ChatCompletionStreamusesBaseURL+"/v1/chat/completions".This means any per-request path override injected via
BifrostContextapplies to synchronous calls but not to streaming calls, which can break integrations that rely onGetPathFromContextfor routing.Updating the streaming functions to mirror the non-streaming path construction will make behavior consistent:
Proposed fix
func (provider *CerebrasProvider) TextCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostTextCompletionRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { @@ - return openai.HandleOpenAITextCompletionStreaming( - ctx, - provider.client, - provider.networkConfig.BaseURL+"/v1/completions", + return openai.HandleOpenAITextCompletionStreaming( + ctx, + provider.client, + provider.networkConfig.BaseURL+providerUtils.GetPathFromContext(ctx, "/v1/completions"), @@ func (provider *CerebrasProvider) ChatCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, key schemas.Key, request *schemas.BifrostChatRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) { @@ - return openai.HandleOpenAIChatCompletionStreaming( - ctx, - provider.client, - provider.networkConfig.BaseURL+"/v1/chat/completions", + return openai.HandleOpenAIChatCompletionStreaming( + ctx, + provider.client, + provider.networkConfig.BaseURL+providerUtils.GetPathFromContext(ctx, "/v1/chat/completions"),Also applies to: 97-117, 119-133, 139-161
🤖 Fix all issues with AI agents
In @core/internal/testutil/speech_synthesis.go:
- Around line 16-17: The unused parameter ctx in both RunSpeechSynthesisTest and
RunSpeechSynthesisAdvancedTest causes a Go compile error; to fix it with no
runtime change, rename the parameter from ctx to _ in both function signatures
so the incoming *schemas.BifrostContext is ignored; leave the existing creation
of requestCtx := context.Background() and other code paths unchanged (this is
the minimal change to compile without threading the BifrostContext into
SpeechRequest calls).
🧹 Nitpick comments (22)
core/providers/utils/utils.go (1)
835-874: Pointer-based BifrostContext in streaming utilities improves safety and matches new span-finalization flowSwitching
Send*ResponsesChunk,ProcessAndSend*helpers, andcompleteDeferredSpanto accept*schemas.BifrostContext(and invokingpostHookRunner(ctx, ...)directly) removes the prior by-value copies of BifrostContext, which is safer given the embedded mutexes and Once fields. The checks forBifrostContextKeyStreamEndIndicatorbefore callingcompleteDeferredSpannow line up with providers setting this flag viactx.SetValue, and the finalizer invocation usingcontext.WithValue(ctx, BifrostContextKeySpanID, spanID)remains correct assuming*BifrostContextimplementscontext.Context. Overall, the streaming + tracing pipeline remains logically intact with a cleaner context model.Also applies to: 881-931, 939-981, 988-1033, 1284-1325, 1428-1513
core/providers/anthropic/anthropic.go (1)
123-251: Anthropic: BifrostContext wiring, Files/Batch/CountTokens look correct, with a small URL-building inconsistencyThe switch to
*schemas.BifrostContextinbuildRequestURL,completeRequest, list models, text/chat/response completions, batch operations, file APIs, andCountTokensconsistently passesctxinto helpers likeGetRequestPath,SetExtraHeaders,MakeRequestWithContext, andCheckContextAndGetRequestBody, so behavior is preserved while enabling Bifrost-specific context values. New File APIs (upload/list/retrieve/delete/content) and Batch helpers use Anthropic’s expected endpoints, headers (includinganthropic-versionandanthropic-betawhere needed), and standardizedHandleProviderResponse+NewUnsupportedOperationErrorpatterns, with pagination handled viaNewSerialListHelperwhere appropriate.One minor inconsistency: some endpoints (
BatchCancel,BatchResults,FileDelete,FileContent) still construct URLs withprovider.networkConfig.BaseURL + "/..."instead of going throughbuildRequestURL(...). That means they won’t respect per-request path overrides or custom provider path mappings the way other methods do. If you intend all Anthropic endpoints to be configurable in the same way, consider routing these throughbuildRequestURLas well for consistency.Also applies to: 176-231, 237-374, 693-742, 1033-1329, 1547-1576, 1579-1783, 1785-2057, 2059-2125
core/mcp/agent_test.go (1)
19-19: Update MockLLMCaller to use BifrostContext for type consistency.The MockLLMCaller methods still accept
context.Context, but the test closures now pass*schemas.BifrostContext(lines 223-225, 283-285, 397-399, 445-447). While BifrostContext likely implements the context.Context interface (allowing this to work), the mock should be updated for type safety and consistency with the real BifrostLLMCaller interface.🔎 Proposed update to MockLLMCaller signatures
-func (m *MockLLMCaller) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) { +func (m *MockLLMCaller) ChatCompletionRequest(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) { if m.chatCallCount >= len(m.chatResponses) { return nil, &schemas.BifrostError{ IsBifrostError: false, Error: &schemas.ErrorField{ Message: "no more mock chat responses available", }, } } response := m.chatResponses[m.chatCallCount] m.chatCallCount++ return response, nil } -func (m *MockLLMCaller) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { +func (m *MockLLMCaller) ResponsesRequest(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { if m.responsesCallCount >= len(m.responsesResponses) { return nil, &schemas.BifrostError{ IsBifrostError: false, Error: &schemas.ErrorField{ Message: "no more mock responses api responses available", }, } } response := m.responsesResponses[m.responsesCallCount] m.responsesCallCount++ return response, nil }Also applies to: 34-34
core/providers/bedrock/bedrock_test.go (2)
680-680: Add space after comma for Go formatting conventions.Missing whitespace after the comma in the
NewBifrostContextcall. This pattern appears in multiple locations throughout the file.🔎 Suggested formatting fix
- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
708-708: Apply consistent spacing after commas.The same formatting issue (missing space after comma) appears on lines 708, 1778, 1989, and 2016. Apply the same fix to maintain consistent Go formatting throughout the test file.
Also applies to: 1778-1778, 1989-1989, 2016-2016
core/providers/bedrock/chat.go (1)
50-50: Migrate ToBifrostChatResponse parameter to use BifrostContext for consistency.
ToBedrockChatCompletionRequest(line 15) and other bedrock response converters (ToBifrostResponsesRequest,ToBifrostResponsesResponsein responses.go) now use*schemas.BifrostContext, butToBifrostChatResponsestill acceptscontext.Context. The call site at line 724 in bedrock.go already passesBifrostContext, so this migration aligns with the ongoing refactoring effort and eliminates the inconsistency across the provider.plugins/mocker/plugin_test.go (1)
24-32: Unusedctxin test account mock is acceptable
BaseAccount.GetKeysForProvidercorrectly matches the updated Account interface but does not usectx. That's fine for a test stub; if linters start complaining, consider renaming the parameter to_ *schemas.BifrostContextto make the intent explicit.transports/bifrost-http/integrations/openai.go (1)
38-97: Pre‑request hooks now accept BifrostContext but don’t use it yet
AzureEndpointPreHook,setQueryParamsAndAzureEndpointPreHook, and the variousextract*helpers correctly plumb*schemas.BifrostContextthrough their signatures and into nested hooks, but the value isn’t read yet. That’s fine for now; if you later need per‑request flags (e.g. structured output, tracing, or direct-key hints), you can store them onbifrostCtxinstead offasthttp.RequestCtx.UserValue.Also applies to: 454-479, 976-1013, 1015-1058, 1060-1137, 1139-1221
core/internal/testutil/speech_synthesis_stream.go (1)
17-118: Streaming speech tests correctly construct per‑scenario BifrostContexts (but ignore the incomingctx)Both streaming helpers now create a fresh
*schemas.BifrostContextviaschemas.NewBifrostContext(context.Background(), schemas.NoDeadline)for each scenario and use that withSpeechStreamRequest, which is consistent with the new API. Thectx *schemas.BifrostContextparameter on the helpers is currently unused; if the testutil API doesn’t rely on it elsewhere, consider either wiring it into the per‑scenario context creation or removing it to avoid confusion.Also applies to: 251-255, 309-313, 455-466
transports/bifrost-http/integrations/anthropic.go (2)
227-235: Redundant pointer dereference in context value access.
ctxis already*schemas.BifrostContext, so(*ctx).Value(...)can be simplified toctx.Value(...). The dereference is unnecessary and reduces readability.This pattern appears in multiple places in this file (lines 229, 323, 394, 436, 472, 508, 732, 777, 818, 855).
🔎 Proposed fix
func shouldUsePassthrough(ctx *schemas.BifrostContext, provider schemas.ModelProvider, model string, deployment string) bool { isClaudeCode := false - if userAgent, ok := (*ctx).Value(schemas.BifrostContextKeyUserAgent).(string); ok { + if userAgent, ok := ctx.Value(schemas.BifrostContextKeyUserAgent).(string); ok { if strings.Contains(userAgent, "claude-cli") { isClaudeCode = true } } return isClaudeCode && isClaudeModel(model, deployment, string(provider)) }
317-326: Same redundant dereference pattern in BatchRequestConverter.Line 323 uses
(*ctx).Value(bifrostContextKeyProvider)- should bectx.Value(...).🔎 Proposed fix
- if provider, ok = (*ctx).Value(bifrostContextKeyProvider).(schemas.ModelProvider); ok && provider != schemas.Anthropic { + if provider, ok = ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider); ok && provider != schemas.Anthropic {core/schemas/provider.go (1)
314-366: Provider/PostHookRunner API now explicitly depends on BifrostContextUpdating
PostHookRunnerand allProvidermethods to take*BifrostContextaligns the public provider API with the new context abstraction (needed for things like per‑request values and stream end markers). This is a clean, consistent change, just be aware it’s a breaking interface update for any out‑of‑tree Provider or PostHookRunner implementations.core/providers/openai/openai.go (2)
78-205: Non‑streaming OpenAI provider methods are cleanly migrated to BifrostContextAll primary OpenAIProvider entry points (list models, text/chat, responses, embedding, speech, transcription, count tokens, file, and batch ops) now take
*schemas.BifrostContextand passctxthrough toproviderUtils.*and HTTP helpers without altering behavior. UsingBifrostContexthere gives hooks and internal plumbing access to per‑request values while remaining compatible with helpers that still acceptcontext.Context.The one notable inconsistency is
HandleOpenAIResponsesRequest, which still takesctx context.Contextwhile its siblings use*schemas.BifrostContext. That’s functionally fine (BifrostContext should satisfy the interface) but you may want to align its signature for API regularity and to allow direct use of BifrostContext-specific helpers inside it if needed later.Also applies to: 216-314, 613-715, 1101-1204, 1460-1665, 1916-2044, 2278-2378, 2382-3342
319-608: Stream end signalling via BifrostContext is sensible; consider clarifying context reuse assumptionsThe streaming helpers now set
ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true)on terminal paths (error chunks and final completion/usage chunks) before invokingProcessAndSendResponse/ProcessAndSendBifrostError. That’s a good fit for PostHookRunner’s new*BifrostContextsignature and lets hooks reliably detect the final event.This does, however, rely on the implicit contract that each API call uses a fresh
*BifrostContext(or that any consumers of this flag treat it as strictly per‑stream and ignore stale values). If there’s any chance a single BifrostContext instance is reused across multiple requests, it might be worth either clearing this key at the start of each streaming call or constructing a per‑stream child context to avoid leaking atrueflag into subsequent operations.Also applies to: 720-1094, 1207-1455, 1669-1910, 2047-2273
core/internal/testutil/structured_outputs.go (1)
203-355: Streaming and Responses structured‑output tests use BifrostContext consistentlyAll four helpers (
RunStructuredOutputChatStreamTest,RunStructuredOutputResponsesTest,RunStructuredOutputResponsesStreamTest) now:
- Take
ctx *schemas.BifrostContextand pass it (or areqCtxderived by settingBifrostContextKeyExtraHeadersfor Claude) into the bifrost client calls.- Continue to use
context.WithTimeout(ctx, …)only for local test timeouts, not as the actual request context, preserving prior behavior.- Keep retry configs and validation logic unchanged, now just riding on BifrostContext instead of plain context.Context.
This is a straightforward, correct migration, and the extra‑header pattern is consistent across chat and responses tests.
Also applies to: 359-488, 491-712
core/bifrost.go (1)
278-380: Tracer reliance on BifrostContext values is consistent but tightly coupled
ListModelsRequestnow setsBifrostContextKeyTraceronctxbefore callingexecuteRequestWithRetries, andexecuteRequestWithRetriesexpects to read a non-nilschemas.Tracerfrom that key. Current call sites (this method and the request-worker paths) all set the key first, so behavior matches the previous global-tracer usage, but any future caller ofexecuteRequestWithRetriesmust remember to populate this key or will see a “tracer not found in context” error.Consider documenting this contract near
executeRequestWithRetriesor guarding with a fallback toschemas.DefaultTracer()to make the function safer to reuse.Also applies to: 2885-2956
core/providers/gemini/gemini.go (1)
21-23: Consider centralizing the response-format context key
BifrostContextKeyResponseFormatis defined locally in this file and used inSpeechviactx.SetValueto driveToBifrostSpeechResponse(ctx). If you expect other providers or transports to read or set this key, it may be worth moving the constant intocore/schemas/bifrost.goalongside the otherBifrostContextKeydefinitions to avoid duplication and mismatched literals later.Also applies to: 969-1015
core/providers/azure/azure.go (1)
902-1150: AzureSpeechStream: consider marking stream end on[DONE]sentinel as well
SpeechStreamnow usesctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true)when:
- Emitting a JSON error decoded from the stream (Line 1076).
- Sending the final
SpeechStreamResponseTypeDonechunk (Line 1140).However, when the SSE payload is exactly
[DONE], the goroutine returns immediately without setting the end-indicator and without emitting a final “done” response (Lines 1055–1061). If any downstream code relies onBifrostContextKeyStreamEndIndicatorto know the stream has fully terminated (especially for non‑[DONE] providers), it may miss this early‑exit path.Consider also setting the end-indicator in the
[DONE]branch, for example:Possible tweak to cover the `[DONE]` path
- if bytes.Equal(audioData, []byte("[DONE]")) { - return - } + if bytes.Equal(audioData, []byte("[DONE]")) { + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + return + }Please double‑check how
BifrostContextKeyStreamEndIndicatoris consumed in the HTTP/router layer for Azure speech streams to confirm whether this additional marker is desirable, and adjust if tests around stream completion expect it.core/providers/bedrock/bedrock.go (1)
681-751: Bedrock ChatCompletion/ChatCompletionStream: BifrostContext integration is good; minor consistency nit on error end-flag
ChatCompletionandChatCompletionStreamnow pass*schemas.BifrostContextinto the Bedrock converters (ToBedrockChatCompletionRequest(ctx, ...)),completeRequest, and the eventstream decoder. Stream state (usage, finishReason, IDs, etc.) is unchanged.In
ChatCompletionStreamyou mark the stream end viactx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true):
- When
ToBifrostChatCompletionStreamreports aBifrostError(Lines 871–881).- Before emitting the final aggregated chunk response (Lines 913–915).
You do not set the flag on low‑level decode errors from
decoder.Decode(Lines 799–808), where you instead log and forward an error response. For symmetry withTextCompletionStreamandResponsesStream, you may want to also set the end‑indicator in that decode‑error path so all terminal cases set the same flag.Optional consistency tweak for the decode-error path
- message, err := decoder.Decode(resp.Body, payloadBuf) - if err != nil { - if err == io.EOF { - // End of stream - this is normal - break - } - provider.logger.Warn(fmt.Sprintf("Error decoding %s EventStream message: %v", providerName, err)) - providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ChatCompletionStreamRequest, providerName, request.Model, provider.logger) - return - } + message, err := decoder.Decode(resp.Body, payloadBuf) + if err != nil { + if err == io.EOF { + // End of stream - this is normal + break + } + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + provider.logger.Warn(fmt.Sprintf("Error decoding %s EventStream message: %v", providerName, err)) + providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ChatCompletionStreamRequest, providerName, request.Model, provider.logger) + return + }Please verify how the router checks
BifrostContextKeyStreamEndIndicatorfor Bedrock chat streams; if it treats any channel close as sufficient, this tweak is optional, but aligning behavior across all error/EOF cases can simplify reasoning and debugging.Also applies to: 755-918
transports/bifrost-http/lib/ctx.go (1)
80-91: ConvertToBifrostContext migration looks correct; consider updating docsThe switch to
schemas.NewBifrostContextWithCancel(ctx)plusbifrostCtx.SetValue(...)for all header- and key-derived values keeps behavior consistent while centralizing on*schemas.BifrostContext. Key usages (BifrostContextKeyRequestID, virtual/api keys, cache keys/TTL/threshold/type, extra headers, Maxim tags, direct key, fallback flag) all map cleanly to the new context type.The function comment above still describes returning
*context.Context, which is now outdated; you may want to update it to mention*schemas.BifrostContextand that values are stored viaSetValue.Also applies to: 157-205, 207-213, 223-225, 227-234, 238-259, 264-287, 290-307, 313-337, 339-342, 348-356, 391-400, 403-405, 407-407
transports/bifrost-http/server/server.go (1)
86-87: Server ctx migration to BifrostContext is coherent; small consistency tweak optionalUsing
*schemas.BifrostContextas the server’s root ctx and wiring it through Bootstrap, model listing, and route registration aligns with the new client API, andRefetchModelsForProvidercorrectly wraps the incomingcontext.Contextin aBifrostContextbefore callingListModelsRequest.For helpers like
ExecuteChatMCPTool/ExecuteResponsesMCPTool, you might optionally standardize onNewBifrostContextWithCancelanddefer cancel()rather than callingNewBifrostContextwithout any explicit cancellation, to keep lifecycle management uniform across these entry points.Also applies to: 498-507, 932-935, 1185-1188, 1243-1247, 1266-1267, 1292-1293, 1318-1319
core/providers/bedrock/utils.go (1)
48-60: Bedrock parameter converters now correctly use BifrostContext; minor robustness tweak optionalSwitching
convertChatParametersand the structured-output helpers to take a*schemas.BifrostContextand recordBifrostContextKeyStructuredOutputToolNameviactx.SetValue(...)preserves the previous behavior while aligning with the new context type. Tool insertion and forcedToolChoicewiring remain unchanged.To harden against accidental nil usage in the future, you might optionally guard the
SetValuecalls:if ctx != nil { ctx.SetValue(schemas.BifrostContextKeyStructuredOutputToolName, toolName) }but this isn’t strictly necessary if all call sites guarantee a non-nil ctx.
Also applies to: 182-185, 678-721, 735-757
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
bd2418a to
a061a86
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 (4)
plugins/jsonparser/plugin_test.go (1)
153-181: Inconsistent context usage:bifrost.Initreceivescontext.Contextinstead of*schemas.BifrostContext.Line 154 uses plain
context.Background()which is then passed tobifrost.Initat line 173. However, based on the PR's migration pattern,bifrost.Initshould receive a*schemas.BifrostContext. This creates inconsistency withTestJsonParserPluginEndToEndwhich properly uses aBifrostContext.🔎 Suggested fix
func TestJsonParserPluginPerRequest(t *testing.T) { - ctx := context.Background() + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) // Check if OpenAI API key is set if os.Getenv("OPENAI_API_KEY") == "" { t.Skip("OPENAI_API_KEY is not set, skipping per-request test")core/bifrost_test.go (1)
698-711: Incomplete migration: several UpdateProvider tests still usecontext.Background()directly.Multiple tests in this file still use plain
context.Background()when callingbifrost.Init:
TestUpdateProvider_UpdateNonExistentProvider(line 703)TestUpdateProvider_UpdateInactiveProvider(line 729)TestUpdateProvider_MultipleProviderUpdates(line 771)TestUpdateProvider_ConcurrentProviderUpdates(line 831)TestUpdateProvider_ProviderSliceIntegrity(lines 884, 937)This is inconsistent with
TestUpdateProvider_SuccessfulUpdate(line 656) which correctly usesschemas.NewBifrostContext.🔎 Suggested fixes (apply to each occurrence)
- ctx := context.Background() + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)Also applies to: 725-737, 764-778, 826-838, 878-892, 933-945
core/providers/openai/openai.go (1)
1122-1204: Fix inconsistent context type in HandleOpenAIResponsesRequest.The
HandleOpenAIResponsesRequestfunction signature still usescontext.Contextinstead of*schemas.BifrostContext. This is inconsistent with:
- The calling function
Responses(line 1101) which passes*schemas.BifrostContext- All other
HandleOpenAI*Requestfunctions in this file which use*schemas.BifrostContext- The broader migration pattern across the codebase
This will cause compilation errors or type mismatches.
🔎 Proposed fix
func HandleOpenAIResponsesRequest( - ctx context.Context, + ctx *schemas.BifrostContext, client *fasthttp.Client, url string, request *schemas.BifrostResponsesRequest,core/bifrost.go (1)
385-518: Guard against nil ctx in ListAllModels to avoid panics
ListAllModelsusesctx.Done()in the worker goroutines but never normalizes a nilctx, unlike most other public APIs that fall back tobifrost.ctx. CallingListAllModels(nil, req)will panic.Consider aligning this with the rest of the API by defaulting a nil
ctxtobifrost.ctxup front.Proposed fix
func (bifrost *Bifrost) ListAllModels(ctx *schemas.BifrostContext, request *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) { - if request == nil { + if ctx == nil { + ctx = bifrost.ctx + } + + if request == nil { request = &schemas.BifrostListModelsRequest{} }
🧹 Nitpick comments (10)
core/internal/testutil/setup.go (1)
34-48: Potential type mismatch betweengetBifrostandSetupTest.
SetupTestcreates a*schemas.BifrostContext(line 52) and passes it togetBifrost(line 53), butgetBifrostdeclares its parameter ascontext.Context(line 34). SinceBifrostContextimplements thecontext.Contextinterface, this will compile, but it creates an inconsistency with the PR's migration goal.Consider updating
getBifrostto accept*schemas.BifrostContextfor consistency:🔎 Suggested change
-func getBifrost(ctx context.Context) (*bifrost.Bifrost, error) { +func getBifrost(ctx *schemas.BifrostContext) (*bifrost.Bifrost, error) {Also applies to: 51-59
core/providers/bedrock/bedrock_test.go (1)
680-680: Fix formatting: missing space after comma in NewBifrostContext calls.Multiple
NewBifrostContextcalls are missing a space after the comma, making the code less readable.🔎 Proposed formatting fix
- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)Apply this formatting fix at lines 680, 708, 1778, 1989, and 2016.
Also applies to: 708-708, 1778-1778, 1989-1989, 2016-2016
core/internal/testutil/structured_outputs.go (1)
103-134: Avoid leaking Anthropic extra headers across calls by deriving a per‑request context
reqCtx := ctxfollowed byreqCtx.SetValue(BifrostContextKeyExtraHeaders, extraHeaders)mutates the shared*schemas.BifrostContext. If the samectxis reused across multiple tests or operations (especially witht.Parallel), Anthropic headers set forclaudemodels can inadvertently affect subsequent non‑claudecalls.Consider deriving a per‑request context instead of mutating the caller’s context (e.g., a
NewBifrostContextWithValuehelper or equivalent) so header scope is confined to the current operation.Also applies to: 219-247, 375-431, 512-563
transports/bifrost-http/integrations/openai.go (1)
38-97: Verify direct‑key propagation from Azure HTTP layer into BifrostContext
AzureEndpointPreHookand the batch/file pre‑callbacks still record direct keys on thefasthttp.RequestCtx(viaSetUserValue), while core key‑selection logic now checksctx.Value(schemas.BifrostContextKeyDirectKey)on*schemas.BifrostContext.Please double‑check that the HTTP router (or middleware) copies this direct‑key value from the request context into the
BifrostContextbefore invoking core APIs; otherwise, Azure direct‑key flows may bypass the newBifrostContext‑based key selection path.Also applies to: 455-479, 976-1013, 1060-1137, 1140-1221
core/bifrost.go (1)
3794-3915: MakeselectKeyFromProviderForModelnil‑safe for future callers
selectKeyFromProviderForModelusesctx.Value(...)without guarding againstctx == nilin theskipKeySelectioncheck:if skipKeySelection, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection ...Today it’s only called with non‑nil
req.Context, so this isn’t an immediate bug, but adding actx != nilguard will prevent accidental panics if the function is reused elsewhere and brings it in line with the checks already used ingetAllSupportedKeys/getKeysForBatchAndFileOps.Proposed tweak
- // Check if key skipping is allowed - if skipKeySelection, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) { - return schemas.Key{}, nil - } + // Check if key skipping is allowed + if ctx != nil { + if skipKeySelection, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skipKeySelection && isKeySkippingAllowed(providerKey) { + return schemas.Key{}, nil + } + }core/providers/anthropic/anthropic.go (3)
502-511: Scanner error path doesn’t currently mark stream end or complete spansFor the happy path and most error conditions you set
BifrostContextKeyStreamEndIndicatorand go throughProcessAndSendResponse/ProcessAndSendBifrostError, which in turn callcompleteDeferredSpan. However, whenscanner.Err()is non-nil you callProcessAndSendErrorwithout setting the stream-end indicator, andProcessAndSendErroritself never finalizes the deferred span. That means streaming spans may be left open (and trace state retained) on lower-level read failures.Consider either:
- Setting
BifrostContextKeyStreamEndIndicatorbeforeProcessAndSendError, and/or- Extending
ProcessAndSendErrorto mirror the stream-end handling used inProcessAndSendBifrostError(includingcompleteDeferredSpanwhen the end indicator is set).The same applies to
HandleAnthropicResponsesStream’sscanner.Err()branch.Also applies to: 665-684
744-785: Responses streaming: BifrostContext integration and end-indicator usage
HandleAnthropicResponsesStreamnow takes*schemas.BifrostContextand correctly:
- Uses
SetExtraHeaderswith the new context type,- Marks
BifrostContextKeyStreamEndIndicatoron empty-body andToBifrostResponsesStreamerror paths,- Marks it again on the final chunk before calling
ProcessAndSendResponse.This matches the tracing pattern used for chat streaming and should integrate cleanly with
completeDeferredSpan. Only thescanner.Err()path (discussed earlier) currently bypasses that end-indicator logic.Also applies to: 790-1031
1313-1419: BatchCancel/BatchResults and Files endpoints could sharebuildRequestURL
BatchCancel,BatchResults,FileDelete, andFileContentbuild request URIs usingprovider.networkConfig.BaseURL + ".../..."instead of the centralizedbuildRequestURL. That’s functionally fine, but it bypasses anyRequestPathOverridesconfigured viaCustomProviderConfigand is slightly inconsistent with the rest of the file wherebuildRequestURLis used.If you want uniform behavior for custom Anthropic endpoints across all operations, consider switching these to
buildRequestURLas well (adding the appropriate path andRequestType), and usingurl.PathEscapefor IDs where applicable.Also applies to: 1421-1526
core/providers/utils/utils.go (1)
987-1033: ProcessAndSendError still bypasses span completion
ProcessAndSendErrornow also takes*BifrostContextand wraps scanner errors into aBifrostError, but unlikeProcessAndSendBifrostErrorit doesn’t look atBifrostContextKeyStreamEndIndicatoror callcompleteDeferredSpan. Given it’s used in streaming readers (e.g., scanner.Err in provider streams), you may want to consider aligning it with the other two helpers so tracing is properly finalized even when the underlying scanner fails.transports/bifrost-http/integrations/router.go (1)
857-859: Minor: Remove extraneous blank line.Extra blank line at line 858 can be removed for consistency.
🔎 Proposed fix
func (g *GenericRouter) handleFileRequest(ctx *fasthttp.RequestCtx, config RouteConfig, req interface{}, fileReq *FileRequest, bifrostCtx *schemas.BifrostContext) { - - var response interface{}
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
2 times, most recently
from
ebe3f5a to
5d1fcf4
Compare
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
core/internal/testutil/cross_provider_scenarios.go (1)
3-11: Critical: Missing context import for context.WithTimeout usage.The file removes the
contextimport but still usescontext.WithTimeoutat lines 568 and 633. This will cause a compilation error.🐛 Fix: Add context import
import ( + "context" "encoding/json" "fmt" "strings" "testing"core/providers/bedrock/chat.go (1)
50-50: *Update ToBifrostChatResponse to use schemas.BifrostContext for consistency with ToBedrockChatCompletionRequest.
ToBifrostChatResponseacceptscontext.ContextwhileToBedrockChatCompletionRequestuses*schemas.BifrostContext. Beyond this asymmetry within Bedrock,ToBifrostChatResponseis also unique—no other provider's response converter takes a context parameter at all. The method usesctx.Value()to access structured output configuration, which works with the interface type, but the signature should be explicit about*schemas.BifrostContextto align with the request converter and the broader pattern of migrating to BifrostContext throughout the codebase.core/providers/openai/openai.go (1)
1122-1133: Inconsistent context type:HandleOpenAIResponsesRequeststill usescontext.Context.This function signature uses
context.Contextwhile all otherHandle*functions in this file have been migrated to*schemas.BifrostContext. This breaks consistency and will cause type mismatches when callingproviderUtils.ShouldSendBackRawRequest(ctx, ...)at line 1194, which likely expects*schemas.BifrostContext.🐛 Proposed fix
// HandleOpenAIResponsesRequest handles a responses request to OpenAI's API. func HandleOpenAIResponsesRequest( - ctx context.Context, + ctx *schemas.BifrostContext, client *fasthttp.Client, url string, request *schemas.BifrostResponsesRequest,core/bifrost.go (1)
2353-2461: Document BifrostContext isolation contract to prevent concurrent request metadata interleavingThe
handleRequestfunction mutates*BifrostContextviaSetValuecalls forFallbackIndex,RequestID,FallbackRequestID, andSpanID. WhileSetValueis thread-safe via mutex, concurrent requests reusing the same*BifrostContextinstance (particularly through theif ctx == nil { ctx = bifrost.ctx }fallback in public API methods) will have their metadata writes interleave, corrupting span relationships and request tracking.This is safe only if callers guarantee one
BifrostContextper logical request. The codebase should enforce or document this invariant explicitly. Consider adding a note to the context type documentation clarifying thatBifrostContextinstances must not be shared across concurrent requests, and optionally extend theblockRestrictedWritesmechanism to the request boundary to catch violations early.
🤖 Fix all issues with AI agents
In @core/internal/testutil/speech_synthesis_stream.go:
- Line 17: RunSpeechSynthesisStreamTest accepts a ctx but builds new contexts
with context.Background(), breaking propagation; update the function
(RunSpeechSynthesisStreamTest) to use the passed ctx when calling client methods
and creating derived contexts (e.g., replace context.Background() with ctx or
ctxWithCancel := contexts derived from ctx), ensuring cancellation, deadlines,
and values propagate; apply the same replacement for every place in the function
where context.Background() is used (the spots creating fresh contexts during
stream setup/tear-down) so the passed ctx is actually used rather than
discarded.
- Line 252: RunSpeechSynthesisStreamAdvancedTest currently ignores the passed-in
ctx parameter and instead creates a new context locally; update the function to
use the provided ctx (only create a new context if the incoming ctx is nil) so
that context propagation and cancellation work correctly. Specifically, in
RunSpeechSynthesisStreamAdvancedTest replace the local creation of a fresh
context (the variable that was being used around the calls at the advanced test
flow) with the existing ctx parameter, and if ctx == nil then initialize it the
same way existing tests do (create a new context or call the helper used
elsewhere) before using it.
- Line 455: The multi-voice test is creating a fresh context with
schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) instead of
using the incoming ctx parameter; update the code to propagate the provided ctx
into the call (i.e., replace context.Background() with ctx) so the requestCtx
uses the test's context and any cancellations/deadlines are respected; locate
the occurrence of NewBifrostContext in the multi-voice test and make the same
change applied previously at the other location.
In @core/schemas/context.go:
- Around line 88-94: The doc comment for NewBifrostContextWithCancel still
refers to "PluginContext"; update the comment text to reference "BifrostContext"
instead—e.g., change "convenience wrapper around NewPluginContext" to
"convenience wrapper around NewBifrostContext" and ensure other mentions in that
comment use BifrostContext; leave the function implementation
(NewBifrostContextWithCancel and its call to NewBifrostContext) unchanged.
In @examples/plugins/hello-world-wasm-typescript/assembly/index.ts:
- Around line 73-87: The extractJsonBool function yields false positives because
it searches for the substring "true" within an arbitrary proximity instead of
validating the actual JSON value; change extractJsonBool to locate the key token
'"' + key + '"' then find the colon ':' after that key, skip any whitespace,
ensure the next non-whitespace character is not a quote (to avoid string
values), then read the next literal token and return true only if it exactly
equals "true" and false if it equals "false"; this removes the 10-char heuristic
and correctly distinguishes boolean literals from quoted strings or other
fields.
- Around line 128-184: The three hook functions (http_intercept, pre_hook,
post_hook) use unsafe manual JSON extraction/building (via extractJsonObject,
addToJsonObject, extractJsonBool, readString, writeString) and lack
validation/error handling; fix by validating readString output and results of
extractJsonObject/extractJsonBool at the start of each function, and return a
clear error payload (e.g., set error/hook_error/has_error or short_circuit
fields and writeString that error) when parsing fails instead of continuing;
stop constructing JSON via raw concatenation—either use a safe JSON
builder/serializer helper or ensure values are properly escaped before
insertion; make the hardcoded context values ('123','789','456') configurable or
document them and move to constants; add unit tests for edge cases (escaped
characters, nested objects, empty/null fields) and update comments to document
accepted input format and failure behavior.
- Around line 28-68: The manual JSON helpers (extractJsonObject,
extractJsonBool, addToJsonObject) are fragile and incorrect (they mis-handle
keys containing truthy substrings, whitespace, string contexts, escaped
quotes/backslashes, and nullable values); fix by removing the ad-hoc parsing and
either (a) use a proper AssemblyScript JSON library (e.g., json-as with
JSON.Box<T> or assemblyscript-json) to parse/modify objects, or (b) reimplement
the functions as a correct JSON state machine that tracks string/escape context
and nested braces and properly serializes/escapes values; update all call sites
to use the library parsing/boxing APIs or the new robust functions (search for
extractJsonObject, extractJsonBool, addToJsonObject to locate spots to change).
🧹 Nitpick comments (14)
examples/plugins/hello-world-wasm-typescript/assembly/types.ts (1)
1-10: Intentional minimization aligns with the BifrostContext migration.The decision to strip type definitions and rely on direct JSON manipulation in
index.tsis well-documented and appropriate for stub WASM runtime compatibility. This aligns with the broader stack migration toschemas.BifrostContext.Consider adding a direct file reference for discoverability:
📝 Optional: Add explicit file reference
// This file is intentionally minimal since we use direct JSON manipulation // in index.ts for maximum compatibility with the stub WASM runtime. +// See ./index.ts for the actual type handling implementation.examples/plugins/hello-world-wasm-rust/src/lib.rs (1)
62-71: Potential panic on UTF-8 string boundary.When
input_strcontains multi-byte UTF-8 characters, slicing by byte index (input_str[start..end]) may panic ifstartorendfall in the middle of a character. The column number from serde errors is a byte offset, butsaturating_sub(50)and+ 50adjustments may land on non-character boundaries.Consider using
char_indices()or thefloor_char_boundarymethod (nightly) for safe slicing, or wrap in a fallback:♻️ Suggested safer approach
let error_context = if let Some(col) = extract_column(&e.to_string()) { - let start = col.saturating_sub(50); - let end = (col + 50).min(input_str.len()); - format!(" | context: ...{}...", &input_str[start..end]) + let start = col.saturating_sub(50); + let end = (col + 50).min(input_str.len()); + // Safe slice: find valid UTF-8 boundaries + let safe_start = input_str.floor_char_boundary(start); + let safe_end = input_str.ceil_char_boundary(end); + if safe_start < safe_end { + format!(" | context: ...{}...", &input_str[safe_start..safe_end]) + } else { + String::new() + } } else { String::new() };Alternatively, for stable Rust, use
.get(start..end).unwrap_or("")which returnsNoneon invalid boundaries:input_str.get(start..end).map_or(String::new(), |s| format!(" | context: ...{}...", s))core/internal/testutil/setup.go (1)
34-34: Consider updatinggetBifrostsignature for consistency.While
getBifroststill acceptscontext.Context(line 34) and works correctly because*schemas.BifrostContextimplements thecontext.Contextinterface, updating it to accept*schemas.BifrostContextwould improve consistency with the broader migration pattern across the codebase.♻️ Optional signature update
-func getBifrost(ctx context.Context) (*bifrost.Bifrost, error) { +func getBifrost(ctx *schemas.BifrostContext) (*bifrost.Bifrost, error) {core/bifrost_test.go (1)
703-710: Consider simplifying context usage in tests rather than standardizing onNewBifrostContext.Lines 656 and 703 show inconsistent context handling: line 656 pre-converts with
schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)while line 703 passescontext.Background()directly. SinceInit()accepts thecontext.Contextinterface and internally converts any context viaschemas.NewBifrostContextWithCancel()at line 136 of core/bifrost.go, pre-converting in tests creates redundancy. The simpler approach—passingcontext.Background()directly—is clearer and avoids unnecessary double-conversion.transports/bifrost-http/server/server.go (1)
498-501: Consider adding context cleanup.Both
ExecuteChatMCPToolandExecuteResponsesMCPToolcreate a newBifrostContextbut don't explicitly cancel it. While these are likely short-lived operations, it's generally good practice to clean up contexts:func (s *BifrostHTTPServer) ExecuteChatMCPTool(ctx context.Context, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) { bifrostCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) + defer bifrostCtx.Cancel() return s.Client.ExecuteChatMCPTool(bifrostCtx, toolCall) }However, if the parent
ctxis already properly managed (e.g., request context), and the operations are synchronous and fast, the current pattern may be acceptable.Also applies to: 504-507
core/bifrost.go (2)
278-380: *Public API methods now uniformly accept BifrostContext; watch nil‑ctx fallback semanticsAll top‑level request methods (ListModels*, Text/Chat/Responses, Embedding/Speech/Transcription, Batch*, File*) now take *schemas.BifrostContext and in some cases fall back to bifrost.ctx when ctx is nil. Logic and error handling otherwise match the prior structure.
Two minor points to keep in mind:
- Only some methods defensively default nil ctx to bifrost.ctx (e.g., ListModelsRequest, ChatCompletionRequest, ResponsesRequest, Batch*/File*). Others (notably Text/Embedding/Speech/Transcription) assume a non‑nil ctx. If external callers might pass nil, consider either documenting that requirement or normalizing consistently.
- When nil fallback is used under concurrency, all such calls will share the same BifrostContext instance; because SetValue is mutating, per‑request keys (request ID, retry count, span ID, etc.) will interleave. If you ever expect concurrent calls with nil ctx, it’s safer to always create a per‑request BifrostContext in the caller and pass that down.
Overall the migration itself is correct, but I’d recommend explicitly stating or enforcing that callers must pass a per‑request *BifrostContext instead of relying on the shared bifrost.ctx except for very controlled internal paths.
Also applies to: 382-518, 520-1057, 1059-1437
3684-3722: Key selection helpers now fully leverage BifrostContext; behavior preserved with extra override hooksgetAllSupportedKeys, getKeysForBatchAndFileOps, and selectKeyFromProviderForModel now all:
- Accept *schemas.BifrostContext and first look for an explicitly injected schemas.Key via BifrostContextKeyDirectKey.
- Use ctx.Value(BifrostContextKeySkipKeySelection) and BifrostContextKeyAPIKeyName to support key skipping and per‑name overrides.
- Fall back to account.GetKeysForProvider(ctx, providerKey) and apply existing filtering logic.
WeightedRandomKeySelector’s signature was updated to use *BifrostContext but the implementation still ignores ctx, which is fine and keeps it usable as a default KeySelector.
These changes look correct and nicely centralize context‑driven overrides.
Also applies to: 3727-3792, 3795-3915, 3917-3939
transports/bifrost-http/integrations/router.go (2)
481-483: Unnecessary pointer dereference when calling Value().At line 481,
(*bifrostCtx).Value(...)uses explicit pointer dereference, but this is inconsistent with other usages in this file where methods are called directly onbifrostCtx(e.g.,bifrostCtx.SetValue(...)). Go allows calling methods on pointer receivers without explicit dereference.Suggested fix
- if sendRawRequestBody, ok := (*bifrostCtx).Value(schemas.BifrostContextKeyUseRawRequestBody).(bool); ok && sendRawRequestBody { + if sendRawRequestBody, ok := bifrostCtx.Value(schemas.BifrostContextKeyUseRawRequestBody).(bool); ok && sendRawRequestBody {
857-858: Minor: trailing whitespace on empty line.Line 858 appears to have trailing whitespace after the function signature which should be cleaned up for consistency.
transports/bifrost-http/integrations/bedrock.go (2)
203-268: Provider extraction viactx.Valuerelies on PreCallback invariants
BatchRequestConverterand batch list/retrieve/cancel converters use:provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider)This will panic if the context key is ever missing or of the wrong type (e.g., if a new route is wired without the matching
PreCallback). Given you already set the provider in the correspondingPreCallbacks, this is fine today but somewhat brittle.Consider a small defensive pattern to decouple converter safety from routing wiring:
providerVal := ctx.Value(bifrostContextKeyProvider) provider, ok := providerVal.(schemas.ModelProvider) if !ok { return nil, errors.New("provider not set in context") }And similarly for the other batch converters.
Also applies to: 271-281, 302-312, 313-323, 344-354, 355-365, 386-396, 397-407
532-568: Safer provider lookup when rewriting Bedrock job ARN
extractBedrockJobArnFromPathcurrently assumes the provider context key is always present and correctly typed:if (*bifrostCtx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) != schemas.Bedrock { decodedJobArn = strings.Replace(decodedJobArn, "...", "", 1) }Two nitpicks here:
- The explicit
(*bifrostCtx)dereference is unnecessary.- The unchecked type assertion will panic if the key isn’t set.
You can simplify and harden this without changing behavior:
- if (*bifrostCtx).Value(bifrostContextKeyProvider).(schemas.ModelProvider) != schemas.Bedrock { + if provider, ok := bifrostCtx.Value(bifrostContextKeyProvider).(schemas.ModelProvider); ok && provider != schemas.Bedrock { decodedJobArn = strings.Replace(decodedJobArn, "arn:aws:bedrock:us-east-1:444444444444:batch:", "", 1) }This keeps the cross-provider rewrite but makes the helper more robust to future routing changes.
core/providers/anthropic/anthropic.go (1)
379-432: Streaming helpers should mark stream-end on scanner errors as well
HandleAnthropicChatCompletionStreamingandHandleAnthropicResponsesStreamcorrectly setBifrostContextKeyStreamEndIndicatorwhen:
- The provider returns an empty body stream.
ToBifrost*Streamreturns aBifrostError.- The final chunk is emitted.
On the other hand, when
scanner.Err()is non-nil you log and callProcessAndSendError, but don’t mark the stream as ended in the context. If downstream cleanup logic relies onBifrostContextKeyStreamEndIndicator(in addition to channel close), this could leave the flag unset on some failure paths.Consider also setting the indicator in the scanner error branch, e.g.:
if err := scanner.Err(); err != nil { logger.Warn(...) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) providerUtils.ProcessAndSendError(...) }Same pattern for both chat and responses streaming helpers.
Also applies to: 436-452, 502-509, 619-627, 681-684
core/providers/vertex/vertex.go (1)
822-1075: Responses/ResponsesStream: new fallback flag is reasonable but context-dependentFor non-Anthropic and non-Gemini deployments,
Responsescontinues to fall back toChatCompletion, andResponsesStreamnow:ctx.SetValue(schemas.BifrostContextKeyIsResponsesToChatCompletionFallback, true) return provider.ChatCompletionStream(ctx, postHookRunner, key, request.ToChatRequest())This is a sensible way to tag the stream as a responses→chat fallback in the new context model so downstream hooks can distinguish it. Just ensure any consumers of
BifrostContextKeyIsResponsesToChatCompletionFallbackexpect it only on streaming fallbacks (the synchronousResponsesfallback does not currently set it).Also applies to: 1077-1263
core/providers/cohere/cohere.go (1)
341-413: Mark stream-end on scanner errors for Cohere streaming as wellFor both
ChatCompletionStreamandResponsesStream, you now setBifrostContextKeyStreamEndIndicatoron:
- Per-event conversion errors (
ToBifrost*Streamreturning an error).- The final chunk when
isLastChunkis true.However, if the SSE scanner itself errors (
scanner.Err()), you log and callProcessAndSendErrorbut do not set the stream-end indicator. To keep semantics consistent with the new context flag (and with other providers), consider setting it there too:if err := scanner.Err(); err != nil { provider.logger.Warn(...) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) providerUtils.ProcessAndSendError(...) }This mirrors the Anthropic recommendation and makes it easier for higher layers to reliably detect all terminal stream states via the context.
Also applies to: 432-507, 572-647, 656-762
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (11)
transports/bifrost-http/lib/ctx.go (1)
70-71: Outdated documentation: return type description.The doc comment on line 70 states the function returns
*context.Context, but the actual return type is now*schemas.BifrostContext. Update the documentation to reflect the new return type.// Returns: -// - *context.Context: A new cancellable context.Context containing the propagated values +// - *schemas.BifrostContext: A new cancellable BifrostContext containing the propagated values // - context.CancelFunc: Function to cancel the context (should be called when request completes)core/internal/testutil/speech_synthesis.go (3)
113-117: Thread the passed context instead of creating a new one.The function receives a
ctx *schemas.BifrostContextparameter but ignores it, creating a fresh context instead. For consistency with other test utilities and to fix the compile error, thread the passed context through.- requestCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - - speechResponse, bifrostErr := WithSpeechTestRetry(t, speechRetryConfig, retryContext, expectations, "SpeechSynthesis_"+tc.name, func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { - return client.SpeechRequest(requestCtx, request) + speechResponse, bifrostErr := WithSpeechTestRetry(t, speechRetryConfig, retryContext, expectations, "SpeechSynthesis_"+tc.name, func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { + return client.SpeechRequest(ctx, request)
219-223: Thread the passed context instead of creating a new one.Same issue: use the passed
ctxparameter instead of creating a fresh context.- requestCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - - speechResponse, bifrostErr := WithSpeechTestRetry(t, speechRetryConfig, retryContext, expectations, "SpeechSynthesis_HD", func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { - return client.SpeechRequest(requestCtx, request) + speechResponse, bifrostErr := WithSpeechTestRetry(t, speechRetryConfig, retryContext, expectations, "SpeechSynthesis_HD", func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { + return client.SpeechRequest(ctx, request)
301-305: Thread the passed context instead of creating a new one.Same issue in the "AllVoiceOptions" subtest: use the passed
ctxparameter.- requestCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - - speechResponse, bifrostErr := WithSpeechTestRetry(t, voiceSpeechRetryConfig, voiceRetryContext, expectations, "SpeechSynthesis_VoiceType_"+voiceType, func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { - return client.SpeechRequest(requestCtx, request) + speechResponse, bifrostErr := WithSpeechTestRetry(t, voiceSpeechRetryConfig, voiceRetryContext, expectations, "SpeechSynthesis_VoiceType_"+voiceType, func() (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { + return client.SpeechRequest(ctx, request)transports/bifrost-http/integrations/anthropic.go (6)
434-448: Missing nil check on type assertion for provider.At line 436,
ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider)performs a type assertion without checking if the value exists or if the assertion succeeds. If the provider is not set in the context, this will panic.🐛 Proposed fix with nil check
BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { if retrieveReq, ok := req.(*anthropic.AnthropicBatchRetrieveRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } if provider == schemas.Gemini {
470-484: Missing nil check on type assertion for provider in batch cancel.Same issue as batch retrieve - line 472 performs an unchecked type assertion that could panic.
🐛 Proposed fix
BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { if cancelReq, ok := req.(*anthropic.AnthropicBatchCancelRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } if provider == schemas.Gemini {
506-520: Missing nil check on type assertion for provider in batch results.Line 508 has the same unchecked type assertion pattern.
🐛 Proposed fix
BatchRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*BatchRequest, error) { if resultsReq, ok := req.(*anthropic.AnthropicBatchResultsRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } if provider == schemas.Gemini {
775-789: Missing nil check on type assertion for provider in file list.Line 777 has the same unchecked type assertion pattern that could panic.
🐛 Proposed fix
FileRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*FileRequest, error) { if listReq, ok := req.(*anthropic.AnthropicFileListRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } return &FileRequest{
816-831: Missing nil check on type assertion for provider in file retrieve.Line 818 has the same unchecked type assertion pattern.
🐛 Proposed fix
FileRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*FileRequest, error) { if retrieveReq, ok := req.(*anthropic.AnthropicFileRetrieveRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } // Handle file id conversion for Gemini
853-868: Missing nil check on type assertion for provider in file delete.Line 855 has the same unchecked type assertion pattern.
🐛 Proposed fix
FileRequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*FileRequest, error) { if deleteReq, ok := req.(*anthropic.AnthropicFileDeleteRequest); ok { - provider := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + provider, ok := ctx.Value(bifrostContextKeyProvider).(schemas.ModelProvider) + if !ok { + return nil, errors.New("provider not found in context") + } if provider == schemas.Gemini {core/providers/openai/openai.go (1)
1122-1133: Fix inconsistent context type in HandleOpenAIResponsesRequest.Line 1123 still uses
context.Contextinstead of*schemas.BifrostContext, which is inconsistent with all other handler functions in this file and will cause a compilation error when called from line 1107.🔧 Proposed fix
func HandleOpenAIResponsesRequest( - ctx context.Context, + ctx *schemas.BifrostContext, client *fasthttp.Client, url string, request *schemas.BifrostResponsesRequest,
🤖 Fix all issues with AI agents
In @core/internal/testutil/speech_synthesis_stream.go:
- Line 17: RunSpeechSynthesisStreamTest and RunSpeechSynthesisStreamAdvancedTest
are creating new request contexts from context.Background() instead of
propagating the provided ctx, so cancellation, deadlines, and values are lost;
update each place that currently uses context.Background() (the streaming
request context construction inside RunSpeechSynthesisStreamTest and
RunSpeechSynthesisStreamAdvancedTest) to derive from the passed-in ctx by
calling schemas.NewBifrostContext(ctx.Context(), schemas.NoDeadline) (or another
appropriate deadline option) so the parent context is propagated to streaming
requests.
In @core/schemas/mcp.go:
- Line 24: Test implementations that assign to FetchNewRequestIDFunc still use
context.Context; change their function signatures to accept
*schemas.BifrostContext to match the type declared as FetchNewRequestIDFunc in
core/schemas/mcp.go. Locate the functions assigned to FetchNewRequestIDFunc in
the tests (the anonymous or named functions creating request IDs) and update the
parameter type to *schemas.BifrostContext, update any call sites that construct
or pass a context to pass a *schemas.BifrostContext instead, and add or adjust
imports to reference the schemas package where BifrostContext is defined.
In @core/schemas/plugin_wasm.go:
- Around line 1-10: The comment for PluginShortCircuit is inaccurate: it
mentions a Stream/streaming short-circuit even though the wasm/tinygo build of
the struct only contains Response (*BifrostResponse) and Error (*BifrostError);
update the documentation on the PluginShortCircuit type to remove any reference
to Stream/streaming short-circuits and clearly state the struct only supports
Response and Error in this build (mentioning PluginShortCircuit,
BifrostResponse, and BifrostError so the maintainer can find and edit the
correct comment).
In @docs/docs.json:
- Around line 180-187: Add a redirect object to the docs.json "redirects" array
mapping the removed path "plugins/writing-plugin" to the new page (e.g.,
"plugins/writing-go-plugin"); locate the "redirects" array and append an entry
with the old path as the source and the chosen new page as the destination so
requests to "plugins/writing-plugin" are forwarded to the new documentation.
In @docs/plugins/writing-wasm-plugin.mdx:
- Around line 26-35: The docs table lists free as `free (ptr: u32)` but the Rust
plugin implementation exposes `free(ptr: u32, size: u32)`; update the table
entry for `free` to show `(ptr: u32, size: u32)` (or add a short parenthetical
note that some language allocators require size) so the doc matches the Rust
`free` implementation and avoids confusion for plugin authors.
In @transports/bifrost-http/integrations/router.go:
- Around line 191-195: The comment for ResponsesStreamResponseConverter is
outdated and claims the function returns "the list of event types and the list
of streaming formats", but the type signature actually returns a single string
and an interface{} plus error; update the doc comment for
ResponsesStreamResponseConverter to describe the current return values (that it
returns a single event type string and a single streaming-format payload as
interface{}, along with an error) and clarify the ordering/semantics expected
for integration authors so the comment matches the signature.
In @ui/app/clientLayout.tsx:
- Around line 19-23: The comment above the DevProfiler dynamic import is
misleading about bundle exclusion; update it to explain that dynamic(() =>
import(...)) creates a separate code-split chunk and that the DevProfiler is
only excluded from production because the component is never rendered when the
NODE_ENV/dev render guard is false (and will be removed by tree-shaking), not
because dynamic() alone prevents inclusion; reference the DevProfiler symbol and
the render guard that checks NODE_ENV/dev and mention that the dynamic import is
used with ssr: false to avoid server rendering.
🧹 Nitpick comments (25)
docs/plugins/writing-wasm-plugin.mdx (1)
296-304: Go example: Exported function name mismatch.The
initfunction is documented as an export, but the Go implementation usesinit_plugin(line 297) with the//export initdirective. This is intentional becauseinitis a reserved function name in Go. The code correctly handles this, but a brief inline comment explaining this Go-specific workaround would help developers understand why the function is named differently.examples/plugins/hello-world-wasm-typescript/assembly/index.ts (2)
73-87: Fragile boolean extraction logic.The
extractJsonBoolfunction checks if"true"appears within 10 characters after the key pattern. This approach has edge cases:
- A string value like
"has_error": "true_story"would incorrectly returntrue- Different whitespace/formatting could push
truebeyond 10 chars:"has_error" : trueFor a hello-world example this is acceptable, but consider adding a comment noting these limitations for developers who might copy this pattern.
92-103: Edge case inaddToJsonObjectwhen object has whitespace before closing brace.If the input JSON has trailing whitespace like
{"key": "value" }, the function checksjson.endsWith('}')which would be false. The function returns the originaljsonunmodified (line 102), silently failing to add the new key-value pair.For a hello-world example, this is minor, but worth noting for robustness.
plugins/governance/main.go (1)
501-505: Optional: Remove unnecessary nil check for context parameter.The nil check on line 501 appears redundant since the function signature guarantees
ctxis*schemas.BifrostContextand other parts of the function (lines 469-470, 498) usectxwithout nil checks. Consider removing the defensive nil check for consistency.♻️ Proposed simplification
- if result.Decision != DecisionAllow { - if ctx != nil { - if _, ok := ctx.Value(governanceRejectedContextKey).(bool); !ok { - ctx.SetValue(governanceRejectedContextKey, true) - } - } - } + if result.Decision != DecisionAllow { + if _, ok := ctx.Value(governanceRejectedContextKey).(bool); !ok { + ctx.SetValue(governanceRejectedContextKey, true) + } + }examples/plugins/hello-world-wasm-rust/src/lib.rs (1)
219-231: Consider avoiding unnecessarySome()wrapping for error field.The post-hook wraps
input.error.clone()inSome(), which means even whenhas_erroris false, you're still cloning and passing an error value. This works but could be more explicit about intent.Optional: More explicit error handling
let output = PostHookOutput { context, response: Some(input.response.clone()), - error: Some(input.error.clone()), + error: if input.has_error { Some(input.error.clone()) } else { None }, has_error: input.has_error, hook_error: String::new(), };core/providers/bedrock/bedrock_test.go (5)
680-681: Minor formatting: add space after comma.- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
708-708: Minor formatting: add space after comma.- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
1778-1778: Minor formatting: add space after comma.- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
1989-1991: Minor formatting: add space after comma.- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
2016-2018: Minor formatting: add space after comma.- ctx := schemas.NewBifrostContext(context.Background(),schemas.NoDeadline) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)transports/bifrost-http/integrations/openai.go (1)
38-97: Azure direct‑key handling may no longer reach BifrostContext
AzureEndpointPreHookand the composite pre‑hooks still only stash the constructedschemas.Keyinfasthttp.RequestCtxviactx.SetUserValue(string(schemas.BifrostContextKeyDirectKey), key), while core key‑selection now reads direct keys from*schemas.BifrostContextusingctx.Value(schemas.BifrostContextKeyDirectKey).Unless some middleware is copying this user value into the
BifrostContextfor each request, Azure “direct key” flows may silently stop working after this migration. A straightforward hardening would be to also write the same key intobifrostCtx:Example adjustment inside
AzureEndpointPreHook- ctx.SetUserValue(string(schemas.BifrostContextKeyDirectKey), key) + ctx.SetUserValue(string(schemas.BifrostContextKeyDirectKey), key) + if bifrostCtx != nil { + bifrostCtx.SetValue(schemas.BifrostContextKeyDirectKey, key) + }Please double‑check how
BifrostContextKeyDirectKeyis populated for OpenAI/Azure requests across the stack and whetherBifrostContextnow needs to be the single source of truth.Also applies to: 456-479, 976-1013, 1016-1058, 1060-1137, 1139-1221
transports/bifrost-http/integrations/bedrock.go (1)
41-106: Bedrock integrations: BifrostContext usage is coherentThe Bedrock converse/invoke, batch, and S3‑compatible file routes now consistently:
- Take
*schemas.BifrostContextin converters and pre‑callbacks.- Store provider selection and direct keys via
bifrostCtx.SetValue(...).- Use per‑request S3 context keys (bucket/prefix/maxKeys) on the same
BifrostContextinstance for response shaping.This keeps the previous behavior while aligning with the new context model. The one
(*bifrostCtx).Value(...)inextractBedrockJobArnFromPathcan be simplified tobifrostCtx.Value(...), but that’s purely stylistic.Also applies to: 195-333, 381-418, 422-476, 530-568, 587-829
core/bifrost.go (1)
385-518: Consider normalizing nil BifrostContext consistently for ListAllModels and MCP tool APIsMost public APIs either normalize a
nil*schemas.BifrostContexttobifrost.ctxdirectly or rely onhandleRequest/handleStreamRequestdoing so. Two areas still assume a non‑nil context:
ListAllModels(ctx *schemas.BifrostContext, ...)(usesctx.Done()in goroutines).ExecuteChatMCPToolandExecuteResponsesMCPTool(pass ctx straight into MCP manager).For symmetry – and to make accidental
nilcontexts less foot‑gunny for callers migrating from the oldcontext.Contextsignatures – you could add a defensive default similar to the batch/file APIs:Illustrative change for
ListAllModelsand MCP tool executorsfunc (bifrost *Bifrost) ListAllModels(ctx *schemas.BifrostContext, request *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) { + if ctx == nil { + ctx = bifrost.ctx + } ... } func (bifrost *Bifrost) ExecuteChatMCPTool(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, *schemas.BifrostError) { + if ctx == nil { + ctx = bifrost.ctx + } ... } func (bifrost *Bifrost) ExecuteResponsesMCPTool(ctx *schemas.BifrostContext, toolCall *schemas.ResponsesToolMessage) (*schemas.ResponsesMessage, *schemas.BifrostError) { + if ctx == nil { + ctx = bifrost.ctx + } ... }Not mandatory, but would make the new API surface more uniform and a bit more defensive.
Also applies to: 1788-1860
core/providers/anthropic/anthropic.go (2)
436-510: Streaming chat: end‑of‑stream signaling via context is wired correctly.
HandleAnthropicChatCompletionStreamingnow setsBifrostContextKeyStreamEndIndicatoron empty‑body failures, chunk‑level errors, and the final summary chunk before callingProcessAndSend{BifrostError,Response}, which allowscompleteDeferredSpanto close tracing spans exactly once per stream. Scanner error still goes throughProcessAndSendErrorwithout marking the stream as final; if you want spans to close on transport errors as well, consider setting the end‑indicator before callingProcessAndSendError.Also applies to: 619-683
791-865: Streaming responses: context flagging for final chunks mirrors chat behavior.For SSE responses, you set
StreamEndIndicatoron provider‑side empty responses and on the last logical chunk before delegating toProcessAndSendBifrostError/ProcessAndSendResponse, which keeps span finalization semantics aligned with chat streaming. As with chat, scanner errors bypass this flag and only emit viaProcessAndSendError; optionally aligning that path would improve tracing completeness.Also applies to: 965-1013
core/providers/ollama/ollama.go (1)
163-187: Responses→chat fallback flag is a nice addition; consider mirroring for non‑streaming.Setting
BifrostContextKeyIsResponsesToChatCompletionFallbackinResponsesStreambefore delegating toChatCompletionStreamgives downstream hooks a clear signal that this is a responses‑via‑chat fallback. If you want identical observability for non‑streaming flows, you might also set this flag inResponsesbefore callingChatCompletion.core/providers/cohere/cohere.go (1)
341-412: Chat streaming now integrates cleanly with Bifrost tracing.
ChatCompletionStreamusesctx *schemas.BifrostContextand setsStreamEndIndicatorboth when a stream error is converted to aBifrostErrorand on the final chunk before callingProcessAndSend{BifrostError,Response}, which letscompleteDeferredSpanclose spans and propagate TTFT/chunk metrics. Scanner errors still route throughProcessAndSendErrorwithout marking final; aligning that path is a possible future improvement if you need spans to close on transport failure too.Also applies to: 423-512
core/providers/vertex/vertex.go (1)
172-223: Minor nit: defers inside the pagination loop could be tightened.In
listModelsByKey,fasthttp.AcquireRequest/AcquireResponseare paired withdefer‑based releases inside the pagination loop, so requests/responses are only returned to the pool when the function exits. Not a correctness issue, but releasing them at the end of each iteration instead of viadeferwould reduce peak resource usage during multi‑page listings.core/providers/gemini/gemini.go (5)
271-520: Streaming chat completion: consider marking end-of-stream on scanner errors as wellThe
ChatCompletionStream/HandleGeminiChatCompletionStreampath now:
- Uses
*schemas.BifrostContextthroughout.- Sets
BifrostContextKeyStreamEndIndicatoron Gemini API errors, conversion errors, and the final completion chunk.- Properly sends errors and responses via the shared helpers.
One robustness gap: in the
scanner.Err()branch, you log and forward the error viaProcessAndSendErrorbut don’t setBifrostContextKeyStreamEndIndicator. If downstream consumers rely on that flag to know the stream has terminated (regardless of reason), it may be worth setting it there too for consistency.Proposed tweak for the scanner error path
@@ - if err := scanner.Err(); err != nil { - logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) - providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ChatCompletionStreamRequest, providerName, model, logger) - } + if err := scanner.Err(); err != nil { + logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ChatCompletionStreamRequest, providerName, model, logger) + }
525-873: Responses and ResponsesStream: consistent BifrostContext and similar stream-end concernNon-streaming
Responsesand streamingResponsesStreamcorrectly:
- Use
CheckContextAndGetRequestBody(ctx, ...)with the new context type.- Invoke
completeRequest/HandleGeminiResponsesStreamwithctx.- Populate
ExtraFieldsand raw payloads via the context-aware helpers.
HandleGeminiResponsesStreammirrors chat streaming in how it setsBifrostContextKeyStreamEndIndicatorfor Gemini API errors, conversion errors, per-chunk completion, and finalization, but again omits it in thescanner.Err()branch. For symmetry with other termination paths, consider also setting the flag there, analogous to the chat-stream suggestion.Proposed tweak for responses stream scanner error path
@@ - if err := scanner.Err(); err != nil { - logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) - providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ResponsesStreamRequest, providerName, model, logger) - } else { + if err := scanner.Err(); err != nil { + logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.ResponsesStreamRequest, providerName, model, logger) + } else {
969-1015: Speech response format key: consider guarding against nilParamsSetting
BifrostContextKeyResponseFormatfromrequest.Params.ResponseFormatassumesrequest.Paramsis always non-nil. If upstream validation guarantees this, it’s fine; otherwise, this can panic on malformed input.Given this is a hot path, a small defensive check would make it more robust without changing semantics when
Paramsis present.Defensive guard around `request.Params`
- ctx.SetValue(BifrostContextKeyResponseFormat, request.Params.ResponseFormat) + if request.Params != nil { + ctx.SetValue(BifrostContextKeyResponseFormat, request.Params.ResponseFormat) + }If
Paramsis required by your API contract, please ignore this suggestion but consider adding an explicit validation earlier in the flow.
1017-1239: Speech streaming: good BifrostContext usage; mirror stream-end flag on IO errors
SpeechStreamcorrectly:
- Uses
*schemas.BifrostContextfor request-body building and extra headers.- Detects Gemini API errors and sets
BifrostContextKeyStreamEndIndicatorbefore sending aBifrostError.- Marks stream completion via the same key in the final “done” chunk.
As with other streaming helpers, the
scanner.Err()branch currently only logs and callsProcessAndSendError. For consistent stream termination signaling, you may want to setBifrostContextKeyStreamEndIndicatorthere as well.Optional stream-end flag on speech scanner error
- if err := scanner.Err(); err != nil { - provider.logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) - providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.SpeechStreamRequest, providerName, request.Model, provider.logger) - } else { + if err := scanner.Err(); err != nil { + provider.logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.SpeechStreamRequest, providerName, request.Model, provider.logger) + } else {
1242-1523: Transcription and TranscriptionStream: consistent context migration; same stream-end detailThe non-streaming
Transcriptionpath and streamingTranscriptionStreamboth:
- Use
*schemas.BifrostContextend-to-end.- Set
BifrostContextKeyStreamEndIndicatoron Gemini API errors and on the final “done” event.For the
scanner.Err()case inTranscriptionStream, consider also settingBifrostContextKeyStreamEndIndicatorbefore forwarding the error, mirroring the other termination paths.Optional stream-end flag on transcription scanner error
- if err := scanner.Err(); err != nil { - provider.logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) - providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.TranscriptionStreamRequest, providerName, request.Model, provider.logger) - } else { + if err := scanner.Err(); err != nil { + provider.logger.Warn(fmt.Sprintf("Error reading stream: %v", err)) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, schemas.TranscriptionStreamRequest, providerName, request.Model, provider.logger) + } else {core/schemas/context.go (1)
73-79: Optional: clarify silent dropping of reserved-key writes
NewBifrostContextWithValueplusSetValue’s guard mean reserved keys set afterBlockRestrictedWritesis enabled are silently ignored:if bc.blockRestrictedWrites.Load() && slices.Contains(reservedKeys, key) { // we silently drop writes for these reserved keys return }Given that this is observable only via missing values, consider briefly documenting this behavior on
SetValue(orBlockRestrictedWrites) so callers know failures won’t surface as errors or logs.Also applies to: 218-224
core/providers/bedrock/responses.go (1)
1206-1224: Bedrock converters now properly use BifrostContext and support structured-output tools
ToBifrostResponsesRequest,ToBedrockResponsesRequest, andToBifrostResponsesResponsenow accept*schemas.BifrostContextand thread it through to message converters, giving Bedrock conversions access to per-request context.ConvertBedrockMessagesToBifrostMessagesandconvertSingleBedrockMessageToBifrostMessagesusectx.Value(schemas.BifrostContextKeyStructuredOutputToolName)to detect a configured structured-output tool and, when aToolUseblock’s name matches, convert itsInputpayload directly into an assistant text message rather than afunction_callmessage.- For all other tool uses the prior behavior is preserved (emit
ResponsesMessageTypeFunctionCall/FunctionCallOutputmessages), and ctx is handled safely when nil.This is a clean extension of the Bedrock mapping that respects the new BifrostContext without regressing existing tool/streaming behaviors. Just ensure upstream code consistently sets
BifrostContextKeyStructuredOutputToolNameonly for calls that should be treated as structured output.Also applies to: 1436-1456, 1698-1713, 2485-2502, 2631-2643, 2644-2713
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
5d1fcf4 to
e894005
Compare
akshaydeo
force-pushed
the
01-05-plugin_loaders
branch
from
afd603c to
fbfd140
Compare
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
e894005 to
e6f4a43
Compare
akshaydeo
force-pushed
the
01-05-plugin_loaders
branch
from
fbfd140 to
016cde0
Compare
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
2 times, most recently
from
3e6657a to
3450664
Compare
akshaydeo
force-pushed
the
01-05-plugin_loaders
branch
from
016cde0 to
511a030
Compare
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
3450664 to
36f0cca
Compare
akshaydeo
force-pushed
the
01-07-move_to_schemas.bifrostcontext
branch
from
36f0cca to
6010e56
Compare
akshaydeo
force-pushed
the
01-05-plugin_loaders
branch
from
511a030 to
6272071
Compare
Merge activity
|
akshaydeo
changed the base branch from
01-05-plugin_loaders
to
graphite-base/1262
1 participant
Summary
Moves entire internal of Bifrost to use schemas.BifrostContext vs context.Context.
Changes
Type of change
Affected areas
How to test
Describe the steps to validate this change. Include commands and expected outcomes.
If adding new configs or environment variables, document them here.
Screenshots/Recordings
If UI changes, add before/after screenshots or short clips.
Breaking changes
If yes, describe impact and migration instructions.
Related issues
Link related issues and discussions. Example: Closes #123
Security considerations
Note any security implications (auth, secrets, PII, sandboxing, etc.).
Checklist
docs/contributing/README.mdand followed the guidelines