[Bug] Sampling loop: no cap on consecutive final_response validation retries

[strawgate]

Description

When result_type is set and the LLM calls final_response with invalid data, the sampling loop retries by appending a validation error to history (lines 657-679 in run.py). However, the only limit on these retries is the shared max_iterations = 100 constant that also governs the agentic tool loop.

This means a model that consistently produces invalid JSON for a structured output schema will burn through up to 100 LLM calls before failing — each one costing tokens and time.

Current behavior

# run.py:614
for _iteration in range(max_iterations):  # max_iterations = 100
    step = await sample_step_impl(...)
    
    for tool_call in step.tool_calls:
        if tool_call.name == "final_response":
            try:
                validated_result = type_adapter.validate_python(input_data)
                return SamplingResult(...)  # success
            except ValidationError as e:
                # Append error to history and... just loop again
                step.history.append(...)  # "Validation error: ... Please try again"
    
    current_messages = step.history
    # No separate counter for consecutive validation failures
    # Loop continues up to 100 iterations

Expected behavior

There should be a separate, smaller cap for consecutive final_response validation failures (e.g. 3-5). If the LLM fails validation N times in a row, it's unlikely to self-correct and the loop should raise early rather than burning through 95+ more iterations.

This is distinct from the overall max_iterations which appropriately limits the total agentic tool loop (where the LLM may legitimately call many tools before finishing).

Suggested approach:

max_validation_retries = 3  # or make configurable
consecutive_validation_failures = 0

for _iteration in range(max_iterations):
    ...
    if tool_call.name == "final_response":
        try:
            validated_result = type_adapter.validate_python(input_data)
            return SamplingResult(...)
        except ValidationError as e:
            consecutive_validation_failures += 1
            if consecutive_validation_failures > max_validation_retries:
                raise RuntimeError(
                    f"Structured output validation failed {consecutive_validation_failures} "
                    f"times consecutively for type {result_type.__name__}: {e}"
                ) from e
            # append error to history and retry...
    
    # Reset counter when the LLM does something other than fail validation
    # (e.g. calls a different tool successfully)
    if tool_call.name != "final_response":
        consecutive_validation_failures = 0

Impact

Without this cap:

• A complex Pydantic schema that the model can't satisfy will silently burn 100 API calls
• Each retry includes the full conversation history, so token usage grows quadratically
• The eventual error after 100 iterations is "Sampling exceeded maximum iterations (100)" which doesn't indicate that validation was the root cause

Version

fastmcp 3.2.3

[marvin-context-protocol]

Claude finished @strawgate's task in 2m 14s —— View job

---

Triage Summary

Confirmed bug. The sample_impl loop in src/fastmcp/server/sampling/run.py:614-703 shares a single max_iterations = 100 cap for both normal tool-loop iterations and final_response validation retries. There is no separate counter for consecutive validation failures.

Recommendation

Add a consecutive_validation_failures counter inside sample_impl. Reset it whenever a non-final_response-validation iteration occurs (i.e., a successful tool call, or a different tool entirely). Raise a clear ValidationError-based error after a small number of consecutive failures (3 is reasonable). This is a low-risk, self-contained change.

The cap should not be user-configurable at this stage — hardcode it (e.g., max_validation_retries = 3). It can be promoted to a parameter later if there's demand.

Findings

The relevant code is in sample_impl() at run.py:574-703:

# run.py:594-595
max_iterations = 100

# run.py:614
for _iteration in range(max_iterations):
    step = await sample_step_impl(...)

    # Lines 629-679: check for final_response tool call
    if result_type is not None and result_type is not str and step.is_tool_use:
        for tool_call in step.tool_calls:
            if tool_call.name == "final_response":
                try:
                    validated_result = type_adapter.validate_python(input_data)
                    return SamplingResult(...)  # success
                except ValidationError as e:
                    # Just appends error to history and loops again — no counter
                    step.history.append(...)

Key observations:

1. No separate counter: The ValidationError catch block (line 657) appends an error message to the history and falls through to the next loop iteration. There's no tracking of how many times this has happened consecutively.
2. max_iterations = 100 is hardcoded (line 595), not exposed as a parameter on sample() or sample_impl().
3. The error message on exhaustion is misleading: When all 100 iterations are consumed by validation failures, the user sees "Sampling exceeded maximum iterations (100)" (line 703) — which doesn't indicate validation was the root cause.
4. Existing test coverage: test_result_type_validation_error_retries in tests/client/test_sampling_result_types.py:138-226 tests the happy path (one failure then success), but there's no test for repeated consecutive failures.
5. No related issues or PRs were found — this is a novel report.

Detailed Action Plan

1. Add consecutive validation failure tracking to sample_impl in src/fastmcp/server/sampling/run.py

Before the loop (around line 613), add:

max_validation_retries = 3
consecutive_validation_failures = 0

In the ValidationError catch block (lines 657-679), increment the counter and check the cap:

except ValidationError as e:
    consecutive_validation_failures += 1
    if consecutive_validation_failures > max_validation_retries:
        raise RuntimeError(
            f"Structured output validation failed {consecutive_validation_failures} "
            f"times consecutively for type {result_type.__name__}: {e}"
        ) from e
    # existing error-append logic unchanged...
    step.history.append(...)

After the final_response handling block but still inside the tool-call loop (around line 680), reset the counter when a non-final_response tool call is processed successfully. The cleanest place is after the for tool_call in step.tool_calls block. If no final_response validation failure occurred during this iteration, reset:

# After processing all tool calls for this step, if we didn't hit
# a final_response validation failure, reset the counter
if not had_validation_failure_this_step:
    consecutive_validation_failures = 0

Use a boolean flag had_validation_failure_this_step = False set at the top of each iteration and flipped to True inside the except ValidationError block.

2. Add tests in tests/client/test_sampling_result_types.py

Add a test class or method that:

1. Tests consecutive validation failures hit the cap: Create a sampling_handler that always returns invalid data for final_response. Assert that a RuntimeError is raised after exactly max_validation_retries + 1 calls, and that the error message mentions "validation failed" and the type name.

2. Tests that the counter resets after a non-validation step: Create a handler that:

   • Call 1: returns an invalid final_response (failure Add echo server #1)
   • Call 2: returns a call to a regular tool (counter resets)
   • Call 3: returns an invalid final_response (failure Add echo server #1 again, not Add github workflows #2)
   • Call 4: returns a valid final_response (success)

   Assert that this succeeds — confirming the counter was reset by the intermediate tool call.

3. Run the test suite

uv run pytest tests/client/test_sampling_result_types.py -v
uv run pytest -n auto  # full suite
uv run prek run --all-files  # lint + type checks

Related Files

Repository	File	Relevance	Sections
PrefectHQ/fastmcp	src/fastmcp/server/sampling/run.py	Contains sample_impl — the sampling loop where the fix goes	574-703
PrefectHQ/fastmcp	src/fastmcp/server/context.py	Context.sample() — the public API that calls sample_impl	941-1010
PrefectHQ/fastmcp	tests/client/test_sampling_result_types.py	Existing tests for result_type validation retries — new tests go here	138-226

---

Marvin Context Protocol | Type /marvin to interact further

Give us feedback! React with 🚀 if perfect, 👍 if helpful, 👎 if not.

[marvin-context-protocol]

Thanks for the report. This issue goes beyond what our contributor guidelines ask for — we just need a short problem description and an MRE. Please see our contributing guidelines and condense this issue. We'll triage it once it's trimmed down.