Merge pull request #40525 from BerriAI/litellm_fix_mcp_optional_discovery

fix(mcp): respect optional discovery capabilities and quiet unsupported methods

litellm/experimental_mcp_client/client.py

@@ -54,18 +54,22 @@ def missing_streamable_http_client_error() -> ImportError:
     )
 
 
-from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
-from mcp.types import CallToolResult as MCPCallToolResult
 from mcp.types import (
+    METHOD_NOT_FOUND,
     ClientResult,
     GetPromptRequestParams,
     GetPromptResult,
+    ListPromptsResult,
+    ListResourcesResult,
+    ListResourceTemplatesResult,
     Prompt,
     ResourceTemplate,
     ServerNotification,
     ServerRequest,
     TextContent,
 )
+from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
+from mcp.types import CallToolResult as MCPCallToolResult
 from mcp.types import Tool as MCPTool
 from pydantic import AnyUrl
 
@@ -777,8 +781,19 @@ class MCPClient:
         """List available prompts from the server."""
         verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
 
-        async def _list_prompts_operation(session: ClientSession):
-            return await session.list_prompts()
+        async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult:
+            capabilities: Final = session.get_server_capabilities()
+            if capabilities is not None and capabilities.prompts is None:
+                return ListPromptsResult(prompts=[])
+            try:
+                return await session.list_prompts()
+            except McpError as error:
+                if error.error.code != METHOD_NOT_FOUND:
+                    raise
+                verbose_logger.debug(
+                    "MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error
+                )
+                return ListPromptsResult(prompts=[])
 
         try:
             result: Final = await self.run_with_session(_list_prompts_operation)
@@ -854,8 +869,19 @@ class MCPClient:
         """List available resources from the server."""
         verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio")
 
-        async def _list_resources_operation(session: ClientSession):
-            return await session.list_resources()
+        async def _list_resources_operation(session: ClientSession) -> ListResourcesResult:
+            capabilities: Final = session.get_server_capabilities()
+            if capabilities is not None and capabilities.resources is None:
+                return ListResourcesResult(resources=[])
+            try:
+                return await session.list_resources()
+            except McpError as error:
+                if error.error.code != METHOD_NOT_FOUND:
+                    raise
+                verbose_logger.debug(
+                    "MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error
+                )
+                return ListResourcesResult(resources=[])
 
         try:
             result: Final = await self.run_with_session(_list_resources_operation)
@@ -890,8 +916,19 @@ class MCPClient:
         """List available resource templates from the server."""
         verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio")
 
-        async def _list_resource_templates_operation(session: ClientSession):
-            return await session.list_resource_templates()
+        async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult:
+            capabilities: Final = session.get_server_capabilities()
+            if capabilities is not None and capabilities.resources is None:
+                return ListResourceTemplatesResult(resourceTemplates=[])
+            try:
+                return await session.list_resource_templates()
+            except McpError as error:
+                if error.error.code != METHOD_NOT_FOUND:
+                    raise
+                verbose_logger.debug(
+                    "MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error
+                )
+                return ListResourceTemplatesResult(resourceTemplates=[])
 
         try:
             result: Final = await self.run_with_session(_list_resource_templates_operation)

tests/test_litellm/experimental_mcp_client/test_mcp_client.py

@@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import anyio
 import httpx
 import pytest
+import respx
 from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
 from mcp import McpError
 from mcp.client.streamable_http import streamable_http_client
@@ -1686,3 +1687,211 @@ async def test_empty_http_event_stream_uses_the_existing_request_deadline() -> N
                 timeout=3,
             )
     assert isinstance(as_mcp_read_timeout(caught.value), TimeoutError)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list"))
