From d5c9d17d4790aa6f5d8347885cf89f991da59639 Mon Sep 17 00:00:00 2001 From: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:37:56 -0700 Subject: [PATCH] fix(mcp): forward caller's extra headers in true_passthrough auth probe Motivation: _probe_upstream_auth predicts whether a true_passthrough MCP server will accept the caller's request before a stateful session opens, so LiteLLM can return a 401 up front instead of masking the rejection as an empty tool list. The probe only ever sent Accept and Authorization upstream. When an upstream identifies the caller through a different header the admin already configured in server.extra_headers (a gateway-issued platform header, for example), the probe always looked anonymous to it, so it answered 401 for every caller including ones with a valid session, and every connect got the pre-session challenge regardless of whether the caller was actually authorized. Approach: Added _extra_headers_for_probe, which takes the caller's raw request headers and the server's configured extra_headers list and returns just the values for those header names (case-insensitively matched), always excluding authorization since the probe already sets that separately. _probe_upstream_auth gained an extra_headers parameter that merges into the outgoing probe request. The is_true_passthrough branch in _raise_preemptive_401_for_unauthenticated_servers now passes the caller's configured extra headers into the probe, using the raw_headers parameter that function already receives. The is_oauth_passthrough / _check_passthrough_upstream_auth path is untouched, matching the issue's stated scope. Validation: PYTHONPATH=. pytest tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py -v passes all 240 tests (235 existing + 5 new). Reverting only the server.py change (git stash) makes the 5 new/changed tests fail, including one that raises the exact HTTPException: 401 the issue describes, confirming they are a real regression guard rather than a structural check. ruff check, ruff format --check, and basedpyright were run on both changed files: no new basedpyright errors were introduced in server.py (236 before and after), and the small increase in the test file (reportPrivateUsage / reportAny noise from calling internal helpers directly, matching this test file's existing whitebox pattern) stays well under the repo's basedpyright-code-budget.json limits for those rules. Disclosure: dependencies were installed into a local virtualenv rather than through the maturin/editable install this repo normally uses, because of disk-space limits in this environment; litellm was imported directly from source with PYTHONPATH, and the Rust bridge already no-ops gracefully when its compiled extension is absent. No live MCP server was reachable to reproduce the issue end to end; the fix and its tests are demonstrated at the unit level described above. Report: https://github.com/BerriAI/litellm/issues/40099 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5 (via Claude Code) --- .../proxy/_experimental/mcp_server/server.py | 29 ++- .../test_mcp_server_tool_calls_and_headers.py | 175 ++++++++++++++++++ 2 files changed, 203 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2090a0c7421..9de3f16299b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1815,7 +1815,11 @@ if MCP_AVAILABLE: ) }, ) - upstream_status, upstream_www_authenticate = await _probe_upstream_auth(server.url or "", "") + upstream_status, upstream_www_authenticate = await _probe_upstream_auth( + server.url or "", + "", + extra_headers=_extra_headers_for_probe(server, raw_headers), + ) if upstream_status == 401 and upstream_www_authenticate: raise HTTPException( status_code=401, @@ -1851,10 +1855,27 @@ if MCP_AVAILABLE: return None return _get_authorization_header_from_scope(scope) + def _extra_headers_for_probe(server: MCPServer, raw_headers: Mapping[str, str] | None) -> dict[str, str] | None: + """Caller's values for headers this upstream is configured to read + (``server.extra_headers``), minus ``authorization`` which the probe + sets itself. Lets the pre-session probe re-run the same auth path a + real client request would take instead of an anonymous one. + """ + if not server.extra_headers or not raw_headers: + return None + normalized: Final = {k.lower(): v for k, v in raw_headers.items()} + forwarded: Final = { + header: normalized[header.lower()] + for header in server.extra_headers + if header.lower() != "authorization" and header.lower() in normalized + } + return forwarded or None + async def _probe_upstream_auth( url: str, auth_header: str, timeout: float = 5.0, + extra_headers: dict[str, str] | None = None, ) -> tuple[int, str | None]: """JSON-RPC initialize-probe the upstream URL to check whether the token is accepted. @@ -1863,6 +1884,11 @@ if MCP_AVAILABLE: Fails-open with (200, None) on network errors so a transient hiccup does not block valid requests. + ``extra_headers`` carries any additional headers this upstream is + configured to read (``server.extra_headers``) so the probe predicts the + outcome of the same request a real client would send, not an anonymous + one, when identity travels in a header other than ``Authorization``. + Uses the public ``AsyncHTTPHandler.post()`` interface and catches ``httpx.HTTPStatusError`` separately so the 401/403 we want to surface is not swallowed by the broad fail-open ``except Exception`` below. @@ -1887,6 +1913,7 @@ if MCP_AVAILABLE: probe_headers: Final = { "Accept": "application/json, text/event-stream", **({"Authorization": auth_header} if auth_header else {}), + **(extra_headers or {}), } try: resp: Final = await client.post( diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 4cc7794d4ad..aa8d7886b53 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -6918,6 +6918,181 @@ async def test_probe_upstream_auth_fails_open_on_network_error(): assert www_auth is None +@pytest.mark.asyncio +async def test_probe_upstream_auth_merges_extra_headers_into_request(): + """extra_headers must reach the outgoing probe request, so an upstream that + identifies the caller via a configured header (not Authorization) sees the + same request a real client would send.""" + from litellm.proxy._experimental.mcp_server.server import _probe_upstream_auth + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", + return_value=mock_client, + ): + await _probe_upstream_auth("http://upstream/mcp", "", extra_headers={"X-Platform-User": "alice"}) + + _, kwargs = mock_client.post.call_args + assert kwargs["headers"]["X-Platform-User"] == "alice" + assert "Authorization" not in kwargs["headers"] + + +def test_extra_headers_for_probe_filters_to_configured_non_authorization_headers(): + """Only forwards values for header names the admin listed in + server.extra_headers, matched case-insensitively, and always drops + Authorization (the probe sets that itself).""" + from litellm.proxy._experimental.mcp_server.server import _extra_headers_for_probe + + server = MCPServer( + server_id="tp-1", + name="tp_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + extra_headers=["X-Platform-User", "Authorization"], + ) + raw_headers = { + "x-platform-user": "alice", + "authorization": "Bearer caller-key", + "content-type": "application/json", + } + + assert _extra_headers_for_probe(server, raw_headers) == {"X-Platform-User": "alice"} + + +def test_extra_headers_for_probe_returns_none_when_nothing_to_forward(): + from litellm.proxy._experimental.mcp_server.server import _extra_headers_for_probe + + no_extra_headers_server = MCPServer( + server_id="tp-2", + name="tp_server_2", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + extra_headers=None, + ) + assert _extra_headers_for_probe(no_extra_headers_server, {"x-platform-user": "alice"}) is None + + unmatched_server = MCPServer( + server_id="tp-3", + name="tp_server_3", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + extra_headers=["X-Platform-User"], + ) + assert _extra_headers_for_probe(unmatched_server, {"content-type": "application/json"}) is None + + +class TestTruePassthroughProbeForwardsCallerHeaders: + """Regression test for the true_passthrough pre-session probe sending an + anonymous request to upstreams that identify callers via a header other + than Authorization (the header the admin already listed in + server.extra_headers). Before the fix, _probe_upstream_auth never received + that header, so the probe's answer described a caller that doesn't exist.""" + + ALIAS: Final = "tp_platform_header" + + def _server(self, extra_headers: list[str] | None) -> MCPServer: + return MCPServer( + server_id="tp-platform-header", + name=self.ALIAS, + alias=self.ALIAS, + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + extra_headers=extra_headers, + ) + + def _scope(self) -> dict[str, object]: + return {"type": "http", "method": "POST", "path": f"/mcp/{self.ALIAS}", "headers": []} + + @pytest.mark.asyncio + async def test_probe_receives_caller_platform_header_and_does_not_challenge(self): + from litellm.proxy._experimental.mcp_server import server as server_module + + server = self._server(["X-Platform-User"]) + + async def fake_upstream_probe( + url: str, auth_header: str, timeout: float = 5.0, extra_headers: dict[str, str] | None = None + ) -> tuple[int, str | None]: + # Simulates an upstream that answers 401 for anyone it cannot + # identify by the platform header -- exactly the case the issue + # describes, where Authorization alone is not the caller's identity. + if not extra_headers or extra_headers.get("X-Platform-User") != "alice": + return 401, 'Bearer realm="upstream"' + return 200, None + + with ( + patch.object( + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=server, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + side_effect=fake_upstream_probe, + ), + ): + # Must not raise: the caller's platform header reaches the probe, + # so the upstream recognizes them and the pre-session challenge is + # skipped, matching what a real session would get. + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope=self._scope(), + mcp_servers=[self.ALIAS], + oauth2_headers=None, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(), + client_ip=None, + raw_headers={"X-Platform-User": "alice"}, + ) + + @pytest.mark.asyncio + async def test_probe_never_forwards_authorization_via_extra_headers(self): + """Authorization listed in server.extra_headers must not be replayed a + second time through extra_headers; the probe call already sets its own + (empty) auth_header for true_passthrough.""" + from litellm.proxy._experimental.mcp_server import server as server_module + + server = self._server(["Authorization", "X-Platform-User"]) + captured: dict[str, dict[str, str] | None] = {} + + async def fake_upstream_probe( + url: str, auth_header: str, timeout: float = 5.0, extra_headers: dict[str, str] | None = None + ) -> tuple[int, str | None]: + captured["extra_headers"] = extra_headers + return 200, None + + with ( + patch.object( + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=server, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + side_effect=fake_upstream_probe, + ), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope=self._scope(), + mcp_servers=[self.ALIAS], + oauth2_headers=None, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(), + client_ip=None, + raw_headers={"Authorization": "Bearer caller-key", "X-Platform-User": "alice"}, + ) + + assert captured["extra_headers"] == {"X-Platform-User": "alice"} + + def test_get_forwarded_auth_from_scope_extracts_header(): """Returns Authorization value when x-litellm-api-key is also present.""" from litellm.proxy._experimental.mcp_server.server import (