mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
5767a2da0f
* fix(mcp): follow tools/list pagination from upstream servers Adopts BerriAI/litellm#32244 by Jupiter363 onto litellm_internal_staging with merge conflicts resolved * fix(mcp): degrade buggy pagination to partial results and bound the preview walk A repeated nextCursor now returns the tools collected so far instead of discarding every page with a RuntimeError, an empty-string cursor is treated as terminal, load_mcp_tools shares the same pagination walk instead of returning only the first page, and the tools/list preview is bounded by the listing timeout instead of only the per-request timeout times the page cap * fix(mcp): annotate deliberate rebind for the preview timeout scope * fix(mcp): bound the shared pagination walk with an overall listing deadline The per-request session read timeout restarts on every page, so direct SDK callers of list_tools and load_mcp_tools could run up to the page cap with no overall bound. The walk now returns the tools collected so far when max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT) expires * fix(mcp): let a per-server timeout extend the pagination deadline MCPClient carries a per-server timeout that can exceed the global default; list_tools now passes max(self.timeout, MCP_TOOL_LISTING_TIMEOUT) into the shared walk so a deliberately slow server is not silently truncated at the global deadline * fix(mcp): honor per-server timeouts in the preview deadline and test the walk sessionless The preview deadline now extends with the created client's own timeout, and the pagination walk's cap, repeated-cursor, and empty-cursor cases are tested directly against the helper instead of through patched SDK internals * fix(mcp): forward the preview request's per-server timeout to the temporary server model The tools preview built its temporary MCPServer without the request's timeout field, so the client factory always fell back to the global default and a per-server timeout could never extend the preview's listing deadline (or its per-request timeout).
468 lines
17 KiB
Python
468 lines
17 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
|
|
from mcp.types import (
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
ListToolsResult,
|
|
PaginatedRequestParams,
|
|
TextContent,
|
|
)
|
|
from mcp.types import Tool as MCPTool
|
|
|
|
from litellm.experimental_mcp_client.tools import (
|
|
list_tools_with_pagination,
|
|
transform_mcp_tool_to_anthropic_tool,
|
|
_get_function_arguments,
|
|
_normalize_mcp_input_schema,
|
|
call_mcp_tool,
|
|
call_openai_tool,
|
|
load_mcp_tools,
|
|
transform_mcp_tool_to_openai_responses_api_tool,
|
|
transform_mcp_tool_to_openai_tool,
|
|
transform_openai_tool_call_request_to_mcp_tool_call_request,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mcp_tool():
|
|
return MCPTool(
|
|
name="test_tool",
|
|
description="A test tool",
|
|
inputSchema={"type": "object", "properties": {"test": {"type": "string"}}},
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_session():
|
|
session = MagicMock()
|
|
session.list_tools = AsyncMock()
|
|
session.call_tool = AsyncMock()
|
|
return session
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_list_tools_result():
|
|
return ListToolsResult(
|
|
tools=[
|
|
MCPTool(
|
|
name="test_tool",
|
|
description="A test tool",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"test": {"type": "string"}},
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mcp_tool_call_result():
|
|
return CallToolResult(content=[TextContent(type="text", text="test_output")])
|
|
|
|
|
|
def test_transform_mcp_tool_to_openai_tool(mock_mcp_tool):
|
|
openai_tool = transform_mcp_tool_to_openai_tool(mock_mcp_tool)
|
|
assert openai_tool["type"] == "function"
|
|
assert openai_tool["function"]["name"] == "test_tool"
|
|
assert openai_tool["function"]["description"] == "A test tool"
|
|
assert openai_tool["function"]["parameters"] == {
|
|
"type": "object",
|
|
"properties": {"test": {"type": "string"}},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
|
|
def testtransform_openai_tool_call_request_to_mcp_tool_call_request(mock_mcp_tool):
|
|
openai_tool = {
|
|
"function": {"name": "test_tool", "arguments": json.dumps({"test": "value"})}
|
|
}
|
|
mcp_tool_call_request = transform_openai_tool_call_request_to_mcp_tool_call_request(
|
|
openai_tool
|
|
)
|
|
assert mcp_tool_call_request.name == "test_tool"
|
|
assert mcp_tool_call_request.arguments == {"test": "value"}
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_load_mcp_tools_mcp_format(mock_session, mock_list_tools_result):
|
|
mock_session.list_tools.return_value = mock_list_tools_result
|
|
result = await load_mcp_tools(mock_session, format="mcp")
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], MCPTool)
|
|
assert result[0].name == "test_tool"
|
|
mock_session.list_tools.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result):
|
|
mock_session.list_tools.return_value = mock_list_tools_result
|
|
result = await load_mcp_tools(mock_session, format="openai")
|
|
assert len(result) == 1
|
|
assert result[0]["type"] == "function"
|
|
assert result[0]["function"]["name"] == "test_tool"
|
|
mock_session.list_tools.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_load_mcp_tools_follows_pagination(mock_session):
|
|
mock_session.list_tools.side_effect = [
|
|
ListToolsResult(
|
|
tools=[
|
|
MCPTool(name="tool_a", description="a", inputSchema={}),
|
|
MCPTool(name="tool_b", description="b", inputSchema={}),
|
|
],
|
|
nextCursor="page-2",
|
|
),
|
|
ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]),
|
|
]
|
|
result = await load_mcp_tools(mock_session, format="mcp")
|
|
assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"]
|
|
assert mock_session.list_tools.call_count == 2
|
|
second_call_params = mock_session.list_tools.call_args_list[1].kwargs["params"]
|
|
assert isinstance(second_call_params, PaginatedRequestParams)
|
|
assert second_call_params.cursor == "page-2"
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch):
|
|
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2)
|
|
mock_session.list_tools.side_effect = [
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
|
nextCursor="page-2",
|
|
),
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
|
|
nextCursor="page-3",
|
|
),
|
|
ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]),
|
|
]
|
|
result = await list_tools_with_pagination(mock_session)
|
|
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
|
assert mock_session.list_tools.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
|
|
mock_session.list_tools.side_effect = [
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
|
nextCursor="same-cursor",
|
|
),
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_1", description="1", inputSchema={})],
|
|
nextCursor="same-cursor",
|
|
),
|
|
]
|
|
result = await list_tools_with_pagination(mock_session)
|
|
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
|
assert mock_session.list_tools.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session):
|
|
mock_session.list_tools.side_effect = [
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_0", description="0", inputSchema={})],
|
|
nextCursor="",
|
|
),
|
|
]
|
|
result = await list_tools_with_pagination(mock_session)
|
|
assert [tool.name for tool in result] == ["tool_0"]
|
|
mock_session.list_tools.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch):
|
|
import anyio
|
|
|
|
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
|
|
|
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2)
|
|
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.2)
|
|
|
|
async def slow_page(params=None):
|
|
await anyio.sleep(0.15)
|
|
idx = int(params.cursor) if params is not None else 0
|
|
return ListToolsResult(
|
|
tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})],
|
|
nextCursor=str(idx + 1),
|
|
)
|
|
|
|
mock_session.list_tools = slow_page
|
|
result = await list_tools_with_pagination(mock_session)
|
|
|
|
assert [tool.name for tool in result] == ["tool_0"]
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_session, monkeypatch):
|
|
import anyio
|
|
|
|
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
|
|
|
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1)
|
|
monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.1)
|
|
|
|
async def slow_page(params=None):
|
|
await anyio.sleep(0.15)
|
|
idx = int(params.cursor) if params is not None else 0
|
|
tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})]
|
|
if idx == 0:
|
|
return ListToolsResult(tools=tools, nextCursor="1")
|
|
return ListToolsResult(tools=tools)
|
|
|
|
mock_session.list_tools = slow_page
|
|
result = await list_tools_with_pagination(mock_session, listing_deadline=2.0)
|
|
|
|
assert [tool.name for tool in result] == ["tool_0", "tool_1"]
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_load_mcp_tools_openai_format_spans_pages(mock_session):
|
|
mock_session.list_tools.side_effect = [
|
|
ListToolsResult(
|
|
tools=[MCPTool(name="tool_a", description="a", inputSchema={})],
|
|
nextCursor="page-2",
|
|
),
|
|
ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]),
|
|
]
|
|
result = await load_mcp_tools(mock_session, format="openai")
|
|
assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"]
|
|
|
|
|
|
def test_get_function_arguments():
|
|
# Test with string arguments
|
|
function = {"arguments": '{"test": "value"}'}
|
|
result = _get_function_arguments(function)
|
|
assert result == {"test": "value"}
|
|
|
|
# Test with dict arguments
|
|
function = {"arguments": {"test": "value"}}
|
|
result = _get_function_arguments(function)
|
|
assert result == {"test": "value"}
|
|
|
|
# Test with invalid JSON string
|
|
function = {"arguments": "invalid json"}
|
|
result = _get_function_arguments(function)
|
|
assert result == {}
|
|
|
|
# Test with no arguments
|
|
function = {}
|
|
result = _get_function_arguments(function)
|
|
assert result == {}
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_call_openai_tool(mock_session, mock_mcp_tool_call_result):
|
|
mock_session.call_tool.return_value = mock_mcp_tool_call_result
|
|
openai_tool = {
|
|
"function": {"name": "test_tool", "arguments": json.dumps({"test": "value"})}
|
|
}
|
|
result = await call_openai_tool(mock_session, openai_tool)
|
|
print("result of call_openai_tool", result)
|
|
assert result.content[0].text == "test_output"
|
|
mock_session.call_tool.assert_called_once_with(
|
|
name="test_tool", arguments={"test": "value"}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_call_mcp_tool(mock_session, mock_mcp_tool_call_result):
|
|
mock_session.call_tool.return_value = mock_mcp_tool_call_result
|
|
request_params = CallToolRequestParams(
|
|
name="test_tool", arguments={"test": "value"}
|
|
)
|
|
result = await call_mcp_tool(mock_session, request_params)
|
|
print("call_mcp_tool result", result)
|
|
assert result.content[0].text == "test_output"
|
|
mock_session.call_tool.assert_called_once_with(
|
|
name="test_tool", arguments={"test": "value"}
|
|
)
|
|
|
|
|
|
def test_normalize_mcp_input_schema():
|
|
"""Test MCP input schema normalization for OpenAI compatibility."""
|
|
# Test case 1: Empty/None schema should get default structure
|
|
assert _normalize_mcp_input_schema(None) == {
|
|
"type": "object",
|
|
"properties": {},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
assert _normalize_mcp_input_schema({}) == {
|
|
"type": "object",
|
|
"properties": {},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
# Test case 2: Schema with only type should get properties added
|
|
schema_with_type_only = {"type": "object"}
|
|
normalized = _normalize_mcp_input_schema(schema_with_type_only)
|
|
assert normalized == {
|
|
"type": "object",
|
|
"properties": {},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
# Test case 3: Schema missing type should get type added
|
|
schema_missing_type = {"properties": {"param": {"type": "string"}}}
|
|
normalized = _normalize_mcp_input_schema(schema_missing_type)
|
|
assert normalized == {
|
|
"type": "object",
|
|
"properties": {"param": {"type": "string"}},
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
# Test case 4: Complete schema should be preserved with additionalProperties added
|
|
complete_schema = {
|
|
"type": "object",
|
|
"properties": {"param": {"type": "string"}},
|
|
"required": ["param"],
|
|
}
|
|
normalized = _normalize_mcp_input_schema(complete_schema)
|
|
assert normalized == {
|
|
"type": "object",
|
|
"properties": {"param": {"type": "string"}},
|
|
"required": ["param"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
# Test case 5: Schema with existing additionalProperties should be preserved
|
|
schema_with_additional = {
|
|
"type": "object",
|
|
"properties": {"param": {"type": "string"}},
|
|
"additionalProperties": True,
|
|
}
|
|
normalized = _normalize_mcp_input_schema(schema_with_additional)
|
|
assert normalized["additionalProperties"] == True
|
|
|
|
|
|
def test_transform_mcp_tool_to_openai_responses_api_tool():
|
|
"""Test transformation to OpenAI Responses API tool format with schema normalization."""
|
|
# Test case 1: Tool with minimal schema (the problematic case from the error)
|
|
minimal_tool = MCPTool(
|
|
name="GitMCP-fetch_litellm_documentation",
|
|
description="Fetch entire documentation file from GitHub repository",
|
|
inputSchema={"type": "object"}, # This was causing the error
|
|
)
|
|
|
|
openai_tool = transform_mcp_tool_to_openai_responses_api_tool(minimal_tool)
|
|
assert openai_tool["name"] == "GitMCP-fetch_litellm_documentation"
|
|
assert openai_tool["type"] == "function"
|
|
assert openai_tool["strict"] == False
|
|
assert openai_tool["parameters"]["type"] == "object"
|
|
assert openai_tool["parameters"]["properties"] == {}
|
|
assert openai_tool["parameters"]["additionalProperties"] == False
|
|
|
|
# Test case 2: Tool with complete schema
|
|
complete_tool = MCPTool(
|
|
name="test_tool_complete",
|
|
description="A test tool with complete schema",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string", "description": "Search query"}},
|
|
"required": ["query"],
|
|
},
|
|
)
|
|
|
|
openai_tool = transform_mcp_tool_to_openai_responses_api_tool(complete_tool)
|
|
assert openai_tool["parameters"]["type"] == "object"
|
|
assert "query" in openai_tool["parameters"]["properties"]
|
|
assert openai_tool["parameters"]["required"] == ["query"]
|
|
assert openai_tool["parameters"]["additionalProperties"] == False
|
|
|
|
|
|
def test_transform_mcp_tool_to_anthropic_tool():
|
|
"""
|
|
Regression test (LIT-4517): MCP tools must reach /v1/messages in Anthropic's
|
|
own tool shape.
|
|
|
|
Given: An MCP tool
|
|
When: It is transformed for the Anthropic Messages API
|
|
Then: It carries name/description/input_schema, the shape that endpoint
|
|
accepts, rather than an OpenAI function block
|
|
|
|
/v1/messages rejects an OpenAI-shaped tool outright ("Input tag 'function'
|
|
does not match any of the expected tags"), so reusing either OpenAI
|
|
transform here loses every MCP tool.
|
|
"""
|
|
tool = MCPTool(
|
|
name="read_wiki_structure",
|
|
description="Get a list of documentation topics",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"repoName": {"type": "string"}},
|
|
"required": ["repoName"],
|
|
},
|
|
)
|
|
|
|
anthropic_tool = transform_mcp_tool_to_anthropic_tool(tool)
|
|
|
|
assert anthropic_tool["name"] == "read_wiki_structure"
|
|
assert anthropic_tool["description"] == "Get a list of documentation topics"
|
|
assert anthropic_tool["type"] == "custom"
|
|
assert anthropic_tool["input_schema"]["type"] == "object"
|
|
assert "repoName" in anthropic_tool["input_schema"]["properties"]
|
|
assert anthropic_tool["input_schema"]["required"] == ["repoName"]
|
|
assert "function" not in anthropic_tool, "Anthropic tools must not carry an OpenAI function block"
|
|
assert "parameters" not in anthropic_tool, "Anthropic names the schema input_schema, not parameters"
|
|
|
|
|
|
def test_transform_mcp_tool_to_anthropic_tool_normalizes_empty_schema():
|
|
"""A tool with no declared arguments must still present a valid object schema."""
|
|
anthropic_tool = transform_mcp_tool_to_anthropic_tool(
|
|
MCPTool(name="noargs", description=None, inputSchema={})
|
|
)
|
|
|
|
assert anthropic_tool["name"] == "noargs"
|
|
assert anthropic_tool["description"] == ""
|
|
assert anthropic_tool["input_schema"]["type"] == "object"
|
|
assert anthropic_tool["input_schema"]["properties"] == {}
|
|
|
|
|
|
def test_transform_mcp_tool_to_anthropic_tool_strips_keys_anthropic_rejects():
|
|
"""
|
|
Regression test (LIT-4517): an MCP schema with keys Anthropic does not accept
|
|
must be sanitized, so the same tool cannot succeed on /chat/completions and 400
|
|
on /v1/messages.
|
|
|
|
Given: An MCP tool whose inputSchema carries $schema, legacy definitions and oneOf
|
|
When: It is transformed for the Anthropic Messages API
|
|
Then: Only keys in AnthropicInputSchema survive, matching the chat path
|
|
|
|
The chat path runs the schema through the same sanitizer, so before this the two
|
|
routes diverged: a clean-schema server (deepwiki) worked on both, but a server
|
|
with a richer schema would be rejected only on messages.
|
|
"""
|
|
from litellm.types.llms.anthropic import AnthropicInputSchema
|
|
|
|
tool = MCPTool(
|
|
name="rich",
|
|
description="tool with a dirty schema",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {"q": {"type": "string"}},
|
|
"required": ["q"],
|
|
"$schema": "http://json-schema.org/draft-07/schema#",
|
|
"definitions": {"D": {"type": "string"}},
|
|
"oneOf": [{"required": ["q"]}],
|
|
},
|
|
)
|
|
|
|
anthropic_tool = transform_mcp_tool_to_anthropic_tool(tool)
|
|
schema_keys = set(anthropic_tool["input_schema"].keys())
|
|
|
|
assert schema_keys <= set(AnthropicInputSchema.__annotations__.keys()), (
|
|
f"schema must only carry keys Anthropic accepts, got {schema_keys}"
|
|
)
|
|
assert "$schema" not in schema_keys
|
|
assert "definitions" not in schema_keys
|
|
assert "oneOf" not in schema_keys
|
|
assert anthropic_tool["input_schema"]["properties"] == {"q": {"type": "string"}}
|
|
assert anthropic_tool["input_schema"]["required"] == ["q"]
|