Skip to content

fix: propagate Agent(retries=...) to user-provided toolsets#4745

Merged
DouweM merged 16 commits intopydantic:mainfrom
dsfaccini:propagate-agent-retries
Apr 28, 2026
Merged

fix: propagate Agent(retries=...) to user-provided toolsets#4745
DouweM merged 16 commits intopydantic:mainfrom
dsfaccini:propagate-agent-retries

Conversation

@dsfaccini
Copy link
Copy Markdown
Collaborator

FunctionToolset and FastMCPToolset defaulted max_retries=1, ignoring the agent's retries setting. Changed default to None (inherit from agent) and resolve via ToolManager at runtime.

Resolution chain: ToolToolsetAgent (via ToolManager.default_max_retries)

Pre-Review Checklist

  • Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
  • No breaking changes in accordance with the version policy.
  • Linting and type checking pass per make format and make typecheck.
  • PR title is fit for the release changelog.

Pre-Merge Checklist

  • New tests for any fix or new behavior, maintaining 100% coverage.
  • Updated documentation for new features and behaviors, including docstrings for API docs.

🤖 Generated with Claude Code

`FunctionToolset` and `FastMCPToolset` defaulted `max_retries=1`,
ignoring the agent's `retries` setting. Changed default to `None`
(inherit from agent) and resolve via `ToolManager` at runtime.

Resolution chain: Tool → Toolset → Agent (via ToolManager.default_max_retries)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@github-actions github-actions Bot added the size: S Small PR (≤100 weighted lines) label Mar 20, 2026
@dsfaccini dsfaccini added the bug Report that something isn't working, or PR implementing a fix label Mar 20, 2026
devin-ai-integration[bot]

This comment was marked as resolved.

Comment thread tests/test_toolsets.py Outdated
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 5 additional findings in Devin Review.

Open in Devin Review

Comment thread pydantic_ai_slim/pydantic_ai/toolsets/abstract.py Outdated
@Kludex
Copy link
Copy Markdown
Member

Kludex commented Apr 7, 2026

The retries parameter on Agent defaults to 3, right? If that's right, then this is a breaking change - although I think we can get this in anyway.

…_retries as int

- Propagate default_max_retries onto RunContext in for_run_step() so
  prepare functions see the agent's retry count, not the RunContext default
- Resolve None→agent default in FunctionToolset/FastMCPToolset.get_tools()
  so ToolsetTool.max_retries is always int (remove _resolve_max_retries)
- Update docstrings to document None = inherit from agent at runtime
- Add test_prepare_function_sees_agent_max_retries covering ctx propagation

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
devin-ai-integration[bot]

This comment was marked as resolved.

…ries

# Conflicts:
#	pydantic_ai_slim/pydantic_ai/tool_manager.py
#	pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py
#	tests/test_toolsets.py
devin-ai-integration[bot]

This comment was marked as resolved.

… output-retry fix

The `ctx = replace(ctx, max_retries=self.default_max_retries)` line in
`ToolManager.for_run_step` was overriding the `max_result_retries` value
set by `_agent_graph.py` before each `for_run_step` call. This caused
output validators to see `max_retries=1` (the tool default) instead of the
correct `output_retries` value when `output_retries != retries`.

Upstream PR pydantic#4956 already propagates `ctx.max_retries=max_result_retries`
at all three `for_run_step` call sites in `_agent_graph.py`, so the
override in `ToolManager` is now both redundant and harmful on the output
path.

Also update `build_run_context` in tests to accept a `max_retries` param
so unit tests can simulate realistic agent contexts (where `_agent_graph.py`
sets `ctx.max_retries` before calling `for_run_step`).

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 5 additional findings in Devin Review.

Open in Devin Review

Comment thread tests/test_fastmcp.py Outdated
Comment thread tests/test_fastmcp.py Outdated
Comment thread pydantic_ai_slim/pydantic_ai/toolsets/function.py
Comment thread pydantic_ai_slim/pydantic_ai/toolsets/function.py
… in ToolManager

