mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
811 lines
30 KiB
Python
811 lines
30 KiB
Python
from typing import cast
|
|
|
|
from mcp_types import TextContent
|
|
|
|
from fastmcp import Client, Context, FastMCP
|
|
from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
|
|
from fastmcp.server.sampling import SamplingTool
|
|
|
|
|
|
class TestAutomaticToolLoop:
|
|
"""Tests for automatic tool execution loop in ctx.sample()."""
|
|
|
|
async def test_automatic_tool_loop_executes_tools(self):
|
|
"""Test that ctx.sample() automatically executes tool calls."""
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
call_count = 0
|
|
tool_was_called = False
|
|
|
|
def get_weather(city: str) -> str:
|
|
"""Get weather for a city."""
|
|
nonlocal tool_was_called
|
|
tool_was_called = True
|
|
return f"Weather in {city}: sunny, 72°F"
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
# First call: return tool use
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="get_weather",
|
|
input={"city": "Seattle"},
|
|
)
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
# Second call: return final response
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="The weather is sunny!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def weather_assistant(question: str, context: Context) -> str:
|
|
result = await context.sample(
|
|
messages=question,
|
|
tools=[get_weather],
|
|
)
|
|
# Get text from SamplingResult
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool(
|
|
"weather_assistant", {"question": "What's the weather?"}
|
|
)
|
|
|
|
assert tool_was_called
|
|
assert call_count == 2
|
|
assert result.data == "The weather is sunny!"
|
|
|
|
async def test_automatic_tool_loop_multiple_tools(self):
|
|
"""Test that multiple tool calls in one response are all executed."""
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
executed_tools: list[str] = []
|
|
|
|
def tool_a(x: int) -> int:
|
|
"""Tool A."""
|
|
executed_tools.append(f"tool_a({x})")
|
|
return x * 2
|
|
|
|
def tool_b(y: int) -> int:
|
|
"""Tool B."""
|
|
executed_tools.append(f"tool_b({y})")
|
|
return y + 10
|
|
|
|
call_count = 0
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
# Return multiple tool calls
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use", id="call_a", name="tool_a", input={"x": 5}
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use", id="call_b", name="tool_b", input={"y": 3}
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def multi_tool(context: Context) -> str:
|
|
result = await context.sample(messages="Run tools", tools=[tool_a, tool_b])
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("multi_tool", {})
|
|
|
|
assert executed_tools == ["tool_a(5)", "tool_b(3)"]
|
|
assert result.data == "Done!"
|
|
|
|
async def test_automatic_tool_loop_handles_unknown_tool(self):
|
|
"""Test that unknown tool names result in error being passed to LLM."""
|
|
from mcp_types import (
|
|
CreateMessageResultWithTools,
|
|
ToolResultContent,
|
|
ToolUseContent,
|
|
)
|
|
|
|
def known_tool() -> str:
|
|
"""A known tool."""
|
|
return "known result"
|
|
|
|
messages_received: list[list[SamplingMessage]] = []
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
messages_received.append(list(messages))
|
|
|
|
if len(messages_received) == 1:
|
|
# Request unknown tool
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="unknown_tool",
|
|
input={},
|
|
)
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Handled error")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_unknown(context: Context) -> str:
|
|
result = await context.sample(messages="Test", tools=[known_tool])
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_unknown", {})
|
|
|
|
# Check that error was passed back in messages
|
|
assert len(messages_received) == 2
|
|
last_messages = messages_received[1]
|
|
# Find the tool result in list content
|
|
tool_result = None
|
|
for msg in last_messages:
|
|
# Tool results are now in a list
|
|
if isinstance(msg.content, list):
|
|
for item in msg.content:
|
|
if isinstance(item, ToolResultContent):
|
|
tool_result = item
|
|
break
|
|
elif isinstance(msg.content, ToolResultContent):
|
|
tool_result = msg.content
|
|
break
|
|
assert tool_result is not None
|
|
assert tool_result.is_error is True
|
|
# Content is list of TextContent objects
|
|
assert isinstance(tool_result.content[0], TextContent)
|
|
error_text = tool_result.content[0].text
|
|
assert "Unknown tool" in error_text
|
|
assert result.data == "Handled error"
|
|
|
|
async def test_automatic_tool_loop_handles_tool_exception(self):
|
|
"""Test that tool exceptions are caught and passed to LLM as errors."""
|
|
from mcp_types import (
|
|
CreateMessageResultWithTools,
|
|
ToolResultContent,
|
|
ToolUseContent,
|
|
)
|
|
|
|
def failing_tool() -> str:
|
|
"""A tool that raises an exception."""
|
|
raise ValueError("Tool failed intentionally")
|
|
|
|
messages_received: list[list[SamplingMessage]] = []
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
messages_received.append(list(messages))
|
|
|
|
if len(messages_received) == 1:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="failing_tool",
|
|
input={},
|
|
)
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Handled error")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_exception(context: Context) -> str:
|
|
result = await context.sample(messages="Test", tools=[failing_tool])
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_exception", {})
|
|
|
|
# Check that error was passed back
|
|
assert len(messages_received) == 2
|
|
last_messages = messages_received[1]
|
|
# Find the tool result in list content
|
|
tool_result = None
|
|
for msg in last_messages:
|
|
# Tool results are now in a list
|
|
if isinstance(msg.content, list):
|
|
for item in msg.content:
|
|
if isinstance(item, ToolResultContent):
|
|
tool_result = item
|
|
break
|
|
elif isinstance(msg.content, ToolResultContent):
|
|
tool_result = msg.content
|
|
break
|
|
assert tool_result is not None
|
|
assert tool_result.is_error is True
|
|
# Content is list of TextContent objects
|
|
assert isinstance(tool_result.content[0], TextContent)
|
|
error_text = tool_result.content[0].text
|
|
assert "Tool failed intentionally" in error_text
|
|
assert result.data == "Handled error"
|
|
|
|
async def test_concurrent_tool_execution_default_sequential(self):
|
|
"""Test that tools execute sequentially by default."""
|
|
import asyncio
|
|
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
# Ordering is guaranteed structurally (the loop awaits each tool call
|
|
# to completion before starting the next when tool_concurrency is
|
|
# None), so no real delay is needed to prove it - a single
|
|
# `asyncio.sleep(0)` still yields control to the event loop.
|
|
execution_order: list[str] = []
|
|
|
|
async def slow_tool_a(x: int) -> int:
|
|
"""Slow tool A."""
|
|
execution_order.append("tool_a_start")
|
|
await asyncio.sleep(0)
|
|
execution_order.append("tool_a_end")
|
|
return x * 2
|
|
|
|
async def slow_tool_b(y: int) -> int:
|
|
"""Slow tool B."""
|
|
execution_order.append("tool_b_start")
|
|
await asyncio.sleep(0)
|
|
execution_order.append("tool_b_end")
|
|
return y + 10
|
|
|
|
call_count = 0
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_a",
|
|
name="slow_tool_a",
|
|
input={"x": 5},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_b",
|
|
name="slow_tool_b",
|
|
input={"y": 3},
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[slow_tool_a, slow_tool_b],
|
|
# Default: tool_concurrency=None (sequential)
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Done!"
|
|
# Verify sequential execution: tool_a must complete before tool_b starts
|
|
assert execution_order == [
|
|
"tool_a_start",
|
|
"tool_a_end",
|
|
"tool_b_start",
|
|
"tool_b_end",
|
|
]
|
|
|
|
async def test_concurrent_tool_execution_unlimited(self):
|
|
"""Test unlimited parallel tool execution with tool_concurrency=0."""
|
|
import asyncio
|
|
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
# tool_a blocks on an event that only tool_b sets. This is only
|
|
# satisfiable if both tools are genuinely running concurrently: under
|
|
# sequential execution tool_b would never start (tool_a would never
|
|
# finish awaiting it) and the test would fail via the wait_for
|
|
# timeout rather than racing on wall-clock timestamps.
|
|
execution_order: list[str] = []
|
|
tool_b_started = asyncio.Event()
|
|
|
|
async def slow_tool_a(x: int) -> int:
|
|
"""Slow tool A."""
|
|
execution_order.append("tool_a_start")
|
|
await asyncio.wait_for(tool_b_started.wait(), timeout=1.0)
|
|
execution_order.append("tool_a_end")
|
|
return x * 2
|
|
|
|
async def slow_tool_b(y: int) -> int:
|
|
"""Slow tool B."""
|
|
execution_order.append("tool_b_start")
|
|
tool_b_started.set()
|
|
execution_order.append("tool_b_end")
|
|
return y + 10
|
|
|
|
call_count = 0
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_a",
|
|
name="slow_tool_a",
|
|
input={"x": 5},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_b",
|
|
name="slow_tool_b",
|
|
input={"y": 3},
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[slow_tool_a, slow_tool_b],
|
|
tool_concurrency=0, # Unlimited parallel
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Done!"
|
|
# Verify parallel execution: tool_b started and finished entirely
|
|
# inside tool_a's blocked wait, which is only possible if the two
|
|
# tools were running concurrently.
|
|
assert execution_order == [
|
|
"tool_a_start",
|
|
"tool_b_start",
|
|
"tool_b_end",
|
|
"tool_a_end",
|
|
]
|
|
|
|
async def test_concurrent_tool_execution_bounded(self):
|
|
"""Test bounded parallel execution with tool_concurrency=2."""
|
|
import asyncio
|
|
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
# tool_1 and tool_2 each block until *both* have started, which is
|
|
# only possible if two slots are occupied simultaneously (proving
|
|
# concurrency=2 admits two tools at once). tool_3 has no blocking
|
|
# branch, so its appearance in the log tells us when the real
|
|
# semaphore in the implementation let it in - only after a slot
|
|
# frees, i.e. after tool_1 or tool_2 finishes.
|
|
execution_order: list[str] = []
|
|
both_started = asyncio.Event()
|
|
started_names: set[str] = set()
|
|
|
|
async def slow_tool(name: str) -> str:
|
|
"""Generic tool used to observe bounded concurrency."""
|
|
execution_order.append(f"{name}_start")
|
|
if name in ("tool_1", "tool_2"):
|
|
started_names.add(name)
|
|
if {"tool_1", "tool_2"} <= started_names:
|
|
both_started.set()
|
|
await asyncio.wait_for(both_started.wait(), timeout=1.0)
|
|
execution_order.append(f"{name}_end")
|
|
return f"{name} done"
|
|
|
|
call_count = 0
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
# Request 3 tools (with concurrency=2, first 2 run parallel, then 3rd)
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="slow_tool",
|
|
input={"name": "tool_1"},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_2",
|
|
name="slow_tool",
|
|
input={"name": "tool_2"},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_3",
|
|
name="slow_tool",
|
|
input={"name": "tool_3"},
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[slow_tool],
|
|
tool_concurrency=2, # Max 2 concurrent
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Done!"
|
|
# Verify that at most 2 tools run concurrently
|
|
# First 2 tools should start before either ends
|
|
assert execution_order[0] in ["tool_1_start", "tool_2_start"]
|
|
assert execution_order[1] in ["tool_1_start", "tool_2_start"]
|
|
# Third tool should start after at least one of the first two finishes
|
|
tool_3_start_idx = execution_order.index("tool_3_start")
|
|
assert (
|
|
"tool_1_end" in execution_order[:tool_3_start_idx]
|
|
or "tool_2_end" in execution_order[:tool_3_start_idx]
|
|
)
|
|
|
|
async def test_sequential_tool_forces_sequential_execution(self):
|
|
"""Test that sequential=True forces all tools to execute sequentially."""
|
|
import asyncio
|
|
|
|
from mcp_types import CreateMessageResultWithTools, ToolUseContent
|
|
|
|
# A sequential=True tool in the batch forces the whole batch through
|
|
# the plain for-loop path (see run.py's `requires_sequential`), so
|
|
# ordering is guaranteed structurally and no real delay is needed.
|
|
execution_order: list[str] = []
|
|
|
|
async def normal_tool(x: int) -> int:
|
|
"""Normal tool."""
|
|
execution_order.append("normal_start")
|
|
await asyncio.sleep(0)
|
|
execution_order.append("normal_end")
|
|
return x * 2
|
|
|
|
async def sequential_tool(y: int) -> int:
|
|
"""Sequential tool."""
|
|
execution_order.append("sequential_start")
|
|
await asyncio.sleep(0)
|
|
execution_order.append("sequential_end")
|
|
return y + 10
|
|
|
|
call_count = 0
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
if call_count == 1:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="normal_tool",
|
|
input={"x": 5},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_2",
|
|
name="sequential_tool",
|
|
input={"y": 3},
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
# Create tools with sequential=True for one of them
|
|
normal = SamplingTool.from_function(normal_tool, sequential=False)
|
|
sequential = SamplingTool.from_function(sequential_tool, sequential=True)
|
|
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[normal, sequential],
|
|
tool_concurrency=0, # Request unlimited, but sequential tool forces sequential
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Done!"
|
|
# Verify sequential execution: first tool must complete before second starts
|
|
assert execution_order[0] in ["normal_start", "sequential_start"]
|
|
assert execution_order[1] in ["normal_end", "sequential_end"]
|
|
# Ensure the second tool starts after the first ends
|
|
if execution_order[0] == "normal_start":
|
|
assert execution_order[1] == "normal_end"
|
|
assert execution_order[2] == "sequential_start"
|
|
else:
|
|
assert execution_order[1] == "sequential_end"
|
|
assert execution_order[2] == "normal_start"
|
|
|
|
async def test_concurrent_tool_execution_error_handling(self):
|
|
"""Test that errors are captured per-tool in parallel execution."""
|
|
from mcp_types import (
|
|
CreateMessageResultWithTools,
|
|
ToolResultContent,
|
|
ToolUseContent,
|
|
)
|
|
|
|
def good_tool() -> str:
|
|
return "success"
|
|
|
|
def bad_tool() -> str:
|
|
raise ValueError("Tool error")
|
|
|
|
messages_received: list[list[SamplingMessage]] = []
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
messages_received.append(list(messages))
|
|
|
|
if len(messages_received) == 1:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use", id="call_1", name="good_tool", input={}
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use", id="call_2", name="bad_tool", input={}
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Handled errors")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[good_tool, bad_tool],
|
|
tool_concurrency=0, # Parallel execution
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Handled errors"
|
|
# Check that tool results include both success and error
|
|
tool_result_message = messages_received[1][-1]
|
|
assert tool_result_message.role == "user"
|
|
tool_results = cast(list[ToolResultContent], tool_result_message.content)
|
|
assert len(tool_results) == 2
|
|
# One should be success, one should be error
|
|
assert any(not r.is_error for r in tool_results)
|
|
assert any(r.is_error for r in tool_results)
|
|
|
|
async def test_concurrent_tool_result_order_preserved(self):
|
|
"""Test that tool results maintain the same order as tool calls."""
|
|
import asyncio
|
|
|
|
from mcp_types import (
|
|
CreateMessageResultWithTools,
|
|
ToolResultContent,
|
|
ToolUseContent,
|
|
)
|
|
|
|
# Chain events so the tools finish in a different order (2, 3, 1)
|
|
# than they were called (1, 2, 3), without depending on real delays:
|
|
# tool 2 finishes immediately and unblocks tool 3, which finishes and
|
|
# unblocks tool 1. This only resolves if all three run concurrently -
|
|
# under sequential execution tool 1 would deadlock waiting on tool 3,
|
|
# which itself would never have been started yet.
|
|
tool_2_done = asyncio.Event()
|
|
tool_3_done = asyncio.Event()
|
|
|
|
async def tool_with_delay(value: int) -> int:
|
|
"""Tool that finishes out of call order."""
|
|
if value == 1:
|
|
await asyncio.wait_for(tool_3_done.wait(), timeout=1.0)
|
|
elif value == 3:
|
|
await asyncio.wait_for(tool_2_done.wait(), timeout=1.0)
|
|
|
|
if value == 2:
|
|
tool_2_done.set()
|
|
elif value == 3:
|
|
tool_3_done.set()
|
|
return value
|
|
|
|
messages_received: list[list[SamplingMessage]] = []
|
|
|
|
def sampling_handler(
|
|
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
|
|
) -> CreateMessageResultWithTools:
|
|
messages_received.append(list(messages))
|
|
|
|
if len(messages_received) == 1:
|
|
# Call order is 1, 2, 3 but they finish out of order (2, 3, 1)
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_1",
|
|
name="tool_with_delay",
|
|
input={"value": 1},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_2",
|
|
name="tool_with_delay",
|
|
input={"value": 2},
|
|
),
|
|
ToolUseContent(
|
|
type="tool_use",
|
|
id="call_3",
|
|
name="tool_with_delay",
|
|
input={"value": 3},
|
|
),
|
|
],
|
|
model="test-model",
|
|
stop_reason="toolUse",
|
|
)
|
|
else:
|
|
return CreateMessageResultWithTools(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Done!")],
|
|
model="test-model",
|
|
stop_reason="endTurn",
|
|
)
|
|
|
|
mcp = FastMCP(sampling_handler=sampling_handler)
|
|
|
|
@mcp.tool
|
|
async def test_tool(context: Context) -> str:
|
|
result = await context.sample(
|
|
messages="Run tools",
|
|
tools=[tool_with_delay],
|
|
tool_concurrency=0, # Parallel execution
|
|
)
|
|
return result.text or ""
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("test_tool", {})
|
|
|
|
assert result.data == "Done!"
|
|
# Check that results are in the correct order (1, 2, 3) despite finishing order (2, 3, 1)
|
|
tool_result_message = messages_received[1][-1]
|
|
tool_results = cast(list[ToolResultContent], tool_result_message.content)
|
|
assert len(tool_results) == 3
|
|
assert tool_results[0].tool_use_id == "call_1"
|
|
assert tool_results[1].tool_use_id == "call_2"
|
|
assert tool_results[2].tool_use_id == "call_3"
|
|
# Check values are correct
|
|
result_texts = [cast(TextContent, r.content[0]).text for r in tool_results]
|
|
assert result_texts == ["1", "2", "3"]
|