mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-10-11 07:23:19 +02:00
239 lines
7.7 KiB
Python
239 lines
7.7 KiB
Python
import importlib.util
|
|
|
|
import pytest
|
|
|
|
from fastmcp import Client, ClientGroup, FastMCP
|
|
from fastmcp.exceptions import ToolError
|
|
from fastmcp.experimental.client.code_mode import CodeModeClient
|
|
from fastmcp.experimental.transforms.code_mode import GetTags, ListTools, Search
|
|
|
|
requires_monty = pytest.mark.skipif(
|
|
importlib.util.find_spec("pydantic_monty") is None,
|
|
reason="pydantic-monty is required for the real Monty sandbox provider",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def math_server() -> FastMCP:
|
|
mcp = FastMCP("Math")
|
|
|
|
@mcp.tool(tags={"math"})
|
|
def add(x: int, y: int) -> int:
|
|
"""Add two numbers."""
|
|
return x + y
|
|
|
|
@mcp.tool(tags={"math"})
|
|
def multiply(x: int, y: int) -> int:
|
|
"""Multiply two numbers."""
|
|
return x * y
|
|
|
|
@mcp.tool
|
|
def boom() -> None:
|
|
"""Always fails."""
|
|
raise ToolError("deliberate tool failure")
|
|
|
|
return mcp
|
|
|
|
|
|
@pytest.fixture
|
|
def text_server() -> FastMCP:
|
|
mcp = FastMCP("Text")
|
|
|
|
@mcp.tool
|
|
def shout(text: str) -> str:
|
|
"""Uppercase text."""
|
|
return text.upper()
|
|
|
|
return mcp
|
|
|
|
|
|
async def test_list_tools_returns_only_meta_tools(math_server: FastMCP) -> None:
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
tools = await code_mode.list_tools()
|
|
|
|
assert [tool.name for tool in tools] == ["search", "get_schema", "execute"]
|
|
|
|
|
|
async def test_search_and_get_schema_read_the_remote_catalog(
|
|
math_server: FastMCP,
|
|
) -> None:
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
found = await code_mode.call_tool("search", {"query": "multiply numbers"})
|
|
schema = await code_mode.call_tool("get_schema", {"tools": ["add"]})
|
|
|
|
assert "- multiply: Multiply two numbers." in found.data
|
|
assert "### add" in schema.data
|
|
assert "`x` (integer, required)" in schema.data
|
|
|
|
|
|
async def test_discovery_sees_tags_from_remote_meta(math_server: FastMCP) -> None:
|
|
code_mode = CodeModeClient(Client(math_server), discovery_tools=[GetTags()])
|
|
async with code_mode:
|
|
result = await code_mode.call_tool("tags", {})
|
|
|
|
assert "- math (2 tools)" in result.data
|
|
assert "- untagged (1 tool)" in result.data
|
|
|
|
|
|
async def test_discovery_reflects_tools_added_after_connect(
|
|
math_server: FastMCP,
|
|
) -> None:
|
|
code_mode = CodeModeClient(Client(math_server), discovery_tools=[ListTools()])
|
|
async with code_mode:
|
|
math_server.tool(lambda: "late", name="late_tool")
|
|
result = await code_mode.call_tool("list_tools", {})
|
|
|
|
assert "late_tool" in result.data
|
|
|
|
|
|
@requires_monty
|
|
async def test_execute_chains_remote_calls(math_server: FastMCP) -> None:
|
|
code = (
|
|
"a = await call_tool('add', {'x': 3, 'y': 4})\n"
|
|
"return await call_tool('multiply', {'x': a['result'], 'y': 2})"
|
|
)
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
result = await code_mode.call_tool("execute", {"code": code})
|
|
|
|
assert result.data == {"result": 14}
|
|
|
|
|
|
@requires_monty
|
|
async def test_execute_chains_tools_across_a_client_group(
|
|
math_server: FastMCP, text_server: FastMCP
|
|
) -> None:
|
|
group = ClientGroup({"math": Client(math_server), "text": Client(text_server)})
|
|
code = (
|
|
"total = await call_tool('math_add', {'x': 1, 'y': 2})\n"
|
|
"return await call_tool('text_shout', {'text': f\"sum={total['result']}\"})"
|
|
)
|
|
async with CodeModeClient(group) as code_mode:
|
|
result = await code_mode.call_tool("execute", {"code": code})
|
|
|
|
assert result.data == {"result": "SUM=3"}
|
|
|
|
|
|
@requires_monty
|
|
@pytest.mark.parametrize(
|
|
("target", "failing_call", "expected_message"),
|
|
[
|
|
(
|
|
"client",
|
|
"await call_tool('boom', {})",
|
|
"call_tool('boom') failed: deliberate tool failure",
|
|
),
|
|
("group", "await call_tool('math_nope', {})", "Unknown tool: math_nope"),
|
|
],
|
|
ids=["tool-error", "group-unknown-tool"],
|
|
)
|
|
async def test_call_tool_errors_are_catchable_in_the_sandbox(
|
|
math_server: FastMCP, target: str, failing_call: str, expected_message: str
|
|
) -> None:
|
|
wrapped = (
|
|
Client(math_server)
|
|
if target == "client"
|
|
else ClientGroup({"math": Client(math_server)})
|
|
)
|
|
add = "add" if target == "client" else "math_add"
|
|
code = (
|
|
f"total = (await call_tool('{add}', {{'x': 2, 'y': 3}}))['result']\n"
|
|
"caught = None\n"
|
|
"try:\n"
|
|
f" {failing_call}\n"
|
|
"except Exception as exc:\n"
|
|
" caught = str(exc)\n"
|
|
"return {'caught': caught, 'total': total}"
|
|
)
|
|
async with CodeModeClient(wrapped) as code_mode:
|
|
result = await code_mode.call_tool("execute", {"code": code})
|
|
|
|
assert result.data == {"caught": expected_message, "total": 5}
|
|
|
|
|
|
@requires_monty
|
|
async def test_validation_errors_list_valid_parameters(math_server: FastMCP) -> None:
|
|
code = "return await call_tool('add', {'x': 1, 'z': 2})"
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
with pytest.raises(ToolError) as exc_info:
|
|
await code_mode.call_tool("execute", {"code": code})
|
|
|
|
message = str(exc_info.value)
|
|
assert "call_tool('add') failed" in message
|
|
assert "Valid parameters for add (* = required): x*, y*" in message
|
|
|
|
|
|
@requires_monty
|
|
async def test_uncaught_failure_is_an_error_result(math_server: FastMCP) -> None:
|
|
code = "return await call_tool('boom', {})"
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
with pytest.raises(ToolError, match="deliberate tool failure"):
|
|
await code_mode.call_tool("execute", {"code": code})
|
|
result = await code_mode.call_tool(
|
|
"execute", {"code": code}, raise_on_error=False
|
|
)
|
|
|
|
assert result.is_error
|
|
assert result.data is None
|
|
|
|
|
|
@requires_monty
|
|
async def test_max_tool_calls_is_enforced(math_server: FastMCP) -> None:
|
|
code = (
|
|
"for _ in range(3):\n"
|
|
" await call_tool('add', {'x': 1, 'y': 1})\n"
|
|
"return 'done'"
|
|
)
|
|
code_mode = CodeModeClient(Client(math_server), max_tool_calls=2)
|
|
async with code_mode:
|
|
with pytest.raises(ToolError, match="Tool call limit exceeded: at most 2"):
|
|
await code_mode.call_tool("execute", {"code": code})
|
|
|
|
|
|
async def test_unknown_meta_tool_is_an_error(math_server: FastMCP) -> None:
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
with pytest.raises(ToolError, match="Unknown tool: 'add'"):
|
|
await code_mode.call_tool("add", {"x": 1, "y": 2})
|
|
raw = await code_mode.call_tool_mcp("add", {"x": 1, "y": 2})
|
|
|
|
assert raw.is_error
|
|
|
|
|
|
async def test_invalid_meta_tool_arguments_are_an_error_result(
|
|
math_server: FastMCP,
|
|
) -> None:
|
|
async with CodeModeClient(Client(math_server)) as code_mode:
|
|
raw = await code_mode.call_tool_mcp("get_schema", {"tools": "add", "x": 1})
|
|
|
|
assert raw.is_error
|
|
|
|
|
|
async def test_custom_execute_name_and_description(math_server: FastMCP) -> None:
|
|
code_mode = CodeModeClient(
|
|
Client(math_server),
|
|
discovery_tools=[],
|
|
execute_tool_name="run_python",
|
|
execute_description="Run it.",
|
|
)
|
|
tools = await code_mode.list_tools()
|
|
|
|
assert [(tool.name, tool.description) for tool in tools] == [
|
|
("run_python", "Run it.")
|
|
]
|
|
|
|
|
|
def test_discovery_name_colliding_with_execute_is_rejected(
|
|
math_server: FastMCP,
|
|
) -> None:
|
|
with pytest.raises(ValueError, match="collides with execute_tool_name"):
|
|
CodeModeClient(Client(math_server), discovery_tools=[Search(name="execute")])
|
|
|
|
|
|
async def test_context_manager_connects_the_wrapped_client(
|
|
math_server: FastMCP,
|
|
) -> None:
|
|
client = Client(math_server)
|
|
async with CodeModeClient(client) as code_mode:
|
|
assert code_mode.client is client
|
|
assert client.is_connected()
|
|
|
|
assert not client.is_connected()
|