Adding per tool usage limit

docs/tools-advanced.md

@@ -379,7 +379,7 @@ If a tool requires sequential/serial execution, you can pass the [`sequential`][
 Async functions are run on the event loop, while sync functions are offloaded to threads. To get the best performance, _always_ use an async function _unless_ you're doing blocking I/O (and there's no way to use a non-blocking library instead) or CPU-bound work (like `numpy` or `scikit-learn` operations), so that simple functions are not offloaded to threads unnecessarily.
 
 !!! note "Limiting tool executions"
-    You can cap tool executions within a run using [`UsageLimits(tool_calls_limit=...)`](agents.md#usage-limits). The counter increments only after a successful tool invocation. Output tools (used for [structured output](output.md)) are not counted in the `tool_calls` metric.
+    You can cap the total number of tool executions within a run using [`UsageLimits(tool_calls_limit=...)`](agents.md#usage-limits). For finer control, you can limit how many times a *specific* tool can be called by setting the `max_uses` parameter when registering the tool (e.g., `@agent.tool(max_uses=3)` or `Tool(func, max_uses=3)`). Once a tool reaches its `max_uses` limit, it is automatically removed from the available tools for subsequent steps in the run. The `tool_calls` counter increments only after a successful tool invocation. Output tools (used for [structured output](output.md)) are not counted in the `tool_calls` metric.
 
 ## See Also
 

docs/tools.md

@@ -361,7 +361,7 @@ _(This example is complete, it can be run "as is")_
 
 For more tool features and integrations, see:
 
-- [Advanced Tool Features](tools-advanced.md) - Custom schemas, dynamic tools, tool execution and retries
+- [Advanced Tool Features](tools-advanced.md) - Custom schemas, dynamic tools, tool execution, retries, and usage limits
 - [Toolsets](toolsets.md) - Managing collections of tools
 - [Builtin Tools](builtin-tools.md) - Native tools provided by LLM providers
 - [Common Tools](common-tools.md) - Ready-to-use tool implementations

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -95,6 +95,7 @@ class GraphAgentState:
     retries: int = 0
     run_step: int = 0
     run_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4()))
+    tool_usage: dict[str, int] = dataclasses.field(default_factory=dict)
 
     def increment_retries(
         self,
@@ -821,6 +822,7 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         else DEFAULT_INSTRUMENTATION_VERSION,
         run_step=ctx.state.run_step,
         run_id=ctx.state.run_id,
+        tool_usage=ctx.state.tool_usage,
     )
     validation_context = build_validation_context(ctx.deps.validation_context, run_context)
     run_context = replace(run_context, validation_context=validation_context)
@@ -999,6 +1001,17 @@ async def process_tool_calls(  # noqa: C901
         output_final_result.append(final_result)
 
 
+def _projection_count_of_tool_usage(
+    tool_call_counts: defaultdict[str, int], tool_calls: list[_messages.ToolCallPart]
+) -> None:
+    """Populate a count of tool usage based on the provided tool calls for this run step.
+
+    We will use this to make sure the calls do not exceed tool usage limits.
+    """
+    for call in tool_calls:
+        tool_call_counts[call.tool_name] += 1
+
+
 async def _call_tools(
     tool_manager: ToolManager[DepsT],
     tool_calls: list[_messages.ToolCallPart],
@@ -1020,14 +1033,32 @@ async def _call_tools(
         projected_usage.tool_calls += len(tool_calls)
         usage_limits.check_before_tool_call(projected_usage)
 
+    calls_to_run: list[_messages.ToolCallPart] = []
+
+    # For each tool, check how many calls are going to be made
+    tool_call_counts: defaultdict[str, int] = defaultdict(int)
+    _projection_count_of_tool_usage(tool_call_counts, tool_calls)
+
     for call in tool_calls:
-        yield _messages.FunctionToolCallEvent(call)
+        current_tool_use = tool_manager.get_current_use_of_tool(call.tool_name)
+        max_tool_use = tool_manager.get_max_use_of_tool(call.tool_name)
+        if max_tool_use is not None and current_tool_use + tool_call_counts[call.tool_name] > max_tool_use:
+            return_part = _messages.ToolReturnPart(
+                tool_name=call.tool_name,
+                content=f'Tool call limit reached for tool "{call.tool_name}".',
+                tool_call_id=call.tool_call_id,
+            )
+            output_parts.append(return_part)
+            yield _messages.FunctionToolResultEvent(return_part)
+        else:
+            yield _messages.FunctionToolCallEvent(call)
+            calls_to_run.append(call)
 
     with tracer.start_as_current_span(
         'running tools',
         attributes={
-            'tools': [call.tool_name for call in tool_calls],
-            'logfire.msg': f'running {len(tool_calls)} tool{"" if len(tool_calls) == 1 else "s"}',
+            'tools': [call.tool_name for call in calls_to_run],
+            'logfire.msg': f'running {len(calls_to_run)} tool{"" if len(calls_to_run) == 1 else "s"}',
         },
     ):
 
@@ -1061,8 +1092,8 @@ async def handle_call_or_result(
 
                 return _messages.FunctionToolResultEvent(tool_part, content=tool_user_content)
 
-        if tool_manager.should_call_sequentially(tool_calls):
-            for index, call in enumerate(tool_calls):
+        if tool_manager.should_call_sequentially(calls_to_run):
+            for index, call in enumerate(calls_to_run):
                 if event := await handle_call_or_result(
                     _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)),
                     index,
@@ -1075,7 +1106,7 @@ async def handle_call_or_result(
                     _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)),
                     name=call.tool_name,
                 )
-                for call in tool_calls
+                for call in calls_to_run
             ]
 
             pending = tasks
