Skip to content

fix(proxy): defer logging until post-call guardrails complete - #24135

Merged
12 commits merged into
BerriAI:litellm_oss_staging_03_19_2026from
michelligabriele:fix/deferred-guardrail-logging-v2
Mar 20, 2026
Merged

fix(proxy): defer logging until post-call guardrails complete#24135
12 commits merged into
BerriAI:litellm_oss_staging_03_19_2026from
michelligabriele:fix/deferred-guardrail-logging-v2

Conversation

@michelligabriele

Copy link
Copy Markdown
Contributor

guardrail_information is None in StandardLoggingPayload because logging fires before post-call guardrails write to metadata.

Non-streaming: wrapper_async stores a closure instead of calling create_task immediately. The proxy fires it in a try/finally after post_call_success_hook so the SLP is built with guardrail info.

Streaming: a closure on logging_obj is called by CSW.anext at stream end. The closure runs only guardrail hooks (not all callbacks) on the assembled response, then fires both logging handlers. This avoids behavioral changes for non-guardrail callbacks on streaming.

Relevant issues

Replaces #23929

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

Problem

guardrail_information is always None in StandardLoggingPayload when post-call guardrails (e.g. OpenAI Moderation) are configured. This happens because:

  • Non-streaming: asyncio.create_task in wrapper_async (utils.py) fires the logging task before post_call_success_hook runs in base_process_llm_request, so the SLP is built before guardrails write to metadata.
  • Streaming: logging fires at stream exhaustion in CustomStreamWrapper.__anext__ without any guardrail data from the assembled response — post_call_success_hook is never called for streaming early-return routes.

Fix

Two deferral mechanisms — same concept (store a closure, call it at the right time), different execution points.

Non-streaming (utils.py + common_request_processing.py):

  1. _has_post_call_guardrails() checks if any CustomGuardrail with post_call event hook is registered
  2. If true and non-streaming: set logging_obj._defer_async_logging = True
  3. wrapper_async sees the flag → stores closure on logging_obj._enqueue_deferred_logging instead of asyncio.create_task. Sync callbacks fire immediately (unchanged).
  4. base_process_llm_request runs post_call_success_hook (guardrails write to metadata)
  5. finally block calls the stored closure → create_task fires → SLP built with guardrail info

Streaming (common_request_processing.py + streaming_handler.py):

  1. If _has_post_call_guardrails and response is CustomStreamWrapper: attach _on_deferred_stream_complete closure to logging_obj
  2. The closure runs only guardrail hooks — iterates litellm.callbacks, filters for CustomGuardrail instances with post_call event hook, calls their async_post_call_success_hook. This is the same pattern ProxyLogging.post_call_success_hook uses internally, but filtered to guardrails only. Non-guardrail callbacks are not called (avoids behavioral changes for streaming).
  3. CSW.__anext__ at stream end: checks for closure. If set, clears it and calls it via asyncio.create_task. If not set, fires logging directly (original behavior preserved).
  4. Fallthrough safety: if code reaches the inline post_call_success_hook (no early return), the closure is cleared first to prevent double invocation.

Files

File Change
litellm/utils.py _defer_async_logging flag → store closure instead of create_task
litellm/proxy/common_request_processing.py _has_post_call_guardrails() static method, deferral flag, streaming closure (guardrail-only), try/finally
litellm/litellm_core_utils/streaming_handler.py CSW.__anext__ checks for _on_deferred_stream_complete, calls it instead of logging directly
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py 17 tests for both paths
docs/my-website/docs/proxy/guardrails/custom_guardrail.md Document streaming post_call guardrails as audit-only

Tests (17 total)

Detection (7): _has_post_call_guardrails returns correct result for post_call, pre_call, event_hook=None, list event hooks, non-guardrail callbacks, empty callbacks

Non-streaming (3): deferred flag stores and executes closure, sync callbacks fire immediately, regression test without flag

Non-streaming exception (1): deferred logging fires even if guardrail raises HTTPException (try/finally)

Streaming (6): closure defers logging, regression without closure, closure runs only guardrail hooks (not all callbacks), guardrail-modified response flows to logging, exception resilience with guardrail_blocked, transient errors don't set guardrail_blocked, production closure integration test

Krrish Dholakia and others added 5 commits March 18, 2026 19:52
…idebar

- Add 'Contributing to Guardrails' category with links to:
  - Generic Guardrail API (integrate without PR)
  - Adding a New Guardrail Integration tutorial
  - Adding Guardrail Support to Endpoints

- Add 'Team Bring-Your-Own Guardrails' link for team BYOG workflow

These docs existed but were only accessible from the 'LiteLLM AI Gateway'
sidebar. Now they're also accessible when browsing the 'Guardrail Providers'
section.

Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
…ls-docs-143b

docs: add Contributing to Guardrails section to Guardrail Providers sidebar
@vercel

vercel Bot commented Mar 19, 2026

Copy link
Copy Markdown

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

Project Deployment Actions Updated (UTC)
litellm Ready Ready Preview, Comment Mar 19, 2026 6:31pm

Request Review

@codspeed-hq

codspeed-hq Bot commented Mar 19, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 16 untouched benchmarks


