Fix race condition in LLM callback system

[devin-ai-integration]

Summary

Fixes a race condition in the LLM callback system where multiple LLM instances calling set_callbacks concurrently could cause callbacks to be removed before they fire. This was identified in issue #4214, where a test required a sleep(5) workaround to pass.

The fix adds a class-level threading.RLock to synchronize access to global litellm callbacks during call() and acall() methods. The lock ensures that callback registration and the subsequent LLM call execution are atomic, preventing one instance from modifying callbacks while another is using them.

Changes:

• Add _callback_lock: threading.RLock class attribute to LLM
• Wrap callback setup and LLM call execution in both call() and acall() with the lock
• Remove sleep(5) workaround from test_llm_callback_replacement
• Add new test_llm_callback_lock_prevents_race_condition test for concurrent access

Review & Testing Checklist for Human

• Performance impact: The lock is held during the entire LLM call (including network I/O), which serializes all LLM calls across all instances. Verify this is acceptable for multi-agent scenarios where parallelism is expected.
• Lock scope: Consider if the lock scope is too broad - could it just protect the callback registration/cleanup rather than the entire call?
• Real-world validation: Run a multi-agent crew with concurrent LLM calls to verify the fix works in production scenarios and doesn't introduce deadlocks or significant latency.

Recommended test plan:

1. Run the existing test suite to ensure no regressions
2. Create a simple multi-agent crew with 3+ agents making concurrent LLM calls
3. Verify callbacks (like token tracking) work correctly without interference
4. Monitor for any performance degradation compared to before the fix

Notes

• Fixes [BUG] Hidden race condition in LLM callback system causing test failures #4214
• Uses RLock (reentrant lock) to handle recursive calls that occur when retrying with unsupported 'stop' parameter
• Link to Devin run: https://app.devin.ai/sessions/387dcbd2330148399a988256dee3fa6a
• Requested by: João

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment and CI monitoring

[VedantMadane]

The race condition diagnosis is correct, and RLock is the right choice for the recursive retry case. One significant concern worth discussing:

Performance consideration:

The lock is held for the entire duration of the LLM call (including network I/O):
python with LLM._callback_lock: # ... entire completion call happens here

This effectively serializes all LLM calls across all instances in the process. For multi-agent crews that rely on parallel LLM calls, this could be a significant performance regression.

Alternative approaches to consider:

1. Narrow the lock scope - Only protect the callback registration/cleanup, not the actual LLM call:
   `python
   with LLM._callback_lock:
   self.set_callbacks(callbacks)

result = litellm.completion(**params) # Outside lock

with LLM._callback_lock:
# cleanup callbacks
`

2. Per-instance callbacks instead of global litellm callbacks - would require changes to how callbacks are wired but eliminates the global state issue entirely.

3. Copy-on-write pattern for the callback list.

The current fix is correct for correctness, but the performance tradeoff should be explicitly accepted or an alternative considered. The PR description does flag this for review which is good.

Would love to hear thoughts from maintainers on acceptable parallelism tradeoffs.

[devin-ai-integration]

Closing due to inactivity for more than 7 days. Configure here.

[VedantMadane]

/review

[VedantMadane]

Opened follow-up PR #4252 with the same intent as this PR but without a global lock held across the network call (keeps concurrency): #4252