@@ -1092,7 +1123,11 @@ async def handle_call_or_result(
     output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)])
 
     _populate_deferred_calls(
-        tool_calls, deferred_calls_by_index, deferred_metadata_by_index, output_deferred_calls, output_deferred_metadata
+        calls_to_run,
+        deferred_calls_by_index,
+        deferred_metadata_by_index,
+        output_deferred_calls,
+        output_deferred_metadata,
     )
 
 

pydantic_ai_slim/pydantic_ai/_output.py

@@ -972,6 +972,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
                 toolset=self,
                 tool_def=tool_def,
                 max_retries=self.max_retries,
+                max_uses=None,
                 args_validator=self.processors[tool_def.name].validator,
             )
             for tool_def in self._tool_defs

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -48,6 +48,8 @@ class RunContext(Generic[RunContextAgentDepsT]):
     """Instrumentation settings version, if instrumentation is enabled."""
     retries: dict[str, int] = field(default_factory=dict)
     """Number of retries for each tool so far."""
+    tool_usage: dict[str, int] = field(default_factory=dict)
+    """Number of calls for each tool so far."""
     tool_call_id: str | None = None
     """The ID of the tool call."""
     tool_name: str | None = None

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -66,11 +66,22 @@ async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDe
 
     @property
     def tool_defs(self) -> list[ToolDefinition]:
-        """The tool definitions for the tools in this tool manager."""
-        if self.tools is None:
+        """The tool definitions for the tools in this tool manager.
+
+        Tools that have reached their `max_uses` limit are filtered out.
+        """
+        if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
 
-        return [tool.tool_def for tool in self.tools.values()]
+        result: list[ToolDefinition] = []
+        for tool in self.tools.values():
+            # Filter out tools that have reached their max_uses limit
+            if tool.max_uses is not None:
+                current_uses = self.ctx.tool_usage.get(tool.tool_def.name, 0)
+                if current_uses >= tool.max_uses:
+                    continue
+            result.append(tool.tool_def)
+        return result
 
     def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool:
         """Whether to require sequential tool calls for a list of tool calls."""
@@ -161,6 +172,8 @@ async def _call_tool(
                 partial_output=allow_partial,
             )
 
+            self.ctx.tool_usage[name] = self.ctx.tool_usage.get(name, 0) + 1
+
             pyd_allow_partial = 'trailing-strings' if allow_partial else 'off'
             validator = tool.args_validator
             if isinstance(call.args, str):
@@ -274,3 +287,21 @@ async def _call_function_tool(
                 )
 
         return tool_result
