diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index 40264d111d7..c332ac3c964 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -36,7 +36,6 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( AGENT_365_PROD_API_BASE, @@ -686,8 +685,10 @@ def _applicable_guardrails( """Agent 365 guardrails whose sign-in the gateway advertises for ``server``: the ``default_on`` ones, minus those the caller's key or team opted out of once the caller is known. A guardrail only a key or policy selects still enforces at the tool call but never challenges, since the anonymous metadata fetch that - follows a challenge cannot see which key selected it and would advertise the wrong issuer.""" - if server.auth_type == MCPAuth.oauth2 or not server.advertises_gateway_authorization_server: + follows a challenge cannot see which key selected it and would advertise the wrong issuer. Only servers + that leave the caller's top-level ``Authorization`` with the gateway qualify: a forwarded API-key header + travels upstream in its own slot and does not displace the Entra assertion.""" + if not server.keeps_caller_authorization: return () advertised: Final = tuple( callback diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 985d31af997..70ad101dcb0 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -276,10 +276,10 @@ class MCPServer(BaseModel): return self.per_server_oauth_discovery and self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials @property - def advertises_gateway_authorization_server(self) -> bool: - """Whether named discovery should advertise the aggregate gateway authorization server.""" - if self.auth_type == MCPAuth.oauth2: - return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + def keeps_caller_authorization(self) -> bool: + """Whether the caller's top-level ``Authorization`` stays with the gateway: the server neither relays + it upstream nor runs an OAuth mode that fills that slot itself, so a gateway guardrail may consume it + as the caller's own assertion. Forwarding a separate API-key header leaves the slot untouched.""" if self.auth_type not in ( None, MCPAuth.none, @@ -291,9 +291,15 @@ class MCPServer(BaseModel): MCPAuth.aws_sigv4, ): return False - return not any( - header.lower() in ("authorization", "x-api-key", "api-key", "apikey") - for header in (self.extra_headers or ()) + return not any(header.lower() == "authorization" for header in (self.extra_headers or ())) + + @property + def advertises_gateway_authorization_server(self) -> bool: + """Whether named discovery should advertise the aggregate gateway authorization server.""" + if self.auth_type == MCPAuth.oauth2: + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + return self.keeps_caller_authorization and not any( + header.lower() in ("x-api-key", "api-key", "apikey") for header in (self.extra_headers or ()) ) @property diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 25164d4b215..975c4b03545 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7148,7 +7148,7 @@ def agent_365_guardrail(): ) -async def _agent_365_gated_prm(scopes): +async def _agent_365_gated_prm(scopes, extra_headers=None): from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -7168,6 +7168,7 @@ async def _agent_365_gated_prm(scopes): transport=MCPTransport.http, auth_type=MCPAuth.none, scopes=scopes, + extra_headers=extra_headers, ) mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" @@ -7200,6 +7201,18 @@ async def test_agent_365_prm_defaults_scopeless_server_to_the_gateway_app_scope( } +@pytest.mark.asyncio +async def test_agent_365_prm_survives_a_forwarded_upstream_api_key_header(agent_365_guardrail): + """The forwarded ``x-api-key`` is the upstream's credential and rides in its own header, so the caller's + ``Authorization`` still carries the Entra assertion and discovery must keep naming the Entra tenant.""" + response = await _agent_365_gated_prm(scopes=None, extra_headers=["x-api-key"]) + assert jsonable_encoder(response) == { + "authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"], + "resource": "https://litellm.example.com/mcp/tools", + "scopes_supported": ["api://client-xyz/access_as_user"], + } + + @pytest.mark.asyncio async def test_key_selected_agent_365_guardrail_leaves_anonymous_prm_on_the_gateway_issuer(agent_365_guardrail): """A default-off guardrail gates only the keys that select it, so the anonymous discovery fetch must keep diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index aaa472a2400..59973b51054 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8926,7 +8926,7 @@ class TestAgent365ChallengeAtConnect: request=httpx.Request("POST", "https://login.microsoftonline.com/tenant-abc/oauth2/v2.0/token"), ) - def _server(self, scopes: list[str] | None) -> MCPServer: + def _server(self, scopes: list[str] | None, extra_headers: list[str] | None = None) -> MCPServer: return MCPServer( server_id="id-tools", name="tools", @@ -8936,6 +8936,7 @@ class TestAgent365ChallengeAtConnect: transport=MCPTransport.http, auth_type=MCPAuth.none, scopes=scopes, + extra_headers=extra_headers, mcp_info={"server_name": "tools"}, ) @@ -9181,6 +9182,26 @@ class TestAgent365ChallengeAtConnect: assert challenge is not None and challenge.status_code == 401 assert 'error="invalid_token"' in (challenge.headers or {}).get("WWW-Authenticate", "") + @pytest.mark.asyncio + async def test_server_forwarding_an_upstream_api_key_header_is_still_challenged(self, agent_365_guardrail): + """``x-api-key`` travels upstream in its own header and leaves the caller's ``Authorization`` free for + the Entra assertion, so a key-only connect must still be sent to sign in.""" + challenge = await self._connect(self._server(None, extra_headers=["x-api-key"]), None) + + assert challenge is not None and challenge.status_code == 401 + www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "") + assert 'error="invalid_token"' in www_authenticate + assert ( + 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"' + in www_authenticate + ) + + @pytest.mark.asyncio + async def test_server_relaying_the_caller_authorization_is_not_challenged(self, agent_365_guardrail): + """Forwarding ``Authorization`` hands the caller's bearer to the upstream, so the gateway holds no Entra + assertion of its own to exchange and must not advertise a sign-in it cannot consume.""" + assert await self._connect(self._server(None, extra_headers=["Authorization"]), None) is None + @pytest.mark.asyncio async def test_no_registered_guardrail_means_no_challenge(self): assert await self._connect(self._server([self.GATEWAY_SCOPE]), None) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index bf670f62c4c..2a1064ec510 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5092,6 +5092,36 @@ class TestMCPServerManager: ) assert server2.requires_per_user_auth is False + @pytest.mark.parametrize( + "auth_type, extra_headers, keeps_authorization, advertises_gateway", + [ + (MCPAuth.none, None, True, True), + (MCPAuth.api_key, None, True, True), + (MCPAuth.none, ["x-api-key"], True, False), + (MCPAuth.none, ["API-Key"], True, False), + (MCPAuth.none, ["Authorization"], False, False), + (MCPAuth.none, ["x-api-key", "authorization"], False, False), + (MCPAuth.oauth_delegate, None, False, False), + (MCPAuth.true_passthrough, None, False, False), + (MCPAuth.oauth2_token_exchange, None, False, False), + ], + ) + def test_forwarded_api_key_header_keeps_caller_authorization_but_not_gateway_discovery( + self, auth_type, extra_headers, keeps_authorization, advertises_gateway + ): + """A forwarded API-key header is the upstream's own credential and leaves the caller's top-level + ``Authorization`` with the gateway, while still ruling out the gateway's aggregate OAuth discovery.""" + server = MCPServer( + server_id="s", + name="s", + transport=MCPTransport.http, + auth_type=auth_type, + url="http://s.example", + extra_headers=extra_headers, + ) + assert server.keeps_caller_authorization is keeps_authorization + assert server.advertises_gateway_authorization_server is advertises_gateway + @pytest.mark.asyncio async def test_register_openapi_tools_includes_static_headers(self, tmp_path): """Ensure OpenAPI-to-MCP tool calls include server.static_headers (Issue #19341).""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py index ba9253cb415..f73c0e87547 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -1107,6 +1107,15 @@ class TestAgent365AuthorizationServers: def test_leaves_servers_whose_own_auth_mode_owns_sign_in_alone(self, registered_guardrail, server): assert agent_365_authorization_servers(server, None) == () + @pytest.mark.parametrize("header", ["x-api-key", "API-Key", "apikey"]) + def test_forwarded_api_key_header_leaves_authorization_to_entra(self, registered_guardrail, header): + """An upstream API key rides in its own header, so the caller's ``Authorization`` still carries the + Entra assertion and a key-only client must be told where to sign in.""" + server: Final = _mcp_server(MCPAuth.none, scopes=None, extra_headers=[header]) + + assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,) + assert agent_365_scopes_supported(server, None) == ("api://client-xyz/access_as_user",) + def test_dedupes_guardrails_sharing_a_tenant(self, registered_guardrail): twin: Final = _make_guardrail(FakeHandler([])) twin.guardrail_name = "agent-365-twin"