fix(proxy): defer logging until post-call guardrails complete - #24135
Conversation
…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
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes How it works:
Key design decisions addressed from prior review:
Remaining minor concerns:
Confidence Score: 4/5
|
| 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
Last reviewed commit: "fix(proxy): split or..."
| guardrail_data = _check_and_merge_model_level_guardrails( | ||
| data=_captured_data, llm_router=None |
There was a problem hiding this comment.
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 mergingThis 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:
| 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 |
There was a problem hiding this comment.
_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.
c3ff034 to
0a54601
Compare
| from unittest.mock import MagicMock, patch | ||
|
|
||
| import pytest | ||
| from fastapi import HTTPException |
There was a problem hiding this comment.
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:
| 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)
| 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] |
There was a problem hiding this comment.
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 runSince 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"] = TrueThis ensures all configured post-call guardrails contribute to audit logging even when one fails.
0a54601 to
04ef82a
Compare
| 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, | ||
| ) | ||
|
|
There was a problem hiding this comment.
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)
| 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): |
There was a problem hiding this comment.
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.
| llm_router=llm_router, | ||
| ) | ||
|
|
||
| # Defer async logging when post-call guardrails are configured so the |
There was a problem hiding this comment.
_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:
_defer_async_logging = Trueis set onlogging_obj- In
utils.py'swrapper_async, the result IS a streaming response → it exits early before the_defer_async_loggingblock _enqueue_deferred_loggingis never stored onlogging_obj- The
finallyblock sees_enqueue_fn = None→ no-op - Meanwhile,
_on_deferred_stream_completeis set (becauseisinstance(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.
04ef82a to
001501f
Compare
| 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, | ||
| ) |
There was a problem hiding this comment.
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 skippedBecause 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.
|
|
||
| _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: |
There was a problem hiding this comment.
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, ...)
)
...|
@greptileai how would you address these concerns? |
|
Based on the concerns raised, here are specific recommendations: 1. Wrap initialization in
|
…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() |
There was a problem hiding this comment.
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 loadAlthough 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_guardrailNote: 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() |
There was a problem hiding this comment.
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:
| _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: |
There was a problem hiding this comment.
_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.
|
@greptileai how would you fix this last one remaining issue in a way that this PR can raise the score to 4/5? 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. |
|
Replace the 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:
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.
d4857f6
into
BerriAI:litellm_oss_staging_03_19_2026
…uardrail-logging-v2 fix(proxy): defer logging until post-call guardrails complete
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
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays 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)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🐛 Bug Fix
Changes
Problem
guardrail_informationis alwaysNoneinStandardLoggingPayloadwhen post-call guardrails (e.g. OpenAI Moderation) are configured. This happens because:asyncio.create_taskinwrapper_async(utils.py) fires the logging task beforepost_call_success_hookruns inbase_process_llm_request, so the SLP is built before guardrails write to metadata.CustomStreamWrapper.__anext__without any guardrail data from the assembled response —post_call_success_hookis 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):_has_post_call_guardrails()checks if anyCustomGuardrailwithpost_callevent hook is registeredlogging_obj._defer_async_logging = Truewrapper_asyncsees the flag → stores closure onlogging_obj._enqueue_deferred_logginginstead ofasyncio.create_task. Sync callbacks fire immediately (unchanged).base_process_llm_requestrunspost_call_success_hook(guardrails write to metadata)finallyblock calls the stored closure →create_taskfires → SLP built with guardrail infoStreaming (
common_request_processing.py+streaming_handler.py):_has_post_call_guardrailsand response isCustomStreamWrapper: attach_on_deferred_stream_completeclosure tologging_objlitellm.callbacks, filters forCustomGuardrailinstances withpost_callevent hook, calls theirasync_post_call_success_hook. This is the same patternProxyLogging.post_call_success_hookuses internally, but filtered to guardrails only. Non-guardrail callbacks are not called (avoids behavioral changes for streaming).CSW.__anext__at stream end: checks for closure. If set, clears it and calls it viaasyncio.create_task. If not set, fires logging directly (original behavior preserved).post_call_success_hook(no early return), the closure is cleared first to prevent double invocation.Files
litellm/utils.py_defer_async_loggingflag → store closure instead ofcreate_tasklitellm/proxy/common_request_processing.py_has_post_call_guardrails()static method, deferral flag, streaming closure (guardrail-only),try/finallylitellm/litellm_core_utils/streaming_handler.pyCSW.__anext__checks for_on_deferred_stream_complete, calls it instead of logging directlytests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.pydocs/my-website/docs/proxy/guardrails/custom_guardrail.mdpost_callguardrails as audit-onlyTests (17 total)
Detection (7):
_has_post_call_guardrailsreturns correct result for post_call, pre_call, event_hook=None, list event hooks, non-guardrail callbacks, empty callbacksNon-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