Comparing michelligabriele:fix/deferred-guardrail-logging-v2 (573f6b7) with main (81dadb6)

Open in CodSpeed

@greptile-apps

greptile-apps Bot commented Mar 19, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes guardrail_information being None in StandardLoggingPayload for both non-streaming and streaming requests when post-call guardrails (e.g., OpenAI Moderation) are configured. The root cause was that async logging fired via asyncio.create_task before post_call_success_hook had a chance to write guardrail results to metadata.

How it works:

  • Non-streaming (utils.py): When _defer_async_logging is set on the logging object, wrapper_async stores a closure (_enqueue_deferred_logging) instead of immediately calling create_task. base_process_llm_request fires this closure in a try/finally after post_call_success_hook completes, ensuring the SLP is built with guardrail data.
  • Streaming (common_request_processing.py + streaming_handler.py): A closure (_on_deferred_stream_complete) is attached to logging_obj when a post-call guardrail is active and the response is a CustomStreamWrapper. CSW.__anext__ calls this closure at stream end (instead of logging directly), running only guardrail hooks via the extracted static method _run_deferred_stream_guardrails, then firing both async and sync logging handlers with the guardrail-populated response.

Key design decisions addressed from prior review:

  • _run_deferred_stream_guardrails is extracted as a testable static method (not an anonymous closure)
  • Uses the unified_guardrail singleton from proxy/utils.py, consistent with the rest of the proxy
  • Imports the global llm_router at call time to correctly merge model-level guardrails
  • Per-guardrail try/except prevents one failing guardrail from silently skipping subsequent ones
  • _exception_raised flag in finally cleans up orphaned streaming closures on error paths
  • 17 targeted unit tests cover detection, deferral mechanics, merged-data propagation, and exception resilience

Remaining minor concerns:

  • The thread_pool_executor import in _run_deferred_stream_guardrails sits outside the try/finally guard — an import failure there would bypass the finally block and silently drop logging (extremely unlikely but inconsistent with the protection applied to other imports in the same function)
  • test_deferred_logging_fires_on_guardrail_exception manually mirrors the production finally block instead of exercising it through base_process_llm_request, reducing regression protection for the non-streaming exception path

Confidence Score: 4/5

  • Safe to merge with minor test and resilience improvements; the core fix is correct and well-guarded
  • All substantive concerns from the prior review round have been addressed: model-level guardrail merging, per-guardrail exception isolation, UnifiedLLMGuardrails singleton reuse, and the static method extraction for testability. The two remaining concerns are low-severity: one unlikely import-ordering edge case and one test that mirrors rather than exercises the production finally block. The behavioral change (streaming post_call guardrails are now audit-only) accurately reflects what the code has always done and is correctly documented.
  • litellm/proxy/common_request_processing.py — specifically the _run_deferred_stream_guardrails static method; verify the executor import is covered by the try/finally guard before merge

Important Files Changed

Filename Overview
litellm/utils.py Adds _defer_async_logging flag support in wrapper_async: when the flag is set, stores a closure on _enqueue_deferred_logging instead of immediately calling asyncio.create_task. Sync callbacks still fire immediately. Change is minimal, well-guarded, and non-breaking for the default (unflagged) path.
litellm/proxy/common_request_processing.py Core of the PR. Adds _has_post_call_guardrails() static method, the streaming deferred closure mechanism, a try/except/finally wrapper around post-call processing, and the extracted _run_deferred_stream_guardrails static method. Previous review concerns (model-level guardrail merging, per-guardrail exception handling, UnifiedLLMGuardrails singleton reuse, orphaned closure cleanup) have all been addressed. Minor residual concern: the thread_pool_executor import sits outside the try/finally guard in _run_deferred_stream_guardrails.
litellm/litellm_core_utils/streaming_handler.py Minimal, surgical change in CSW.__anext__: checks for _on_deferred_stream_complete closure on logging_obj, clears it before dispatch (preventing double invocation), then either runs the deferred path or falls back to the existing direct logging path. Both async and sync logging handlers are preserved in both paths.
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py 17 tests covering detection, non-streaming deferral, streaming closure behavior, per-guardrail exception resilience, merged guardrail data propagation, and the apply_guardrail unified path. Most streaming tests now call the real _run_deferred_stream_guardrails static method. The non-streaming exception test (test_deferred_logging_fires_on_guardrail_exception) still manually reimplements the production finally block rather than invoking it, reducing regression protection for that path.
docs/my-website/docs/proxy/guardrails/custom_guardrail.md Documents the new streaming post_call guardrail behavior as "audit-only" and updates the capability table to clarify that async_post_call_success_hook for streaming cannot block content delivery. Accurately reflects the implementation.

Sequence Diagram

