litellm/tests/unit/mcp_client/test_tools.py

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.mcp_client.tools import (
    _get_function_arguments,
    _normalize_mcp_input_schema,
    call_mcp_tool,
    call_openai_tool,
    list_tools_with_pagination,
    load_mcp_tools,
    transform_mcp_tool_to_anthropic_tool,
    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()
@pytest.mark.parametrize("require_complete", [False, True])
async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch, require_complete):
    monkeypatch.setattr("litellm.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={})]),
    ]
    if require_complete:
        from mcp import MCPError

        with pytest.raises(MCPError, match="incomplete"):
            await list_tools_with_pagination(mock_session, require_complete=True)
        return
    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()
@pytest.mark.parametrize("require_complete", [False, True])
async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch, require_complete):
    import anyio

    from litellm.mcp_client.tools import list_tools_with_pagination

    monkeypatch.setattr("litellm.mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2)
    monkeypatch.setattr("litellm.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
    if require_complete:
        from mcp import MCPError

        with pytest.raises(MCPError, match="incomplete"):
            await list_tools_with_pagination(mock_session, require_complete=True)
        return
    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.mcp_client.tools import list_tools_with_pagination

    monkeypatch.setattr("litellm.mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1)
    monkeypatch.setattr("litellm.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"]


@pytest.mark.asyncio
async def test_incomplete_discovery_cannot_satisfy_a_complete_listing(mock_session):
    from mcp.shared.exceptions import MCPError

    mock_session.list_tools.return_value = ListToolsResult(
        tools=[MCPTool(name="partial", input_schema={"type": "object"})], next_cursor="repeat"
    )
    with pytest.raises(MCPError, match="incomplete"):
        await list_tools_with_pagination(mock_session, require_complete=True)


@pytest.mark.asyncio
async def test_complete_discovery_preserves_tools_when_completeness_is_required(mock_session):
    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]