From 2e118bd672ee29c3625ec4c05428fd0b39adee97 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 3 Oct 2026 15:16:35 +0000 Subject: [PATCH] fix(mcp): advertise the Agent 365 sign-in only when the guardrail gates a tagless connect Anonymous protected-resource discovery now runs the same should_run_guardrail probe as a keyed connect, so a default_on guardrail whose Mode only has tags and no default advertises the gateway authorization server as the merge base did instead of the Entra issuer and scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/agent_365/agent_365.py | 36 +++++++++---------- .../guardrail_hooks/test_agent_365.py | 20 ++++++++++- 2 files changed, 37 insertions(+), 19 deletions(-) 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 dbc041ec806..ed501b007cd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -480,27 +480,27 @@ class Agent365Guardrail(CustomGuardrail): return str(uuid.uuid4()) def caller_sign_in(self, server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> CallerSignIn | None: - """The Entra sign-in this guardrail requires of callers: only a ``default_on`` guardrail the caller's - key or team has not opted out of, because the anonymous metadata fetch that follows a challenge cannot - see which key selected a guardrail 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.""" + """The Entra sign-in this guardrail requires of callers: only a ``default_on`` guardrail whose mode gates a + tagless MCP connect and that the caller's key or team has not opted out of, because the anonymous + metadata fetch that follows a challenge cannot see which key selected a guardrail 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 (self.default_on and server.keeps_caller_authorization): return None - if user_api_key_auth is not None: - probe: Final[_AdmissionProbe] = { - "metadata": { - "user_api_key_metadata": user_api_key_auth.metadata, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.metadata is a raw dict - "user_api_key_team_metadata": user_api_key_auth.team_metadata, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.team_metadata is a raw dict - } + probe: Final[_AdmissionProbe] = { + "metadata": { + "user_api_key_metadata": user_api_key_auth.metadata if user_api_key_auth else None, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.metadata is a raw dict + "user_api_key_team_metadata": user_api_key_auth.team_metadata if user_api_key_auth else None, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.team_metadata is a raw dict } - if ( - self.should_run_guardrail( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # should_run_guardrail takes an untyped data dict - data=probe, event_type=GuardrailEventHooks.pre_mcp_call - ) - is not True - ): - return None + } + if ( + self.should_run_guardrail( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # should_run_guardrail takes an untyped data dict + data=probe, event_type=GuardrailEventHooks.pre_mcp_call + ) + is not True + ): + return None return CallerSignIn( issuers=(ENTRA_ISSUER_TEMPLATE.format(tenant_id=self.tenant_id),), scopes=tuple(server.scopes) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_agent_365.py index 0d5f5b3ad8b..a225be3cfda 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -39,6 +39,7 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, + Mode, SupportedGuardrailIntegrations, ) from litellm.types.mcp import MCPAuth, MCPTransport @@ -179,6 +180,7 @@ def _make_guardrail( exchanger: StubTokenExchanger | None = None, unreachable_fallback: str = "fail_closed", default_on: bool = True, + event_hook: str | Mode = "pre_mcp_call", ) -> Agent365Guardrail: return Agent365Guardrail( guardrail_name="agent-365-guard", @@ -188,7 +190,7 @@ def _make_guardrail( unreachable_fallback=unreachable_fallback, async_handler=handler, token_exchanger=exchanger if exchanger is not None else StubTokenExchanger(_obo_ok()), - event_hook="pre_mcp_call", + event_hook=event_hook, default_on=default_on, ) @@ -1403,6 +1405,22 @@ class TestCallerSignIn: assert guardrail.caller_sign_in(_server(), UserAPIKeyAuth(api_key="k", user_id="u-1")) is not None assert guardrail.caller_sign_in(_server(), None) is not None + def test_tag_mode_advertises_entra_only_when_it_gates_a_tagless_connect(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + plain_key: Final = UserAPIKeyAuth(api_key="k", user_id="u-1") + tag_only: Final = _make_guardrail(FakeHandler([]), event_hook=Mode(tags={"a365": "pre_mcp_call"})) + assert tag_only.caller_sign_in(_server(), plain_key) is None + assert tag_only.caller_sign_in(_server(), None) is None + with_default: Final = _make_guardrail( + FakeHandler([]), event_hook=Mode(tags={"a365": "pre_mcp_call"}, default="pre_mcp_call") + ) + expected: Final = CallerSignIn( + issuers=("https://login.microsoftonline.com/tenant-abc/v2.0",), + scopes=("api://client-xyz/access_as_user",), + ) + assert with_default.caller_sign_in(_server(), plain_key) == expected + assert with_default.caller_sign_in(_server(), None) == expected + def test_obo_server_with_provider_advertises_both_issuers_and_the_server_scopes(self, monkeypatch): monkeypatch.setenv("JWT_ISSUER", "https://jwt-idp.test") guardrail: Final = _make_guardrail(FakeHandler([]))