sequenceDiagram
    participant Client
    participant BPLLM as base_process_llm_request
    participant WA as wrapper_async (utils.py)
    participant CSW as CustomStreamWrapper.__anext__
    participant PL as ProxyLogging.post_call_success_hook
    participant RDSG as _run_deferred_stream_guardrails
    participant Logger as async_success_handler / success_handler

    Note over BPLLM: _has_post_call_guardrails() → True

    rect rgb(200, 230, 255)
        Note over BPLLM,Logger: Non-streaming path
        BPLLM->>WA: acompletion() with _defer_async_logging=True
        WA-->>BPLLM: result + stores _enqueue_deferred_logging closure
        BPLLM->>PL: post_call_success_hook() (guardrails write to metadata)
        Note over BPLLM: finally block
        BPLLM->>WA: _enqueue_deferred_logging()
        WA->>Logger: asyncio.create_task(async_success_handler) ← SLP built with guardrail_information
    end

    rect rgb(200, 255, 210)
        Note over BPLLM,Logger: Streaming path
        BPLLM->>CSW: attach _on_deferred_stream_complete closure
        BPLLM-->>Client: return StreamingResponse (CSW)
        Client->>CSW: iterate chunks
        CSW-->>Client: yield chunks (stream delivered)
        Note over CSW: stream exhausted
        CSW->>RDSG: asyncio.create_task(_on_deferred_stream_complete)
        RDSG->>PL: _check_and_merge_model_level_guardrails
        loop for each CustomGuardrail with post_call hook
            RDSG->>RDSG: cb.async_post_call_success_hook (or unified_guardrail)
        end
        Note over RDSG: finally block
        RDSG->>Logger: async_success_handler(_response) ← SLP built with guardrail_information
        RDSG->>Logger: executor.submit(success_handler)
    end
Loading

Last reviewed commit: "fix(proxy): split or..."

Comment on lines +1082 to +1083
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Model-level guardrails silently skipped for streaming

_check_and_merge_model_level_guardrails immediately returns data unchanged when llm_router is None:

def _check_and_merge_model_level_guardrails(data, llm_router):
    if llm_router is None:
        return data  # no merging

This means deployments that configure guardrails at the model/router level (via litellm_params.guardrails) will have those guardrails silently skipped in the streaming deferred path — a behavioral regression compared to the non-streaming path.

ProxyLogging.post_call_success_hook (in proxy/utils.py) solves this correctly by importing the global llm_router at call time:

from litellm.proxy.proxy_server import llm_router
_check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)

The streaming closure should do the same:

Suggested change
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=None
from litellm.proxy.proxy_server import (
llm_router as _global_llm_router,
)
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=_global_llm_router
)

if _has_post_call_guardrails and isinstance(
response, CustomStreamWrapper
):
_captured_data = self.data

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 _captured_data is a mutable reference to self.data

_captured_data = self.data captures a reference, not a copy. The line _captured_data["guardrail_to_apply"] = cb mutates the original request data dict. While this pattern mirrors ProxyLogging.post_call_success_hook and works correctly in the current implementation (the closure runs after the response is returned), it leaves a stale guardrail_to_apply key on self.data after the closure finishes and iterates through each guardrail.

Consider using _captured_data = dict(self.data) or documenting that this is intentionally a live reference, to make the mutation intent explicit for future maintainers.

from unittest.mock import MagicMock, patch

import pytest
from fastapi import HTTPException

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 FastAPI import outside proxy/ folder

from fastapi import HTTPException imports FastAPI, which is a proxy-only dependency, into a test file that lives outside litellm/proxy/. The custom instruction for this repo prohibits FastAPI imports outside the proxy/ directory.

Since this test only needs HTTPException as a sentinel exception type that the guardrail raises, you can substitute it with a stdlib-compatible alternative:

Suggested change
from fastapi import HTTPException
from starlette.exceptions import HTTPException

starlette is already a transitive dependency (FastAPI is built on it), and starlette.exceptions.HTTPException is identical to fastapi.HTTPException. Alternatively, you can raise a plain Exception subclass in the test guardrail and skip the FastAPI dependency entirely.

Rule Used: What: Do not allow fastapi imports on files outsid... (source)

Comment on lines +1063 to +1154
async def _on_deferred_stream_complete(
assembled_response, cache_hit
):
from litellm.litellm_core_utils.thread_pool_executor import (
executor,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.proxy.proxy_server import (
llm_router as _global_llm_router,
)
from litellm.proxy.utils import _check_and_merge_model_level_guardrails

# NOTE: This closure runs after all chunks have been
# delivered to the client. Blocking guardrails that
# raise HTTPException cannot prevent content delivery
# for streaming — this is an inherent limitation of
# SSE streaming. The purpose here is to populate
# guardrail_information in the logging payload for
# audit/compliance. Per-chunk filtering should use
# async_post_call_streaming_hook instead.
_response = assembled_response
_unified_guardrail = UnifiedLLMGuardrails()
try:
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=_global_llm_router
)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,
):
continue
guardrail_result = None
if "apply_guardrail" in type(cb).__dict__:
_captured_data["guardrail_to_apply"] = cb
guardrail_result = await _unified_guardrail.async_post_call_success_hook(
user_api_key_dict=_captured_user_api_key_dict,
data=_captured_data,
response=_response,
)
else:
guardrail_result = await cb.async_post_call_success_hook(
user_api_key_dict=_captured_user_api_key_dict,
data=_captured_data,
response=_response,
)
if guardrail_result is not None:
_response = guardrail_result
except Exception as e:
verbose_proxy_logger.exception(
"Error running post-call guardrails on streaming response: %s",
e,
)
if isinstance(e, HTTPException) and hasattr(
_captured_logging_obj, "model_call_details"
):
_captured_logging_obj.model_call_details.setdefault(
"metadata", {}
)["guardrail_blocked"] = True

