extra header forwarding for mcp tools

[akshaydeo]

Summary

Per-request extra headers set via BifrostContextKeyMCPExtraHeaders in a PreMCPHook were not reaching the upstream MCP server for health-check probes (ping and tools/list). The mcp-go client drops request.Header for these internally-generated calls, so headers injected at the CallToolRequest level were silently lost. This PR centralizes all per-request extra header injection onto the transport layer via WithHTTPHeaderFunc / WithHeaderFunc, ensuring headers flow on every outgoing message — including ping, tools/list, and tools/call — while still being filtered by AllowedExtraHeaders.

Changes

• Registered a headerFunc on the StreamableHTTP and SSE transports (in createHTTPConnection, createSSEConnection, and AcquireClientConn) that reads BifrostContextKeyMCPExtraHeaders from the request context and injects only allowlisted headers per MCPClientConfig.AllowedExtraHeaders. This replaces the previous per-call CallToolRequest.Header approach.
• Removed credStore.RequestHeaders calls and CallToolRequest.Header assignments from executeToolInternal (tool manager) and callMCPTool (Starlark code mode), since header injection is now handled uniformly by the transport.
• Fixed runListToolsWithHooks and runPingWithHooks to pass gateCtx (the child context that carries PreMCPHook writes) instead of the outer ctx to the wire calls, so transport headerFunc can see values written during the plugin gate.
• Relaxed ExtractFilteredExtras to accept a plain context.Context instead of *schemas.BifrostContext, enabling it to be called from the transport headerFunc closure.
• Added end-to-end wire-level tests (extraheaders_test.go) using a real httptest streamable-HTTP server that records inbound headers per JSON-RPC method, covering: allowlisted headers reaching ping, tools/list, and tools/call; and non-allowlisted headers being filtered on all requests.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go test ./core/internal/mcptests/... -run TestExtraHeaders -v
go test ./...

The three new tests validate:

• TestExtraHeadersHealthCheckPingReachWire — allowlisted header appears on ping probes; non-allowlisted header never appears on any request.
• TestExtraHeadersHealthCheckListToolsReachWire — same guarantee for tools/list health-check probes when ping is unavailable.
• TestExtraHeadersToolCallReachWire — allowlisted header appears on a normal tools/call; non-allowlisted header is filtered.

Breaking changes

• No

The CallToolRequest.Header field is no longer populated by Bifrost internals, but this is an internal implementation detail with no public API impact. Header forwarding behavior is preserved (and extended to health-check probes).

Security considerations

AllowedExtraHeaders filtering is now enforced at the transport layer for all outgoing MCP requests. Non-allowlisted headers set by plugins are dropped before reaching the wire on every request type, including health-check probes that previously bypassed the per-call header path entirely.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d90363f4-ab31-4754-a7c9-c088f77354df

📥 Commits

Reviewing files that changed from the base of the PR and between 5f40c10 and 411f221.

📒 Files selected for processing (6)

• core/internal/mcptests/extraheaders_test.go
• core/mcp/clientmanager.go
• core/mcp/codemode/starlark/executecode.go
• core/mcp/pluginpipeline.go
• core/mcp/toolmanager.go
• core/mcp/utils/utils.go

🚧 Files skipped from review as they are similar to previous changes (5)

• core/mcp/toolmanager.go
• core/mcp/clientmanager.go
• core/mcp/pluginpipeline.go
• core/mcp/codemode/starlark/executecode.go
• core/mcp/utils/utils.go

---

📝 Walkthrough

Summary by CodeRabbit

• Tests

  • Added end-to-end wire-level regression tests to verify allowlisted MCP extra headers are forwarded correctly and denied headers never reach upstream requests.
  • Covered both HTTP (streamable) and SSE for health checks (ping/list-tools) and tool calls, including fallback behavior when ping isn’t available.

• Refactor

  • Unified per-request extra-header injection so filtering/enforcement is applied consistently across HTTP, SSE, and all MCP operations using the same transport path.

