Convert to asyncio.Lock for async compatibility

- Replace threading.Lock with asyncio.Lock to fix async context issues
- Make incr() method async and use async with lock
- Update all usage.incr() calls to use await
- Fix tool manager to use await usage.incr(RunUsage(tool_calls=1))
- Update agent graph _finish_handling to be async
- Fix streaming tests that were failing with TypeError

This resolves the 16 TypeError failures in streaming tests by ensuring
thread-safe operations work correctly in async contexts.

Author: certainly-param

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -426,7 +426,7 @@ async def stream(
 
         model_response = streamed_response.get()
 
-        self._finish_handling(ctx, model_response)
+        await self._finish_handling(ctx, model_response)
         assert self._result is not None  # this should be set by the previous line
 
     async def _make_request(
@@ -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,7 +1,7 @@
 from __future__ import annotations as _annotations
 
+import asyncio
 import dataclasses
-import threading
 from copy import copy
 from dataclasses import dataclass, fields
 from typing import Annotated, Any
@@ -193,19 +193,19 @@ class RunUsage(UsageBase):
     """Lock to prevent race conditions when incrementing usage from concurrent tool calls."""
 
     @property
-    def _lock(self) -> threading.Lock:
+    def _lock(self) -> asyncio.Lock:
         """Get the lock, creating it if it doesn't exist."""
         if not hasattr(self, '__lock'):
-            self.__lock = threading.Lock()
+            self.__lock = asyncio.Lock()
         return self.__lock
 
-    def incr(self, incr_usage: RunUsage | RequestUsage) -> None:
+    async def incr(self, incr_usage: RunUsage | RequestUsage) -> None:
         """Increment the usage in place.
 
         Args:
             incr_usage: The usage to increment by.
         """
-        with self._lock:
+        async with self._lock:
             if isinstance(incr_usage, RunUsage):
                 self.requests += incr_usage.requests
                 self.tool_calls += incr_usage.tool_calls

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')