try:
asyncio.create_task(
_captured_logging_obj.async_success_handler(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming async logging: %s", e,
)

try:
executor.submit(
_captured_logging_obj.success_handler,
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming sync logging: %s", e,
)

logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[attr-defined]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Streaming closure exits early on first guardrail exception, skipping subsequent guardrails

The entire guardrail for loop is wrapped in a single try/except. If any guardrail in the loop raises (including a non-HTTPException), the except block is entered and the loop is exited — all remaining guardrails are silently skipped. Only the first raised exception sets guardrail_blocked.

try:
    for cb in litellm.callbacks:
        ...
        guardrail_result = await cb.async_post_call_success_hook(...)  # raises → loop exits
        ...
except Exception as e:
    ...  # subsequent guardrails never run

Since this is an audit-only path (streaming content already delivered), a per-guardrail try/except would be more resilient and consistent with how ProxyLogging.post_call_success_hook handles multiple callbacks:

for cb in litellm.callbacks:
    if not isinstance(cb, CustomGuardrail):
        continue
    if not cb.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call):
        continue
    try:
        guardrail_result = None
        if "apply_guardrail" in type(cb).__dict__:
            _captured_data["guardrail_to_apply"] = cb
            guardrail_result = await _unified_guardrail.async_post_call_success_hook(...)
        else:
            guardrail_result = await cb.async_post_call_success_hook(...)
        if guardrail_result is not None:
            _response = guardrail_result
    except Exception as e:
        verbose_proxy_logger.exception("Error in guardrail %s: %s", cb, e)
        if isinstance(e, HTTPException) and hasattr(_captured_logging_obj, "model_call_details"):
            _captured_logging_obj.model_call_details.setdefault("metadata", {})["guardrail_blocked"] = True

This ensures all configured post-call guardrails contribute to audit logging even when one fails.

Comment on lines +354 to +431
logger_called = False

class TrackingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="tracker",
default_on=True,
event_hook=GuardrailEventHooks.post_call,
)

async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any
) -> Any:
nonlocal guardrail_called
guardrail_called = True
return response

class TrackingLogger(CustomLogger):
async def async_post_call_success_hook(
self, user_api_key_dict, data, response
):
nonlocal logger_called
logger_called = True
return response

mock_logging_obj = MagicMock()
mock_logging_obj.model_call_details = {"metadata": {}}

async def track_async_success(*args, **kwargs):
pass

mock_logging_obj.async_success_handler = track_async_success

tracking_guardrail = TrackingGuardrail()
tracking_logger = TrackingLogger()

# Build the closure using the same pattern as production code
_captured_data = {"model": "gpt-4", "metadata": {}}
_captured_user_api_key_dict = UserAPIKeyAuth(api_key="test")
_captured_logging_obj = mock_logging_obj

async def _on_deferred_stream_complete(assembled_response, cache_hit):
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.proxy.utils import _check_and_merge_model_level_guardrails

_response = assembled_response
try:
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=None
)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,
):
continue
guardrail_result = await cb.async_post_call_success_hook(
user_api_key_dict=_captured_user_api_key_dict,
data=_captured_data,
response=_response,
)
if guardrail_result is not None:
_response = guardrail_result
except Exception:
pass

