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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* Fix ruff format

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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 ""),

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"

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",