feat: add MCP tool execution governance with virtual key allow-lists#1940
Conversation
This was referenced Mar 5, 2026
Collaborator
Author
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughSwitch MCP-related APIs from Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Transport
participant Governance
participant MCPManager
participant ToolRuntime
Client->>Transport: HTTP request
Transport->>Governance: HTTPTransportPreHook (may inject header)
Governance->>Transport: header decision (respect DisableAutoToolInject / VK)
Transport->>MCPManager: GetAvailableMCPTools(BifrostContext)
MCPManager->>MCPManager: Append MCPAddedTools to BifrostContext
MCPManager->>ToolRuntime: ExecuteToolCall (per-tool, validated)
ToolRuntime-->>MCPManager: tool result
MCPManager-->>Transport: tools/results
Transport-->>Client: response (or 403 if DecisionMCPToolBlocked)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
Pratham-Mishra04
force-pushed
the
03-05-feat_add_option_to_disable_automatic_mcp_tool_injection_per_request
branch
from
b515efc to
1ee6d71
Compare
Pratham-Mishra04
force-pushed
the
03-05-feat_virtual_key_mcp_configs_now_act_as_an_execution-time_allow-list
branch
2 times, most recently
from
8af3ad3 to
dafd4b6
Compare
Pratham-Mishra04
force-pushed
the
03-05-feat_add_option_to_disable_automatic_mcp_tool_injection_per_request
branch
from
1ee6d71 to
57671fd
Compare
Pratham-Mishra04
force-pushed
the
03-05-feat_virtual_key_mcp_configs_now_act_as_an_execution-time_allow-list
branch
from
dafd4b6 to
df4061c
Compare
Pratham-Mishra04
force-pushed
the
03-05-feat_add_option_to_disable_automatic_mcp_tool_injection_per_request
branch
from
57671fd to
dae2e01
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/mcp/toolmanager.go`:
- Around line 197-203: In GetAvailableTools, avoid dereferencing
tool.Function.Name before nil checks and make tool suppression deterministic:
first check tool.Function != nil && tool.Function.Name != "" and only then call
schemas.AppendToContextList(ctx, schemas.BifrostContextKeyMCPAddedTools,
tool.Function.Name); next, replace the current global gating on
includeCodeModeTools with a per-tool check (e.g., skip adding the tool only if
!includeCodeModeTools && tool.IsCodeMode or equivalent flag on the tool) and
then use seenToolNames[tool.Function.Name] to deduplicate and append to
availableTools; update references to clientTools, includeCodeModeTools,
seenToolNames, availableTools, and schemas.AppendToContextList accordingly.
In `@plugins/governance/main.go`:
- Around line 1202-1228: The current blind single-tool check using
p.store.GetVirtualKey(virtualKeyValue) lets execution proceed (return req, nil,
nil) if the virtual key is missing/inactive even though
evaluateGovernanceRequest already ran, which creates a race that can allow
unauthorized tool execution; change this to fail-closed: when GetVirtualKey
returns !ok || vk == nil || !vk.IsActive, return an MCPPluginShortCircuit
BifrostError (DecisionMCPToolBlocked, 403) denying execution, and keep the
existing handling for len(vk.MCPConfigs) == 0 and isMCPToolAllowedByVK(vk,
toolName) as-is so empty configs or disallowed tools also return the same
short-circuit error.
In `@transports/changelog.md`:
- Around line 7-8: Remove the duplicated changelog entries by keeping only one
instance of the "feat: add option to disable automatic MCP tool injection per
request" entry and one instance of the "fix: preserve original audio filename in
transcription requests" entry; locate the duplicate lines in
transports/changelog.md (the exact strings above) and delete the repeated
occurrences so each change appears only once in the file.
In `@ui/app/workspace/config/views/mcpView.tsx`:
- Line 185: Update the explanatory sentence in the MCP view text (the JSX string
in mcpView.tsx where the paragraph about header-based tool inclusion is
rendered) to explicitly state that using the x-bf-mcp-include-tools header does
not bypass Virtual Key (VK) MCP execution-time allow-lists; tool injection via
the header is still subject to VK MCP allow-list checks and may be blocked by
them. Locate the paragraph around the existing wording "When enabled, MCP tools
are not automatically included..." and replace it with a clarified sentence that
mentions both the header and the VK MCP allow-lists (keep the existing header
code element <code className="text-xs">x-bf-mcp-include-tools</code> intact).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e14b7c87-859a-40e2-95da-57c4be9301a8
📒 Files selected for processing (15)
core/bifrost.gocore/changelog.mdcore/mcp/interface.gocore/mcp/mcp.gocore/mcp/toolmanager.gocore/schemas/bifrost.godocs/features/governance/mcp-tools.mdxplugins/governance/changelog.mdplugins/governance/main.goplugins/governance/resolver.gotransports/bifrost-http/server/plugins.gotransports/bifrost-http/server/server.gotransports/changelog.mdui/app/workspace/config/views/mcpView.tsxui/components/sidebar.tsx
Pratham-Mishra04
force-pushed
the
03-05-feat_add_option_to_disable_automatic_mcp_tool_injection_per_request
branch
from
dae2e01 to
fc9336c
Compare
Pratham-Mishra04
force-pushed
the
03-05-feat_virtual_key_mcp_configs_now_act_as_an_execution-time_allow-list
branch
from
df4061c to
9255034
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
core/mcp/toolmanager.go (1)
193-203:⚠️ Potential issue | 🟠 MajorFix nondeterministic tool availability and incorrect
MCPAddedToolstracking inGetAvailableTools.At Line 200, regular tool append is globally gated by
includeCodeModeTools, so once any code-mode client is encountered, regular tools may be dropped depending on map iteration order.
At Line 199,BifrostContextKeyMCPAddedToolsis updated before confirming the tool was actually added. Also, code-mode tools appended at Lines 212-217 are not tracked.💡 Proposed fix
for clientName, clientTools := range availableToolsPerClient { client := m.clientManager.GetClientByName(clientName) if client == nil { m.logger.Warn("%s Client %s not found, skipping", MCPLogPrefix, clientName) continue } - if client.ExecutionConfig.IsCodeModeClient { + isCodeModeClient := client.ExecutionConfig.IsCodeModeClient + if isCodeModeClient { includeCodeModeTools = true + // Code mode tools are injected from m.codeMode.GetTools() below. + continue } // Add tools from this client, checking for duplicates for _, tool := range clientTools { - if tool.Function != nil && tool.Function.Name != "" && !seenToolNames[tool.Function.Name] { - schemas.AppendToContextList(ctx, schemas.BifrostContextKeyMCPAddedTools, tool.Function.Name) - if !includeCodeModeTools { - availableTools = append(availableTools, tool) - seenToolNames[tool.Function.Name] = true - } - } + if tool.Function == nil || tool.Function.Name == "" { + continue + } + if seenToolNames[tool.Function.Name] { + continue + } + availableTools = append(availableTools, tool) + seenToolNames[tool.Function.Name] = true + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyMCPAddedTools, tool.Function.Name) } } @@ if includeCodeModeTools && m.codeMode != nil { codeModeTools := m.codeMode.GetTools() // Add code mode tools, checking for duplicates for _, tool := range codeModeTools { if tool.Function != nil && tool.Function.Name != "" { if !seenToolNames[tool.Function.Name] { availableTools = append(availableTools, tool) seenToolNames[tool.Function.Name] = true + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyMCPAddedTools, tool.Function.Name) } } } }As per coding guidelines:
core/mcp/*.go: “MCP tool access filtering follows a 4-level hierarchy: Global filter → Client-level filter → Tool-level filter → Per-request filter (HTTP headers). All four levels must agree for a tool to be available.”Also applies to: 212-217
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/mcp/toolmanager.go` around lines 193 - 203, The GetAvailableTools loop currently gates adding regular tools via a global includeCodeModeTools flag and updates schemas.BifrostContextKeyMCPAddedTools before confirming addition, which causes nondeterministic drops and missed tracking; update the logic inside the clientTools iteration (referencing client.ExecutionConfig.IsCodeModeClient, includeCodeModeTools, seenToolNames, availableTools, and schemas.AppendToContextList) so that: evaluate per-tool whether code-mode allows adding it (do not let a single client set includeCodeModeTools globally to veto other tools), only call schemas.AppendToContextList when the tool is actually appended to availableTools or otherwise accepted, mark seenToolNames only upon successful addition, and ensure the same tracking is applied for code-mode tool additions (the code-path that appends code-mode tools must also update seenToolNames and MCPAddedTools).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@plugins/governance/main.go`:
- Around line 405-407: The header presence check currently uses a case-sensitive
map lookup on req.Headers (e.g., _, headerPresent :=
req.Headers["x-bf-mcp-include-tools"]) which can miss differently-cased incoming
headers and cause unintended injection; change the logic in the handler that
decides whether to call p.addMCPIncludeTools (and the similar check around lines
412-415) to perform a case-insensitive lookup — either normalize keys with
http.CanonicalHeaderKey or iterate req.Headers and use strings.EqualFold to
detect "x-bf-mcp-include-tools" — and only call addMCPIncludeTools(nil,
virtualKey) when no matching header is found ignoring case.
---
Duplicate comments:
In `@core/mcp/toolmanager.go`:
- Around line 193-203: The GetAvailableTools loop currently gates adding regular
tools via a global includeCodeModeTools flag and updates
schemas.BifrostContextKeyMCPAddedTools before confirming addition, which causes
nondeterministic drops and missed tracking; update the logic inside the
clientTools iteration (referencing client.ExecutionConfig.IsCodeModeClient,
includeCodeModeTools, seenToolNames, availableTools, and
schemas.AppendToContextList) so that: evaluate per-tool whether code-mode allows
adding it (do not let a single client set includeCodeModeTools globally to veto
other tools), only call schemas.AppendToContextList when the tool is actually
appended to availableTools or otherwise accepted, mark seenToolNames only upon
successful addition, and ensure the same tracking is applied for code-mode tool
additions (the code-path that appends code-mode tools must also update
seenToolNames and MCPAddedTools).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5847046b-32f7-41bf-a6ab-c734e9f6a6a4
📒 Files selected for processing (15)
core/bifrost.gocore/changelog.mdcore/mcp/interface.gocore/mcp/mcp.gocore/mcp/toolmanager.gocore/schemas/bifrost.godocs/features/governance/mcp-tools.mdxplugins/governance/changelog.mdplugins/governance/main.goplugins/governance/resolver.gotransports/bifrost-http/server/plugins.gotransports/bifrost-http/server/server.gotransports/changelog.mdui/app/workspace/config/views/mcpView.tsxui/components/sidebar.tsx
🚧 Files skipped from review as they are similar to previous changes (6)
- ui/app/workspace/config/views/mcpView.tsx
- core/mcp/mcp.go
- plugins/governance/resolver.go
- core/bifrost.go
- core/changelog.md
- plugins/governance/changelog.md
Comment on lines
+405
to
+407
| _, headerPresent := req.Headers["x-bf-mcp-include-tools"] | ||
| if !headerPresent { | ||
| headers, err := p.addMCPIncludeTools(nil, virtualKey) |
Contributor
There was a problem hiding this comment.
Case-sensitive header check can break the “caller already provided header” contract.
req.Headers["x-bf-mcp-include-tools"] is case-sensitive, so X-BF-MCP-Include-Tools won’t be detected and auto-injection may run unexpectedly. Use case-insensitive key detection before injecting.
💡 Suggested fix
- _, headerPresent := req.Headers["x-bf-mcp-include-tools"]
+ headerPresent := false
+ for h := range req.Headers {
+ if strings.EqualFold(h, "x-bf-mcp-include-tools") {
+ headerPresent = true
+ break
+ }
+ }
if !headerPresent {
+ if req.Headers == nil {
+ req.Headers = make(map[string]string)
+ }
headers, err := p.addMCPIncludeTools(nil, virtualKey)
if err != nil {
p.logger.Error("failed to add MCP include tools: %v", err)
return nil, nil
}Also applies to: 412-415
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@plugins/governance/main.go` around lines 405 - 407, The header presence check
currently uses a case-sensitive map lookup on req.Headers (e.g., _,
headerPresent := req.Headers["x-bf-mcp-include-tools"]) which can miss
differently-cased incoming headers and cause unintended injection; change the
logic in the handler that decides whether to call p.addMCPIncludeTools (and the
similar check around lines 412-415) to perform a case-insensitive lookup —
either normalize keys with http.CanonicalHeaderKey or iterate req.Headers and
use strings.EqualFold to detect "x-bf-mcp-include-tools" — and only call
addMCPIncludeTools(nil, virtualKey) when no matching header is found ignoring
case.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR implements MCP tool governance by enforcing virtual key MCP configurations as an execution-time allow-list. When virtual keys have empty MCPConfigs, all MCP tools are denied. When non-empty, each tool is validated against the configured allow-list at both inference time and MCP tool execution.
Changes
*schemas.BifrostContextinstead ofcontext.Contextto enable tool trackingBifrostContextKeyMCPAddedToolscontext key to track which MCP tools are added to requestsPreMCPHookandevaluateGovernanceRequestDisableAutoToolInjectconfiguration option that respects the toggle and skips auto-injection when headers are already set by callersDecisionMCPToolBlockedfor MCP tool governance violationsType of change
Affected areas
How to test
Test MCP tool governance with virtual keys:
New configuration options:
disable_auto_tool_inject: Boolean flag to disable automatic MCP tool injectionMCPConfigs: Array of MCP client configurations that act as allow-listsScreenshots/Recordings
UI changes include updated MCP configuration view with clearer descriptions for the disable auto tool injection toggle and improved sidebar navigation labels.
Breaking changes
Impact: MCP-related function signatures now require
*schemas.BifrostContextinstead ofcontext.Context. Virtual keys with empty MCPConfigs will now deny all MCP tools by default.Migration: Update any custom MCP integrations to use the new context parameter type. Configure MCPConfigs on virtual keys that need MCP tool access.
Related issues
Implements MCP tool governance and execution-time validation for virtual key configurations.
Security considerations
Checklist
docs/contributing/README.mdand followed the guidelines