asyncio.create_task(
_captured_logging_obj.async_success_handler(
_response, cache_hit=cache_hit, start_time=None, end_time=None
)
)
executor.submit(
_captured_logging_obj.success_handler,
_response, cache_hit=cache_hit, start_time=None, end_time=None,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Tests replicate production closure instead of exercising it

test_closure_runs_only_guardrail_hooks, test_production_closure_integration, and test_apply_guardrail_path_uses_unified_guardrail each manually build their own _on_deferred_stream_complete closure that mirrors the production code in common_request_processing.py. They never invoke the actual production closure — they test hand-crafted copies.

This violates Mock Test Integrity: if someone later changes the production closure (e.g. adds a new guardrail dispatch path, changes the exception-handling logic, or reorders async/sync logging calls), these tests will still pass because they are testing a snapshot copy of the old logic.

The tests should exercise the actual production code path. For example, the streaming integration tests could call ProxyBaseLLMRequestProcessing with a real (mocked-I/O) CSW and assert on side-effects, or the closure could be extracted to a named, importable helper so tests can call the real thing.

This same issue applies at:

  • test_production_closure_integration (~line 588 in the test file)
  • test_apply_guardrail_path_uses_unified_guardrail (~line 700 in the test file)

Rule Used: # Code Review Rule: Mock Test Integrity

What:... (source)

Comment on lines +1063 to +1160
async def _on_deferred_stream_complete(
assembled_response, cache_hit
):
from litellm.litellm_core_utils.thread_pool_executor import (
executor,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.proxy.proxy_server import (
llm_router as _global_llm_router,
)
from litellm.proxy.utils import _check_and_merge_model_level_guardrails

# NOTE: This closure runs after all chunks have been
# delivered to the client. Blocking guardrails that
# raise HTTPException cannot prevent content delivery
# for streaming — this is an inherent limitation of
# SSE streaming. The purpose here is to populate
# guardrail_information in the logging payload for
# audit/compliance. Per-chunk filtering should use
# async_post_call_streaming_hook instead.
_response = assembled_response
_unified_guardrail = UnifiedLLMGuardrails()
guardrail_data = _check_and_merge_model_level_guardrails(
data=_captured_data, llm_router=_global_llm_router
)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,
):
continue
try:
guardrail_result = None
if "apply_guardrail" in type(cb).__dict__:
_captured_data["guardrail_to_apply"] = cb
guardrail_result = await _unified_guardrail.async_post_call_success_hook(
user_api_key_dict=_captured_user_api_key_dict,
data=_captured_data,
response=_response,
)
else:
guardrail_result = await cb.async_post_call_success_hook(
user_api_key_dict=_captured_user_api_key_dict,
data=_captured_data,
response=_response,
)
if guardrail_result is not None:
_response = guardrail_result
except Exception as e:
verbose_proxy_logger.exception(
"Error running post-call guardrail %s on streaming response: %s",
getattr(cb, "guardrail_name", type(cb).__name__),
e,
)
if isinstance(e, HTTPException) and hasattr(
_captured_logging_obj, "model_call_details"
):
_captured_logging_obj.model_call_details.setdefault(
"metadata", {}
)["guardrail_blocked"] = True

try:
asyncio.create_task(
_captured_logging_obj.async_success_handler(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming async logging: %s", e,
)

try:
executor.submit(
_captured_logging_obj.success_handler,
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming sync logging: %s", e,
)

logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[attr-defined]

if route_type == "allm_passthrough_route":
# Check if response is an async generator
if self._is_streaming_response(response):
if asyncio.iscoroutine(response):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 60-line inline closure is hard to test and maintain

The _on_deferred_stream_complete async closure (~60 lines) is defined inline inside base_process_llm_request. Because it is a closure (not a named, importable function), tests cannot import and call the production implementation directly — they are forced to write hand-copies of it (as seen in test_production_closure_integration etc.), which defeats regression testing.

Consider extracting this to a static or module-level helper, e.g.:

@staticmethod
async def _run_deferred_stream_guardrails(
    captured_data: dict,
    captured_user_api_key_dict: UserAPIKeyAuth,
    captured_logging_obj: Any,
    assembled_response: Any,
    cache_hit: Any,
) -> None:
    ...

The closure can then simply delegate to it:

async def _on_deferred_stream_complete(assembled_response, cache_hit):
    await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
        _captured_data, _captured_user_api_key_dict, _captured_logging_obj,
        assembled_response, cache_hit,
    )

This makes the logic importable and directly testable, and prevents the test-copy drift already present in this PR.

Comment on lines 928 to +931
llm_router=llm_router,
)

# Defer async logging when post-call guardrails are configured so the

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 _defer_async_logging set before checking _is_streaming_response

_defer_async_logging = True is set when _is_streaming_request returns False. However, _is_streaming_response(response) is only evaluated after the LLM API call completes. If the response turns out to be a CustomStreamWrapper even though the request was not marked as streaming (e.g. some provider wraps non-streaming calls in a stream), the following occurs:

  1. _defer_async_logging = True is set on logging_obj
  2. In utils.py's wrapper_async, the result IS a streaming response → it exits early before the _defer_async_logging block
  3. _enqueue_deferred_logging is never stored on logging_obj
  4. The finally block sees _enqueue_fn = None → no-op
  5. Meanwhile, _on_deferred_stream_complete is set (because isinstance(response, CustomStreamWrapper) is True), so CSW defers logging correctly

The async logging from the _client_async_logging_helper path in wrapper_async is silently dropped for this edge case. The CSW path fires _on_deferred_stream_complete instead, which calls async_success_handler directly, so logging is not fully lost — but the _client_async_logging_helper path (which may include fallback/router-level callbacks) is skipped.

Adding not self._is_streaming_response(response) to the flag guard would make this explicit, though determining response type before the API call is infeasible. At minimum a clarifying comment noting this intentional reliance on wrapper_async's early-return behaviour would help future maintainers.

guardrail_information is None in StandardLoggingPayload because logging
fires before post-call guardrails write to metadata.

Non-streaming: wrapper_async stores a closure instead of calling
create_task immediately. The proxy fires it in a try/finally after
post_call_success_hook so the SLP is built with guardrail info.

