From 0cade620df2e7bf8ef459747d08dc6fcd1ac2218 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 21:21:55 -0700 Subject: [PATCH] fix(mcp): probe legacy delegates with only the caller's token, as v1.102 did The probe is the check on a delegate caller's own token, so sending the server's static headers with it let a static Authorization stand in for that token. On a delegate server that also sets a static Authorization, any junk bearer passed the probe and reached the upstream with the operator's token The probe goes back to sending only the caller's token for every server, as v1.102 and main do. This drops the static header probing added for the earlier review, so an upstream that needs a companion static header next to the token gets the same probe it got on v1.102 --- .../proxy/mcp_server/mcp_server_manager.py | 6 -- litellm/proxy/mcp_server/server.py | 42 ++------------ .../test_mcp_server_tool_calls_and_headers.py | 58 ++++--------------- 3 files changed, 16 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 9f00a3f2138..c772e2e9d09 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -4176,12 +4176,6 @@ class MCPServerManager: referenced: Final = collect_env_var_references(strings=static_headers.values()) return bool(referenced & user_var_names) - async def resolved_static_headers( - self, server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None - ) -> dict[str, str] | None: - """``server.static_headers`` interpolated exactly as a tool call sends them.""" - return await self._resolve_static_headers_with_env_vars(server, user_api_key_auth) - async def _resolve_static_headers_with_env_vars( self, server: MCPServer, diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index 0fb1e6bbd3e..2105d875b58 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -2104,7 +2104,6 @@ if MCP_AVAILABLE: url: str, auth_header: str, timeout: float = 5.0, - headers: Mapping[str, str] | None = None, ) -> tuple[int, str | None]: """JSON-RPC initialize-probe the upstream URL to check whether the token is accepted. @@ -2137,7 +2136,6 @@ if MCP_AVAILABLE: probe_headers: Final = { "Accept": "application/json, text/event-stream", **({"Authorization": auth_header} if auth_header else {}), - **(headers or {}), } try: resp: Final = await client.post( @@ -2156,30 +2154,6 @@ if MCP_AVAILABLE: verbose_logger.debug("_probe_upstream_auth: probe to %s failed (%s), allowing request through", url, exc) return 200, None - async def _probe_with_outbound_headers( - server: MCPServer, - auth_header: str, - caller_headers: Mapping[str, str], - user_api_key_auth: UserAPIKeyAuth | None, - ) -> tuple[int, str | None]: - """Probe with the headers a real tool call sends: the caller's prepared headers with the server's - resolved static headers on top, so a static ``Authorization`` wins here too. When those cannot be - resolved yet, such as a per-user variable the caller has not set, the probe fails open like a network - error and leaves the error to the call.""" - try: - static_headers: Final = await global_mcp_server_manager.resolved_static_headers(server, user_api_key_auth) - except Exception as exc: # noqa: BLE001 # the probe fails open on anything it cannot reproduce - verbose_logger.debug("Skipping upstream auth probe for %s (%s)", server.name, type(exc).__name__) - return 200, None - outbound: Final = merge_mcp_headers(extra_headers=caller_headers, static_headers=static_headers) or {} - authorization: Final = next( - (value for name, value in reversed(outbound.items()) if name.lower() == "authorization"), auth_header - ) - companion_headers: Final = {name: value for name, value in outbound.items() if name.lower() != "authorization"} - if not companion_headers: - return await _probe_upstream_auth(server.url or "", authorization) - return await _probe_upstream_auth(server.url or "", authorization, headers=companion_headers) - async def _check_passthrough_upstream_auth( scope: Scope, user_api_key_auth: UserAPIKeyAuth | None, @@ -2257,13 +2231,8 @@ if MCP_AVAILABLE: ) for srv, (auth, extra) in prepared_headers ) - probe_targets: Final[tuple[tuple[MCPServer, str, str, Mapping[str, str]], ...]] = tuple( - ( - srv, - auth_header, - srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name, - headers or {}, - ) + probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple( + (srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name) for srv, headers in probe_headers if ( auth_header := next( @@ -2275,12 +2244,9 @@ if MCP_AVAILABLE: return probe_results: Final = await asyncio.gather( - *[ - _probe_with_outbound_headers(srv, auth_header, headers, user_api_key_auth) - for srv, auth_header, _, headers in probe_targets - ] + *[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets] ) - for (_srv, _, challenge_server_name, _), (probe_status, _) in zip(probe_targets, probe_results): + for (_srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results): if probe_status == 401: # Token is missing or expired: keep pass-through clients on the # protected-resource discovery flow so they re-authorize against diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index a56a51d5d87..621a7b65005 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -7437,51 +7437,13 @@ async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_ @pytest.mark.asyncio @pytest.mark.usefixtures("httpx_transport") -@pytest.mark.parametrize( - "static_headers, expected_authorization", - [ - ({"X-Tenant": "tenant-1"}, "Bearer expired-upstream-token"), - ({"X-Tenant": "tenant-1", "Authorization": "Bearer server-wide-token"}, "Bearer server-wide-token"), - ], - ids=["companion-header", "static-authorization-wins"], -) -async def test_legacy_delegate_probe_sends_the_headers_a_tool_call_sends( - static_headers: dict[str, str], expected_authorization: str -) -> None: - """An upstream that needs a static header next to the token sees it in the probe too, so a token the - upstream accepts on real calls is not rejected before the session opens.""" - import respx - - server: Final = _delegate_auth_mcp_server().model_copy(update={"static_headers": static_headers}) - with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - respx.mock as upstream, - ): - probe: Final = upstream.post(server.url).mock( - side_effect=lambda request: httpx.Response(200 if request.headers.get("X-Tenant") == "tenant-1" else 401) - ) - assert await _expired_delegate_token_preflight(["delegate_test"]) is None - - assert probe.call_count == 1 - assert probe.calls[0].request.headers["Authorization"] == expected_authorization - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("httpx_transport") -async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_user_values() -> None: - """A static header built from a per-user value the caller has not set cannot be reproduced, so the - probe lets the request through and the tool call reports the missing value.""" +async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_authorization() -> None: + """The probe is what checks a delegate caller's own token, so a static Authorization the server sends on + tool calls cannot stand in for it: a junk token still gets the re-authorization challenge.""" import respx server: Final = _delegate_auth_mcp_server().model_copy( - update={ - "static_headers": {"X-Tenant-Token": "${TENANT_TOKEN}"}, - "env_vars": [{"name": "TENANT_TOKEN", "value": "", "scope": "user"}], - } + update={"static_headers": {"Authorization": "Bearer operator-token"}} ) with ( _patch_delegate_resolver(server, "delegate_test"), @@ -7491,12 +7453,16 @@ async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_ ), respx.mock as upstream, ): - probe: Final = upstream.post(server.url).respond( - 401, headers={"www-authenticate": 'Bearer error="invalid_token"'} + upstream.post(server.url).mock( + side_effect=lambda request: httpx.Response( + 200 if request.headers["Authorization"] == "Bearer operator-token" else 401 + ) + ) + outcome: Final = await asyncio.gather( + _expired_delegate_token_preflight(["delegate_test"]), return_exceptions=True ) - assert await _expired_delegate_token_preflight(["delegate_test"]) is None - assert probe.call_count == 0 + assert getattr(outcome[0], "status_code", None) == 401 @pytest.mark.asyncio