Walkthrough

Extra MCP headers are shifted from per-call CallToolRequest fields into transport-level per-request header functions. ExtractFilteredExtras is generalized to context.Context, the plugin pipeline is fixed to pass gateCtx to wire calls, and end-to-end tests verify allowlisted headers reach the upstream server while denied headers are filtered.

Changes

MCP Extra Headers via Transport Injection

Layer / File(s)	Summary
ExtractFilteredExtras generalized to context.Context core/mcp/utils/utils.go	Adds context import and changes ExtractFilteredExtras parameter from *schemas.BifrostContext to context.Context, allowing the function to be called from any transport-carrying site.
Per-request header functions in HTTP/SSE transports core/mcp/clientmanager.go	Adds WithHTTPHeaderFunc to the ephemeral AcquireClientConn HTTP transport, replaces static WithHTTPHeaders with a combined static+dynamic WithHTTPHeaderFunc in createHTTPConnection, and replaces static WithHeaders with a WithHeaderFunc in createSSEConnection, all using ExtractFilteredExtras.
Gate context propagated to wire calls in plugin pipeline core/mcp/pluginpipeline.go	Switches retrieveExternalToolsDetailed in runListToolsWithHooks and conn.Ping in runPingWithHooks from the outer ctx to gateCtx, so PreMCPHook extra header writes are visible to transport header functions.
Remove per-call header construction from tool and Starlark paths core/mcp/toolmanager.go, core/mcp/codemode/starlark/executecode.go	Removes credStore.RequestHeaders calls and the Header: reqHeaders field from mcp.CallToolRequest construction in both executeToolInternal and callMCPTool, delegating header injection entirely to the transport layer.
End-to-end wire-level extra header tests core/internal/mcptests/extraheaders_test.go	Adds a thread-safe recording HTTP/SSE test server, an inject plugin setting allowlisted and denied headers, client configurations, and six tests (three HTTP, three SSE) asserting allowlisted headers reach the upstream for ping, tools/list fallback, and tools/call, while denied headers never appear.

Sequence Diagram(s)

sequenceDiagram
  participant Plugin as extraHeaderInjectPlugin
  participant Pipeline as pluginpipeline
  participant Transport as clientmanager (HeaderFunc)
  participant ExtractFilteredExtras as utils.ExtractFilteredExtras
  participant UpstreamMCP as Upstream MCP Server

  Plugin->>Pipeline: PreMCPHook sets MCPExtraHeaders in gateCtx
  Pipeline->>Transport: wire call with gateCtx (ping / tools/list / tools/call)
  Transport->>ExtractFilteredExtras: gateCtx + AllowedExtraHeaders
  ExtractFilteredExtras-->>Transport: filtered http.Header (allowlisted only)
  Transport->>UpstreamMCP: HTTP request with allowlisted header injected
  UpstreamMCP-->>Transport: response

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• maximhq/bifrost#3656: Introduces MCPCredentialStore and the original ExtractFilteredExtras for per-request header filtering, which this PR extends by switching to context.Context and wiring into transport-level header funcs.
• maximhq/bifrost#3702: Also modifies clientmanager.go and removes per-request CallToolRequest header construction in executecode.go, overlapping directly with this PR's centralization of header injection.
• maximhq/bifrost#3794: Shifts Starlark tool calls to RunWithPluginPipeline, which is the same gate-context pipeline this PR fixes to pass gateCtx to wire calls.

Suggested reviewers

• danpiths
• roroghost17

🐇 A header lost, a header found,
Through transport layers, tightly bound.
No more per-call reqHeaders set—
The gateCtx carries what you get!
Allowlisted headers hop the wire,
Denied ones never sneak through higher. 🎉

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'extra header forwarding for mcp tools' clearly and concisely summarizes the main change: implementing extra header forwarding for MCP tool calls.
Description check	✅ Passed	The PR description is comprehensive and well-structured, covering all key template sections: summary, detailed changes, type of change, affected areas, testing instructions, breaking changes, security considerations, and mostly complete checklist.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-20-extra_header_forwarding_for_mcp_tools

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[akshaydeo]

• adds mcp cleanup support for bedrock models #4573
• extra header forwarding for mcp tools #4572 👈 (View in Graphite)
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the change moves header injection from per-call to transport-level and fixes missing context propagation for health-check probes; no regressions are introduced on the tool-call path.

The two-part fix (transport headerFunc registration + gateCtx propagation to wire calls) is mechanically correct: PreHook writes are local to gateCtx, so the previous code that passed the outer ctx to conn.Ping and retrieveExternalToolsDetailed was provably unable to see those writes. The new code correctly passes gateCtx. The tool-call path in exec.go was already using the correct context and continues to do so. AllowedExtraHeaders filtering is preserved at the transport layer for all request types. Six wire-level tests exercise both transports across all three request types and verify the deny-list is enforced everywhere.

No files require special attention; all changes are straightforward and well-tested.

Important Files Changed

Filename	Overview
core/internal/mcptests/extraheaders_test.go	New wire-level tests covering allowlisted/denied headers on HTTP and SSE transports across ping, tools/list, and tools/call
core/mcp/clientmanager.go	Registers a per-request headerFunc on StreamableHTTP and SSE transports in createHTTPConnection, createSSEConnection, and AcquireClientConn to inject allowlisted extra headers on every outgoing message
core/mcp/pluginpipeline.go	Fixes runListToolsWithHooks and runPingWithHooks to pass gateCtx (carrying PreMCPHook writes) instead of outer ctx to the wire call
core/mcp/toolmanager.go	Removes per-call CallToolRequest.Header assignment; header injection now handled uniformly by transport headerFunc
core/mcp/codemode/starlark/executecode.go	Removes credStore.RequestHeaders call and CallToolRequest.Header assignment from callMCPTool; header forwarding is now at the transport layer
core/mcp/utils/utils.go	ExtractFilteredExtras signature relaxed to accept plain context.Context; stale doc comment still references the removed RequestHeaders code path

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Plugin as PreMCPHook Plugin
    participant Gate as RunWithPluginPipeline
    participant Wire as Wire Op (ping/list/call)
    participant HeaderFunc as Transport headerFunc
    participant MCP as Upstream MCP Server

    Plugin->>Gate: SetValue(BifrostContextKeyMCPExtraHeaders, headers) on gateCtx
    Gate->>Wire: op(preReq) — passes gateCtx to wire call
    Wire->>HeaderFunc: headerFunc(gateCtx) called per outgoing request
    HeaderFunc->>HeaderFunc: ExtractFilteredExtras(gateCtx, config) → filter by AllowedExtraHeaders
    HeaderFunc-->>Wire: map[string]string (allowlisted headers only)
    Wire->>MCP: HTTP request with injected headers (ping / tools/list / tools/call)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Plugin as PreMCPHook Plugin
    participant Gate as RunWithPluginPipeline
    participant Wire as Wire Op (ping/list/call)
    participant HeaderFunc as Transport headerFunc
    participant MCP as Upstream MCP Server

    Plugin->>Gate: SetValue(BifrostContextKeyMCPExtraHeaders, headers) on gateCtx
    Gate->>Wire: op(preReq) — passes gateCtx to wire call
    Wire->>HeaderFunc: headerFunc(gateCtx) called per outgoing request
    HeaderFunc->>HeaderFunc: ExtractFilteredExtras(gateCtx, config) → filter by AllowedExtraHeaders
    HeaderFunc-->>Wire: map[string]string (allowlisted headers only)
    Wire->>MCP: HTTP request with injected headers (ping / tools/list / tools/call)

Loading

_{Reviews (2): Last reviewed commit: "extra header forwarding for mcp tools" | Re-trigger Greptile}

[akshaydeo]

Merge activity

• Jun 20, 12:05 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 20, 12:06 PM UTC: @akshaydeo merged this pull request with Graphite.