Allow pre_mcp_call guardrail hooks to mutate outbound MCP headers

[noahnistler]

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 17, 2026 9:08pm

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing Evernorth:feature/pre-mcp-call-hook-header-mutation (9f443a3) with main (ef9cc33)}

[greptile-apps]

Greptile Summary

This PR extends the pre_mcp_call guardrail hook pipeline to allow hooks to inject custom HTTP headers into outbound MCP server requests. It also propagates JWT claims onto UserAPIKeyAuth so hooks can access them when deciding which headers to inject (e.g., signing a downstream JWT).

Key changes:

• pre_call_tool_check now returns a Dict[str, Any] (previously returned nothing) containing arguments if modified by a hook and extra_headers if injected by a hook. This also incidentally fixes a pre-existing bug where hook-modified tool arguments were silently discarded because the caller never used the return value.
• _call_regular_mcp_tool gains a hook_extra_headers parameter; these headers are merged last, giving them the highest priority over oauth2_headers, extra_headers config, and static_headers.
• OpenAPI-backed servers (spec_path set) reject hook header injection early with HTTP 400 before any background tasks are scheduled.
• _convert_mcp_hook_response_to_kwargs in ProxyLogging now also extracts extra_headers from hook responses.
• UserAPIKeyAuth gains an optional jwt_claims: Optional[Dict] = None field; JWT claims are now propagated in the virtual-key fast path and both branches of the standard JWT auth path.
• A comprehensive 701-line mock-only test suite is included covering all new paths and backward-compatibility scenarios.

Minor concerns:

• jwt_claims is typed as Optional[Dict] rather than Optional[Dict[str, Any]].
• TestHookHeaderMergePriority tests use broad except Exception: pass blocks that could mask unrelated failures.
• The modified_kwargs["arguments"] access in pre_call_tool_check is implicitly safe but creates a fragile coupling to the internal behaviour of _convert_mcp_hook_response_to_kwargs.

Confidence Score: 4/5

• Safe to merge; changes are additive and backward compatible with only minor style concerns.
• All production code changes are additive: new optional parameters with defaults, a new optional field on a Pydantic model, and a new return value from a function whose return was previously unused. The OpenAPI server guard prevents misuse. The comprehensive mock test suite covers all critical paths. The score is 4 rather than 5 due to the implicit coupling around modified_kwargs["arguments"] access and the test quality issue with broad exception swallowing in TestHookHeaderMergePriority.
• litellm/proxy/_experimental/mcp_server/mcp_server_manager.py — particularly the pre_call_tool_check return value construction and the header-merge ordering in _call_regular_mcp_tool.

Important Files Changed

Filename	Overview
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py	Core change: pre_call_tool_check now returns a Dict[str, Any] containing optional arguments and extra_headers from hook responses. call_tool captures this dict and propagates it downstream; _call_regular_mcp_tool gains a new hook_extra_headers parameter merged last (highest priority) into outbound headers. OpenAPI-backed servers reject hook header injection early via an HTTP 400. One minor defensive-coding concern on the modified_kwargs["arguments"] access pattern.
litellm/proxy/utils.py	Small additive change to _convert_mcp_hook_response_to_kwargs: now extracts extra_headers from the hook response dict in addition to modified_arguments. Backward compatible — no extra_headers in the response simply leaves the key absent.
litellm/proxy/_types.py	Adds optional jwt_claims: Optional[Dict] = None field to UserAPIKeyAuth. Backward compatible. Minor typing concern: should be Optional[Dict[str, Any]] for consistency with actual usage.
litellm/proxy/auth/user_api_key_auth.py	JWT claims are now propagated to UserAPIKeyAuth in three paths: virtual-key fast path (post _resolve_jwt_to_virtual_key), proxy-admin standard JWT path, and non-admin standard JWT path. All changes are additive and backward compatible.
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py	701-line mock-only test suite covering all major paths: _convert_mcp_hook_response_to_kwargs, pre_call_tool_check return values, call_tool header/argument propagation, OpenAPI rejection, JWT claims field. Tests use only mocks — no real network calls. Minor quality concern: TestHookHeaderMergePriority tests use broad except Exception: pass which can mask unrelated failures.

