From 803f8a69f07ccc9720065e1c79d2e336e6402edd Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 11 Sep 2026 02:56:34 +0000 Subject: [PATCH] test(mcp): assert the no-challenge outcome and register openapi tools through the registry api Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/test_mcp_server.py | 36 ++++++++++--------- .../mcp_server/test_mcp_server_manager.py | 12 +++---- .../guardrail_hooks/test_agent_365.py | 4 ++- 3 files changed, 29 insertions(+), 23 deletions(-) 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 4c77825232b..4fd9c923709 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 @@ -8867,7 +8867,7 @@ class TestAgent365ChallengeAtConnect: litellm.callbacks, guardrail, require_self=False ) - async def _connect(self, server: MCPServer, oauth2_headers: dict[str, str] | None) -> None: + async def _connect(self, server: MCPServer, oauth2_headers: dict[str, str] | None) -> HTTPException | None: from litellm.proxy._experimental.mcp_server import server as server_module with ( @@ -8878,36 +8878,40 @@ class TestAgent365ChallengeAtConnect: server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[]) ), ): - await server_module._raise_preemptive_401_for_unauthenticated_servers( - scope={"type": "http", "method": "POST", "path": "/mcp/tools", "headers": []}, - mcp_servers=["tools"], - oauth2_headers=oauth2_headers, - mcp_server_auth_headers=None, - user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), - client_ip=None, - ) + try: + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": "/mcp/tools", "headers": []}, + mcp_servers=["tools"], + oauth2_headers=oauth2_headers, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + except HTTPException as challenge: + return challenge + return None @pytest.mark.asyncio async def test_no_bearer_gets_the_discovery_challenge(self, agent_365_guardrail): - with pytest.raises(HTTPException) as exc: - await self._connect(self._server([self.GATEWAY_SCOPE]), None) + challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), None) - assert exc.value.status_code == 401 - www_authenticate = (exc.value.headers or {}).get("WWW-Authenticate", "") + 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_bearer_present_connects(self, agent_365_guardrail): - await self._connect(self._server([self.GATEWAY_SCOPE]), {"Authorization": "Bearer entra-user-token"}) + bearer = {"Authorization": "Bearer entra-user-token"} + 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): - await self._connect(self._server(None), None) + assert await self._connect(self._server(None), None) is None @pytest.mark.asyncio async def test_no_registered_guardrail_means_no_challenge(self): - await self._connect(self._server([self.GATEWAY_SCOPE]), None) + assert await self._connect(self._server([self.GATEWAY_SCOPE]), None) is None def _make_obo_server(alias: str) -> MCPServer: 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 ba7b76eae91..b32787d4947 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 @@ -6508,7 +6508,6 @@ class TestMCPServerManager: @pytest.mark.parametrize("add_prefix", [True, False]) async def test_openapi_listing_records_tool_metadata_for_pre_call_hooks(self, add_prefix): from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry - from litellm.types.mcp_server.tool_registry import MCPTool as RegistryTool server = MCPServer( server_id="petstore-id", @@ -6519,14 +6518,15 @@ class TestMCPServerManager: spec_path="https://example.com/petstore.yaml", ) schema = {"type": "object", "properties": {"petId": {"type": "integer"}}} - registered = RegistryTool( - name="petstore-get_pet", description="Fetch a pet", input_schema=schema, handler=lambda: None - ) manager = MCPServerManager() manager._create_mcp_client = AsyncMock(return_value=AsyncMock()) - - with patch.dict(global_mcp_tool_registry.tools, {"petstore-get_pet": registered}, clear=True): + global_mcp_tool_registry.register_tool( + name="petstore-get_pet", description="Fetch a pet", input_schema=schema, handler=lambda: None + ) + try: listed = await manager._get_tools_from_server(server=server, add_prefix=add_prefix) + finally: + global_mcp_tool_registry.unregister_tools_with_prefix("petstore-") assert [t.name for t in listed] == ["petstore-get_pet" if add_prefix else "get_pet"] for spelling in ("get_pet", "petstore-get_pet"): 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 b928ccad4c6..c82f830cd84 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 @@ -991,7 +991,9 @@ class TestAgent365AuthorizationServers: api_key="sk-guarded", user_id="u-2", metadata={"guardrails": ["agent-365-guard"]} ) try: - with patch("litellm.proxy.proxy_server.premium_user", True): + with patch( # test-quality-ok: key-selected guardrails read the proxy server premium global, no injection seam + "litellm.proxy.proxy_server.premium_user", True + ): assert agent_365_authorization_servers(server, plain_key) == () assert agent_365_authorization_servers(server, guarded_key) == (ENTRA_ISSUER,) assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,)