- ToolsetTool.max_retries widened to int|None; ToolManager resolves None to
  ctx.max_retries at execution time via new _resolve_max_retries helper,
  applied at _build_tool_context, _run_execute_hooks, _raw_execute,
  validate_tool_call
- FunctionToolset.get_tools passes unresolved value to FunctionToolsetTool;
  keeps local resolution only for the prepare_tool_def ctx so prepare
  functions still see the resolved agent-default
- FastMCPToolset.get_tools and tool_for_tool_def pass self.max_retries
  directly, dropping the hardcoded 1 fallback so durable execution paths
  inherit the agent's retries via ToolManager
- _AgentFunctionToolset: widen max_retries to int | None = None for
  parent-class consistency
- temporal _ToolInfo.max_retries mirrors ToolsetTool to int | None
- tests: new test_toolset_tool_max_retries_none_resolved_by_tool_manager
  and test_fastmcp_tool_for_tool_def_uses_toolset_max_retries; rewrite
  test_fastmcp.py::test_get_tools snapshot with direct field assertions

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@github-actions github-actions Bot added size: M Medium PR (101-500 weighted lines) and removed size: S Small PR (≤100 weighted lines) labels Apr 14, 2026
devin-ai-integration[bot]

This comment was marked as resolved.

`ToolManager._resolve_max_retries` and `FunctionToolset.get_tools` were
falling back to `ctx.max_retries`, but `RunContext.max_retries` is populated
from `max_result_retries` (the output-validation retry count), not the tool
retry count. When a toolset/tool has no explicit `max_retries` and the agent
is built with `Agent(retries=..., output_retries=...)` (divergent values),
tools and prepare functions saw the output retry count instead of the tool
retry count.

- ToolManager._resolve_max_retries falls back to self.default_max_retries,
  which holds _max_tool_retries (set by Agent at construction)
- FunctionToolset.get_tools resolves the prepare ctx max_retries from
  ctx.tool_manager.default_max_retries (propagated via build_run_context),
  with ctx.max_retries as the fallback for standalone toolset use
- tests: new regression tests exercise the divergent retries/output_retries
  scenario for both the execute path and the prepare path

Addresses devin r3076977400, r3076977506, r3076977694.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

Root cause of prior workarounds: the graph set `ctx.max_retries =
max_result_retries` before calling `for_run_step`, so regular toolsets'
`get_tools` saw the output retry count. Fix at the source: pass
`tool_manager.default_max_retries`; have `OutputToolset.for_run_step`
snapshot from `self.max_retries` instead of `ctx.max_retries`.

Consequences:
- Revert `ToolsetTool.max_retries` to `int`; each toolset resolves
  `None` at `get_tools` time (tool override ?? toolset default ??
  ctx.max_retries).
- Drop `ToolManager._resolve_max_retries` and the patch-and-restore
  in `for_run_step`; use `tool.max_retries` directly.
- `CombinedToolset._output_prep_ctx` now handles only the one case
  the graph fix can't: output-tool `prepare_output_tools` needs
  `ctx.max_retries` = output retries.
- Collapse `FunctionToolset.get_tools` three-level fallback to two.
- Strengthen new tests with tool-call attempt counts.
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 7 additional findings in Devin Review.

Open in Devin Review

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

github-actions[bot]

This comment was marked as resolved.

…n tests

Addresses auto-review feedback on pydantic#4745.

- Introduce `AbstractToolset._adjusted_ctx_for_get_tools(ctx)` so each
  toolset can customize its own `get_tools` context instead of
  `CombinedToolset` hard-coding knowledge of `OutputToolset` and
  manually unwrapping `WrapperToolset` layers.
  - `OutputToolset` overrides to inject `self.max_retries` into
    `ctx.max_retries` for `prepare_output_tools`.
  - `WrapperToolset` delegates to the wrapped toolset.
- Drop the `else 1` fallbacks in `OutputToolset.for_run_step` /
  `get_tools`: Agent always sets `max_retries` before the run, so
  assert that invariant instead of masking it.
