Fix RunUsage.tool_calls being undercounted due to race condition when running tools in parallel

[certainly-param]

• 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

NOTE: While working on this fix, I noticed something interesting about the lock implementation. Since PydanticAI typically uses a single shared RunUsage object per agent run (the ctx.state.usage), I was curious about whether we could optimize the lock granularity. I ran some quick benchmarks and found that using context-based locks (where all tool calls in the same agent run share the same lock) could give about 26-29% better performance. The current instance-level approach works fine though. But, it might be worth exploring this optimization in the future. Current implementation should handle the concurrent tool execution issue nicely :)

Resolves #3120

[DouweM]

@certainly-param Considering the performance difference you call out, I'd prefer to use a lock only where we know we are running incr in parallel, meaning specifically inside the ToolManager.

We also definitely need a test that is failing on main, but will succeed after this PR. I'm still a bit surprised this issue exists at all, because as I understand it the usage. tool_calls += 1 line is run in parallel asyncio tasks, not threads, and asyncio tasks should not have race condition issues as they're not truly concurrent (or whatever the more precise word is for the difference between thread and event loop concurrency/parallelism).

[certainly-param]

@certainly-param Considering the performance difference you call out, I'd prefer to use a lock only where we know we are running incr in parallel, meaning specifically inside the ToolManager.

We also definitely need a test that is failing on main, but will succeed after this PR. I'm still a bit surprised this issue exists at all, because as I understand it the usage. tool_calls += 1 line is run in parallel asyncio tasks, not threads, and asyncio tasks should not have race condition issues as they're not truly concurrent (or whatever the more precise word is for the difference between thread and event loop concurrency/parallelism).

You're absolutely right that asyncio isn't truly concurrent like threads, but race conditions can still happen here. The issue is that usage.tool_calls += 1 looks like one operation, but it's actually three: read the current value, add one to it, and write it back. When you have multiple tools running in parallel and they hit an await point (like when actually calling the tool), the event loop can switch between tasks mid-operation. So imagine Task 1 reads tool_calls = 0, then the event loop switches to Task 2 which also reads 0, and both write back 1 instead of 2. Even though we're in a single-threaded event loop, that task switching at await points creates the exact same race condition you'd see with threads. It caught me by surprise too when I first dug into it 😄

So, I took your feedback moved the lock to exactly where the parallel execution happens instead of putting it in every incr() call. So now there's an asyncio.Lock() created right in _call_tools() where we know tools are running in parallel, and I use a ContextVar to pass it down to the tool manager. This way, it only locks when tools are actually executing in parallel, sequential calls don't touch the lock at all, so zero performance overhead there. I also added a test that really exercises this: 20 iterations with 10 parallel tools each, with multiple await points to maximize task switching. Without the fix, this test would fail intermittently on main (classic race condition behavior), but with the fix it's rock solid. Everything passes at 100% coverage, and RunUsage stays nice and clean without any lock logic cluttering it up.

[DouweM]

No need to include this comment, or the next identical one

[DouweM]

I think this could be a cached_property on ToolManager, and we wouldn't need to touch agent_graph at all

[DouweM]

Duplicated with if branch

[DouweM]

This looks like a breaking change we shouldn't make

[DouweM]

As mentioned above, this would be a breaking change, so I'd rather ensure at the call site where we call incr that RunUsage.requests == 0.

[DouweM]

When you have multiple tools running in parallel and they hit an await point (like when actually calling the tool), the event loop can switch between tasks mid-operation. So imagine Task 1 reads tool_calls = 0, then the event loop switches to Task 2 which also reads 0, and both write back 1 instead of 2. Even though we're in a single-threaded event loop, that task switching at await points creates the exact same race condition you'd see with threads. It caught me by surprise too when I first dug into it 😄

@certainly-param Are you sure? I thought that switches only happen at await points, so not in the middle of the += operation.

Either way, if we have a test that used to fail and now succeeds, we're good.

[certainly-param]

