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>

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:

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"):

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,)