+@pytest.mark.parametrize(
+    "outcome",
+    (
+        "absent",
+        "other_capability",
+        "supported",
+        "method_not_found",
+        "internal_error",
+        "unauthorized",
+        "timeout",
+        "initialize_not_found",
+    ),
+)
+async def test_optional_discovery_capabilities_and_errors(
+    method: str, outcome: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    import logging
+    from unittest.mock import Mock
+
+    from mcp.types import JSONRPCRequest
+
+    capability: Final = "prompts" if method == "prompts/list" else "resources"
+    field: Final = {
+        "prompts/list": "prompts",
+        "resources/list": "resources",
+        "resources/templates/list": "resourceTemplates",
+    }[method]
+    advertised: Final = "resources" if capability == "prompts" else "prompts"
+    entry: Final = {
+        "prompts/list": {"name": "example"},
+        "resources/list": {"name": "example", "uri": "test://example"},
+        "resources/templates/list": {"name": "example", "uriTemplate": "test://{name}"},
+    }[method]
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        payload: Final = JSONRPCMessage.model_validate_json(request.content).root
+        if not isinstance(payload, JSONRPCRequest):
+            return httpx.Response(202)
+        if outcome == "initialize_not_found":
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": payload.id,
+                    "error": {"code": -32601, "message": "Initialization rejected"},
+                },
+            )
+        if payload.method == "initialize":
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": payload.id,
+                    "result": {
+                        "protocolVersion": LATEST_PROTOCOL_VERSION,
+                        "capabilities": {}
+                        if outcome == "absent"
+                        else {advertised if outcome == "other_capability" else capability: {}},
+                        "serverInfo": {"name": "discovery", "version": "1"},
+                    },
+                },
+            )
+        if outcome == "timeout":
+            raise httpx.ReadTimeout("Optional list timed out", request=request)
+        if outcome == "unauthorized":
+            return httpx.Response(401)
+        if outcome in ("method_not_found", "internal_error", "absent", "other_capability"):
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": payload.id,
+                    "error": {
+                        "code": -32603 if outcome == "internal_error" else -32601,
+                        "message": "Optional list rejected",
+                    },
+                },
+            )
+        return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry]}})
+
+    responder: Final = Mock(side_effect=respond)
+    caplog.set_level(logging.DEBUG, logger="LiteLLM")
+    with respx.mock(base_url="https://example.com") as router:
+        router.route().mock(side_effect=responder)
+        client: Final = MCPClient(server_url="https://example.com/mcp")
+        operation: Final = {
+            "prompts/list": client.list_prompts,
+            "resources/list": client.list_resources,
+            "resources/templates/list": client.list_resource_templates,
+        }[method]
+        result: Final = await operation()
+
+    requests: Final = tuple(
+        JSONRPCMessage.model_validate_json(call.args[0].content).root
+        for call in responder.call_args_list
+        if call.args[0].method == "POST"
+    )
+    assert sum(isinstance(request, JSONRPCRequest) and request.method == method for request in requests) == (
+        0 if outcome in ("absent", "other_capability", "initialize_not_found") else 1
+    )
+    assert [item.name for item in result] == (["example"] if outcome == "supported" else [])
+    failures: Final = tuple(
+        record for record in caplog.records if record.name == "LiteLLM" and record.levelno >= logging.WARNING
+    )
+    if outcome in ("internal_error", "unauthorized", "timeout", "initialize_not_found"):
+        assert any(record.levelno == logging.ERROR and "failed" in record.message for record in failures)
+    else:
+        assert failures == ()
+    if outcome == "method_not_found":
+        assert any(
+            record.levelno == logging.DEBUG and "Optional list rejected" in record.message for record in caplog.records
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("supports_first", (True, False))
+async def test_optional_discovery_uses_each_sessions_capabilities(supports_first: bool) -> None:
+    from unittest.mock import Mock
+    from mcp.types import JSONRPCRequest
+
+    capabilities: Final = iter(({"resources": {}}, {}) if supports_first else ({}, {"resources": {}}))
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        payload: Final = JSONRPCMessage.model_validate_json(request.content).root
+        if not isinstance(payload, JSONRPCRequest):
+            return httpx.Response(202)
+        result: Final = (
+            {
+                "protocolVersion": LATEST_PROTOCOL_VERSION,
+                "capabilities": next(capabilities),
+                "serverInfo": {"name": "changing", "version": "1"},
+            }
+            if payload.method == "initialize"
+            else {"resources": [{"name": "example", "uri": "test://example"}]}
+        )
+        return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result})
+
+    responder: Final = Mock(side_effect=respond)
+    with respx.mock(base_url="https://example.com") as router:
+        router.route().mock(side_effect=responder)
+        client: Final = MCPClient(server_url="https://example.com/mcp")
+        first: Final = await client.list_resources()
+        second: Final = await client.list_resources()
+
+    assert [item.name for item in first] == (["example"] if supports_first else [])
+    assert [item.name for item in second] == ([] if supports_first else ["example"])
+    requests: Final = tuple(
+        JSONRPCMessage.model_validate_json(call.args[0].content).root
+        for call in responder.call_args_list
+        if call.args[0].method == "POST"
+    )
+    assert sum(isinstance(request, JSONRPCRequest) and request.method == "resources/list" for request in requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list"))
+async def test_optional_discovery_preserves_cancellation(method: str) -> None:
+    from mcp.types import JSONRPCRequest
+
+    ready: Final = asyncio.Event()
+    pending: Final = asyncio.Event()
+
+    async def respond(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        payload: Final = JSONRPCMessage.model_validate_json(request.content).root
+        if not isinstance(payload, JSONRPCRequest):
+            return httpx.Response(202)
+        if payload.method == "initialize":
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": payload.id,
+                    "result": {
+                        "protocolVersion": LATEST_PROTOCOL_VERSION,
+                        "capabilities": {"resources": {}, "prompts": {}},
+                        "serverInfo": {"name": "pending", "version": "1"},
+                    },
+                },
+            )
+        ready.set()
+        await pending.wait()
+        return httpx.Response(202)
+
+    with respx.mock(base_url="https://example.com") as router:
+        router.route().mock(side_effect=respond)
+        client: Final = MCPClient(server_url="https://example.com/mcp")
+        operation: Final = {
+            "prompts/list": client.list_prompts,
+            "resources/list": client.list_resources,
+            "resources/templates/list": client.list_resource_templates,
+        }[method]
+        task: Final = asyncio.create_task(operation())
+        try:
+            await asyncio.wait_for(ready.wait(), timeout=3)
+        finally:
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await asyncio.wait_for(task, timeout=3)