diff --git a/docs/servers/sampling.mdx b/docs/servers/sampling.mdx index f15d50aa3..2212890d0 100644 --- a/docs/servers/sampling.mdx +++ b/docs/servers/sampling.mdx @@ -243,7 +243,21 @@ async def research(question: str, ctx: Context) -> str: return result.text or "" ``` -The LLM sees each function's signature and docstring, using this information to decide when and how to call them. Tool errors are caught and sent back to the LLM, allowing it to recover gracefully. An internal safety limit prevents infinite loops. +The LLM sees each function's signature and docstring, using this information to decide when and how to call them. Tool errors are caught and sent back to the LLM, allowing it to recover gracefully. + +### Iteration Limits + +The `max_sampling_rounds` parameter controls the maximum number of tool-calling iterations (defaults to 100). This prevents infinite loops if the LLM repeatedly calls tools without converging on a final response. + +```python +result = await ctx.sample( + messages=question, + tools=[search], + max_sampling_rounds=50 # Allow up to 50 tool-calling iterations +) +``` + +When the limit is exceeded, a `RuntimeError` is raised. ### Custom Tool Definitions diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 0fe505b16..890d54c56 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -748,6 +748,7 @@ class Context: tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, result_type: type[ResultT], mask_error_details: bool | None = None, + max_sampling_rounds: int = 100, ) -> SamplingResult[ResultT]: """Overload: With result_type, returns SamplingResult[ResultT].""" @@ -763,6 +764,7 @@ class Context: tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, result_type: None = None, mask_error_details: bool | None = None, + max_sampling_rounds: int = 100, ) -> SamplingResult[str]: """Overload: Without result_type, returns SamplingResult[str].""" @@ -777,6 +779,7 @@ class Context: tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, result_type: type[ResultT] | None = None, mask_error_details: bool | None = None, + max_sampling_rounds: int = 100, ) -> SamplingResult[ResultT] | SamplingResult[str]: """ Send a sampling request to the client and await the response. @@ -807,6 +810,9 @@ 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. + max_sampling_rounds: Maximum number of tool-calling iterations before + stopping. Defaults to 100. Prevents infinite loops if the LLM + repeatedly calls tools without converging on a final response. Returns: SamplingResult[T] containing: @@ -824,6 +830,7 @@ class Context: tools=tools, result_type=result_type, mask_error_details=mask_error_details, + max_sampling_rounds=max_sampling_rounds, ) @overload diff --git a/src/fastmcp/server/sampling/run.py b/src/fastmcp/server/sampling/run.py index f5c73afef..fa897c8a3 100644 --- a/src/fastmcp/server/sampling/run.py +++ b/src/fastmcp/server/sampling/run.py @@ -523,6 +523,7 @@ async def sample_impl( tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, result_type: type[ResultT] | None = None, mask_error_details: bool | None = None, + max_sampling_rounds: int = 100, ) -> SamplingResult[ResultT]: """Implementation of Context.sample(). @@ -531,7 +532,7 @@ async def sample_impl( provides a final text response. """ # Safety limit to prevent infinite loops - max_iterations = 100 + max_iterations = max_sampling_rounds # Convert tools to SamplingTools sampling_tools = prepare_tools(tools) diff --git a/tests/server/sampling/test_max_sampling_rounds.py b/tests/server/sampling/test_max_sampling_rounds.py new file mode 100644 index 000000000..23f36d00d --- /dev/null +++ b/tests/server/sampling/test_max_sampling_rounds.py @@ -0,0 +1,161 @@ +"""Tests for max_sampling_rounds parameter.""" + +import pytest +from mcp.types import CreateMessageResultWithTools, TextContent, ToolUseContent + +from fastmcp import Client, Context, FastMCP +from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams +from fastmcp.exceptions import ToolError + + +async def test_max_sampling_rounds_default(): + """Test that default max_sampling_rounds is 100.""" + call_count = 0 + + def loop_tool() -> str: + """A tool that the LLM will keep calling.""" + nonlocal call_count + call_count += 1 + return "keep going" + + # Handler that always returns tool use to create an infinite loop + def sampling_handler( + messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext + ) -> CreateMessageResultWithTools: + # Always request tool use to create infinite loop + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="loop_tool", + input={}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def infinite_tool(context: Context) -> str: + """Tool that never stops calling itself.""" + result = await context.sample( + "Keep calling loop_tool", + tools=[loop_tool], + ) + return result.text or "" + + async with Client(mcp) as client: + with pytest.raises(ToolError, match="Sampling exceeded maximum iterations \\(100\\)"): + await client.call_tool("infinite_tool", {}) + + # Verify it actually made 100 iterations + assert call_count == 100 + + +async def test_max_sampling_rounds_custom(): + """Test that custom max_sampling_rounds works.""" + call_count = 0 + + def loop_tool() -> str: + """A tool that the LLM will keep calling.""" + nonlocal call_count + call_count += 1 + return "keep going" + + # Handler that always returns tool use + def sampling_handler( + messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext + ) -> CreateMessageResultWithTools: + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="loop_tool", + input={}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def limited_tool(context: Context) -> str: + """Tool with custom max rounds.""" + result = await context.sample( + "Keep calling loop_tool", + tools=[loop_tool], + max_sampling_rounds=5, + ) + return result.text or "" + + async with Client(mcp) as client: + with pytest.raises(ToolError, match="Sampling exceeded maximum iterations \\(5\\)"): + await client.call_tool("limited_tool", {}) + + # Verify it only made 5 iterations + assert call_count == 5 + + +async def test_max_sampling_rounds_completes_normally(): + """Test that sampling completes normally when rounds don't exceed limit.""" + call_count = 0 + + def helper_tool() -> str: + """A helper tool.""" + return "done" + + # Handler that returns tool use once, then text response + def sampling_handler( + messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext + ) -> CreateMessageResultWithTools: + nonlocal call_count + call_count += 1 + + if call_count == 1: + # First call: request tool use + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="call_1", + name="helper_tool", + input={}, + ) + ], + model="test-model", + stopReason="toolUse", + ) + else: + # Second call: return text response + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text="All done!")], + model="test-model", + stopReason="endTurn", + ) + + mcp = FastMCP(sampling_handler=sampling_handler) + + @mcp.tool + async def normal_tool(context: Context) -> str: + """Tool that completes normally.""" + result = await context.sample( + "Use helper_tool once", + tools=[helper_tool], + max_sampling_rounds=10, + ) + return result.text or "" + + async with Client(mcp) as client: + result = await client.call_tool("normal_tool", {}) + assert result.data == "All done!" + assert call_count == 2 # Should only take 2 iterations