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 cb6ff9a5355..40264d111d7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -60,6 +60,10 @@ GATEWAY_SCOPE_TEMPLATE: Final = "api://{client_id}/access_as_user" _GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset( {"invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"} ) +# Entra reports a malformed or unverifiable assertion as ``invalid_client`` too; only its AADSTS50027xx +# (InvalidJwtToken) sub-codes tell that apart from a bad gateway secret. +_INVALID_ASSERTION_AADSTS_PREFIX: Final = "50027" +_AADSTS_CODES_ADAPTER: Final = TypeAdapter(tuple[int, ...]) _MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool") _TOOL_INPUT_SCHEMA_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBO_CACHE_MAX_ENTRIES: Final = 1000 @@ -76,6 +80,13 @@ def _parse_expires_in(raw: object) -> float: return _DEFAULT_TOKEN_TTL_SECONDS +def _parse_aadsts_codes(raw: object) -> tuple[int, ...]: + try: + return _AADSTS_CODES_ADAPTER.validate_python(raw) + except ValidationError: + return () + + def _parse_tool_input_schema(raw: object) -> Mapping[str, object] | None: try: return _TOOL_INPUT_SCHEMA_ADAPTER.validate_python(raw) @@ -117,11 +128,20 @@ class _BlockedDetail(TypedDict): class Agent365TokenExchangeError(Exception): - def __init__(self, status_code: int, error_code: str, description: str) -> None: + def __init__(self, status_code: int, error_code: str, description: str, aadsts_codes: tuple[int, ...] = ()) -> None: super().__init__(f"{error_code}: {description}") self.status_code = status_code self.error_code = error_code self.description = description + self.aadsts_codes = aadsts_codes + + @property + def gateway_owned(self) -> bool: + """Whether the gateway's own client credentials, scope or resource were refused, as opposed to the + caller's assertion. The caller cannot fix a gateway-owned rejection by signing in again.""" + if self.error_code not in _GATEWAY_OWNED_TOKEN_ERRORS: + return False + return not any(str(code).startswith(_INVALID_ASSERTION_AADSTS_PREFIX) for code in self.aadsts_codes) class Agent365MalformedResponseError(Exception): @@ -219,7 +239,7 @@ class Agent365Guardrail(CustomGuardrail): try: obo_token: Final = await self._get_obo_token(assertion) except Agent365TokenExchangeError as exc: - if exc.error_code in _GATEWAY_OWNED_TOKEN_ERRORS: + if exc.gateway_owned: return self._handle_unavailable( data=data, tool_name=tool_name, @@ -446,7 +466,7 @@ class Agent365Guardrail(CustomGuardrail): try: await self._get_obo_token(assertion) except Agent365TokenExchangeError as exc: - return exc.error_code not in _GATEWAY_OWNED_TOKEN_ERRORS + return not exc.gateway_owned except (Agent365ThrottledError, Agent365MalformedResponseError, httpx.HTTPError, LitellmTimeout, TimeoutError): return False return False @@ -492,6 +512,7 @@ class Agent365Guardrail(CustomGuardrail): status_code=response.status_code, error_code=str(body.get("error", "invalid_grant")), description=str(body.get("error_description", ""))[:512], + aadsts_codes=_parse_aadsts_codes(body.get("error_codes")), ) if "access_token" not in body: raise Agent365MalformedResponseError("the Entra token endpoint returned no access_token") 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 511433837c0..aaa472a2400 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 @@ -9066,13 +9066,24 @@ class TestAgent365ChallengeAtConnect: assert await self._connect(self._server([self.GATEWAY_SCOPE]), self.ENTRA_BEARER) is None @pytest.mark.asyncio - async def test_assertion_entra_refuses_to_exchange_is_challenged(self, agent_365_guardrail): + @pytest.mark.parametrize( + "entra_body", + [ + {"error": "invalid_grant", "error_description": "AADSTS700084: The refresh token was issued..."}, + { + "error": "invalid_client", + "error_description": "AADSTS5002723: Invalid JWT token.", + "error_codes": [5002723], + }, + ], + ids=["expired-or-wrong-audience", "forged-reported-as-invalid_client"], + ) + async def test_assertion_entra_refuses_to_exchange_is_challenged(self, agent_365_guardrail, entra_body): """An expired, wrong-audience, or forged Entra token looks like a valid one. Only the OBO exchange can tell, and its verdict must arrive at connect, where WWW-Authenticate reaches the client, - rather than inside every tools/call JSON-RPC error.""" - agent_365_guardrail.async_handler.post.return_value = self._entra_response( - 400, {"error": "invalid_grant", "error_description": "AADSTS700084: The refresh token was issued..."} - ) + rather than inside every tools/call JSON-RPC error. Entra files a forged assertion under + ``invalid_client`` with an AADSTS50027xx sub-code, which must not read as a gateway secret problem.""" + agent_365_guardrail.async_handler.post.return_value = self._entra_response(400, entra_body) challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), self.ENTRA_BEARER) @@ -9088,7 +9099,16 @@ class TestAgent365ChallengeAtConnect: @pytest.mark.parametrize( "entra_outcome", [ - {"return_value": _entra_response(401, {"error": "invalid_client"})}, + { + "return_value": _entra_response( + 401, + { + "error": "invalid_client", + "error_description": "AADSTS7000215: Invalid client secret", + "error_codes": [7000215], + }, + ) + }, {"return_value": _entra_response(503, {"error": "temporarily_unavailable"})}, {"side_effect": httpx.ConnectError("dns")}, ], 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 6466bd4549d..ba9253cb415 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 @@ -631,6 +631,30 @@ class TestUnreachableFallback: assert error_code in info["guardrail_response"]["reason"] assert "client_secret" in info["guardrail_response"]["reason"] + @pytest.mark.asyncio + @pytest.mark.parametrize("aadsts_code", [5002710, 5002723], ids=["malformed-header", "no-kid"]) + async def test_malformed_assertion_reported_as_invalid_client_is_a_caller_401(self, aadsts_code: int): + """Entra answers ``invalid_client`` for a forged or garbled assertion (AADSTS50027xx) exactly as for a + bad gateway secret; the sub-code is what says the caller, not the gateway, has to fix it.""" + handler: Final = FakeHandler( + [ + _response( + 401, + { + "error": "invalid_client", + "error_description": f"AADSTS{aadsts_code}: Invalid JWT token.", + "error_codes": [aadsts_code], + }, + ) + ] + ) + guardrail: Final = _make_guardrail(handler) + data: Final = _mcp_data() + with pytest.raises(HTTPException) as exc_info: + await _run(guardrail, data) + assert exc_info.value.status_code == 401 + assert "client_secret" not in _guardrail_info(data)["guardrail_response"]["reason"] + @pytest.mark.asyncio async def test_gateway_credential_rejection_follows_fail_open(self): handler: Final = FakeHandler(