scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(mcp): follow tools/list pagination from upstream servers

Browse source 
Adopts BerriAI/litellm#32244 by Jupiter363 onto litellm_internal_staging
with merge conflicts resolved
This commit is contained in:
Yucheng Zhu 2026-09-01 11:50:32 -07:00
parent 8a4ba78869
commit 772dba1d11
5 changed files with 272 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -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.

44
litellm/experimental_mcp_client/client.py
View file

@ -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

6
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

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

117
tests/mcp_tests/test_mcp_client_unit.py
View file

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

117
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -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