Compare commits

...

1 commit

Author SHA1 Message Date
claude[bot]
af786cef46 Add max_sampling_rounds parameter to sample() method
Add configurable max_sampling_rounds parameter to ctx.sample() to allow
users to control 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.

Co-authored-by: Bill Easton <strawgate@users.noreply.github.com>
2026-01-29 01:45:27 +00:00
4 changed files with 185 additions and 2 deletions

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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