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