fix(proxy): eliminate race condition in streaming guardrail_information logging

[michelligabriele]

asyncio.create_task in CSW.anext scheduled the deferred logging callback as an independent task that raced with unified_guardrail's end-of-stream block. For short-stream providers (Vertex AI, Azure, Anthropic), the logging fired before guardrail_information was written, causing post_call guardrail entries to be missing from StandardLoggingPayload.

Move the deferred callback trigger from CSW.anext to ProxyLogging.async_post_call_streaming_iterator_hook (after the full streaming pipeline completes). CSW now stores the assembled response args; the outer consumer fires the callback after all guardrail end-of-stream blocks finish. Also skip apply_guardrail guardrails in _run_deferred_stream_guardrails to eliminate duplicate API calls.

Relevant issues

Fixes guardrail_information missing post_call entries for non-OpenAI streaming providers (Vertex AI/Gemini, Azure, Anthropic). Related to #23910 and #24135.

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🐛 Bug Fix

Changes

litellm/litellm_core_utils/streaming_handler.py

• CSW.anext no longer calls asyncio.create_task for the deferred callback. Instead, it stores (assembled_response, cache_hit) on logging_obj._deferred_stream_complete_args for the outer consumer to pick up.

litellm/proxy/utils.py

• Added ProxyLogging._fire_deferred_stream_logging() static method that retrieves the stored callback and args from logging_obj, then fires asyncio.create_task.
• Called after the async for loop in async_post_call_streaming_iterator_hook, ensuring all guardrail end-of-stream blocks (including unified_guardrail) have fully completed before logging fires.

litellm/proxy/common_request_processing.py

• _run_deferred_stream_guardrails now skips guardrails that define apply_guardrail — these already ran via unified_guardrail's streaming iterator end-of-stream block. Running them again would duplicate the guardrail API call (e.g., double OpenAI Moderation charges).
• Updated docstring and closure comment to reflect the new flow.
• Removed unused unified_guardrail import.

tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py

• Updated test_streaming_closure_defers_logging → test_streaming_stores_deferred_args: verifies CSW stores args instead of calling the closure.
• Updated test_apply_guardrail_path_uses_unified_guardrail → test_apply_guardrail_skipped_in_deferred_path: verifies apply_guardrail guardrails are skipped in the deferred path.
• Updated integration tests to use _fire_deferred_stream_logging after CSW iteration.
• Added TestFireDeferredStreamLogging class with 4 new tests: fires callback with stored args, no-op when no args, no-op when no logging_obj, short-stream guardrail info populated.
• 25 tests total, all passing.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 25, 2026 10:15pm

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing michelligabriele:fix/streaming-guardrail-info-race-condition (f3774a8) with main (90b850e)}

[greptile-apps]

Greptile Summary

This PR fixes a real race condition in streaming guardrail logging. CSW.__anext__ previously called asyncio.create_task at stream-end, which raced with unified_guardrail's end-of-stream block (where guardrail_information is written). For short-stream providers (Vertex AI, Azure, Anthropic), logging would fire before the guardrail metadata was populated. The fix delays the create_task call until after the full guardrail pipeline has drained, by storing the assembled-response args on logging_obj and having ProxyLogging.async_post_call_streaming_iterator_hook fire the deferred callback after the async for loop. Additionally, apply_guardrail callbacks are now correctly skipped in _run_deferred_stream_guardrails because they already executed in unified_guardrail's streaming iterator, eliminating duplicate API calls (e.g. double OpenAI Moderation charges).

Key changes:

• CSW.__anext__: stores (assembled_response, cache_hit) on logging_obj._deferred_stream_complete_args instead of immediately scheduling create_task
• ProxyLogging.async_post_call_streaming_iterator_hook: calls new _fire_deferred_stream_logging after the async for loop, guaranteeing all guardrail end-of-stream blocks are complete
• _run_deferred_stream_guardrails: skips apply_guardrail callbacks (they already ran via the streaming iterator) to prevent duplicate guardrail API calls
• Tests updated and extended: 25 tests total, with a new TestFireDeferredStreamLogging class

One concern: the _fire_deferred_stream_logging call in async_post_call_streaming_iterator_hook sits after the async for loop with no try/finally guard. If a guardrail's end-of-stream block raises after CSW has already stored its args, the deferred logging callback will be silently dropped. Wrapping the loop in try/finally would make logging resilient to such failures.

Confidence Score: 4/5

• Safe to merge after addressing the missing try/finally guard in async_post_call_streaming_iterator_hook.
• The race-condition fix is correct and well-motivated. The new store-then-fire pattern is cleaner than the previous create_task approach and the 25-test suite covers the happy path thoroughly. The main outstanding concern is the P1 logic issue: without a try/finally, a guardrail end-of-stream exception after CSW stores args will silently drop deferred logging. That scenario is unlikely in practice (audit-only guardrails rarely throw), but it's a real regression from the previous behaviour where the old create_task would still fire. The missing model-level guardrail merge test is a P2 gap. Neither concern is blocking-critical, so score is 4 rather than 3.
• litellm/proxy/utils.py (async_post_call_streaming_iterator_hook — missing try/finally around _fire_deferred_stream_logging)

Important Files Changed

Filename	Overview
litellm/litellm_core_utils/streaming_handler.py	CSW.anext now stores (assembled_response, cache_hit) on logging_obj._deferred_stream_complete_args instead of scheduling asyncio.create_task directly — eliminates the race with unified_guardrail's end-of-stream block. Change is minimal and correct.
litellm/proxy/utils.py	New _fire_deferred_stream_logging static method and its call site after the async for loop in async_post_call_streaming_iterator_hook. The call site has no try/finally, so if a guardrail end-of-stream block raises (after CSW has already stored args), the deferred logging will be silently dropped.
litellm/proxy/common_request_processing.py	_run_deferred_stream_guardrails now skips apply_guardrail callbacks (they already ran via unified_guardrail's streaming iterator) to avoid duplicate API calls. Unused unified_guardrail import removed. Logic change is well-reasoned.
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py	Tests updated to reflect new store-then-fire pattern; new TestFireDeferredStreamLogging class adds good coverage. However, test_apply_guardrail_path_receives_merged_guardrail_data was removed without a replacement — no test now verifies that model-level guardrails are properly merged for the apply_guardrail path in the streaming iterator.

Sequence Diagram

sequenceDiagram
    participant Client
    participant RouteHandler
    participant PLHook as ProxyLogging.async_post_call_streaming_iterator_hook
    participant UG as unified_guardrail.async_post_call_streaming_iterator_hook
    participant CSW as CustomStreamWrapper.__anext__
    participant DL as ProxyLogging._fire_deferred_stream_logging
    participant DCB as _on_deferred_stream_complete closure
    participant Logger as logging_obj handlers

    RouteHandler->>PLHook: iterate stream
    PLHook->>UG: chain generator (for apply_guardrail callbacks)
    UG->>CSW: chain generator

    loop For each chunk
        Client->>PLHook: __anext__()
        PLHook->>UG: __anext__()
        UG->>CSW: __anext__()
        CSW-->>UG: chunk
        UG-->>PLHook: chunk
        PLHook-->>Client: yield chunk
    end

    Note over CSW: All chunks consumed → StopAsyncIteration path
    CSW->>CSW: Store (assembled_response, cache_hit) in logging_obj._deferred_stream_complete_args
    CSW-->>UG: StopAsyncIteration

    Note over UG: End-of-stream block runs
    UG->>UG: Run apply_guardrail post-call logic
    UG->>UG: Write guardrail_information to request_data
    UG-->>PLHook: StopAsyncIteration

    Note over PLHook: async for loop exits
    PLHook->>DL: _fire_deferred_stream_logging(request_data)
    DL->>DL: Retrieve _on_deferred_stream_complete + _deferred_stream_complete_args
    DL->>DCB: asyncio.create_task(_deferred_cb(*args))

    Note over DCB: Runs non-apply_guardrail guardrails
    DCB->>Logger: async_success_handler(assembled_response)
    DCB->>Logger: success_handler(assembled_response)

Loading

Comments Outside Diff (1)

1. litellm/proxy/utils.py, line 2234-2242 (link)

   Deferred logging silently dropped on guardrail exception

   _fire_deferred_stream_logging is called after the async for loop with no try/finally guard. By the time this code runs, CSW has already stored _deferred_stream_complete_args (it does so just before raising StopAsyncIteration). If a guardrail's end-of-stream block raises an exception after CSW stores the args — propagating out of async for chunk in current_response — execution never reaches _fire_deferred_stream_logging, and the deferred logging callback is silently lost.

   While unified_guardrail's streaming end-of-stream block is audit-only and unlikely to raise, other custom guardrails could. Adding a try/finally makes logging fire regardless:

           # Actually iterate through the chained async generator and yield chunks
           try:
               async for chunk in current_response:
                   yield chunk
           finally:
               # Fire deferred logging AFTER all guardrail end-of-stream blocks
               # completed.  unified_guardrail writes guardrail_information during
               # its end-of-stream block (inside current_response), so by the time
               # we reach this point the metadata is fully populated.
               ProxyLogging._fire_deferred_stream_logging(request_data)

   Note: yield inside a try/finally is valid for async generators in Python 3.10+, which is already required by this codebase.

_{Reviews (1): Last reviewed commit: "fix(proxy): eliminate race condition in ..." | Re-trigger Greptile}

[greptile-apps]

Removed test leaves model-level guardrail merging uncovered

test_apply_guardrail_path_receives_merged_guardrail_data was deleted without a replacement. That test verified a specific, previously-flagged correctness property: when an apply_guardrail callback is configured at the model level (not the global level), the merged guardrail data — produced by _check_and_merge_model_level_guardrails — was forwarded to UnifiedLLMGuardrails.async_post_call_success_hook, so that the inner should_run_guardrail re-check inside UnifiedLLMGuardrails could see the model-level guardrails.

With this PR, apply_guardrail callbacks now run exclusively via unified_guardrail.async_post_call_streaming_iterator_hook (called from ProxyLogging.async_post_call_streaming_iterator_hook) using the raw request_data dict, which does not go through _check_and_merge_model_level_guardrails. The original concern — a default_on=False guardrail configured only at the model level being silently skipped — now applies to the streaming-iterator path and there is no test guarding it.

Consider adding a test that configures a model-level-only apply_guardrail guardrail and asserts it runs (or at minimum asserts that async_post_call_streaming_iterator_hook receives data that includes the model-level guardrail list).

Rule Used: What: Flag any modifications to existing tests and... (source)