- Strengthen new max_retries tests with `capture_run_messages()` +
  message-trace snapshots per the testing guidelines.

Known follow-up tracked in pydantic#5180 (FastMCP `tool_for_tool_def` fallback).
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 5 additional findings in Devin Review.

Open in Devin Review

…d_ctx_for_get_tools

Agent always sets `OutputToolset.max_retries` before the run, so the
`return ctx` branch is unreachable — coverage failed on it. Replace
the guard with the same `assert self.max_retries is not None` used
in `for_run_step` and `get_tools`, and unconditionally `replace(ctx)`.
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 5 additional findings in Devin Review.

Open in Devin Review

…ction

Replaces the `_adjusted_ctx_for_get_tools` hook on `AbstractToolset`
(and its `WrapperToolset`/`OutputToolset` participants) with an optional
`PreparedToolset.max_retries` kwarg set at agent wiring time.

Addresses @DouweM's review of the prior refactor: don't pollute the
abstract toolset hierarchy with a hook that exists solely for one caller
(`prepare_output_tools`); decide at construction time where both the
wrapper and the retry limit are known.

Changes:

- `AbstractToolset`, `WrapperToolset`, `CombinedToolset`: reverted to
  upstream — no hook, no delegation, `CombinedToolset.get_tools` passes
  `ctx` through to children.
- `PreparedToolset`: new optional `max_retries: int | None = None` field.
  When set, `ctx.max_retries` is overridden for the `prepare_func` call
  only. Default `None` preserves existing behavior.
- `OutputToolset`: removed the hook override.
- `Agent._get_toolset`: when wrapping an `OutputToolset` with
  `prepare_output_tools`, pass `max_retries=output_toolset.max_retries`
  so the callback observes the output retry limit.

Semantically matches pydantic#5075's direction of "ctx.max_retries always
reflects the enforcement scope" — inside `prepare_output_tools` that
scope is the output toolset, whose `max_retries` is the per-tool default.
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 4 additional findings in Devin Review.

Open in Devin Review

Comment on lines +21 to +25
max_retries: int | None = None
"""If set, overrides `ctx.max_retries` passed to `prepare_func`. Used by the agent when
wrapping `OutputToolset` so `prepare_output_tools` sees the output retry limit rather
than the graph-level tool retry default.
"""
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's remove the whole output tools related fix because it was meant to address a bug in the prepare output tool hook but we've decided to deprecate that soon and replace it with dedicated output hooks from #4859 (v2: check consequences for https://github.com/pydantic/pydantic-ai-notes/blob/main/david/v2-cards/)

Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 0 new potential issues.

View 5 additional findings in Devin Review.

Open in Devin Review

@DouweM DouweM merged commit 3387ac3 into pydantic:main Apr 28, 2026
47 checks passed
DouweM added a commit that referenced this pull request Apr 28, 2026
PR #4745 (now in main) split tool retries from output retries so the
`tool_manager.ctx.max_retries` is the **tool** retry budget. Two follow-ups:

- `execute_output_tool_call` builds the validator context using
  `OutputToolset.max_retries` (= agent `max_result_retries`), not
  `tool_manager.ctx.max_retries`. Validators now see the right budget on the
  output-tool path; the text-output path already used `_build_output_run_context`
  with the correct value.
- `prepare_tools` capability hook now receives the tool-manager run context
  (per-tool `ctx.retries` populated by `for_run_step`) instead of a fresh
  `build_run_context` that lost the retry counts. Restores parity with the
  pre-migration `PreparedToolset` wrapping path.
