diff --git a/core/bifrost.go b/core/bifrost.go index 2c53f83d7fa..3d88c7b427c 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -6852,7 +6852,6 @@ func (p *PluginPipeline) RunPreRequestHooks(ctx *schemas.BifrostContext, req *sc return } ctx.BlockRestrictedWrites() - defer ctx.UnblockRestrictedWrites() for _, plugin := range p.llmPlugins { pluginName := plugin.GetName() p.logger.Debug("running pre-request hook for plugin %s", pluginName) @@ -6876,6 +6875,20 @@ func (p *PluginPipeline) RunPreRequestHooks(ctx *schemas.BifrostContext, req *sc } p.tracer.EndSpan(handle, schemas.SpanStatusOk, "") } + ctx.UnblockRestrictedWrites() + + // Commit the routing-rule key pin. A matched routing rule writes the pinned key ID to the + // non-reserved BifrostContextKeyRoutingPinnedAPIKeyID during the blocked phase above — a + // direct write to the reserved BifrostContextKeyAPIKeyID would have been silently dropped. + // Core is the sole writer of the reserved key, so normalize the routing pin into it here, + // after unblocking, so key selection reads a single canonical pin. A non-empty routing pin + // overrides a caller-supplied pin: the routing rule is authoritative server-side policy and + // has typically already rewritten provider/model for this request. + if pin, ok := ctx.Value(schemas.BifrostContextKeyRoutingPinnedAPIKeyID).(string); ok { + if pin = strings.TrimSpace(pin); pin != "" { + ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, pin) + } + } } // RunPostLLMHooks executes PostHooks in reverse order for the plugins whose PreLLMHook ran. diff --git a/core/bifrost_test.go b/core/bifrost_test.go index 5023684c7d0..e7d53c21754 100644 --- a/core/bifrost_test.go +++ b/core/bifrost_test.go @@ -2756,3 +2756,73 @@ func TestFilterKeysByID(t *testing.T) { } }) } + +// fakeRoutingPlugin is a minimal LLMPlugin whose PreRequestHook writes a routing key pin to the +// non-reserved BifrostContextKeyRoutingPinnedAPIKeyID, mirroring what the governance routing +// engine does. It exists to exercise the commit step in PluginPipeline.RunPreRequestHooks. +type fakeRoutingPlugin struct { + name string + pinKeyID string // written to BifrostContextKeyRoutingPinnedAPIKeyID when non-empty +} + +func (f *fakeRoutingPlugin) GetName() string { return f.name } +func (f *fakeRoutingPlugin) Cleanup() error { return nil } +func (f *fakeRoutingPlugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + if f.pinKeyID != "" { + // A direct write to the reserved BifrostContextKeyAPIKeyID here would be dropped by the + // restricted-write block; routing must use the non-reserved key. + ctx.SetValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID, f.pinKeyID) + } + return nil +} +func (f *fakeRoutingPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { + return req, nil, nil +} +func (f *fakeRoutingPlugin) PostLLMHook(ctx *schemas.BifrostContext, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + return resp, bifrostErr, nil +} + +func newRoutingCommitPipeline(plugins ...schemas.LLMPlugin) *PluginPipeline { + return &PluginPipeline{ + logger: NewDefaultLogger(schemas.LogLevelError), + tracer: &schemas.NoOpTracer{}, + llmPlugins: plugins, + } +} + +// TestRunPreRequestHooks_CommitsRoutingPinnedKey verifies that the pinned key a routing rule +// writes to the non-reserved BifrostContextKeyRoutingPinnedAPIKeyID (during the blocked +// PreRequestHook phase) is committed by core into the reserved BifrostContextKeyAPIKeyID that +// key selection reads — and that the routing pin's precedence over a caller-supplied pin holds. +func TestRunPreRequestHooks_CommitsRoutingPinnedKey(t *testing.T) { + const pinned = "routing-pinned-key-id" + + t.Run("routing pin is committed to reserved api-key-id", func(t *testing.T) { + p := newRoutingCommitPipeline(&fakeRoutingPlugin{name: "gov", pinKeyID: pinned}) + ctx := schemas.NewBifrostContext(context.Background(), time.Now()) + p.RunPreRequestHooks(ctx, &schemas.BifrostRequest{}) + if got, _ := ctx.Value(schemas.BifrostContextKeyAPIKeyID).(string); got != pinned { + t.Fatalf("APIKeyID = %q, want %q", got, pinned) + } + }) + + t.Run("routing pin overrides a caller-supplied api-key-id", func(t *testing.T) { + p := newRoutingCommitPipeline(&fakeRoutingPlugin{name: "gov", pinKeyID: pinned}) + ctx := schemas.NewBifrostContext(context.Background(), time.Now()) + ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, "caller-pin") + p.RunPreRequestHooks(ctx, &schemas.BifrostRequest{}) + if got, _ := ctx.Value(schemas.BifrostContextKeyAPIKeyID).(string); got != pinned { + t.Fatalf("APIKeyID = %q, want %q (routing pin must override caller pin)", got, pinned) + } + }) + + t.Run("caller api-key-id preserved when no routing pin", func(t *testing.T) { + p := newRoutingCommitPipeline(&fakeRoutingPlugin{name: "noop"}) + ctx := schemas.NewBifrostContext(context.Background(), time.Now()) + ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, "caller-pin") + p.RunPreRequestHooks(ctx, &schemas.BifrostRequest{}) + if got, _ := ctx.Value(schemas.BifrostContextKeyAPIKeyID).(string); got != "caller-pin" { + t.Fatalf("APIKeyID = %q, want %q (no routing pin must not clobber caller pin)", got, "caller-pin") + } + }) +} diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 57f64891d25..b206b52da6d 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -222,6 +222,7 @@ const ( BifrostContextKeyGovernanceCustomerNames BifrostContextKey = "bifrost-governance-customer-names" // []string (display names, aligned with customer-ids; set by enterprise governance plugin - DO NOT SET THIS MANUALLY) BifrostContextKeyGovernanceRoutingRuleID BifrostContextKey = "bifrost-governance-routing-rule-id" // string (to store the routing rule ID (set by bifrost governance plugin - DO NOT SET THIS MANUALLY)) BifrostContextKeyGovernanceRoutingRuleName BifrostContextKey = "bifrost-governance-routing-rule-name" // string (to store the routing rule name (set by bifrost governance plugin - DO NOT SET THIS MANUALLY)) + BifrostContextKeyRoutingPinnedAPIKeyID BifrostContextKey = "bifrost-routing-pinned-api-key-id" // string (provider key ID pinned by a matched routing rule target; resolved against the configured key pool during key selection and takes precedence over a caller-supplied pin (set by bifrost governance plugin - DO NOT SET THIS MANUALLY)) BifrostContextKeySelectedPromptName BifrostContextKey = "bifrost-selected-prompt-name" // string (display name of the selected prompt (set by prompts plugin - DO NOT SET THIS MANUALLY)) BifrostContextKeySelectedPromptVersion BifrostContextKey = "bifrost-selected-prompt-version" // string (numeric version as string, e.g. "3" (set by prompts plugin - DO NOT SET THIS MANUALLY)) BifrostContextKeySelectedPromptID BifrostContextKey = "bifrost-selected-prompt-id" // string (id of the selected prompt (set by prompts plugin - DO NOT SET THIS MANUALLY)) @@ -1555,8 +1556,8 @@ func (r *BifrostMCPResponse) PopulateExtraFields(mcpRequestType MCPRequestType, // BifrostResponseExtraFields contains additional fields in a response. type BifrostResponseExtraFields struct { - RequestType RequestType `json:"request_type"` - RoutingInfo RoutingInfo `json:"routing_info"` + RequestType RequestType `json:"request_type"` + RoutingInfo RoutingInfo `json:"routing_info"` // Deprecated: use RoutingInfo.Provider. Still populated for backward // compatibility; new consumers should read from RoutingInfo. Provider ModelProvider `json:"provider,omitempty"` @@ -1569,9 +1570,9 @@ type BifrostResponseExtraFields struct { // matched (i.e. RoutingInfo.ResolvedKeyAlias != nil), otherwise // RoutingInfo.Model. Still populated for backward compatibility; new // consumers should read from RoutingInfo. - ResolvedModelUsed string `json:"resolved_model_used,omitempty"` - Latency int64 `json:"latency"` // in milliseconds (for streaming responses this will be each chunk latency, and the last chunk latency will be the total latency) - ChunkIndex int `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses + ResolvedModelUsed string `json:"resolved_model_used,omitempty"` + Latency int64 `json:"latency"` // in milliseconds (for streaming responses this will be each chunk latency, and the last chunk latency will be the total latency) + ChunkIndex int `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses RawRequest interface{} `json:"raw_request,omitempty"` RawResponse interface{} `json:"raw_response,omitempty"` CacheDebug *BifrostCacheDebug `json:"cache_debug,omitempty"` diff --git a/plugins/governance/main.go b/plugins/governance/main.go index c30a8a31858..ba0fbf9375b 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -774,9 +774,13 @@ func (p *GovernancePlugin) applyRoutingRules(ctx *schemas.BifrostContext, req *s req.SetFallbacks(resolvedFallbacks) } - // Pin specific API key by ID if the routing rule specifies one + // Pin specific API key by ID if the routing rule specifies one. This uses a dedicated, + // non-reserved context key (not BifrostContextKeyAPIKeyID): routing runs inside + // PreRequestHook, where core blocks writes to reserved key-selection keys, so a write to + // the caller-pin key would be silently dropped. Key selection reads this routing pin first + // and resolves it against the configured key pool. if decision.KeyID != "" { - ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, decision.KeyID) + ctx.SetValue(schemas.BifrostContextKeyRoutingPinnedAPIKeyID, decision.KeyID) } p.logger.Debug("[Governance] Applied routing decision: provider=%s, model=%s, keyID=%s, fallbacks=%v", decision.Provider, decision.Model, decision.KeyID, decision.Fallbacks) diff --git a/plugins/governance/routing_test.go b/plugins/governance/routing_test.go index 231db1c861d..19b8b30a41d 100644 --- a/plugins/governance/routing_test.go +++ b/plugins/governance/routing_test.go @@ -334,13 +334,37 @@ func TestEvaluateRoutingRules_MultiTargetDeterministicWithPinnedKey(t *testing.T // KeyID must be propagated through the routing decision. assert.Equal(t, pinnedKeyID, decision.KeyID) - // Simulate the propagation step performed by governance/main.go so that we can - // assert the pinned key_id is visible in the BifrostContext. - if decision.KeyID != "" { - bgCtx.SetValue(schemas.BifrostContextKeyAPIKeyID, decision.KeyID) - } - ctxKeyID, _ := bgCtx.Value(schemas.BifrostContextKeyAPIKeyID).(string) - assert.Equal(t, pinnedKeyID, ctxKeyID) + // Now exercise the REAL propagation path through applyRoutingRules, under the same + // restricted-write block that core's RunPreRequestHooks installs around every + // PreRequestHook (core/bifrost.go: ctx.BlockRestrictedWrites + ctx.WithPluginScope). + // This is what production actually does — the routing pin lands on the dedicated, + // non-reserved BifrostContextKeyRoutingPinnedAPIKeyID (a write to the reserved + // BifrostContextKeyAPIKeyID would be silently dropped during this phase). Key selection + // (selectKeyFromProviderForModelWithPool) reads the pinned key back from this context. + plugin := &GovernancePlugin{ + logger: NewMockLogger(), + store: store, + engine: engine, + } + req := &schemas.BifrostRequest{ + RequestType: schemas.ChatCompletionRequest, + ChatRequest: &schemas.BifrostChatRequest{Provider: schemas.OpenAI, Model: "gpt-4o"}, + } + + root := schemas.NewBifrostContext(context.Background(), time.Now()) + root.BlockRestrictedWrites() + pluginName := PluginName + scoped := root.WithPluginScope(&pluginName) + + appliedDecision, err := plugin.applyRoutingRules(scoped, req, nil) + require.NoError(t, err) + require.NotNil(t, appliedDecision) + assert.Equal(t, pinnedKeyID, appliedDecision.KeyID) + + // The pinned key_id must be readable from the root context that key selection consults. + ctxKeyID, _ := root.Value(schemas.BifrostContextKeyRoutingPinnedAPIKeyID).(string) + assert.Equal(t, pinnedKeyID, ctxKeyID, + "routing-rule pinned key_id must reach BifrostContextKeyRoutingPinnedAPIKeyID that selectKeyFromProviderForModelWithPool reads") } // TestEvaluateRoutingRules_ScopePrecedence tests virtual_key scope takes precedence over global