diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 8c7f4557992..421a3c84d75 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/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) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index b07e5876e8b..713330a8280 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/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)