- `test_output_toolset_call_tool_raises` sets `OutputToolset.max_retries`
  explicitly — `get_tools` now asserts it (also #4745).
DouweM added a commit that referenced this pull request Apr 28, 2026
…ols()` time

Wraps the function and output toolsets in `Agent._get_toolset` with `PreparedToolset`s
that dispatch the capability hook chain. The filtered/modified tool defs now flow into
`ToolManager.tools` (so hallucinated calls to filtered tools fail with unknown-tool
errors) and into the model's `ModelRequestParameters` simultaneously, instead of only
the latter as before.

Previously the dispatch happened in `_prepare_request_parameters` after `for_run_step`
had already built `tool_manager.tools`, so the hook only changed what the model saw —
it didn't affect tool execution lookup. Reported by Devin on PR #4859.

The output-side dispatch closure overrides `ctx.max_retries` to the agent's
`max_result_retries` to preserve the contract that `prepare_output_tools` sees the
output retry budget (#4745).

Adds a regression test asserting that a tool removed via `prepare_tools` is unreachable
even when the model hallucinates a call to it.
DouweM added a commit that referenced this pull request Apr 29, 2026
…s additive

Reverts the breaking-change part of the prepare-tools split so existing capabilities
that override `prepare_tools` (released for ~2 weeks) keep getting the full tool
set — function + output — as documented and shipped. `prepare_output_tools` is now
a purely additive hook for output-tool-specific filtering with `ctx.max_retries`
reflecting the output retry budget (the original motivation, see #4745).

Implementation:
- `PrepareTools` capability is back to overriding `get_wrapper_toolset` (matches main),
  so the agent-level `prepare_tools=` arg sees function tools only — unchanged from
  release. The `prepare_tools` capability hook itself now dispatches via a
  `PreparedToolset` wrap on the **combined** toolset, so it sees both function and
  output tool defs and the result still flows into `ToolManager.tools`.
- `PrepareOutputTools` keeps its `prepare_output_tools` hook implementation; the
  hook is dispatched via a `PreparedToolset` wrap on the output toolset specifically
  (with `ctx.max_retries` overridden to the output budget). `Agent(prepare_output_tools=...)`
  injects a `PrepareOutputTools` capability mirroring `prepare_tools=`.

Order: `prepare_output_tools` runs first (innermost — only sees output tools),
then `prepare_tools` sees the merged list (outermost — fires after output prep).

Also updates the `prepare_tools` / `prepare_output_tools` docstrings and the
`docs/capabilities.md` / `docs/hooks.md` sections to reflect the additive shape, and
flips the `TestPrepareToolsHook` snapshot test to assert that the hook sees both
function and output tool kinds.
Alex-Resch pushed a commit to Alex-Resch/pydantic-ai that referenced this pull request Apr 29, 2026
dsfaccini added a commit to dsfaccini/pydantic-ai that referenced this pull request Apr 29, 2026
…(), document tool retries

- Fix `_agent_graph.py`'s `allows_none` empty-response retry path that still
  referenced `increment_retries` / `max_result_retries` after the rename, so
  agents with `output_type=str | None` no longer raise `AttributeError` when
  an output validator triggers a retry. (Devin's r3155859175.)

- Add `output_retries` keyword to `Agent.override()` and the AbstractAgent /
  WrapperAgent / dbos / prefect / temporal `override()` signatures, wired via
  a new `_override_output_retries` `ContextVar`. The previous warning when
  `spec={'output_retries': N}` was passed to `override()` is replaced by
  honoring the value. Precedence is `override > run kwarg > spec > agent
  default`, matching `model` / `deps` / `instructions` so test fixtures that
  wrap call sites in `agent.override(output_retries=N)` keep their override
  effective even when production code passes its own `run(output_retries=N)`.
  (DouweM's r3076333408.)

- Add a "How tool retries are enforced" subsection to `docs/agent.md` covering
  per-tool counters, the three configuration levels (per-tool / per-toolset /
  agent-wide), what `ctx.max_retries` and `ctx.retry` show inside a tool, and
  the inheritance behavior for user-provided toolsets after pydantic#4745. (DouweM's
  r3083353737.)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Report that something isn't working, or PR implementing a fix size: M Medium PR (101-500 weighted lines)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Agent(retries=...) not propagated to user-provided toolsets

3 participants