From f26c8fad3e65dc7f441d2ced7700afb49c8325ad Mon Sep 17 00:00:00 2001 From: Bill Easton Date: Sun, 12 Apr 2026 11:52:16 -0500 Subject: [PATCH] fix: retry when LLM returns text instead of calling final_response (#3850) * Retry when LLM returns text instead of calling final_response tool Instead of raising RuntimeError immediately when the LLM returns a text response instead of calling the `final_response` tool for structured output, retry up to 3 times with an explicit nudge message asking the model to use the tool. This mirrors the existing retry behavior for validation errors but with a separate, smaller cap. Fixes #3847 Co-Authored-By: Claude Opus 4.6 (1M context) * Add tests for text response retry logic Tests cover: - Text response followed by successful final_response (retry works) - Text response exceeding max retries (raises RuntimeError) - Nudge message appended to history on retry - No retry when result_type is None (text is valid) Addresses review feedback from PR review tool (v1 flagged missing tests as high severity). Co-Authored-By: Claude Opus 4.6 (1M context) * Slim down text response retry tests Remove test_nudge_message_in_history (implementation detail). Reduce boilerplate in remaining 3 tests. Co-Authored-By: Claude Opus 4.6 (1M context) * Fix ruff format Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/fastmcp/server/sampling/run.py | 30 ++++- tests/client/test_sampling_result_types.py | 110 +++++++++++++++++++ tests/server/tasks/test_task_return_types.py | 5 +- 3 files changed, 139 insertions(+), 6 deletions(-) diff --git a/src/fastmcp/server/sampling/run.py b/src/fastmcp/server/sampling/run.py index 4ea31f0b6..de5e28780 100644 --- a/src/fastmcp/server/sampling/run.py +++ b/src/fastmcp/server/sampling/run.py @@ -49,6 +49,9 @@ ResultT = TypeVar("ResultT") # Simplified tool choice type - just the mode string instead of the full MCP object ToolChoiceOption = Literal["auto", "required", "none"] +# How many times we retry when the LLM returns text instead of calling final_response +_MAX_TEXT_RESPONSE_RETRIES = 3 + @dataclass class SamplingResult(Generic[ResultT]): @@ -611,6 +614,8 @@ async def sample_impl( # Convert messages for the loop current_messages: str | Sequence[str | SamplingMessage] = messages + text_response_retries = 0 + for _iteration in range(max_iterations): step = await sample_step_impl( context, @@ -682,11 +687,28 @@ async def sample_impl( if not step.is_tool_use: # For structured output, the LLM must use the final_response tool if result_type is not None and result_type is not str: - raise RuntimeError( - f"Expected structured output of type {result_type.__name__}, " - "but the LLM returned a text response instead of calling " - "the final_response tool." + text_response_retries += 1 + if text_response_retries > _MAX_TEXT_RESPONSE_RETRIES: + raise RuntimeError( + f"Expected structured output of type {result_type.__name__}, " + "but the LLM returned a text response instead of calling " + f"the final_response tool ({text_response_retries} attempts)." + ) + # Nudge the LLM to use the tool + step.history.append( + SamplingMessage( + role="user", + content=TextContent( + type="text", + text=( + "You must call the `final_response` tool to provide " + "your answer. Do not respond with text — use the tool." + ), + ), + ) ) + current_messages = step.history + continue return SamplingResult( text=step.text, result=cast(ResultT, step.text if step.text else ""), diff --git a/tests/client/test_sampling_result_types.py b/tests/client/test_sampling_result_types.py index 73d6fbddc..1aaccdf4a 100644 --- a/tests/client/test_sampling_result_types.py +++ b/tests/client/test_sampling_result_types.py @@ -1,3 +1,4 @@ +import pytest from mcp.types import TextContent from fastmcp import Client, Context, FastMCP @@ -440,3 +441,112 @@ class TestSampleStep: result = await client.call_tool("test_step", {}) assert result.data == "ok" + + +class TestTextResponseRetry: + """Tests for retry logic when LLM returns text instead of calling final_response.""" + + @staticmethod + def _text_reply(text: str = "some text"): + from mcp.types import CreateMessageResultWithTools + + return CreateMessageResultWithTools( + role="assistant", + content=[TextContent(type="text", text=text)], + model="m", + stopReason="endTurn", + ) + + @staticmethod + def _tool_reply(value: int): + from mcp.types import CreateMessageResultWithTools, ToolUseContent + + return CreateMessageResultWithTools( + role="assistant", + content=[ + ToolUseContent( + type="tool_use", + id="c1", + name="final_response", + input={"value": value}, + ) + ], + model="m", + stopReason="toolUse", + ) + + async def test_text_response_then_success(self): + """Text on first call, final_response on second -- verify call_count == 2.""" + from pydantic import BaseModel + + class R(BaseModel): + value: int + + call_count = 0 + + def handler(messages, params, ctx): + nonlocal call_count + call_count += 1 + return self._text_reply() if call_count == 1 else self._tool_reply(42) + + mcp = FastMCP(sampling_handler=handler) + + @mcp.tool + async def t(context: Context) -> str: + return str((await context.sample(messages="q", result_type=R)).result.value) + + async with Client(mcp) as client: + result = await client.call_tool("t", {}) + + assert call_count == 2 + assert result.data == "42" + + async def test_text_response_exceeds_max_retries(self): + """Always text, never tool -- verify error after _MAX_TEXT_RESPONSE_RETRIES+1 calls.""" + from pydantic import BaseModel + + from fastmcp.exceptions import ToolError + from fastmcp.server.sampling.run import _MAX_TEXT_RESPONSE_RETRIES + + class R(BaseModel): + value: int + + call_count = 0 + + def handler(messages, params, ctx): + nonlocal call_count + call_count += 1 + return self._text_reply() + + mcp = FastMCP(sampling_handler=handler) + + @mcp.tool + async def t(context: Context) -> str: + return str((await context.sample(messages="q", result_type=R)).result) + + async with Client(mcp) as client: + with pytest.raises(ToolError, match="attempts"): + await client.call_tool("t", {}) + + assert call_count == _MAX_TEXT_RESPONSE_RETRIES + 1 + + async def test_no_retry_when_result_type_is_none(self): + """Text response with no result_type -- single call, normal return.""" + call_count = 0 + + def handler(messages, params, ctx): + nonlocal call_count + call_count += 1 + return self._text_reply("hello") + + mcp = FastMCP(sampling_handler=handler) + + @mcp.tool + async def t(context: Context) -> str: + return (await context.sample(messages="q")).text or "" + + async with Client(mcp) as client: + result = await client.call_tool("t", {}) + + assert call_count == 1 + assert result.data == "hello" diff --git a/tests/server/tasks/test_task_return_types.py b/tests/server/tasks/test_task_return_types.py index f29f1e19b..967d41ee6 100644 --- a/tests/server/tasks/test_task_return_types.py +++ b/tests/server/tasks/test_task_return_types.py @@ -259,8 +259,9 @@ async def binary_type_server(): ( "return_bytes", type(None), - lambda r: r.data is None - and any("Hello bytes!" in c.text for c in r.content), + lambda r: ( + r.data is None and any("Hello bytes!" in c.text for c in r.content) + ), ), ( "return_uuid",