Thanks for the detailed review! I've made all the changes you suggested.

Removed the comments about request counting, moved the lock into ToolManager as a cached property (much cleaner this way - no more context var juggling), and simplified the RunUsage.incr() logic.

About the test_google.py diffs - those _identifier fields got added from inline-snapshot updates in an earlier commit. Happy to do an interactive rebase to clean them up if you want, or we can just squash them away when merging.

For test_openai_responses.py - yeah, there's a behavior change here. Before, requests was incremented right when streaming started, now it happens when the response finishes (in _finish_handling). This was needed to avoid double-counting since response.usage also includes requests. But it does mean run.usage() shows requests=0 during the stream now instead of requests=1.

Not sure if that's a problem for anyone? If we need to keep the old behavior, I could increment requests at the start but then skip that field when calling incr(response.usage).

On the race condition - you're right that asyncio doesn't have traditional threading issues. The problem is that += isn't atomic. It's roughly temp = x; temp = temp + 1; x = temp, and when two tasks do this concurrently they can interleave:

Task A reads tool_calls=0
Task A awaits (yields control)
Task B reads tool_calls=0  <- still 0!
Task B increments, writes tool_calls=1
Task A resumes, increments, writes tool_calls=1  <- lost update

The test I added (test_race_condition_parallel_tool_calls) fails pretty reliably without the lock, which is how I caught this.

Anyway, all tests are passing now. Let me know what you think about the test_google.py cleanup and the usage timing question!

[DouweM]

This looks like a broken merge conflict resolution! Please remove it from the diff to make sure we don't accidentally merge this into main.

[DouweM]

As mentioned above, this would be a breaking change, so I'd rather ensure at the call site where we call incr that RunUsage.requests == 0.

[DouweM]

Is this needed?

[DouweM]

@certainly-param Thank you!

[phemmer]

I thought that switches only happen at await points, so not in the middle of the += operation.

This is true

The problem is that += isn't atomic. It's roughly temp = x; temp = temp + 1; x = temp, and when two tasks do this concurrently they can interleave:

This is not possible. There must be something within that flow to suspend the current task. Without doing so, the active task will continue to run. The asyncio design is extremely explicit on this. If tasks could be interrupted at any time, it could cause havoc with async applications.

I tried running the test case without the fix applied. I could not get any failures. I even bumped for iteration in range(20): up to 2000 and for i in range(10) to 1000. The issue could not be reproduced. I suspect there's something else going on in your setup, or you have a bug in your python or something.

[DouweM]

@phemmer Yeah I'm confused as well; I couldn't and can't reproduce the issue with the old code. @certainly-param wrote "I tested with 3 tools and sometimes get tool_calls=1, sometimes 2, rarely 3." and that it's fixed now, so I merged this mostly off the back of that, but I'm thinking I'll revert this until we have a failure multiple people can reproduce.

[certainly-param]

Looking deeper at the code, I think I found the actual issue that was causing those results. In _tool_manager.py, the usage.tool_calls += 1 increment is inside the try block:

try:
tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
usage.tool_calls += 1 # Only happens on success
except ToolRetryError as e:
# Tool fails and retries, but increment never happened
raise e

This means when a tool fails with ToolRetryError and gets retried, it only gets counted once even though it was called twice. This could explain the intermittent undercounting I was seeing in my tests - it would depend on which tools fail and retry.
So the lock I added doesn't actually solve this retry counting issue. Should I update the fix to move the increment outside the try block so retries are properly counted? That seems like it would address the real root cause of the undercounting behavior I observed.

[DouweM]

@certainly-param Only counting successful calls was a conscious decision in #2633 and #2978, and changing it now could be considered a breaking change if people were relying on the existing behavior. I suppose at the time our thinking was that if your goal is to give your models a "budget" in terms of tool usage, an unsuccessful call (possibly because of arg validation) does not use of the budget. But of course there are cases where a call, even if unsuccessful did use up resources...

@tradeqvest What do you think?