Streaming: a closure on logging_obj is called by CSW.__anext__ at
stream end. The closure runs only guardrail hooks (not all callbacks)
on the assembled response, then fires both logging handlers. This
avoids behavioral changes for non-guardrail callbacks on streaming.
Comment on lines +1362 to +1376
if "apply_guardrail" in type(cb).__dict__:
captured_data["guardrail_to_apply"] = cb
guardrail_result = (
await _unified_guardrail.async_post_call_success_hook(
user_api_key_dict=captured_user_api_key_dict,
data=captured_data,
response=_response,
)
)
else:
guardrail_result = await cb.async_post_call_success_hook(
user_api_key_dict=captured_user_api_key_dict,
data=captured_data,
response=_response,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Model-level non-default guardrails silently skipped at execution

guardrail_data (the model-merged dict) is only used for the outer should_run_guardrail gate, but both execution branches pass the original captured_data to the actual hook calls.

UnifiedLLMGuardrails.async_post_call_success_hook re-runs should_run_guardrail(data=data, ...) internally (at unified_guardrail.py lines 211–216):

if guardrail_to_apply.should_run_guardrail(
    data=data, event_type=GuardrailEventHooks.post_call
) is not True:
    return   # silently skipped

Because captured_data does not contain the model-level guardrails that were merged into guardrail_data, this second should_run_guardrail call returns False for any default_on=False guardrail configured via litellm_params.guardrails, and the guardrail is silently dropped at execution time even though it passed the outer gate.

The non-streaming ProxyLogging.post_call_success_hook avoids this by reassigning data after the merge so the same merged dict is used for both gating and execution. The fix is to pass guardrail_data (not captured_data) to all execution calls, and set guardrail_data["guardrail_to_apply"] = cb on the merged dict:

guardrail_data = _check_and_merge_model_level_guardrails(
    data=captured_data, llm_router=_global_llm_router
)
for cb in litellm.callbacks:
    ...
    try:
        guardrail_result = None
        if "apply_guardrail" in type(cb).__dict__:
            guardrail_data["guardrail_to_apply"] = cb  # use merged dict
            guardrail_result = await _unified_guardrail.async_post_call_success_hook(
                data=guardrail_data,                   # pass merged dict
                ...
            )
        else:
            guardrail_result = await cb.async_post_call_success_hook(
                data=guardrail_data,                   # pass merged dict
                ...
            )

Use the merged guardrail_data dict (from _check_and_merge_model_level_guardrails)
for hook invocations in _run_deferred_stream_guardrails, instead of the original
captured_data. This ensures model-level non-default guardrails are visible to
inner should_run_guardrail re-checks inside UnifiedLLMGuardrails.

Rewrite three hand-crafted closure tests to exercise the production
_run_deferred_stream_guardrails exception-handling path. Add three new tests
that use deep-copy mocks to prove hooks receive the merged dict.
Comment on lines +1346 to +1360

_response = assembled_response
_unified_guardrail = UnifiedLLMGuardrails()
guardrail_data = _check_and_merge_model_level_guardrails(
data=captured_data, llm_router=_global_llm_router
)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,
):
continue
try:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Logging silently dropped if guardrail initialization raises

_run_deferred_stream_guardrails calls UnifiedLLMGuardrails() and _check_and_merge_model_level_guardrails(...) before the loop and before the create_task logging calls. If either of these raises (e.g., a transient import error, a router issue, or a future __init__ change), the function exits with an unhandled exception and the asyncio.create_task / executor.submit calls at the bottom are never reached. Logging for that streaming response is silently dropped — the event loop will log a task exception warning but no structured proxy log entry is written.

This contrasts with the per-guardrail try/except inside the loop (which correctly continues logging even if a single guardrail fails) and with the finally pattern used for the non-streaming path.

@staticmethod
async def _run_deferred_stream_guardrails(
    ...
) -> None:
    from litellm.litellm_core_utils.thread_pool_executor import executor
    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
        UnifiedLLMGuardrails,
    )
    from litellm.proxy.proxy_server import llm_router as _global_llm_router
    from litellm.proxy.utils import _check_and_merge_model_level_guardrails

    _response = assembled_response
    try:
        _unified_guardrail = UnifiedLLMGuardrails()
        guardrail_data = _check_and_merge_model_level_guardrails(
            data=captured_data, llm_router=_global_llm_router
        )
        for cb in litellm.callbacks:
            ...  # per-guardrail try/except already present
    except Exception as e:
        verbose_proxy_logger.exception(
            "Error during guardrail initialization/dispatch for streaming response: %s", e
        )
    # always fire logging regardless of guardrail errors
    try:
        asyncio.create_task(
            captured_logging_obj.async_success_handler(_response, ...)
        )
    ...

@michelligabriele

Copy link
Copy Markdown
Contributor Author

@greptileai how would you address these concerns?
The non-streaming deferral path is solid and well-tested. The streaming path correctly handles per-guardrail errors and model-level merging, and the static method extraction improves testability. However, an unguarded initialization section in _run_deferred_stream_guardrails (before the per-guardrail try/except loop) can cause silent logging loss if UnifiedLLMGuardrails() or _check_and_merge_model_level_guardrails() throws. Additionally, a fastapi.HTTPException import in the test file (previously flagged) remains unresolved. The overall complexity of the two-mechanism deferral approach is high, and the streaming path has edge cases that are difficult to test end-to-end without a running proxy.
Pay close attention to litellm/proxy/common_request_processing.py — specifically the initialization block at the top of _run_deferred_stream_guardrails (lines 1340–1360) which can drop logging on transient errors.

@greptile-apps

greptile-apps Bot commented Mar 19, 2026

Copy link
Copy Markdown
Contributor

Based on the concerns raised, here are specific recommendations:

1. Wrap initialization in _run_deferred_stream_guardrails with try/finally

