Fix race condition in RunUsage.incr() when running tools in parallel

- Add asyncio.Lock to RunUsage class to prevent race conditions
- Make incr() method async and use lock for thread-safe increments
- Update all calls to usage.incr() to use await
- Replace direct tool_calls += 1 with await usage.incr(RunUsage(tool_calls=1))
- Fixes issue where concurrent tool calls could cause undercounted tool_calls
- Maintains backward compatibility with synchronous __add__ method

Resolves #3120

Author: certainly-param

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -408,7 +408,7 @@ async def stream(
             message_history, model_settings, model_request_parameters, run_context
         ) as streamed_response:
             self._did_stream = True
-            ctx.state.usage.requests += 1
+            await ctx.state.usage.incr(_usage.RunUsage(requests=1))
             agent_stream = result.AgentStream[DepsT, T](
                 _raw_stream_response=streamed_response,
                 _output_schema=ctx.deps.output_schema,
@@ -437,9 +437,9 @@ async def _make_request(
 
         model_settings, model_request_parameters, message_history, _ = await self._prepare_request(ctx)
         model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
-        ctx.state.usage.requests += 1
+        await ctx.state.usage.incr(_usage.RunUsage(requests=1))
 
-        return self._finish_handling(ctx, model_response)
+        return await self._finish_handling(ctx, model_response)
 
     async def _prepare_request(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
@@ -475,19 +475,19 @@ async def _prepare_request(
             usage = deepcopy(usage)
 
             counted_usage = await ctx.deps.model.count_tokens(message_history, model_settings, model_request_parameters)
-            usage.incr(counted_usage)
+            await usage.incr(counted_usage)
 
         ctx.deps.usage_limits.check_before_request(usage)
 
         return model_settings, model_request_parameters, message_history, run_context
 
-    def _finish_handling(
+    async def _finish_handling(
         self,
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
         response: _messages.ModelResponse,
     ) -> CallToolsNode[DepsT, NodeRunEndT]:
         # Update usage
-        ctx.state.usage.incr(response.usage)
+        await ctx.state.usage.incr(response.usage)
         if ctx.deps.usage_limits:  # pragma: no branch
             ctx.deps.usage_limits.check_tokens(ctx.state.usage)
 

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -234,7 +234,7 @@ async def _call_function_tool(
         ) as span:
             try:
                 tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
-                usage.tool_calls += 1
+                await usage.incr(RunUsage(tool_calls=1))
 
             except ToolRetryError as e:
                 part = e.tool_retry

pydantic_ai_slim/pydantic_ai/usage.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import asyncio
 import dataclasses
 from copy import copy
 from dataclasses import dataclass, fields
@@ -189,24 +190,34 @@ class RunUsage(UsageBase):
     details: dict[str, int] = dataclasses.field(default_factory=dict)
     """Any extra details returned by the model."""
 
-    def incr(self, incr_usage: RunUsage | RequestUsage) -> None:
+    _lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock)
+    """Lock to prevent race conditions when incrementing usage from concurrent tool calls."""
+
+    async def incr(self, incr_usage: RunUsage | RequestUsage) -> None:
         """Increment the usage in place.
 
         Args:
             incr_usage: The usage to increment by.
         """
-        if isinstance(incr_usage, RunUsage):
-            self.requests += incr_usage.requests
-            self.tool_calls += incr_usage.tool_calls
-        return _incr_usage_tokens(self, incr_usage)
+        async with self._lock:
+            if isinstance(incr_usage, RunUsage):
+                self.requests += incr_usage.requests
+                self.tool_calls += incr_usage.tool_calls
+            return _incr_usage_tokens(self, incr_usage)
 
     def __add__(self, other: RunUsage | RequestUsage) -> RunUsage:
         """Add two RunUsages together.
 
         This is provided so it's trivial to sum usage information from multiple runs.
         """
         new_usage = copy(self)
-        new_usage.incr(other)
+        # Note: We can't use await here since __add__ must be synchronous
+        # But __add__ creates a new object, so there's no race condition
+        # The race condition only happens when modifying the same object concurrently
+        if isinstance(other, RunUsage):
+            new_usage.requests += other.requests
+            new_usage.tool_calls += other.tool_calls
+        _incr_usage_tokens(new_usage, other)
         return new_usage
 
 

tests/test_usage_limits.py

@@ -200,9 +200,9 @@ async def test_multi_agent_usage_sync():
     controller_agent = Agent(TestModel())
 
     @controller_agent.tool
-    def delegate_to_other_agent(ctx: RunContext[None], sentence: str) -> int:
+    async def delegate_to_other_agent(ctx: RunContext[None], sentence: str) -> int:
         new_usage = RunUsage(requests=5, input_tokens=2, output_tokens=3)
-        ctx.usage.incr(new_usage)
+        await ctx.usage.incr(new_usage)
         return 0
 
     result = await controller_agent.run('foobar')