From 886220375f9a13fe616a32521085d6fb896baf41 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 03:15:59 +0000 Subject: [PATCH] fix(guardrails): agent 365 sign-in for scopeless servers and gateway credential errors Scopeless Agent 365 gated servers now advertise api:///access_as_user instead of staying silent, so a client can still sign in. Entra rejecting the gateway's own credentials (invalid_client, unauthorized_client, invalid_scope, invalid_resource) follows unreachable_fallback rather than telling the caller to sign in again with a 401 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 7 ++- .../guardrail_hooks/agent_365/agent_365.py | 58 ++++++++++++++----- .../mcp_server/test_discoverable_endpoints.py | 8 ++- .../mcp_server/test_mcp_server.py | 7 ++- .../guardrail_hooks/test_agent_365.py | 53 +++++++++++++++-- 5 files changed, 110 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 4ecd5de6121..c912c69bf6a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -82,7 +82,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -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_scopes_supported, +) from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod @@ -2485,7 +2488,7 @@ async def _build_oauth_protected_resource_response( agent_365_metadata: Final[_ProtectedResourceMetadata] = { "authorization_servers": agent_365_issuers, "resource": resource_url, - "scopes_supported": tuple(mcp_server.scopes or ()), + "scopes_supported": agent_365_scopes_supported(mcp_server, None), } return agent_365_metadata 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 5813240f01b..3cc2394a18e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -57,6 +57,10 @@ ENTRA_ISSUER_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/v2 EVALUATE_PATH: Final = "/agents/tool-evaluation/evaluate" MCP_SESSION_ID_HEADER: Final = "mcp-session-id" DEFENDER_STATUS_EVALUATED: Final = "Evaluated" +GATEWAY_SCOPE_TEMPLATE: Final = "api://{client_id}/access_as_user" +_GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset( + {"invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"} +) _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 @@ -212,6 +216,15 @@ 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: + return self._handle_unavailable( + data=data, + tool_name=tool_name, + reason=( + f"Entra rejected the gateway's own Agent 365 credentials ({exc.error_code}); " + "check the guardrail's client_id, client_secret and resource_app_id" + ), + ) self._handle_caller_fault( data=data, tool_name=tool_name, @@ -627,23 +640,42 @@ def _applies_to_caller(guardrail: Agent365Guardrail, user_api_key_auth: "UserAPI return guardrail.should_run_guardrail(data=probe, event_type=GuardrailEventHooks.pre_mcp_call) -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. - - Empty unless the admin advertised the server's ``scopes`` (the audience the client requests), the gateway - would otherwise own sign-in for the server, and an Agent 365 guardrail applies: every registered one for the - anonymous discovery fetch, otherwise those the caller's key, team, or policies select. - """ - if not server.scopes or server.auth_type == MCPAuth.oauth2 or not server.advertises_gateway_authorization_server: +def _applicable_guardrails( + server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None" +) -> tuple[Agent365Guardrail, ...]: + """Agent 365 guardrails that gate ``server`` for this caller: every registered one for the anonymous + discovery fetch, otherwise those the caller's key, team, or policies select. Empty when the gateway + does not own sign-in for the server.""" + if server.auth_type == MCPAuth.oauth2 or not server.advertises_gateway_authorization_server: return () registered: Final = tuple( callback for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(Agent365Guardrail) if isinstance(callback, Agent365Guardrail) ) - applicable: Final = ( - registered - if user_api_key_auth is None - else tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) + if user_api_key_auth is None: + return registered + return tuple(g for g in registered if _applies_to_caller(g, user_api_key_auth)) + + +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( + dict.fromkeys( + ENTRA_ISSUER_TEMPLATE.format(tenant_id=g.tenant_id) + for g in _applicable_guardrails(server, user_api_key_auth) + ) + ) + + +def agent_365_scopes_supported(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]: + """Scopes the client requests from Entra for ``server``: the admin's ``scopes`` when set, otherwise the + ``access_as_user`` scope of each gating guardrail's gateway app registration (``api://``).""" + if server.scopes: + return tuple(server.scopes) + return tuple( + dict.fromkeys( + GATEWAY_SCOPE_TEMPLATE.format(client_id=g.client_id) + for g in _applicable_guardrails(server, user_api_key_auth) + ) ) - return tuple(dict.fromkeys(ENTRA_ISSUER_TEMPLATE.format(tenant_id=g.tenant_id) for g in applicable)) 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 a9ac0d528a7..a24568c43e8 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 @@ -7157,9 +7157,13 @@ async def test_agent_365_gated_server_prm_names_the_entra_tenant(agent_365_guard @pytest.mark.asyncio -async def test_agent_365_prm_falls_back_to_gateway_without_scopes(agent_365_guardrail): +async def test_agent_365_prm_defaults_scopeless_server_to_the_gateway_app_scope(agent_365_guardrail): response = await _agent_365_gated_prm(scopes=None) - assert response["authorization_servers"] == ["https://litellm.example.com/mcp"] + 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"], + } def _token_request(headers): 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 4fd9c923709..3db642bcc12 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 @@ -8906,8 +8906,11 @@ class TestAgent365ChallengeAtConnect: assert await self._connect(self._server([self.GATEWAY_SCOPE]), bearer) is None @pytest.mark.asyncio - async def test_server_without_advertised_scopes_is_not_challenged(self, agent_365_guardrail): - assert await self._connect(self._server(None), None) is None + async def test_scopeless_server_is_still_challenged(self, agent_365_guardrail): + challenge = await self._connect(self._server(None), None) + + 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_no_registered_guardrail_means_no_challenge(self): 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 c82f830cd84..c9373de8f82 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 @@ -20,7 +20,10 @@ from litellm.proxy.guardrails.guardrail_hooks.agent_365 import ( guardrail_initializer_registry, initialize_guardrail, ) -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_scopes_supported, +) from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, @@ -590,6 +593,40 @@ class TestUnreachableFallback: assert exc_info.value.status_code == 401 assert "invalid_grant" in exc_info.value.detail["message"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error_code", ["invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"] + ) + async def test_gateway_credential_rejection_is_unavailable_not_a_caller_401(self, error_code: str): + handler: Final = FakeHandler( + [_response(401, {"error": error_code, "error_description": "AADSTS7000215: invalid client secret"})] + ) + 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 == 503 + assert exc_info.value.headers is None or "WWW-Authenticate" not in exc_info.value.headers + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Unavailable" + assert error_code in info["guardrail_response"]["reason"] + assert "client_secret" in info["guardrail_response"]["reason"] + + @pytest.mark.asyncio + async def test_gateway_credential_rejection_follows_fail_open(self): + handler: Final = FakeHandler( + [_response(401, {"error": "invalid_client", "error_description": "AADSTS7000215: invalid client secret"})] + ) + guardrail: Final = _make_guardrail(handler, unreachable_fallback="fail_open") + data: Final = _mcp_data() + result: Final = await _run(guardrail, data) + assert result is data + info: Final = _guardrail_info(data) + assert info["guardrail_status"] == "guardrail_failed_to_respond" + assert info["guardrail_response"]["verdict"] == "Unscanned" + assert "invalid_client" in info["guardrail_response"]["reason"] + @pytest.mark.asyncio async def test_obo_endpoint_5xx_fail_open(self): handler: Final = FakeHandler([_response(503, text="entra down")]) @@ -948,9 +985,17 @@ class TestAgent365AuthorizationServers: _mcp_server(MCPAuth.api_key, scopes=[GATEWAY_SCOPE], auth_value="k"), None ) == (ENTRA_ISSUER,) - def test_silent_without_advertised_scopes(self, registered_guardrail): - assert agent_365_authorization_servers(_mcp_server(scopes=None), None) == () - assert agent_365_authorization_servers(_mcp_server(scopes=[]), None) == () + @pytest.mark.parametrize("scopes", [None, []], ids=["unset", "empty"]) + def test_scopeless_server_signs_in_with_the_gateway_app_scope(self, registered_guardrail, scopes): + server: Final = _mcp_server(scopes=scopes) + assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,) + assert agent_365_scopes_supported(server, None) == ("api://client-xyz/access_as_user",) + + def test_admin_scopes_override_the_default_gateway_scope(self, registered_guardrail): + assert agent_365_scopes_supported(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (GATEWAY_SCOPE,) + + def test_no_default_scope_when_no_guardrail_gates_the_server(self): + assert agent_365_scopes_supported(_mcp_server(scopes=None), None) == () def test_silent_when_no_guardrail_is_registered(self): assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == ()