fix: propagate ToolOutput max_retries to OutputToolset#4685
fix: propagate ToolOutput max_retries to OutputToolset#4685alvinttang wants to merge 1 commit intopydantic:mainfrom
Conversation
`ToolOutput(Foo, max_retries=3)` accepts a `max_retries` parameter, but `OutputToolset.build()` never extracted it. The build method read `name`, `description`, and `strict` from `ToolOutput` instances but skipped `max_retries`, so `OutputToolset.__init__` always got the default of 1. Extract `max_retries` from `ToolOutput` alongside the other fields and pass it through to the `cls(...)` constructor call. When multiple `ToolOutput` instances specify `max_retries`, take the maximum value. Fixes pydantic#4678 Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
| kwargs: dict[str, Any] = dict(processors=processors, tool_defs=tool_defs) | ||
| if max_retries is not None: | ||
| kwargs['max_retries'] = max_retries | ||
| return cls(**kwargs) |
There was a problem hiding this comment.
🔴 Agent unconditionally overwrites max_retries, making ToolOutput propagation ineffective
The max_retries value propagated from ToolOutput during OutputToolset.build() is immediately overwritten by the Agent. At pydantic_ai_slim/pydantic_ai/agent/__init__.py:380, the agent does self._output_toolset.max_retries = self._max_result_retries (which defaults to the agent's retries=1), and similarly at agent/__init__.py:632. This means ToolOutput(Foo, max_retries=3) has no effect when used through the Agent — the max_retries is always replaced by the agent-level default. The tests only pass because they call OutputToolset.build() directly, bypassing the Agent entirely, so they don't detect this issue.
Prompt for agents
The fix needs to coordinate changes across two files:
1. In pydantic_ai_slim/pydantic_ai/agent/__init__.py, lines 378-380 and 629-632: The agent should only override max_retries on the output toolset when the user has NOT specified per-ToolOutput max_retries. One approach is to only override when the OutputToolset's max_retries is still the default (1), or to add a flag indicating whether max_retries was explicitly set from ToolOutput. Another approach is to have the agent check whether any ToolOutput in the output spec provided a max_retries, and if so, skip the override.
2. The tests in tests/test_agent_output_schemas.py (lines 707-734) should test through the public Agent API instead of calling OutputToolset.build() directly. For example, create an Agent with ToolOutput(Foo, max_retries=3) and verify the actual retry behavior, or at minimum verify that agent._output_toolset.max_retries reflects the ToolOutput setting.
Was this helpful? React with 👍 or 👎 to provide feedback.
| name = output.name | ||
| description = output.description | ||
| strict = output.strict | ||
| if output.max_retries is not None: | ||
| max_retries = max(max_retries or 0, output.max_retries) |
There was a problem hiding this comment.
🚩 All output tools share a single max_retries (the max across all ToolOutputs)
The PR takes the maximum max_retries across all ToolOutput instances and applies it uniformly to the entire OutputToolset (line 901, used at line 972 in get_tools()). This means if a user specifies ToolOutput(Foo, max_retries=2) and ToolOutput(Bar, max_retries=10), both output tools get max_retries=10. Per-tool granularity would require storing max_retries per processor and passing it individually in get_tools(). This is a design limitation worth noting, though supporting per-tool retries would require deeper changes to the OutputToolset/ToolsetTool architecture.
(Refers to lines 889-901)
Was this helpful? React with 👍 or 👎 to provide feedback.
| async def test_tool_output_max_retries_passed_through(): | ||
| """Test that ToolOutput(max_retries=N) is propagated to OutputToolset. | ||
|
|
||
| Regression test for https://github.com/pydantic/pydantic-ai/issues/4678 | ||
| """ | ||
| from pydantic_ai._output import OutputToolset | ||
|
|
||
| toolset = OutputToolset.build([ToolOutput(Foo, max_retries=3)]) | ||
| assert toolset is not None | ||
| assert toolset.max_retries == 3 | ||
|
|
||
|
|
||
| async def test_tool_output_max_retries_default(): | ||
| """Test that OutputToolset uses default max_retries=1 when not specified.""" | ||
| from pydantic_ai._output import OutputToolset | ||
|
|
||
| toolset = OutputToolset.build([ToolOutput(Foo)]) | ||
| assert toolset is not None | ||
| assert toolset.max_retries == 1 | ||
|
|
||
|
|
||
| async def test_tool_output_max_retries_multiple_takes_max(): | ||
| """Test that with multiple ToolOutputs, the highest max_retries wins.""" | ||
| from pydantic_ai._output import OutputToolset | ||
|
|
||
| toolset = OutputToolset.build([ToolOutput(Foo, max_retries=2), ToolOutput(Bar, max_retries=5)]) | ||
| assert toolset is not None | ||
| assert toolset.max_retries == 5 |
There was a problem hiding this comment.
🚩 Tests bypass the Agent and test internal OutputToolset.build() directly
All three new tests import from pydantic_ai._output (a private module prefixed with _) and call OutputToolset.build() directly rather than testing through the public Agent API. Per tests/AGENTS.md rule:177, tests should go through public APIs. More critically, this testing approach masks the fact that the feature doesn't work end-to-end (as reported in BUG-0001), since the Agent overwrites max_retries at pydantic_ai_slim/pydantic_ai/agent/__init__.py:380.
Was this helpful? React with 👍 or 👎 to provide feedback.
|
Closing in favor of #4687 |
Summary
ToolOutput(Foo, max_retries=3)accepts amax_retriesparameter, butOutputToolset.build()never extracted it — the toolset always got the defaultmax_retries=1max_retriesfromToolOutputinstances inOutputToolset.build()and pass it to the constructorToolOutputinstances specifymax_retries, take the maximum valueFixes #4678
Test plan
test_tool_output_max_retries_passed_through— verifiesToolOutput(Foo, max_retries=3)results intoolset.max_retries == 3test_tool_output_max_retries_default— verifies default remains1whenmax_retriesis not specifiedtest_tool_output_max_retries_multiple_takes_max— verifies multipleToolOutputs take the highestmax_retries🤖 Generated with Claude Code