diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e4bc9bdd952..887a4aec033 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -79,7 +79,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils -from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import agent_365_authorization_servers +from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import ( + agent_365_authorization_servers, + agent_365_subject_token_present, +) from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, get_chain_id_from_headers, @@ -4136,13 +4139,13 @@ if MCP_AVAILABLE: # so the client discovers the IdP, SSOs, and retries with a subject token, which LiteLLM # then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the # header lost, so the discovery flow needs this pre-emptive challenge. Servers gated by an - # Agent 365 guardrail (OBO to the evaluate API) get the same challenge. - if ( - server - and not oauth2_headers - and ( - server.auth_type == MCPAuth.oauth2_token_exchange - or agent_365_authorization_servers(server, user_api_key_auth) + # Agent 365 guardrail (OBO to the evaluate API) get the same challenge, also when the only + # bearer is the LiteLLM key itself, which admits the caller but is not an exchangeable subject. + if server and ( + (server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers) + or ( + not agent_365_subject_token_present(oauth2_headers) + and agent_365_authorization_servers(server, user_api_key_auth) ) ): from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph 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 3cc2394a18e..24f1643d89a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -201,8 +201,8 @@ class Agent365Guardrail(CustomGuardrail): return data tool_name: Final = str(data.get("mcp_tool_name") or "") - assertion: Final = data.get("incoming_bearer_token") - if not isinstance(assertion, str) or assertion.count(".") != 2: + assertion: Final = entra_assertion(data.get("incoming_bearer_token")) + if assertion is None: self._handle_caller_fault( data=data, tool_name=tool_name, @@ -658,6 +658,20 @@ def _applicable_guardrails( return tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) +def entra_assertion(value: object) -> str | None: + """``value`` when it is a compact JWS, the only bearer shape the OBO exchange accepts as its assertion. + A LiteLLM virtual key, session bearer, or opaque upstream token in ``Authorization`` yields ``None``.""" + return value if isinstance(value, str) and value.count(".") == 2 else None + + +def agent_365_subject_token_present(oauth2_headers: Mapping[str, str] | None) -> bool: + """Whether the request's ``Authorization`` carries an Entra assertion the guardrail can exchange.""" + authorization: Final = (oauth2_headers or {}).get("Authorization", "") + if not authorization.lower().startswith("bearer "): + return False + return entra_assertion(authorization[len("bearer ") :].strip()) is not None + + def agent_365_authorization_servers(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]: """Entra issuers an MCP client signs in with before calling ``server`` through an Agent 365 guardrail.""" return tuple( 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 7d4159aade8..b599fbadb50 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 @@ -8996,10 +8996,23 @@ class TestAgent365ChallengeAtConnect: assert 'resource_metadata="/.well-known/oauth-protected-resource/mcp/tools"' in www_authenticate @pytest.mark.asyncio - async def test_bearer_present_connects(self, agent_365_guardrail): - bearer = {"Authorization": "Bearer entra-user-token"} + async def test_entra_assertion_present_connects(self, agent_365_guardrail): + bearer = {"Authorization": "Bearer eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1LTEifQ.c2ln"} assert await self._connect(self._server([self.GATEWAY_SCOPE]), bearer) is None + @pytest.mark.asyncio + async def test_litellm_key_in_authorization_is_still_challenged(self, agent_365_guardrail): + """A LiteLLM virtual key admits the caller but is no Entra assertion, so the tools/call would fail + 401 inside JSON-RPC with the WWW-Authenticate header lost. The connect must challenge instead.""" + challenge = await self._connect( + self._server([self.GATEWAY_SCOPE]), {"Authorization": "Bearer sk-litellm-virtual-key"} + ) + + 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="/.well-known/oauth-protected-resource/mcp/tools"' in www_authenticate + @pytest.mark.asyncio async def test_scopeless_server_is_still_challenged(self, agent_365_guardrail): challenge = await self._connect(self._server(None), None)