diff --git a/litellm/constants.py b/litellm/constants.py index b9bea466247..3d0beb22d9f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -130,6 +130,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_TOOL_LISTING_MAX_PAGES: Final = 1000 # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index f0a1bff8fdc..cebe7b13fe9 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -39,6 +39,7 @@ from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import ( GetPromptRequestParams, GetPromptResult, + PaginatedRequestParams, Prompt, ResourceTemplate, TextContent, @@ -47,7 +48,7 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_MAX_PAGES from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -603,17 +604,46 @@ class MCPClient: """ verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_tools_operation(session: ClientSession): - return await session.list_tools() + async def _list_tools_operation(session: ClientSession) -> list[MCPTool]: + tools: list[MCPTool] = [] + cursor: Optional[str] = None + pages_fetched = 0 + seen_cursors: set[str] = set() + + while True: + result = ( + await session.list_tools() + if cursor is None + else await session.list_tools(params=PaginatedRequestParams(cursor=cursor)) + ) + pages_fetched += 1 + tools.extend(result.tools) + + next_cursor = getattr(result, "nextCursor", None) + if not isinstance(next_cursor, str): + return tools + if next_cursor in seen_cursors: + raise RuntimeError( + f"MCP server returned a repeated tools/list cursor while listing tools: {next_cursor}" + ) + if pages_fetched >= MCP_TOOL_LISTING_MAX_PAGES: + verbose_logger.warning( + "MCP server tools/list pagination exceeded the maximum " + f"of {MCP_TOOL_LISTING_MAX_PAGES} pages while listing tools; " + f"returning {len(tools)} tools collected so far" + ) + return tools + seen_cursors.add(next_cursor) + cursor = next_cursor try: - result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count: Final = len(result.tools) - tool_names: Final = [tool.name for tool in result.tools] + tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) + tool_count: Final = len(tools) + tool_names: Final = [tool.name for tool in tools] verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) - return result.tools + return tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..3ca6b6c5f90 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1402,11 +1402,7 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - async def _list_tools_session_operation(session): - return await session.list_tools() - - list_tools_response: Final = await client.run_with_session(_list_tools_session_operation) - list_tools_result: Final[list[MCPTool]] = list_tools_response.tools + list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True) model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index aadaadd510e..1887d0fe94a 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport -from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult +from mcp.types import CallToolResult as MCPCallToolResult +from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Tool as MCPTool def test_mcp_client_uses_configurable_default_timeout(): @@ -185,6 +187,119 @@ class TestMCPClientUnitTests: mock_session_instance.initialize.assert_called_once() mock_session_instance.list_tools.assert_called_once() + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") + @patch.object(mcp_client_module, "ClientSession") + async def test_list_tools_follows_next_cursor_until_exhausted( + self, + mock_session_class, + mock_transport, + ): + """Test listing tools follows MCP pagination cursors until exhausted.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + first_page_tools = [ + MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100) + ] + second_page_tool = MCPTool( + name="tool_100", + description="Tool 100", + inputSchema={}, + ) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult(tools=first_page_tools, nextCursor="page-2"), + ListToolsResult(tools=[second_page_tool]), + ] + + 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 isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") + @patch.object(mcp_client_module, "ClientSession") + async def test_list_tools_stops_when_pagination_reaches_page_cap( + self, + mock_session_class, + mock_transport, + monkeypatch, + ): + """Test listing tools returns accumulated tools if an upstream keeps returning new cursors.""" + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 2, raising=False) + + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})], + nextCursor="page-2", + ), + ListToolsResult( + tools=[MCPTool(name="tool_1", description="Tool 1", inputSchema={})], + nextCursor="page-3", + ), + ListToolsResult( + tools=[MCPTool(name="tool_2", description="Tool 2", inputSchema={})], + ), + ] + + client = MCPClient("http://example.com") + result = await client.list_tools(raise_on_error=True) + + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + assert mock_session_instance.list_tools.call_count == 2 + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") + @patch.object(mcp_client_module, "ClientSession") + async def test_list_tools_raises_on_repeated_next_cursor( + self, + mock_session_class, + mock_transport, + ): + """Test listing tools fails if an upstream repeats a cursor.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + mock_session_instance.list_tools.side_effect = [ + ListToolsResult(tools=[], nextCursor="same-cursor"), + ListToolsResult(tools=[], nextCursor="same-cursor"), + ] + + client = MCPClient("http://example.com") + with pytest.raises(RuntimeError, match="repeated tools/list cursor"): + await client.list_tools(raise_on_error=True) + + assert mock_session_instance.list_tools.call_count == 2 + @pytest.mark.asyncio @patch.object(mcp_client_module, "streamable_http_client") @patch.object(mcp_client_module, "ClientSession") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..3cc798ded9e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -952,6 +952,123 @@ class TestListToolsRestAPI: assert scope_inputs == [session_auth] assert reload_calls == [] + async def test_single_server_response_includes_paginated_upstream_tools( + self, + monkeypatch, + ): + """The REST tools/list path should include tools beyond the upstream first page.""" + import litellm.experimental_mcp_client.client as mcp_client_module + from mcp.types import ListToolsResult, PaginatedRequestParams + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + stub_server = MCPServer( + server_id="server-1", + name="stub", + server_name="stub", + alias="stub", + url="https://example.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "stub"}, + ) + stub_server.available_on_public_internet = True + + mock_transport_ctx = AsyncMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) + mock_transport_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "streamable_http_client", + MagicMock(return_value=mock_transport_ctx), + raising=False, + ) + + mock_session_ctx = AsyncMock() + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=None) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[ + MCPTool( + name="first_page_tool", + description="First page tool", + inputSchema={}, + ) + ], + nextCursor="page-2", + ), + ListToolsResult( + tools=[ + MCPTool( + name="second_page_tool", + description="Second page tool", + inputSchema={}, + ) + ] + ), + ] + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + mock_session_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "ClientSession", + MagicMock(return_value=mock_session_ctx), + raising=False, + ) + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "filter_server_ids_by_ip_with_info", + lambda server_ids, client_ip: (server_ids, 0), + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert set(result.keys()) == {"tools", "error", "message"} + assert [tool.name for tool in result["tools"]] == [ + "first_page_tool", + "second_page_tool", + ] + 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 isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; a non-admin passing it stays filtered so the REST endpoint can't be used