The initialization block (lines ~1340-1360) can silently drop logging if UnifiedLLMGuardrails() or _check_and_merge_model_level_guardrails() throws. Wrap it:

@staticmethod
async def _run_deferred_stream_guardrails(...):
    try:
        # Initialization block
        unified_guardrails = UnifiedLLMGuardrails()
        guardrail_data = _check_and_merge_model_level_guardrails(
            captured_data, captured_logging_obj.llm_router
        )
        
        # ... rest of initialization ...
        
        # Per-guardrail loop with existing try/except
        for callback in litellm.callbacks:
            # ... existing per-guardrail error handling ...
    finally:
        # Always fire logging, even if initialization fails
        asyncio.create_task(
            captured_logging_obj.async_success_handler(...)
        )
        asyncio.create_task(
            captured_logging_obj.async_success_handler_streaming_only(...)
        )

This ensures logging fires even if guardrail setup throws, matching the non-streaming try/finally pattern.

2. Remove fastapi import from test file

test_deferred_guardrail_logging.py:20 imports from fastapi import HTTPException. Per rule 7382761b, fastapi imports are forbidden outside proxy/:

from litellm.proxy._types import ProxyException

# Replace HTTPException usage with ProxyException
raise ProxyException(
    message="Content blocked",
    type="invalid_request",
    param=None,
    code=400
)

If you need HTTP status codes for test assertions, use litellm.proxy._types.ProxyException or check the underlying exception type without importing fastapi directly.

3. Simplify with unified deferral flag

The two-mechanism approach (non-streaming flag + streaming closure) adds cognitive load. Consider a single mechanism:

  • Set logging_obj._defer_logging = True for both paths when _has_post_call_guardrails()
  • Non-streaming: wrapper_async stores closure (current behavior)
  • Streaming: CSW.__anext__ checks the same flag and calls _run_deferred_stream_guardrails directly (no closure attachment)

This reduces state management — one flag instead of flag + closure attribute.

4. Add integration test with real proxy instance

Streaming edge cases (especially the closure→CSW→guardrail→logging chain) are hard to unit test. Add a single end-to-end test:

@pytest.mark.asyncio
async def test_streaming_deferred_logging_e2e(proxy_server_fixture):
    """Integration test: streaming request with post-call guardrail
    verifies guardrail_information appears in logged payload."""
    # Use a real proxy with a test guardrail
    # Make a streaming request
    # Assert guardrail_information is present in the logged payload

This catches issues that mock-based tests miss (e.g., closure firing twice, race conditions).


Priority:

  1. Issue Give me consistent exceptions  #1 (unguarded init) — critical, can cause silent logging loss in production
  2. Issue Enable model / call timeouts #2 (fastapi import) — violates project rules, easy fix
  3. Issue Guarantee format of exceptions #3 (complexity) — nice-to-have refactor, lower priority
  4. Issue Add support for replicate [code in ticket] #4 (integration test) — reduces risk but non-blocking for merge

The first two are blockers; the last two are improvements for future robustness.

…imports

Wrap _run_deferred_stream_guardrails initialization (UnifiedLLMGuardrails
constructor and _check_and_merge_model_level_guardrails) in try/finally
so logging always fires even if init throws. Prevents silent logging loss
on transient errors.

Move fastapi.HTTPException import from module-level to local test-function
scope. Add test_logging_fires_even_if_guardrail_init_raises to verify the
try/finally guard.

_response = assembled_response
try:
_unified_guardrail = UnifiedLLMGuardrails()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 New UnifiedLLMGuardrails() instance per deferred streaming call

_run_deferred_stream_guardrails instantiates a fresh UnifiedLLMGuardrails() on every invocation, while the rest of the proxy (proxy/utils.py line 141) uses a module-level singleton:

# proxy/utils.py (existing pattern)
unified_guardrail = UnifiedLLMGuardrails()  # created once at module load

Although UnifiedLLMGuardrails.__init__ is lightweight (just stores kwargs and calls super()), creating a new object per deferred call is inconsistent with the established singleton pattern used by post_call_success_hook. Consider importing and reusing the existing singleton to keep the behaviour consistent:

from litellm.proxy.utils import unified_guardrail as _unified_guardrail

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

# Defer async logging when post-call guardrails are configured so the
# StandardLoggingPayload is built after guardrails write to metadata.
# Cache the result to avoid scanning litellm.callbacks twice.
_has_post_call_guardrails = self._has_post_call_guardrails()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Local variable shadows static method name

_has_post_call_guardrails is both the name of the static method on the class and the local variable that stores its result. Within the rest of base_process_llm_request, any attempt to call self._has_post_call_guardrails() again would inadvertently call the boolean True/False instead, raising a TypeError.

Consider using a more distinct name for the cached result to prevent this footgun for future maintainers:

Suggested change
_has_post_call_guardrails = self._has_post_call_guardrails()
_post_call_guardrails_active = self._has_post_call_guardrails()

Then update the two downstream uses at lines 946 and 1060 accordingly.

…variable

Reuse the module-level unified_guardrail singleton from proxy/utils.py
in _run_deferred_stream_guardrails instead of creating a new instance
per call, matching the pattern used by post_call_success_hook.