+
+    def get_max_use_of_tool(self, tool_name: str) -> int | None:
+        """Get the maximum number of uses allowed for a given tool, or `None` if unlimited."""
+        if self.tools is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
+
+        tool = self.tools.get(tool_name, None)
+        if tool is None:
+            return None
+
+        return tool.max_uses
+
+    def get_current_use_of_tool(self, tool_name: str) -> int:
+        """Get the current number of uses of a given tool."""
+        if self.ctx is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
+
+        return self.ctx.tool_usage.get(tool_name, 0)

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -1024,6 +1024,7 @@ def tool(
         name: str | None = None,
         description: str | None = None,
         retries: int | None = None,
+        max_uses: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -1042,6 +1043,7 @@ def tool(
         name: str | None = None,
         description: str | None = None,
         retries: int | None = None,
+        max_uses: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -1086,6 +1088,7 @@ async def spam(ctx: RunContext[str], y: float) -> float:
             description: The description of the tool, defaults to the function docstring.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
+            max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited).
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
@@ -1111,6 +1114,7 @@ def tool_decorator(
                 name=name,
                 description=description,
                 retries=retries,
+                max_uses=max_uses,
                 prepare=prepare,
                 docstring_format=docstring_format,
                 require_parameter_descriptions=require_parameter_descriptions,
@@ -1135,6 +1139,7 @@ def tool_plain(
         name: str | None = None,
         description: str | None = None,
         retries: int | None = None,
+        max_uses: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -1153,6 +1158,7 @@ def tool_plain(
         name: str | None = None,
         description: str | None = None,
         retries: int | None = None,
+        max_uses: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -1197,6 +1203,7 @@ async def spam(ctx: RunContext[str]) -> float:
             description: The description of the tool, defaults to the function docstring.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
+            max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited).
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
                 tool from a given step. This is useful if you want to customise a tool at call time,
                 or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
@@ -1220,6 +1227,7 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams
                 name=name,
                 description=description,
                 retries=retries,
+                max_uses=max_uses,
                 prepare=prepare,
                 docstring_format=docstring_format,
                 require_parameter_descriptions=require_parameter_descriptions,

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -590,6 +590,7 @@ def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[Any]:
             toolset=self,
             tool_def=tool_def,
             max_retries=self.max_retries,
+            max_uses=None,
             args_validator=TOOL_SCHEMA_VALIDATOR,
         )
 

pydantic_ai_slim/pydantic_ai/tools.py

@@ -264,6 +264,7 @@ class Tool(Generic[ToolAgentDepsT]):
     function: ToolFuncEither[ToolAgentDepsT]
     takes_ctx: bool
     max_retries: int | None
+    max_uses: int | None
     name: str
     description: str | None
     prepare: ToolPrepareFunc[ToolAgentDepsT] | None
@@ -286,6 +287,7 @@ def __init__(
         *,
         takes_ctx: bool | None = None,
         max_retries: int | None = None,
+        max_uses: int | None = None,
         name: str | None = None,
         description: str | None = None,
         prepare: ToolPrepareFunc[ToolAgentDepsT] | None = None,
@@ -337,6 +339,7 @@ async def prep_my_tool(
             takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] first argument,
                 this is inferred if unset.
             max_retries: Maximum number of retries allowed for this tool, set to the agent default if `None`.
+            max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited).
             name: Name of the tool, inferred from the function if `None`.
             description: Description of the tool, inferred from the function if `None`.
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
@@ -364,6 +367,7 @@ async def prep_my_tool(
         )
         self.takes_ctx = self.function_schema.takes_ctx
         self.max_retries = max_retries
+        self.max_uses = max_uses
         self.name = name or function.__name__
         self.description = description or self.function_schema.description
         self.prepare = prepare

pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

@@ -52,6 +52,8 @@ class ToolsetTool(Generic[AgentDepsT]):
     """The tool definition for this tool, including the name, description, and parameters."""
     max_retries: int
     """The maximum number of retries to attempt if the tool call fails."""
+    max_uses: int | None
+    """The maximum number of uses allowed for this tool."""
     args_validator: SchemaValidator | SchemaValidatorProt
     """The Pydantic Core validator for the tool's arguments.
 

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -77,6 +77,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
                     toolset=tool_toolset,
                     tool_def=tool.tool_def,
                     max_retries=tool.max_retries,
+                    max_uses=tool.max_uses,
                     args_validator=tool.args_validator,
                     source_toolset=toolset,
                     source_tool=tool,

pydantic_ai_slim/pydantic_ai/toolsets/external.py

@@ -36,6 +36,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
                 toolset=self,
                 tool_def=replace(tool_def, kind='external'),
                 max_retries=0,
+                max_uses=None,
                 args_validator=TOOL_SCHEMA_VALIDATOR,
             )
             for tool_def in self.tool_defs

pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py

@@ -170,6 +170,7 @@ def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]
             tool_def=tool_def,
             toolset=self,
             max_retries=self.max_retries,
+            max_uses=None,
             args_validator=TOOL_SCHEMA_VALIDATOR,
         )