diff --git a/litellm/mcp_client/client.py b/litellm/mcp_client/client.py index 0461f1f06ff..5516cc95af0 100644 --- a/litellm/mcp_client/client.py +++ b/litellm/mcp_client/client.py @@ -90,7 +90,7 @@ from mcp.types import ( 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, TypeAdapter +from pydantic import AnyUrl, Field, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -245,6 +245,14 @@ _ListPage = TypeVar("_ListPage", ListPromptsResult, ListResourcesResult, ListRes _ListItem = TypeVar("_ListItem") +class _LenientListPromptsResult(ListPromptsResult): + prompts: list[Prompt] = Field(default_factory=list) # mutable-ok: MCP result model owns the decoded list + + +class _LenientListResourcesResult(ListResourcesResult): + resources: list[Resource] = Field(default_factory=list) # mutable-ok: MCP result model owns the decoded list + + async def _run_bounded_cleanup(operation: Callable[[], Awaitable[TSessionResult]], deadline: float) -> TSessionResult: async def run() -> TSessionResult: with anyio.fail_after(max(0, deadline - anyio.current_time()), shield=True): @@ -1324,7 +1332,8 @@ class MCPClient: return ListPromptsResult(prompts=[]) try: items, freshness = await self._list_optional_pages( - lambda params: session.list_prompts(params=params), lambda page: page.prompts + lambda params: session.send_request(ListPromptsRequest(params=params), _LenientListPromptsResult), + lambda page: page.prompts, ) return ListPromptsResult(prompts=items, ttl_ms=freshness.ttl_ms, cache_scope=freshness.cache_scope) except MCPError as error: @@ -1437,7 +1446,10 @@ class MCPClient: return ListResourcesResult(resources=[]) try: items, freshness = await self._list_optional_pages( - lambda params: session.list_resources(params=params), lambda page: page.resources + lambda params: session.send_request( + ListResourcesRequest(params=params), _LenientListResourcesResult + ), + lambda page: page.resources, ) return ListResourcesResult(resources=items, ttl_ms=freshness.ttl_ms, cache_scope=freshness.cache_scope) except MCPError as error: diff --git a/litellm/mcp_client/tools.py b/litellm/mcp_client/tools.py index 4f204b3629e..dc92a8b650d 100644 --- a/litellm/mcp_client/tools.py +++ b/litellm/mcp_client/tools.py @@ -3,14 +3,14 @@ from typing import Final, Literal import anyio from mcp import ClientSession, MCPError -from mcp.types import INTERNAL_ERROR, PaginatedRequestParams +from mcp.types import INTERNAL_ERROR, ListToolsRequest, ListToolsResult, PaginatedRequestParams from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition -from pydantic import TypeAdapter +from pydantic import Field, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -24,6 +24,10 @@ from litellm.types.utils import ChatCompletionMessageToolCall _DECODED_JSON: Final = TypeAdapter(object) +class _LenientListToolsResult(ListToolsResult): + tools: list[MCPTool] = Field(default_factory=list) # mutable-ok: MCP result model owns the decoded list + + ######################################################## # List MCP Tool functions ######################################################## @@ -125,10 +129,18 @@ async def list_tools_with_pagination( with anyio.move_on_after(effective_deadline): for _ in range(MCP_TOOL_LISTING_MAX_PAGES): - result = ( - await session.list_tools() - if cursor is None - else await session.list_tools(params=PaginatedRequestParams(cursor=cursor)) + params = None if cursor is None else PaginatedRequestParams(cursor=cursor) + result = await session.send_request( + ListToolsRequest(params=params), + _LenientListToolsResult, + ) + # ClientSession.list_tools() normally performs this post-processing + # after decoding: it filters invalid x-mcp-header annotations and + # records the derived header/output-schema state used by tools/call. + # Keep those semantics when injecting the lenient result model. + result = session._absorb_tool_listing( # pyright: ignore[reportPrivateUsage] # no public equivalent + result, + complete=cursor is None and result.next_cursor is None, ) tools.extend(result.tools) diff --git a/tests/unit/mcp_client/test_mcp_client.py b/tests/unit/mcp_client/test_mcp_client.py index 866a29b70e9..a2df2ca1d7e 100644 --- a/tests/unit/mcp_client/test_mcp_client.py +++ b/tests/unit/mcp_client/test_mcp_client.py @@ -80,6 +80,33 @@ from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_ _JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "result_type"), + [ + ("list_prompts", mcp_client_module._LenientListPromptsResult), + ("list_resources", mcp_client_module._LenientListResourcesResult), + ], +) +async def test_catalog_list_tolerates_missing_list_field(method_name: str, result_type: type[BaseModel]) -> None: + client = MCPClient(server_url="http://example.com/mcp", transport_type="http") + session = AsyncMock(spec=ClientSession) + session.server_capabilities = None + + async def send_request(_request, actual_result_type): + assert actual_result_type is result_type + return actual_result_type.model_validate({}) + + async def run_optional_discovery(operation): + return await operation(session) + + session.send_request.side_effect = send_request + client._run_optional_discovery = AsyncMock(side_effect=run_optional_discovery) + + assert await getattr(client, method_name)() == [] + session.send_request.assert_awaited_once() + + @pytest.mark.asyncio async def test_tool_continuation_reaches_modern_upstream() -> None: from mcp.types import ElicitResult, TextContent diff --git a/tests/unit/mcp_client/test_tools.py b/tests/unit/mcp_client/test_tools.py index ca3cf29bd90..faec7315406 100644 --- a/tests/unit/mcp_client/test_tools.py +++ b/tests/unit/mcp_client/test_tools.py @@ -39,6 +39,14 @@ def mock_session(): session = MagicMock() session.list_tools = AsyncMock() session.call_tool = AsyncMock() + + async def send_request(request, _result_type): + if request.params is not None: + return await session.list_tools(params=request.params) + return await session.list_tools() + + session.send_request = AsyncMock(side_effect=send_request) + session._absorb_tool_listing.side_effect = lambda result, **_kwargs: result return session @@ -496,3 +504,14 @@ async def test_complete_discovery_preserves_tools_when_completeness_is_required( tool = MCPTool(name="complete", input_schema={"type": "object"}) mock_session.list_tools.return_value = ListToolsResult(tools=[tool]) assert await list_tools_with_pagination(mock_session, require_complete=True) == [tool] + + +@pytest.mark.asyncio +async def test_list_tools_tolerates_missing_tools_field(mock_session): + async def send_request(_request, result_type): + return result_type.model_validate({}) + + mock_session.send_request.side_effect = send_request + + assert await list_tools_with_pagination(mock_session) == [] + mock_session.send_request.assert_awaited_once() diff --git a/tests/unit/proxy/mcp_server/test_mcp_client_unit.py b/tests/unit/proxy/mcp_server/test_mcp_client_unit.py index 1f804f3b425..110fdfac12d 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_client_unit.py +++ b/tests/unit/proxy/mcp_server/test_mcp_client_unit.py @@ -186,16 +186,15 @@ class TestMCPClientUnitTests: }, ) ] - mock_result = MagicMock() - mock_result.tools = mock_tools - mock_session_instance.list_tools.return_value = mock_result + mock_session_instance.send_request.return_value = ListToolsResult(tools=mock_tools) + mock_session_instance._absorb_tool_listing = MagicMock(side_effect=lambda result, **_kwargs: result) client = MCPClient("http://example.com") result = await client.list_tools() assert result == mock_tools mock_session_instance.initialize.assert_called_once() - mock_session_instance.list_tools.assert_called_once() + mock_session_instance.send_request.assert_awaited_once() @pytest.mark.asyncio @patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py @@ -229,17 +228,18 @@ class TestMCPClientUnitTests: description="Tool 100", inputSchema={}, ) - mock_session_instance.list_tools.side_effect = [ + mock_session_instance.send_request.side_effect = [ ListToolsResult(tools=first_page_tools, nextCursor="page-2"), ListToolsResult(tools=[second_page_tool]), ] + mock_session_instance._absorb_tool_listing = MagicMock(side_effect=lambda result, **_kwargs: result) client = MCPClient("http://example.com") result = await client.list_tools() assert result == [*first_page_tools, second_page_tool] - assert mock_session_instance.list_tools.call_count == 2 - second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert mock_session_instance.send_request.await_count == 2 + second_call_params = mock_session_instance.send_request.await_args_list[1].args[0].params assert isinstance(second_call_params, PaginatedRequestParams) assert second_call_params.cursor == "page-2" @@ -267,19 +267,20 @@ class TestMCPClientUnitTests: ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) - mock_session_instance.list_tools.side_effect = [ + mock_session_instance.send_request.side_effect = [ ListToolsResult( tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})], nextCursor="page-2", ), RuntimeError("transient upstream failure"), ] + mock_session_instance._absorb_tool_listing = MagicMock(side_effect=lambda result, **_kwargs: result) client = MCPClient("http://example.com") result = await client.list_tools() assert result == [] - assert mock_session_instance.list_tools.call_count == 2 + assert mock_session_instance.send_request.await_count == 2 @pytest.mark.asyncio @patch.object(mcp_client_module, "streamable_http_client") diff --git a/tests/unit/proxy/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/mcp_server/test_rest_endpoints.py index fdb1ff3cd09..30bb4de615b 100644 --- a/tests/unit/proxy/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_rest_endpoints.py @@ -1616,7 +1616,7 @@ class TestListToolsRestAPI: server_info=Implementation(name="stub", version="1"), ) ) - mock_session_instance.list_tools.side_effect = [ + mock_session_instance.send_request.side_effect = [ ListToolsResult( tools=[ MCPTool( @@ -1637,6 +1637,7 @@ class TestListToolsRestAPI: ] ), ] + mock_session_instance._absorb_tool_listing = MagicMock(side_effect=lambda result, **_kwargs: result) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_session_ctx.__aexit__ = AsyncMock(return_value=None) monkeypatch.setattr( @@ -1685,8 +1686,8 @@ class TestListToolsRestAPI: assert result["error"] is None assert result["message"] == "Successfully retrieved tools" - assert mock_session_instance.list_tools.call_count == 2 - second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert mock_session_instance.send_request.await_count == 2 + second_call_params = mock_session_instance.send_request.await_args_list[1].args[0].params assert isinstance(second_call_params, PaginatedRequestParams) assert second_call_params.cursor == "page-2"