mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
Add concurrent tool execution with sequential flag (#3022)
Co-authored-by: Bill Easton <strawgate@users.noreply.github.com> Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
This commit is contained in:
parent
939cf5fcf2
commit
5bab188106
5 changed files with 650 additions and 53 deletions
|
|
@ -289,6 +289,45 @@ def search(query: str) -> str:
|
|||
|
||||
`ToolError` messages always pass through to the LLM, making it the escape hatch for errors you want the LLM to see and handle.
|
||||
|
||||
### Concurrent Tool Execution
|
||||
|
||||
By default, tools execute sequentially — one at a time, in order. When your tools are independent (no shared state between them), you can execute them in parallel with `tool_concurrency`:
|
||||
|
||||
```python
|
||||
result = await ctx.sample(
|
||||
messages="Research these three topics",
|
||||
tools=[search, fetch_url],
|
||||
tool_concurrency=0, # Unlimited parallel execution
|
||||
)
|
||||
```
|
||||
|
||||
The `tool_concurrency` parameter controls how many tools run at once:
|
||||
|
||||
- **`None`** (default): Sequential execution
|
||||
- **`0`**: Unlimited parallel execution
|
||||
- **`N > 0`**: Execute at most N tools concurrently
|
||||
|
||||
For tools that must not run concurrently (file writes, shared state mutations, etc.), mark them as `sequential` when creating the `SamplingTool`:
|
||||
|
||||
```python
|
||||
from fastmcp.server.sampling import SamplingTool
|
||||
|
||||
db_writer = SamplingTool.from_function(
|
||||
write_to_db,
|
||||
sequential=True, # Forces all tools in the batch to run sequentially
|
||||
)
|
||||
|
||||
result = await ctx.sample(
|
||||
messages="Process this data",
|
||||
tools=[search, db_writer],
|
||||
tool_concurrency=0, # Would be parallel, but db_writer forces sequential
|
||||
)
|
||||
```
|
||||
|
||||
<Note>
|
||||
When any tool in a batch has `sequential=True`, the entire batch executes sequentially regardless of `tool_concurrency`. This is a conservative guarantee — if one tool needs ordering, all tools in that batch respect it.
|
||||
</Note>
|
||||
|
||||
### Client Requirements
|
||||
|
||||
<Note>
|
||||
|
|
@ -463,6 +502,10 @@ tool_result = ToolResultContent(
|
|||
If True, mask detailed error messages from tool execution. When None (default), uses the global `settings.mask_error_details` value. Tools can raise `ToolError` to bypass masking and provide specific error messages to the LLM.
|
||||
</ResponseField>
|
||||
|
||||
<ResponseField name="tool_concurrency" type="int | None" default="None">
|
||||
Controls parallel execution of tools. `None` (default) for sequential, `0` for unlimited parallel, or a positive integer for bounded concurrency. If any tool has `sequential=True`, all tools execute sequentially regardless.
|
||||
</ResponseField>
|
||||
|
||||
</Expandable>
|
||||
|
||||
<Expandable title="Response">
|
||||
|
|
@ -511,6 +554,10 @@ tool_result = ToolResultContent(
|
|||
<ResponseField name="mask_error_details" type="bool | None" default="None">
|
||||
If True, mask detailed error messages from tool execution.
|
||||
</ResponseField>
|
||||
|
||||
<ResponseField name="tool_concurrency" type="int | None" default="None">
|
||||
Controls parallel execution of tools. `None` (default) for sequential, `0` for unlimited parallel, or a positive integer for bounded concurrency.
|
||||
</ResponseField>
|
||||
</Expandable>
|
||||
|
||||
<Expandable title="Response">
|
||||
|
|
|
|||
|
|
@ -754,6 +754,7 @@ class Context:
|
|||
tool_choice: ToolChoiceOption | str | None = None,
|
||||
execute_tools: bool = True,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SampleStep:
|
||||
"""
|
||||
Make a single LLM sampling call.
|
||||
|
|
@ -777,6 +778,12 @@ class Context:
|
|||
mask_error_details: If True, mask detailed error messages from tool
|
||||
execution. When None (default), uses the global settings value.
|
||||
Tools can raise ToolError to bypass masking.
|
||||
tool_concurrency: Controls parallel execution of tools:
|
||||
- None (default): Sequential execution (one at a time)
|
||||
- 0: Unlimited parallel execution
|
||||
- N > 0: Execute at most N tools concurrently
|
||||
If any tool has sequential=True, all tools execute sequentially
|
||||
regardless of this setting.
|
||||
|
||||
Returns:
|
||||
SampleStep containing:
|
||||
|
|
@ -810,6 +817,7 @@ class Context:
|
|||
tool_choice=tool_choice,
|
||||
auto_execute_tools=execute_tools,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
@ -824,6 +832,7 @@ class Context:
|
|||
tools: Sequence[SamplingTool | Callable[..., Any]] | None = None,
|
||||
result_type: type[ResultT],
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SamplingResult[ResultT]:
|
||||
"""Overload: With result_type, returns SamplingResult[ResultT]."""
|
||||
|
||||
|
|
@ -839,6 +848,7 @@ class Context:
|
|||
tools: Sequence[SamplingTool | Callable[..., Any]] | None = None,
|
||||
result_type: None = None,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SamplingResult[str]:
|
||||
"""Overload: Without result_type, returns SamplingResult[str]."""
|
||||
|
||||
|
|
@ -853,6 +863,7 @@ class Context:
|
|||
tools: Sequence[SamplingTool | Callable[..., Any]] | None = None,
|
||||
result_type: type[ResultT] | None = None,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SamplingResult[ResultT] | SamplingResult[str]:
|
||||
"""
|
||||
Send a sampling request to the client and await the response.
|
||||
|
|
@ -883,6 +894,12 @@ class Context:
|
|||
mask_error_details: If True, mask detailed error messages from tool
|
||||
execution. When None (default), uses the global settings value.
|
||||
Tools can raise ToolError to bypass masking.
|
||||
tool_concurrency: Controls parallel execution of tools:
|
||||
- None (default): Sequential execution (one at a time)
|
||||
- 0: Unlimited parallel execution
|
||||
- N > 0: Execute at most N tools concurrently
|
||||
If any tool has sequential=True, all tools execute sequentially
|
||||
regardless of this setting.
|
||||
|
||||
Returns:
|
||||
SamplingResult[T] containing:
|
||||
|
|
@ -906,6 +923,7 @@ class Context:
|
|||
tools=tools,
|
||||
result_type=result_type,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from collections.abc import Callable, Sequence
|
|||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Generic, Literal, cast
|
||||
|
||||
import anyio
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
CreateMessageResult,
|
||||
|
|
@ -31,6 +32,7 @@ from typing_extensions import TypeVar
|
|||
from fastmcp import settings
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.sampling.sampling_tool import SamplingTool
|
||||
from fastmcp.utilities.async_utils import gather
|
||||
from fastmcp.utilities.json_schema import compress_schema
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
from fastmcp.utilities.types import get_cached_typeadapter
|
||||
|
|
@ -239,6 +241,7 @@ async def execute_tools(
|
|||
tool_calls: list[ToolUseContent],
|
||||
tool_map: dict[str, SamplingTool],
|
||||
mask_error_details: bool = False,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> list[ToolResultContent]:
|
||||
"""Execute tool calls and return results.
|
||||
|
||||
|
|
@ -249,66 +252,96 @@ async def execute_tools(
|
|||
When masked, only generic error messages are returned to the LLM.
|
||||
Tools can explicitly raise ToolError to bypass masking when they want
|
||||
to provide specific error messages to the LLM.
|
||||
tool_concurrency: Controls parallel execution of tools:
|
||||
- None (default): Sequential execution (one at a time)
|
||||
- 0: Unlimited parallel execution
|
||||
- N > 0: Execute at most N tools concurrently
|
||||
If any tool has sequential=True, all tools execute sequentially
|
||||
regardless of this setting.
|
||||
|
||||
Returns:
|
||||
List of tool result content blocks.
|
||||
List of tool result content blocks in the same order as tool_calls.
|
||||
"""
|
||||
tool_results: list[ToolResultContent] = []
|
||||
if tool_concurrency is not None and tool_concurrency < 0:
|
||||
raise ValueError(
|
||||
f"tool_concurrency must be None, 0 (unlimited), or a positive integer, "
|
||||
f"got {tool_concurrency}"
|
||||
)
|
||||
|
||||
for tool_use in tool_calls:
|
||||
async def _execute_single_tool(tool_use: ToolUseContent) -> ToolResultContent:
|
||||
"""Execute a single tool and return its result."""
|
||||
tool = tool_map.get(tool_use.name)
|
||||
if tool is None:
|
||||
tool_results.append(
|
||||
ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"Error: Unknown tool '{tool_use.name}'",
|
||||
)
|
||||
],
|
||||
isError=True,
|
||||
)
|
||||
return ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[
|
||||
TextContent(
|
||||
type="text",
|
||||
text=f"Error: Unknown tool '{tool_use.name}'",
|
||||
)
|
||||
],
|
||||
isError=True,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
result_value = await tool.run(tool_use.input)
|
||||
tool_results.append(
|
||||
ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=str(result_value))],
|
||||
)
|
||||
)
|
||||
except ToolError as e:
|
||||
# ToolError is the escape hatch - always pass message through
|
||||
logger.exception(f"Error calling sampling tool '{tool_use.name}'")
|
||||
tool_results.append(
|
||||
ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=str(e))],
|
||||
isError=True,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
# Generic exceptions - mask based on setting
|
||||
logger.exception(f"Error calling sampling tool '{tool_use.name}'")
|
||||
if mask_error_details:
|
||||
error_text = f"Error executing tool '{tool_use.name}'"
|
||||
else:
|
||||
error_text = f"Error executing tool '{tool_use.name}': {e}"
|
||||
tool_results.append(
|
||||
ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=error_text)],
|
||||
isError=True,
|
||||
)
|
||||
)
|
||||
|
||||
return tool_results
|
||||
try:
|
||||
result_value = await tool.run(tool_use.input)
|
||||
return ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=str(result_value))],
|
||||
)
|
||||
except ToolError as e:
|
||||
# ToolError is the escape hatch - always pass message through
|
||||
logger.exception(f"Error calling sampling tool '{tool_use.name}'")
|
||||
return ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=str(e))],
|
||||
isError=True,
|
||||
)
|
||||
except Exception as e:
|
||||
# Generic exceptions - mask based on setting
|
||||
logger.exception(f"Error calling sampling tool '{tool_use.name}'")
|
||||
if mask_error_details:
|
||||
error_text = f"Error executing tool '{tool_use.name}'"
|
||||
else:
|
||||
error_text = f"Error executing tool '{tool_use.name}': {e}"
|
||||
return ToolResultContent(
|
||||
type="tool_result",
|
||||
toolUseId=tool_use.id,
|
||||
content=[TextContent(type="text", text=error_text)],
|
||||
isError=True,
|
||||
)
|
||||
|
||||
# Check if any tool requires sequential execution
|
||||
requires_sequential = any(
|
||||
tool.sequential
|
||||
for tool_use in tool_calls
|
||||
if (tool := tool_map.get(tool_use.name)) is not None
|
||||
)
|
||||
|
||||
# Execute sequentially if required or if concurrency is None (default)
|
||||
if tool_concurrency is None or requires_sequential:
|
||||
tool_results: list[ToolResultContent] = []
|
||||
for tool_use in tool_calls:
|
||||
result = await _execute_single_tool(tool_use)
|
||||
tool_results.append(result)
|
||||
return tool_results
|
||||
|
||||
# Execute in parallel
|
||||
if tool_concurrency == 0:
|
||||
# Unlimited parallel execution
|
||||
return await gather(*[_execute_single_tool(tc) for tc in tool_calls])
|
||||
else:
|
||||
# Bounded parallel execution with semaphore
|
||||
semaphore = anyio.Semaphore(tool_concurrency)
|
||||
|
||||
async def bounded_execute(tool_use: ToolUseContent) -> ToolResultContent:
|
||||
async with semaphore:
|
||||
return await _execute_single_tool(tool_use)
|
||||
|
||||
return await gather(*[bounded_execute(tc) for tc in tool_calls])
|
||||
|
||||
|
||||
# --- Helper functions for sampling ---
|
||||
|
|
@ -412,6 +445,7 @@ async def sample_step_impl(
|
|||
tool_choice: ToolChoiceOption | str | None = None,
|
||||
auto_execute_tools: bool = True,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SampleStep:
|
||||
"""Implementation of Context.sample_step().
|
||||
|
||||
|
|
@ -498,7 +532,10 @@ async def sample_step_impl(
|
|||
else settings.mask_error_details
|
||||
)
|
||||
tool_results: list[ToolResultContent] = await execute_tools(
|
||||
step_tool_calls, tool_map, mask_error_details=effective_mask
|
||||
step_tool_calls,
|
||||
tool_map,
|
||||
mask_error_details=effective_mask,
|
||||
tool_concurrency=tool_concurrency,
|
||||
)
|
||||
|
||||
if tool_results:
|
||||
|
|
@ -523,6 +560,7 @@ async def sample_impl(
|
|||
tools: Sequence[SamplingTool | Callable[..., Any]] | None = None,
|
||||
result_type: type[ResultT] | None = None,
|
||||
mask_error_details: bool | None = None,
|
||||
tool_concurrency: int | None = None,
|
||||
) -> SamplingResult[ResultT]:
|
||||
"""Implementation of Context.sample().
|
||||
|
||||
|
|
@ -561,6 +599,7 @@ async def sample_impl(
|
|||
tools=sampling_tools,
|
||||
tool_choice=tool_choice,
|
||||
mask_error_details=mask_error_details,
|
||||
tool_concurrency=tool_concurrency,
|
||||
)
|
||||
|
||||
# Check for final_response tool call for structured output
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ class SamplingTool(FastMCPBaseModel):
|
|||
description: str | None = None
|
||||
parameters: dict[str, Any]
|
||||
fn: Callable[..., Any]
|
||||
sequential: bool = False
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
|
@ -79,6 +80,7 @@ class SamplingTool(FastMCPBaseModel):
|
|||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
sequential: bool = False,
|
||||
) -> SamplingTool:
|
||||
"""Create a SamplingTool from a function.
|
||||
|
||||
|
|
@ -89,6 +91,10 @@ class SamplingTool(FastMCPBaseModel):
|
|||
fn: The function to create a tool from.
|
||||
name: Optional name override. Defaults to the function's name.
|
||||
description: Optional description override. Defaults to the function's docstring.
|
||||
sequential: If True, this tool requires sequential execution and prevents
|
||||
parallel execution of all tools in the batch. Set to True for tools
|
||||
with shared state, file writes, or other operations that cannot run
|
||||
concurrently. Defaults to False.
|
||||
|
||||
Returns:
|
||||
A SamplingTool wrapping the function.
|
||||
|
|
@ -106,4 +112,5 @@ class SamplingTool(FastMCPBaseModel):
|
|||
description=description or parsed.description,
|
||||
parameters=parsed.input_schema,
|
||||
fn=parsed.fn,
|
||||
sequential=sequential,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -563,6 +563,492 @@ class TestAutomaticToolLoop:
|
|||
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
|
||||
import time
|
||||
|
||||
from mcp.types import CreateMessageResultWithTools, ToolUseContent
|
||||
|
||||
execution_order: list[tuple[str, float]] = []
|
||||
|
||||
async def slow_tool_a(x: int) -> int:
|
||||
"""Slow tool A."""
|
||||
start = time.time()
|
||||
execution_order.append(("tool_a_start", start))
|
||||
await asyncio.sleep(0.1)
|
||||
execution_order.append(("tool_a_end", time.time()))
|
||||
return x * 2
|
||||
|
||||
async def slow_tool_b(y: int) -> int:
|
||||
"""Slow tool B."""
|
||||
start = time.time()
|
||||
execution_order.append(("tool_b_start", start))
|
||||
await asyncio.sleep(0.1)
|
||||
execution_order.append(("tool_b_end", time.time()))
|
||||
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",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Done!")],
|
||||
model="test-model",
|
||||
stopReason="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
|
||||
events = [e[0] for e in execution_order]
|
||||
assert events == ["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
|
||||
import time
|
||||
|
||||
from mcp.types import CreateMessageResultWithTools, ToolUseContent
|
||||
|
||||
execution_times: dict[str, dict[str, float]] = {}
|
||||
|
||||
async def slow_tool_a(x: int) -> int:
|
||||
"""Slow tool A."""
|
||||
execution_times["tool_a"] = {"start": time.time()}
|
||||
await asyncio.sleep(0.1)
|
||||
execution_times["tool_a"]["end"] = time.time()
|
||||
return x * 2
|
||||
|
||||
async def slow_tool_b(y: int) -> int:
|
||||
"""Slow tool B."""
|
||||
execution_times["tool_b"] = {"start": time.time()}
|
||||
await asyncio.sleep(0.1)
|
||||
execution_times["tool_b"]["end"] = time.time()
|
||||
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",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Done!")],
|
||||
model="test-model",
|
||||
stopReason="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: both tools should overlap in time
|
||||
assert "tool_a" in execution_times
|
||||
assert "tool_b" in execution_times
|
||||
# tool_b should start before tool_a finishes (overlap)
|
||||
assert execution_times["tool_b"]["start"] < execution_times["tool_a"]["end"]
|
||||
|
||||
async def test_concurrent_tool_execution_bounded(self):
|
||||
"""Test bounded parallel execution with tool_concurrency=2."""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
from mcp.types import CreateMessageResultWithTools, ToolUseContent
|
||||
|
||||
execution_order: list[tuple[str, float]] = []
|
||||
|
||||
async def slow_tool(name: str, duration: float = 0.1) -> str:
|
||||
"""Generic slow tool."""
|
||||
execution_order.append((f"{name}_start", time.time()))
|
||||
await asyncio.sleep(duration)
|
||||
execution_order.append((f"{name}_end", time.time()))
|
||||
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", "duration": 0.1},
|
||||
),
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id="call_2",
|
||||
name="slow_tool",
|
||||
input={"name": "tool_2", "duration": 0.1},
|
||||
),
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id="call_3",
|
||||
name="slow_tool",
|
||||
input={"name": "tool_3", "duration": 0.05},
|
||||
),
|
||||
],
|
||||
model="test-model",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Done!")],
|
||||
model="test-model",
|
||||
stopReason="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
|
||||
events = [e[0] for e in execution_order]
|
||||
# First 2 tools should start before either ends
|
||||
assert events[0] in ["tool_1_start", "tool_2_start"]
|
||||
assert events[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 = events.index("tool_3_start")
|
||||
assert (
|
||||
"tool_1_end" in events[:tool_3_start_idx]
|
||||
or "tool_2_end" in events[: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
|
||||
import time
|
||||
|
||||
from mcp.types import CreateMessageResultWithTools, ToolUseContent
|
||||
|
||||
execution_order: list[tuple[str, float]] = []
|
||||
|
||||
async def normal_tool(x: int) -> int:
|
||||
"""Normal tool."""
|
||||
execution_order.append(("normal_start", time.time()))
|
||||
await asyncio.sleep(0.05)
|
||||
execution_order.append(("normal_end", time.time()))
|
||||
return x * 2
|
||||
|
||||
async def sequential_tool(y: int) -> int:
|
||||
"""Sequential tool."""
|
||||
execution_order.append(("sequential_start", time.time()))
|
||||
await asyncio.sleep(0.05)
|
||||
execution_order.append(("sequential_end", time.time()))
|
||||
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",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Done!")],
|
||||
model="test-model",
|
||||
stopReason="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
|
||||
events = [e[0] for e in execution_order]
|
||||
assert events[0] in ["normal_start", "sequential_start"]
|
||||
assert events[1] in ["normal_end", "sequential_end"]
|
||||
# Ensure the second tool starts after the first ends
|
||||
if events[0] == "normal_start":
|
||||
assert events[1] == "normal_end"
|
||||
assert events[2] == "sequential_start"
|
||||
else:
|
||||
assert events[1] == "sequential_end"
|
||||
assert events[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",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Handled errors")],
|
||||
model="test-model",
|
||||
stopReason="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.isError for r in tool_results)
|
||||
assert any(r.isError 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,
|
||||
)
|
||||
|
||||
async def tool_with_delay(value: int, delay: float) -> int:
|
||||
"""Tool that takes variable time."""
|
||||
await asyncio.sleep(delay)
|
||||
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:
|
||||
# Tools with different delays - later tools finish first
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id="call_1",
|
||||
name="tool_with_delay",
|
||||
input={"value": 1, "delay": 0.15},
|
||||
),
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id="call_2",
|
||||
name="tool_with_delay",
|
||||
input={"value": 2, "delay": 0.05},
|
||||
),
|
||||
ToolUseContent(
|
||||
type="tool_use",
|
||||
id="call_3",
|
||||
name="tool_with_delay",
|
||||
input={"value": 3, "delay": 0.1},
|
||||
),
|
||||
],
|
||||
model="test-model",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
else:
|
||||
return CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[TextContent(type="text", text="Done!")],
|
||||
model="test-model",
|
||||
stopReason="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].toolUseId == "call_1"
|
||||
assert tool_results[1].toolUseId == "call_2"
|
||||
assert tool_results[2].toolUseId == "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"]
|
||||
|
||||
|
||||
class TestSamplingResultType:
|
||||
"""Tests for result_type parameter (structured output)."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue