scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 353

Merge e81a958e90 into e9b1c11098

Browse source
This commit is contained in:
Parafee41 2026-11-06 19:10:49 -08:00 • committed by GitHub
commit fe8c687001
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 93 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/mcp_client/client.py
View file

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

24
litellm/mcp_client/tools.py
View file

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

27
tests/unit/mcp_client/test_mcp_client.py
View file

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

19
tests/unit/mcp_client/test_tools.py
View file

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

19
tests/unit/proxy/mcp_server/test_mcp_client_unit.py
View file

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

7
tests/unit/proxy/mcp_server/test_rest_endpoints.py
View file

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