Sequence Diagram

sequenceDiagram
    participant Client
    participant MCPServerManager
    participant ProxyLogging
    participant GuardrailHook
    participant MCPServer

    Client->>MCPServerManager: call_tool(name, arguments, proxy_logging_obj)
    MCPServerManager->>MCPServerManager: _get_mcp_server_from_tool_name()

    alt proxy_logging_obj present
        MCPServerManager->>MCPServerManager: pre_call_tool_check()
        MCPServerManager->>ProxyLogging: pre_call_hook(data, call_type)
        ProxyLogging->>GuardrailHook: custom hook invoked
        GuardrailHook-->>ProxyLogging: {modified_arguments?, extra_headers?}
        ProxyLogging-->>MCPServerManager: hook response dict
        MCPServerManager->>ProxyLogging: _convert_mcp_hook_response_to_kwargs()
        ProxyLogging-->>MCPServerManager: modified_kwargs
        MCPServerManager-->>MCPServerManager: hook_result = {arguments?, extra_headers?}
    end

    alt mcp_server.spec_path AND hook_result.extra_headers
        MCPServerManager-->>Client: HTTP 400 (OpenAPI servers unsupported)
    end

    alt Regular MCP server (no spec_path)
        MCPServerManager->>MCPServerManager: _call_regular_mcp_tool(hook_extra_headers=...)
        Note over MCPServerManager: Merge order: oauth2 → extra_headers config → static_headers → hook_extra_headers (highest priority)
        MCPServerManager->>MCPServer: call_tool(arguments, extra_headers merged)
        MCPServer-->>MCPServerManager: CallToolResult
        MCPServerManager-->>Client: CallToolResult
    else OpenAPI server (spec_path set, no hook headers)
        MCPServerManager->>MCPServerManager: _call_openapi_tool_handler()
        MCPServerManager-->>Client: CallToolResult
    end

Loading

Comments Outside Diff (3)

1. tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py, line 757-775 (link)

   Silent exception swallowing masks test failures

   The broad except Exception: pass blocks in all three TestHookHeaderMergePriority tests silently swallow any exception raised during _call_regular_mcp_tool. If the code under test raises an unexpected error (e.g., an AttributeError, TypeError, or KeyError unrelated to the test's intent), the test would still pass because captured_extra_headers["value"] was already set by fake_create_mcp_client before the exception occurred.

   This is also true for test_no_hook_headers_preserves_existing_behavior and test_hook_headers_merge_with_oauth2.

   Consider only catching expected exceptions, or asserting on the captured headers before the try/except:

   try:
       await manager._call_regular_mcp_tool(...)
   except (AttributeError, KeyError):
       # Expected — _call_regular_mcp_tool may fail downstream of client creation
       pass
   
   headers = captured_extra_headers.get("value", {})
   assert headers["Authorization"] == "Bearer hook-signed-jwt"

2. litellm/proxy/_types.py, line 2471 (link)

   Loose typing on jwt_claims field

   Optional[Dict] is unparameterized; prefer Optional[Dict[str, Any]] to match how jwt_claims dicts are used throughout the codebase (mixed string keys and arbitrary values like sub, iss, groups, exp).

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

3. litellm/proxy/_experimental/mcp_server/mcp_server_manager.py, line 1994-1995 (link)

   modified_kwargs["arguments"] could be None if the key is absent

   modified_kwargs.get("arguments") returns None when "arguments" is not in modified_kwargs, but the subsequent modified_kwargs["arguments"] would then set hook_result["arguments"] = None (which is falsy but not the same as "hook did not modify arguments").

   In practice this path is safe because _convert_mcp_hook_response_to_kwargs always starts with original_kwargs.copy() (which contains "arguments"), so the key is always present. However, the implicit coupling is fragile — if the helper is ever refactored, this line could silently pass None arguments into the tool call.

   A safer pattern would be:

   new_args = modified_kwargs.get("arguments")
   if new_args is not None and new_args != arguments:
       hook_result["arguments"] = new_args

_{Last reviewed commit: 9f443a3}

[noahnistler]

@greptileai

[noahnistler]

@greptileai

[noahnistler]

@greptileai