Rename local variable _has_post_call_guardrails to
_post_call_guardrails_active to avoid shadowing the static method name.
hidden_params.get("fastest_response_batch_completion", None),
hidden_params.get("additional_headers", {}) or {},
)
try:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 _on_deferred_stream_complete not cleaned up in exception path

The try block now wraps the entire post-response processing section. If an exception is raised after _on_deferred_stream_complete is stored on logging_obj (e.g., inside post_call_response_headers_hook), the finally block only cleans up _enqueue_deferred_logging. The streaming closure remains attached to logging_obj.

In practice the CSW is never returned to the caller on this exception path, so the closure never fires — meaning this request produces no structured log entry. The error path itself is logged by _handle_llm_api_exception, but the deferred logging for the successful LLM response is silently dropped.

Adding a cleanup line to the finally block ensures no orphaned state and makes the contract explicit:

finally:
    # Clean up any stale streaming closure (set if exception occurred
    # after closure was attached but before the streaming early-return).
    if getattr(logging_obj, "_on_deferred_stream_complete", None):
        logging_obj._on_deferred_stream_complete = None

    _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None)
    ...

…g closure

Move non-essential lazy imports (llm_router, _check_and_merge,
unified_guardrail) inside the try block of _run_deferred_stream_guardrails
so that import failures are caught and the finally block still fires
logging. Only executor stays outside since the finally block needs it.

Add _on_deferred_stream_complete orphan cleanup in the finally block of
base_process_llm_request. If an exception propagates after the streaming
closure is stored but before a StreamingResponse is returned, the closure
is orphaned (CSW never consumes the stream). Detect this via
sys.exc_info() and fire logging directly to prevent silent loss.
@michelligabriele

Copy link
Copy Markdown
Contributor Author

@greptileai how would you fix this last one remaining issue in a way that this PR can raise the score to 4/5?
One issue found:

The sys.exc_info()[1] check in the finally block (used to detect an orphaned streaming closure on the exception path) can produce false positives if base_process_llm_request is called from within an except block in the caller's scope. In that case the _on_deferred_stream_complete closure would be prematurely cleared and logging would fire with the raw CustomStreamWrapper (not the assembled response), followed by a second logging call from CSW's normal path. Using an explicit boolean sentinel (_exception_raised) would be more reliable.
Confidence Score: 3/5
The PR solves a real bug but the sys.exc_info() check introduces a fragile edge case that could cause double-logging or logging with an unassembled response in certain caller contexts.
The overall approach is sound and well-tested (17 tests, including regression tests). Many issues raised in prior review iterations have been addressed (per-guardrail try/except, static method extraction, singleton reuse, variable shadowing fix). However, the sys.exc_info()[1] check for detecting orphaned streaming closures in the finally block is fragile: Python propagates exception context from outer except blocks into called functions' finally scopes, so a false-positive could prematurely clear the deferred closure, causing logging to fire with the raw CustomStreamWrapper and then again (correctly) from CSW at stream end.
litellm/proxy/common_request_processing.py — specifically the sys.exc_info()[1] check in the finally block at line 1172.

@greptile-apps

greptile-apps Bot commented Mar 19, 2026

Copy link
Copy Markdown
Contributor

Replace the sys.exc_info()[1] check with an explicit boolean sentinel:

async def base_process_llm_request(...):
    _exception_raised = False  # Add at function start
    
    try:
        # ... existing code ...
        
    except Exception as e:
        _exception_raised = True  # Set before re-raising
        raise
    
    finally:
        # Replace sys.exc_info() check with boolean sentinel:
        if _exception_raised and hasattr(logging_obj, "_on_deferred_stream_complete"):
            # Clear orphaned streaming closure (stream won't complete after exception)
            delattr(logging_obj, "_on_deferred_stream_complete")
        
        # ... rest of finally block ...

Why this fixes the issue:

  • _exception_raised is function-scoped, immune to outer exception context
  • Only set True when exception actually occurs in this function
  • Prevents false positives when called from caller's except block
  • Eliminates risk of premature closure clearing → no double-logging or unassembled response logging

This simple change makes exception detection reliable and should raise the confidence score to 4/5.

…ection

Replace sys.exc_info()[1] check with an explicit _exception_raised boolean
sentinel. The flag is function-scoped, immune to outer exception context,
and only set when an exception actually occurs in base_process_llm_request.
This prevents false positives when called from a caller's except block.
Split the single try/except in the _exception_raised cleanup path into
separate try blocks for asyncio.create_task and executor.submit, matching
the pattern used in _run_deferred_stream_guardrails. If create_task
raises, sync logging via executor.submit still fires.
@ghost
ghost changed the base branch from main to litellm_oss_staging_03_19_2026 March 20, 2026 00:55
@ghost
ghost merged commit d4857f6 into BerriAI:litellm_oss_staging_03_19_2026 Mar 20, 2026
38 of 39 checks passed
@micaelmalta

micaelmalta commented Apr 6, 2026

Copy link
Copy Markdown

This PR breaks cost and logging for Bedrock Anthropic in some conditions
#23150
#20179

fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
…uardrail-logging-v2

fix(proxy): defer logging until post-call guardrails complete
This pull request was closed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants