From 1ee84a3314016ecfc9b85b20b05b0d5aa2cdb131 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Thu, 11 Jun 2026 19:59:11 +0530 Subject: [PATCH 1/3] fix: governance mcp tool filter fixes --- plugins/governance/main.go | 98 +++++++-- .../governance/prunemcpincludetools_test.go | 187 ++++++++++++++++++ 2 files changed, 271 insertions(+), 14 deletions(-) create mode 100644 plugins/governance/prunemcpincludetools_test.go diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 73b2116198e..bde7a5b945a 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -428,14 +428,16 @@ func (p *GovernancePlugin) runPreRequestRouting(ctx *schemas.BifrostContext, vir return modelIn, err } + // A caller-provided include-tools list can only narrow the virtual key's + // tool grant, never expand it — prune entries the key does not allow. + callerProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) + p.cfgMutex.RLock() autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject p.cfgMutex.RUnlock() - if !autoInjectDisabled { - if existing := ctx.Value(schemas.MCPContextKeyIncludeTools); existing == nil { - if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { - ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) - } + if !callerProvided && !autoInjectDisabled { + if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) } } } @@ -781,13 +783,19 @@ func (p *GovernancePlugin) applyRoutingRules(ctx *schemas.BifrostContext, req *s // directly; callers store it via ctx.SetValue(schemas.MCPContextKeyIncludeTools, ...). VK-specific // MCP configs take precedence over AllowOnAllVirtualKeys clients. func (p *GovernancePlugin) computeMCPIncludeTools(virtualKey *configstoreTables.TableVirtualKey) []string { - executeOnlyTools := make([]string, 0) - - // Build a lookup of AllowOnAllVirtualKeys clients: clientID -> clientName var allowAllVKsClients map[string]string if p.inMemoryStore != nil { allowAllVKsClients = p.inMemoryStore.GetMCPClientsAllowingAllVirtualKeys() } + return p.computeMCPIncludeToolsWith(virtualKey, allowAllVKsClients) +} + +// computeMCPIncludeToolsWith is the computeMCPIncludeTools variant taking a pre-fetched +// AllowOnAllVirtualKeys map (clientID → clientName), so callers that make multiple +// grant decisions per request can evaluate them all against one consistent snapshot. +func (p *GovernancePlugin) computeMCPIncludeToolsWith(virtualKey *configstoreTables.TableVirtualKey, allowAllVKsClients map[string]string) []string { + executeOnlyTools := make([]string, 0) + if allowAllVKsClients == nil { allowAllVKsClients = make(map[string]string) } @@ -826,6 +834,67 @@ func (p *GovernancePlugin) computeMCPIncludeTools(virtualKey *configstoreTables. return executeOnlyTools } +// pruneMCPIncludeToolsFromContext narrows a caller-provided include-tools list (stamped on ctx +// from the x-bf-mcp-include-tools header in lib/ctx.go) down to the tools the virtual key +// allows, and writes the pruned list back to ctx. Returns true when a caller list was present, +// regardless of how many entries survived. Entries the key does not grant are dropped; a +// "client-*" wildcard is kept only when the key itself is unrestricted for that client, +// otherwise it is replaced by the key's specific grants for that client (passing the wildcard +// through would read downstream as "all tools of this client"). +func (p *GovernancePlugin) pruneMCPIncludeToolsFromContext(ctx *schemas.BifrostContext, virtualKey *configstoreTables.TableVirtualKey) bool { + existing := ctx.Value(schemas.MCPContextKeyIncludeTools) + if existing == nil { + return false + } + requested, _ := existing.([]string) + + // Fetch the AllowOnAllVirtualKeys snapshot once so the wildcard checks (via vkSet) + // and the per-tool checks (via isMCPToolAllowedByVKWith) can't observe different + // states across a concurrent config reload. + var allowAllClients map[string]string + if p.inMemoryStore != nil { + allowAllClients = p.inMemoryStore.GetMCPClientsAllowingAllVirtualKeys() + } + + vkTools := p.computeMCPIncludeToolsWith(virtualKey, allowAllClients) + vkSet := make(map[string]struct{}, len(vkTools)) + for _, tool := range vkTools { + vkSet[tool] = struct{}{} + } + + pruned := make([]string, 0, len(requested)) + seen := make(map[string]struct{}, len(requested)) + add := func(tool string) { + if _, dup := seen[tool]; !dup { + seen[tool] = struct{}{} + pruned = append(pruned, tool) + } + } + for _, pattern := range requested { + if pattern == "" { + continue + } + if clientName, isWildcard := strings.CutSuffix(pattern, "-*"); isWildcard { + if _, ok := vkSet[pattern]; ok { + add(pattern) + continue + } + for _, tool := range vkTools { + if strings.HasPrefix(tool, clientName+"-") { + add(tool) + } + } + continue + } + if p.isMCPToolAllowedByVKWith(virtualKey, pattern, allowAllClients) { + add(pattern) + } + } + + ctx.SetValue(schemas.MCPContextKeyIncludeTools, pruned) + return true +} + // EvaluateGovernanceRequest is a common function that handles virtual key validation // and governance evaluation logic. It returns the evaluation result and a BifrostError // if the request should be rejected, or nil if allowed. @@ -1152,15 +1221,16 @@ func (p *GovernancePlugin) PreRequestHook(ctx *schemas.BifrostContext, req *sche return err } + // A caller-provided include-tools list can only narrow the virtual key's + // tool grant, never expand it — prune entries the key does not allow. + callerProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) + p.cfgMutex.RLock() autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject p.cfgMutex.RUnlock() - if !autoInjectDisabled { - // Don't overwrite a caller-provided include-tools value (set via header in lib/ctx.go). - if existing := ctx.Value(schemas.MCPContextKeyIncludeTools); existing == nil { - if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { - ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) - } + if !callerProvided && !autoInjectDisabled { + if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) } } } diff --git a/plugins/governance/prunemcpincludetools_test.go b/plugins/governance/prunemcpincludetools_test.go new file mode 100644 index 00000000000..d8491461600 --- /dev/null +++ b/plugins/governance/prunemcpincludetools_test.go @@ -0,0 +1,187 @@ +package governance + +import ( + "context" + "testing" + + "github.com/maximhq/bifrost/core/schemas" + configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newCtxWithIncludeTools returns a BifrostContext pre-stamped with a caller-provided +// include-tools list, mirroring what lib/ctx.go does for the x-bf-mcp-include-tools header. +func newCtxWithIncludeTools(tools []string) *schemas.BifrostContext { + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) + return ctx +} + +// includeToolsFromCtx reads back the (possibly pruned) include-tools list from ctx. +func includeToolsFromCtx(t *testing.T, ctx *schemas.BifrostContext) []string { + t.Helper() + value := ctx.Value(schemas.MCPContextKeyIncludeTools) + require.NotNil(t, value, "include-tools ctx value should be set") + tools, ok := value.([]string) + require.True(t, ok, "include-tools ctx value should be a []string") + return tools +} + +// No caller-provided list on ctx → returns false and leaves ctx untouched. +func TestPruneMCPIncludeTools_NoCallerList(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects"}) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + + assert.False(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Nil(t, ctx.Value(schemas.MCPContextKeyIncludeTools), + "ctx should remain unset when the caller provided no list") +} + +// A tool the VK does not grant is dropped; the result is an empty (deny-all) list. +func TestPruneMCPIncludeTools_DisallowedToolDropped(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects"}) + ctx := newCtxWithIncludeTools([]string{"sentry-search_tools"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Empty(t, includeToolsFromCtx(t, ctx), + "a tool outside the VK grant must be pruned, leaving a deny-all list") +} + +// A tool the VK explicitly grants survives pruning. +func TestPruneMCPIncludeTools_GrantedToolKept(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects", "search_issues"}) + ctx := newCtxWithIncludeTools([]string{"sentry-find_projects", "sentry-search_tools"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-find_projects"}, includeToolsFromCtx(t, ctx), + "granted entries survive, ungranted entries are dropped") +} + +// A specific tool requested under an unrestricted ("*") VK grant survives — the +// header narrows within the wildcard grant. +func TestPruneMCPIncludeTools_SpecificToolUnderUnrestrictedGrant(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"*"}) + ctx := newCtxWithIncludeTools([]string{"sentry-search_tools"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-search_tools"}, includeToolsFromCtx(t, ctx), + "specific request should be allowed by the client's unrestricted grant") +} + +// A caller wildcard is kept verbatim only when the VK itself is unrestricted for that client. +func TestPruneMCPIncludeTools_WildcardKeptWhenVKUnrestricted(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"*"}) + ctx := newCtxWithIncludeTools([]string{"sentry-*"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-*"}, includeToolsFromCtx(t, ctx)) +} + +// A caller wildcard against a specific VK grant is narrowed to the grant's entries — +// passing the wildcard through would read downstream as "all tools of this client". +func TestPruneMCPIncludeTools_WildcardNarrowedToSpecificGrants(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects", "search_issues"}) + ctx := newCtxWithIncludeTools([]string{"sentry-*"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-find_projects", "sentry-search_issues"}, includeToolsFromCtx(t, ctx)) +} + +// A caller wildcard for a client the VK does not grant at all yields nothing. +func TestPruneMCPIncludeTools_WildcardForUngrantedClientDropped(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects"}) + ctx := newCtxWithIncludeTools([]string{"github-*"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Empty(t, includeToolsFromCtx(t, ctx)) +} + +// An empty header value (parsed as [""] by lib/ctx.go) prunes to a deny-all list, +// letting callers suppress tool injection for a request. +func TestPruneMCPIncludeTools_EmptyHeaderOptOut(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"*"}) + ctx := newCtxWithIncludeTools([]string{""}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Empty(t, includeToolsFromCtx(t, ctx)) +} + +// Wildcard expansion plus an overlapping specific request must not produce duplicates. +func TestPruneMCPIncludeTools_DedupAcrossWildcardAndSpecific(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"find_projects", "search_issues"}) + ctx := newCtxWithIncludeTools([]string{"sentry-*", "sentry-find_projects"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-find_projects", "sentry-search_issues"}, includeToolsFromCtx(t, ctx)) +} + +// AllowOnAllVirtualKeys client with no explicit VK config: both specific and wildcard +// requests survive (the implicit grant is client-wide). +func TestPruneMCPIncludeTools_AllowOnAllVirtualKeysClient(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{ + allowAllClients: map[string]string{"client-1": "youtube"}, + }) + vk := buildVKNoMCPConfigs() + ctx := newCtxWithIncludeTools([]string{"youtube-search", "youtube-*", "github-list_repos"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"youtube-search", "youtube-*"}, includeToolsFromCtx(t, ctx), + "AllowOnAllVirtualKeys grants the whole client; other clients are still pruned") +} + +// An explicit empty VK config (deny-all) overrides the client's AllowOnAllVirtualKeys flag. +func TestPruneMCPIncludeTools_ExplicitEmptyConfigOverridesAllowAll(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{ + allowAllClients: map[string]string{"client-1": "youtube"}, + }) + vk := buildVKWithMCPConfigs("client-1", "youtube", []string{}) + ctx := newCtxWithIncludeTools([]string{"youtube-search", "youtube-*"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Empty(t, includeToolsFromCtx(t, ctx)) +} + +// Pruning spans multiple VK clients independently. +func TestPruneMCPIncludeTools_MultipleClients(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := &configstoreTables.TableVirtualKey{ + ID: "vk-multi", + Name: "test-vk-multi", + MCPConfigs: []configstoreTables.TableVirtualKeyMCPConfig{ + { + MCPClient: configstoreTables.TableMCPClient{ClientID: "client-1", Name: "sentry"}, + ToolsToExecute: []string{"find_projects"}, + }, + { + MCPClient: configstoreTables.TableMCPClient{ClientID: "client-2", Name: "github"}, + ToolsToExecute: []string{"*"}, + }, + }, + } + ctx := newCtxWithIncludeTools([]string{"sentry-find_projects", "sentry-search_issues", "github-list_repos", "github-*"}) + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Equal(t, []string{"sentry-find_projects", "github-list_repos", "github-*"}, includeToolsFromCtx(t, ctx)) +} + +// A ctx value of the wrong type fails closed: treated as a present-but-empty caller list. +func TestPruneMCPIncludeTools_WrongTypeFailsClosed(t *testing.T) { + p := newPluginWithInMemoryStore(&mockInMemoryStore{}) + vk := buildVKWithMCPConfigs("client-1", "sentry", []string{"*"}) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.MCPContextKeyIncludeTools, "sentry-search_tools") + + assert.True(t, p.pruneMCPIncludeToolsFromContext(ctx, vk)) + assert.Empty(t, includeToolsFromCtx(t, ctx), + "a malformed ctx value must prune to deny-all, not pass through") +} From 25f7aaad63e994598daf0444c71a591c2cd686ac Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Thu, 11 Jun 2026 20:47:35 +0530 Subject: [PATCH 2/3] fix: mcp allowed clients governance fixes --- plugins/governance/main.go | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/plugins/governance/main.go b/plugins/governance/main.go index bde7a5b945a..c30a8a31858 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -430,12 +430,16 @@ func (p *GovernancePlugin) runPreRequestRouting(ctx *schemas.BifrostContext, vir // A caller-provided include-tools list can only narrow the virtual key's // tool grant, never expand it — prune entries the key does not allow. - callerProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) + includeToolsProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) p.cfgMutex.RLock() autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject p.cfgMutex.RUnlock() - if !callerProvided && !autoInjectDisabled { + // An include-clients filter opts the request into tool injection even when + // auto-injection is disabled (see ParseAndAddToolsToRequest in core/mcp), so + // the key's allowlist must be stamped on every path where injection can run. + includeClientsPresent := ctx.Value(schemas.MCPContextKeyIncludeClients) != nil + if !includeToolsProvided && (!autoInjectDisabled || includeClientsPresent) { if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) } @@ -1223,12 +1227,16 @@ func (p *GovernancePlugin) PreRequestHook(ctx *schemas.BifrostContext, req *sche // A caller-provided include-tools list can only narrow the virtual key's // tool grant, never expand it — prune entries the key does not allow. - callerProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) + includeToolsProvided := p.pruneMCPIncludeToolsFromContext(ctx, virtualKey) p.cfgMutex.RLock() autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject p.cfgMutex.RUnlock() - if !callerProvided && !autoInjectDisabled { + // An include-clients filter opts the request into tool injection even when + // auto-injection is disabled (see ParseAndAddToolsToRequest in core/mcp), so + // the key's allowlist must be stamped on every path where injection can run. + includeClientsPresent := ctx.Value(schemas.MCPContextKeyIncludeClients) != nil + if !includeToolsProvided && (!autoInjectDisabled || includeClientsPresent) { if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) } From a5a841d7e81b6f836b964c5cdc508297092f263c Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Thu, 11 Jun 2026 20:59:00 +0530 Subject: [PATCH 3/3] tests: mcp tools filtering tests for governance --- plugins/governance/prerequesthookmcp_test.go | 408 +++++++++++++++++++ 1 file changed, 408 insertions(+) create mode 100644 plugins/governance/prerequesthookmcp_test.go diff --git a/plugins/governance/prerequesthookmcp_test.go b/plugins/governance/prerequesthookmcp_test.go new file mode 100644 index 00000000000..112ce95df91 --- /dev/null +++ b/plugins/governance/prerequesthookmcp_test.go @@ -0,0 +1,408 @@ +// This suite covers PreRequestHook's MCP include-tools stamping: how a caller-provided +// x-bf-mcp-include-tools list is pruned against the virtual key's tool grant, and when +// the grant itself is stamped onto the context for downstream injection. The rules under +// test, for both the normal request path and the large-payload branch: +// +// - A caller include-tools list is always pruned to the grant (it can only narrow, +// never expand), regardless of the DisableAutoToolInject setting. +// - When no caller list is present, the grant is stamped if auto-injection is enabled +// OR an include-clients filter is present — the latter because any explicit MCP +// filter opts the request into injection downstream even when auto-inject is off. +// - With no filters and auto-injection disabled, nothing is stamped and no injection +// occurs. +// - Deny-all is always an explicit empty list; an unset key means "no filtering" and +// would expose every available tool. +package governance + +import ( + "context" + "fmt" + "testing" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore" + configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const mcpTestVKValue = "sk-bf-mcp-test" + +// buildVKForMCPStamping returns an active VK with an openai provider config (so load +// balancing has a provider pool) and an explicit sentry MCP config granting the given +// tools. Passing nil yields a VK with no MCP configs at all, which is semantically +// different from passing an empty slice (an explicit deny-all config for the client). +func buildVKForMCPStamping(tools []string) *configstoreTables.TableVirtualKey { + vk := buildVirtualKeyWithProviders( + "vk-mcp-stamp", + mcpTestVKValue, + "mcp-stamp-vk", + []configstoreTables.TableVirtualKeyProviderConfig{ + buildProviderConfig("openai", []string{"*"}), + }, + ) + if tools != nil { + vk.MCPConfigs = []configstoreTables.TableVirtualKeyMCPConfig{ + { + MCPClient: configstoreTables.TableMCPClient{ClientID: "client-1", Name: "sentry"}, + ToolsToExecute: tools, + }, + } + } + return vk +} + +// newPluginForMCPStamping builds a governance plugin around a single VK with the +// given DisableAutoToolInject setting. +func newPluginForMCPStamping(t *testing.T, vk *configstoreTables.TableVirtualKey, disableAutoToolInject bool) *GovernancePlugin { + t.Helper() + logger := NewMockLogger() + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + VirtualKeys: []configstoreTables.TableVirtualKey{*vk}, + }, nil) + require.NoError(t, err) + + plugin, err := InitFromStore(context.Background(), &Config{ + IsVkMandatory: boolPtr(false), + DisableAutoToolInject: boolPtr(disableAutoToolInject), + }, logger, store, nil, nil, nil, nil) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, plugin.Cleanup()) }) + return plugin +} + +// newPreRequestCtx returns a ctx carrying the test VK plus optional caller MCP filters, +// mirroring what lib/ctx.go stamps from the x-bf-mcp-include-tools and +// x-bf-mcp-include-clients request headers. +func newPreRequestCtx(includeTools, includeClients []string) *schemas.BifrostContext { + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyVirtualKey, mcpTestVKValue) + if includeTools != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeTools, includeTools) + } + if includeClients != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeClients, includeClients) + } + return ctx +} + +// newChatRequest returns a chat request with the provider already resolved, so +// PreRequestHook's load-balancing step is a no-op and only MCP stamping is exercised. +func newChatRequest() *schemas.BifrostRequest { + return &schemas.BifrostRequest{ + RequestType: schemas.ChatCompletionRequest, + ChatRequest: &schemas.BifrostChatRequest{Provider: schemas.OpenAI, Model: "gpt-4o"}, + } +} + +// stampedIncludeTools runs PreRequestHook and returns the resulting include-tools ctx +// value (nil when nothing was stamped). +func stampedIncludeTools(t *testing.T, p *GovernancePlugin, ctx *schemas.BifrostContext, req *schemas.BifrostRequest) []string { + t.Helper() + require.NoError(t, p.PreRequestHook(ctx, req)) + value := ctx.Value(schemas.MCPContextKeyIncludeTools) + if value == nil { + return nil + } + tools, ok := value.([]string) + require.True(t, ok, "include-tools ctx value should be a []string") + return tools +} + +// ============================================================================ +// Decision matrix: include-tools present/absent × include-clients present/absent +// × DisableAutoToolInject on/off +// ============================================================================ + +// Baseline auto-injection: no caller filters, auto-injection enabled. Governance stamps +// the key's full tool grant (every granted tool, client-prefixed) onto the context, and +// downstream injection attaches exactly these tools to the outgoing request. +func TestPreRequestHookMCP_AutoInjectOn_NoFilters_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx := newPreRequestCtx(nil, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools) +} + +// The caller filters by client only and sends no include-tools list. Governance still +// stamps the full grant: the grant is the tool-level ceiling, and the client-level +// narrowing is applied downstream where the include-clients filter intersects with it. +func TestPreRequestHookMCP_AutoInjectOn_IncludeClientsOnly_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools) +} + +// The caller sends an include-tools list with one granted and one ungranted tool. +// Governance prunes the list in place instead of replacing it: the granted entry +// survives, the ungranted entry is dropped, and the full grant is NOT stamped over +// the caller's narrower selection. +func TestPreRequestHookMCP_AutoInjectOn_IncludeToolsPresent_Prunes(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx := newPreRequestCtx([]string{"sentry-tool_a", "sentry-tool_c"}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a"}, tools, + "granted entry survives, ungranted entry is pruned, grant is not re-expanded") +} + +// With auto-injection disabled and no caller filters, the request has not opted into +// MCP tools in any way: governance must leave the include-tools key unset so downstream +// injection sees neither filter and skips entirely. This is the contract of the +// DisableAutoToolInject setting. +func TestPreRequestHookMCP_AutoInjectOff_NoFilters_StampsNothing(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx := newPreRequestCtx(nil, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Nil(t, tools, "no filters + auto-inject disabled must leave include-tools unset") +} + +// With auto-injection disabled, an include-clients filter still opts the request into +// injection downstream (any explicit MCP filter counts as opt-in). The grant must +// therefore be stamped even though auto-inject is off — left unset, the request would +// be injected with every tool of the included client, bypassing the key's grant. +func TestPreRequestHookMCP_AutoInjectOff_IncludeClientsOnly_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools, + "include-clients triggers injection downstream, so the VK ceiling must be stamped") +} + +// Pruning is independent of the auto-inject setting: a caller include-tools list is +// narrowed against the grant even when auto-injection is disabled, because the list +// itself opts the request into injection downstream. +func TestPreRequestHookMCP_AutoInjectOff_IncludeToolsPresent_Prunes(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx := newPreRequestCtx([]string{"sentry-tool_b", "sentry-tool_c"}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_b"}, tools) +} + +// When the caller sends both filters, the pruned include-tools list is the final value: +// the presence of include-clients must not cause the full grant to overwrite the +// caller's narrower selection. Verified under both DisableAutoToolInject values. +func TestPreRequestHookMCP_BothFilters_PrunedListWins(t *testing.T) { + for _, disabled := range []bool{false, true} { + t.Run(fmt.Sprintf("disableAutoToolInject=%v", disabled), func(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), disabled) + ctx := newPreRequestCtx([]string{"sentry-tool_a"}, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a"}, tools, + "pruned caller list must not be overwritten by the grant") + }) + } +} + +// ============================================================================ +// Grant-shape edge cases +// ============================================================================ + +// A key with no MCP configs at all has an empty effective grant. When include-clients +// opts the request into injection, governance must stamp an explicit empty list +// (deny-all) rather than leave the key unset — an unset key reads downstream as +// "no filtering" and would inject every available tool of the included client. +func TestPreRequestHookMCP_NoGrants_IncludeClients_StampsDenyAll(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping(nil), true) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + require.NotNil(t, tools, "deny-all must be an empty list, not an unset key") + assert.Empty(t, tools) +} + +// Same deny-all outcome as the no-configs case, but through a different path in +// computeMCPIncludeTools: here the key has an explicit MCP config for the client whose +// tools list is empty, exercising the ToolsToExecute.IsEmpty guard that skips the +// client without emitting any grant entries. +func TestPreRequestHookMCP_ExplicitEmptyGrant_IncludeClients_StampsDenyAll(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{}), true) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + require.NotNil(t, tools, "deny-all must be an empty list, not an unset key") + assert.Empty(t, tools) +} + +// An unrestricted ("*") grant is stamped as the client-scoped wildcard "sentry-*", +// which downstream filtering reads as "all tools of this client" — the grant never +// expands to other clients. +func TestPreRequestHookMCP_UnrestrictedGrant_StampsWildcard(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"*"}), false) + ctx := newPreRequestCtx(nil, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-*"}, tools) +} + +// The caller requests a tool of a client the key has no grants for. Pruning drops it +// and stamps an explicit empty list, so the request cannot reach another client's +// tools just by naming them in the header. +func TestPreRequestHookMCP_IncludeToolsForUngrantedClient_PrunesToDenyAll(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a"}), false) + ctx := newPreRequestCtx([]string{"github-list_repos"}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + require.NotNil(t, tools) + assert.Empty(t, tools) +} + +// An empty x-bf-mcp-include-tools header value reaches ctx as [""] (see lib/ctx.go). +// Pruning drops the empty entry and stamps deny-all: a caller can suppress tool +// injection for a single request, but cannot gain access through the empty value. +func TestPreRequestHookMCP_EmptyIncludeToolsHeader_DenyAll(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx := newPreRequestCtx([]string{""}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + require.NotNil(t, tools) + assert.Empty(t, tools) +} + +// ============================================================================ +// Paths that must NOT stamp +// ============================================================================ + +// Without a virtual key on ctx (and no routing rules configured), PreRequestHook +// returns before any MCP handling: there is no grant to enforce, so caller filters +// pass through untouched for downstream layers to interpret. +func TestPreRequestHookMCP_NoVirtualKey_NoStamping(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a"}), false) + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.MCPContextKeyIncludeClients, []string{"sentry"}) + + require.NoError(t, p.PreRequestHook(ctx, newChatRequest())) + assert.Nil(t, ctx.Value(schemas.MCPContextKeyIncludeTools)) +} + +// An inactive key is treated as absent: PreRequestHook returns before MCP handling +// and stamps nothing. +func TestPreRequestHookMCP_InactiveVK_NoStamping(t *testing.T) { + vk := buildVKForMCPStamping([]string{"tool_a"}) + inactive := false + vk.IsActive = &inactive + p := newPluginForMCPStamping(t, vk, false) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + + require.NoError(t, p.PreRequestHook(ctx, newChatRequest())) + assert.Nil(t, ctx.Value(schemas.MCPContextKeyIncludeTools)) +} + +// Passthrough request types skip governance entirely, including MCP stamping. +func TestPreRequestHookMCP_PassthroughRequest_NoStamping(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a"}), false) + ctx := newPreRequestCtx(nil, []string{"sentry"}) + req := &schemas.BifrostRequest{RequestType: schemas.PassthroughRequest} + + require.NoError(t, p.PreRequestHook(ctx, req)) + assert.Nil(t, ctx.Value(schemas.MCPContextKeyIncludeTools)) +} + +// ============================================================================ +// Large-payload branch (runPreRequestRouting) applies the same stamping rules +// ============================================================================ + +// newLargePayloadCtx returns a ctx that routes PreRequestHook through its large-payload +// branch: the body streams to the provider unparsed, so the model comes from +// LargePayloadMetadata instead of the request. The metadata pointer is returned so +// tests can assert the routed model is propagated (the streaming body rewriter +// consumes metadata.Model when rewriting the body prefix). +func newLargePayloadCtx(includeTools, includeClients []string) (*schemas.BifrostContext, *schemas.LargePayloadMetadata) { + ctx := newPreRequestCtx(includeTools, includeClients) + metadata := &schemas.LargePayloadMetadata{Model: "openai/gpt-4o"} + ctx.SetValue(schemas.BifrostContextKeyLargePayloadMetadata, metadata) + return ctx, metadata +} + +// Large-payload counterpart of the include-clients opt-in case: with auto-injection +// disabled, the grant must still be stamped because include-clients triggers injection +// downstream. The provider-prefixed model must survive routing unchanged. +func TestPreRequestHookMCP_LargePayload_AutoInjectOff_IncludeClients_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx, metadata := newLargePayloadCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +} + +// Large-payload counterpart of the baseline auto-injection case: no caller filters, +// auto-injection enabled, full grant stamped. +func TestPreRequestHookMCP_LargePayload_AutoInjectOn_NoFilters_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx, metadata := newLargePayloadCtx(nil, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +} + +// Large-payload counterpart of the include-clients-only case with auto-injection +// enabled: the full grant is stamped as the tool-level ceiling. +func TestPreRequestHookMCP_LargePayload_AutoInjectOn_IncludeClientsOnly_StampsGrant(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx, metadata := newLargePayloadCtx(nil, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a", "sentry-tool_b"}, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +} + +// Large-payload pruning with auto-injection disabled: the caller's include-tools list +// is narrowed against the grant — the toggle never disables grant enforcement. +func TestPreRequestHookMCP_LargePayload_AutoInjectOff_IncludeToolsPresent_Prunes(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx, metadata := newLargePayloadCtx([]string{"sentry-tool_a", "sentry-tool_c"}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a"}, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +} + +// Large-payload pruning with auto-injection enabled: granted entries survive, +// ungranted entries are dropped, and the grant is not re-stamped over the caller's +// narrower selection. +func TestPreRequestHookMCP_LargePayload_IncludeToolsPresent_Prunes(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), false) + ctx, metadata := newLargePayloadCtx([]string{"sentry-tool_a", "sentry-tool_c"}, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a"}, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +} + +// Large-payload counterpart of the both-filters case: the pruned include-tools list is +// the final value and is not overwritten by the grant. Verified under both +// DisableAutoToolInject values. +func TestPreRequestHookMCP_LargePayload_BothFilters_PrunedListWins(t *testing.T) { + for _, disabled := range []bool{false, true} { + t.Run(fmt.Sprintf("disableAutoToolInject=%v", disabled), func(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), disabled) + ctx, metadata := newLargePayloadCtx([]string{"sentry-tool_a"}, []string{"sentry"}) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Equal(t, []string{"sentry-tool_a"}, tools, + "pruned caller list must not be overwritten by the grant") + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") + }) + } +} + +// Large-payload counterpart of the disabled-toggle baseline: no filters and +// auto-injection disabled leaves the include-tools key unset, so downstream injection +// is skipped entirely. +func TestPreRequestHookMCP_LargePayload_AutoInjectOff_NoFilters_StampsNothing(t *testing.T) { + p := newPluginForMCPStamping(t, buildVKForMCPStamping([]string{"tool_a", "tool_b"}), true) + ctx, metadata := newLargePayloadCtx(nil, nil) + + tools := stampedIncludeTools(t, p, ctx, newChatRequest()) + assert.Nil(t, tools) + assert.Equal(t, "openai/gpt-4o", metadata.Model, "provider-prefixed model must survive routing unchanged") +}