From 5bab18810641ed2322a5c3ff4f72d077407ced49 Mon Sep 17 00:00:00 2001 From: Bill Easton Date: Mon, 9 Feb 2026 19:43:53 -0600 Subject: [PATCH] Add concurrent tool execution with sequential flag (#3022) Co-authored-by: Bill Easton Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> --- docs/servers/sampling.mdx | 47 ++ src/fastmcp/server/context.py | 18 + src/fastmcp/server/sampling/run.py | 145 ++++-- src/fastmcp/server/sampling/sampling_tool.py | 7 + tests/client/test_sampling.py | 486 +++++++++++++++++++ 5 files changed, 650 insertions(+), 53 deletions(-) diff --git a/docs/servers/sampling.mdx b/docs/servers/sampling.mdx index f15d50aa3..8ea479eb0 100644 --- a/docs/servers/sampling.mdx +++ b/docs/servers/sampling.mdx @@ -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 +) +``` + + +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. + + ### Client Requirements @@ -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. + + 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. + + @@ -511,6 +554,10 @@ tool_result = ToolResultContent( If True, mask detailed error messages from tool execution. + + + Controls parallel execution of tools. `None` (default) for sequential, `0` for unlimited parallel, or a positive integer for bounded concurrency. + diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index a9245fa2b..31e759c98 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -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 diff --git a/src/fastmcp/server/sampling/run.py b/src/fastmcp/server/sampling/run.py index 729968916..c9aa94a76 100644 --- a/src/fastmcp/server/sampling/run.py +++ b/src/fastmcp/server/sampling/run.py @@ -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 diff --git a/src/fastmcp/server/sampling/sampling_tool.py b/src/fastmcp/server/sampling/sampling_tool.py index 106c55fc6..877be71c5 100644 --- a/src/fastmcp/server/sampling/sampling_tool.py +++ b/src/fastmcp/server/sampling/sampling_tool.py @@ -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, ) diff --git a/tests/client/test_sampling.py b/tests/client/test_sampling.py index e379d6707..e5a45adc7 100644 --- a/tests/client/test_sampling.py +++ b/tests/client/test_sampling.py @@ -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)."""