Skip to content

add forward_headers passthrough to remote::model-context-protocol#5257

Draft
skamenan7 wants to merge 4 commits intollamastack:mainfrom
skamenan7:feat/5152-mcp-tool-passthrough
Draft

add forward_headers passthrough to remote::model-context-protocol#5257
skamenan7 wants to merge 4 commits intollamastack:mainfrom
skamenan7:feat/5152-mcp-tool-passthrough

Conversation

@skamenan7
Copy link
Contributor

Note: this branch is stacked on #5134 (inference + safety passthrough, not yet merged). The only new code here is in the last commit — providers/remote/tool_runtime/model_context_protocol/ and tests/integration/tool_runtime/test_passthrough_mcp.py. For reviews for this PR please focus there; the rest is carried from #5134 pending merge.

Adds forward_headers and extra_blocked_headers to MCPProviderConfig, wiring per-request header forwarding into list_runtime_tools and invoke_tool. This lets deployers map keys from X-LlamaStack-Provider-Data to outbound HTTP headers so request-scoped auth tokens (MaaS API keys, tenant IDs, etc.) reach the downstream MCP server without the caller passing them via authorization= on every tool call.

Follows the same forward_headers pattern introduced for inference and safety passthrough in #5134. Authorization-mapped values are split out and passed via the authorization= param — prepare_mcp_headers() rejects Authorization in the headers dict directly, so it flows through the dedicated param instead.

What changed

  • MCPProviderConfig: added forward_headers: dict[str, str] | None and extra_blocked_headers: list[str] with config-time validation via validate_forward_headers_config() from providers/utils/forward_headers.py
  • MCPProviderDataValidator: added model_config = ConfigDict(extra="allow") so deployer-defined keys survive Pydantic parsing (key names are operator-configured at deploy time and can't be declared as typed fields)
  • ModelContextProtocolImpl: new _get_forwarded_headers_and_auth() reads the allowlist from provider data, splits Authorization for the authorization= param, returns non-auth headers separately. Both list_runtime_tools and invoke_tool merge forwarded headers with the legacy mcp_headers URI-keyed path (kept for backward compat). Explicit authorization= from the caller wins over forwarded auth.

Config example:

providers:
  tool_runtime:
  - provider_type: remote::model-context-protocol
    config:
      forward_headers:
        maas_api_token: Authorization   # bare token, "Bearer " prepended by prepare_mcp_headers
        tenant_id: X-Tenant-ID
        team_id: X-Team-ID

Test plan

Unit/integration tests — tests covering config validation, header forwarding, auth splitting, default-deny enforcement, missing-key soft-skip, and wiring through list_runtime_tools and invoke_tool:

uv run pytest tests/integration/tool_runtime/test_passthrough_mcp.py -v

Local testing with mock MCP server — ran a mock MCP server that captures all inbound headers and exposes them at /debug/last-headers. Started llama-stack with the forward_headers config above, then verified via the agent API that:

  • maas_api_token from X-LlamaStack-Provider-Data arrived as Authorization: Bearer <token> on the downstream server
  • tenant_id arrived as X-Tenant-ID
  • Unlisted keys (secret_internal) did not appear in any downstream header (default-deny confirmed)
  • Requests with missing or no provider data worked without crashing

Note: invoke_tool and list_runtime_tools are internal methods called by the agents layer and not exposed as HTTP endpoints (changed in upstream PRs #4997 and #5246), so full e2e requires a model server via the agent API.

Checklist

  • forward_headers / extra_blocked_headers optional with backward-compatible defaults — configs without them parse correctly
  • Default-deny: unlisted keys are silently dropped, nothing escapes the allowlist
  • Three-state config tested: None, empty dict, populated dict
  • MCPProviderDataValidator uses extra="allow" (intentional — deployer key names can't be pre-declared)
  • MCPProviderConfig uses extra="forbid"
  • Reuses shared build_forwarded_headers() and validate_forward_headers_config() from providers/utils/forward_headers.py — no duplication
  • Existing mcp_headers URI-keyed path preserved for backward compat
  • No breaking changes

Summary by Sourcery

Add configurable per-request header forwarding support to remote passthrough providers, including MCP tool runtime, inference, and safety, using a shared forward_headers utility with stricter validation and default-deny behavior.

New Features:

  • Allow MCP remote tool runtime to forward selected request-scoped headers and auth tokens from provider data to downstream MCP servers via forward_headers configuration.
  • Enable inference and safety passthrough providers to forward whitelisted headers from X-LlamaStack-Provider-Data using a common forward_headers policy and extra_blocked_headers for operator overrides.

Enhancements:

  • Refine passthrough inference auth handling so the OpenAI client relies solely on composed request headers, avoiding sentinel API keys and ensuring static credentials override forwarded Authorization values.
  • Harden header forwarding with centralized validation, case-insensitive blocking of security-sensitive and operator-defined headers, value sanitization, and size limits.
  • Relax provider-data validators for passthrough and MCP so deployer-defined keys are preserved while keeping provider configs strict via extra='forbid'.
  • Document forward_headers and extra_blocked_headers options for MCP tool runtime, inference passthrough, and safety passthrough providers.

Tests:

  • Add comprehensive unit tests for the shared forward_headers utilities, passthrough inference adapter behavior, safety passthrough config and headers, and MCP provider config and wiring through list_runtime_tools and invoke_tool.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant