From 8b24d4c24fb1dc2268a667bfd8591b67fff76b55 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 01:24:52 +0000 Subject: [PATCH 1/5] fix(proxy): log blocked streaming guardrail responses as failures, not success Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 55 ++++++++++-- .../test_post_call_failure_hook.py | 52 ++++++++++++ .../proxy_logging/test_streaming_hooks.py | 83 +++++++++++++++++++ 3 files changed, 184 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e611984bf4c..2a62704c6da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -85,7 +85,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException +from litellm.exceptions import ( + GuardrailRaisedException, + RejectedRequestError, + SensitiveDataRouteException, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -2991,6 +2995,7 @@ class ProxyLogging: - Authentication Errors from user_api_key_auth - HTTP HTTPException (rate limit errors) - ProxyException (guardrail blocks, budget / rate-limit errors) + - GuardrailRaisedException (guardrail blocks / guardrail failures) """ ######################################################### @@ -3005,9 +3010,9 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance(original_exception, (HTTPException, ProxyException)) or ( - error_type == ProxyErrorTypes.auth_error - ) + return isinstance( + original_exception, (HTTPException, ProxyException, GuardrailRaisedException) + ) or (error_type == ProxyErrorTypes.auth_error) async def _handle_logging_proxy_only_error( self, @@ -3564,7 +3569,7 @@ class ProxyLogging: except (GeneratorExit, asyncio.CancelledError): raise except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3639,7 +3644,7 @@ class ProxyLogging: except (GeneratorExit, asyncio.CancelledError): raise except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) raise # Fire deferred logging AFTER all guardrail end-of-stream blocks @@ -3735,6 +3740,44 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) + @staticmethod + def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None: + """Discard the deferred stream-complete dispatch when the stream ends in + a failure (e.g. an end-of-stream guardrail block raising out of the + callback chain). The deferred dispatch is the success logging path — + firing it here would record the blocked request as a success callback + and a ``status=success`` spend row before the outer generator's + ``post_call_failure_hook`` writes the failure row. The CSW shape parks + ``(assembled ModelResponse, cache_hit)``; record its partial usage so + the failure row bills what the stream consumed instead of zero. The + native /v1/messages and responses shapes park ``(coroutine,)`` and + still need the flush (no success row is produced without it), so they + keep the existing fire behaviour. + """ + logging_obj: Final = request_data.get("litellm_logging_obj") + if logging_obj is None: + return + _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( + logging_obj, "_on_deferred_stream_complete", None + ) + _args: Final[tuple[object, ...] | None] = getattr( + logging_obj, "_deferred_stream_complete_args", None + ) + if _deferred_cb is None or _args is None: + return + assembled: Final = _args[0] + if not isinstance(assembled, ModelResponse): + ProxyLogging._fire_deferred_stream_logging(request_data) + return + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + usage: Final[Usage | None] = getattr(assembled, "usage", None) + if isinstance(usage, Usage): + logging_obj.record_partial_usage_for_failure( + usage, + logging_obj._response_cost_calculator(result=assembled) or 0.0, + ) + async def _arelease_max_parallel_requests_on_disconnect( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 8b2b6b4c6ca..49c63a91b3e 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -11,6 +11,7 @@ import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import AlertType, ProxyErrorTypes from litellm.proxy.utils import ProxyLogging @@ -47,12 +48,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging): error_type=ProxyErrorTypes.auth_error, route="/chat/completions", ), + "guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error( + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + route="/chat/completions", + ), } assert snapshot == { "no_route": False, "non_llm_route": False, "http_on_llm_route": True, "auth_short_circuit": True, + "guardrail_raised_on_llm_route": True, } @@ -318,3 +324,49 @@ async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes route=route, ) assert request_data["call_type"] == route + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A ``GuardrailRaisedException`` on an LLM route must reach the logging + object's ``async_failure_handler`` so custom loggers see a ``failure`` + status - without this, guardrail blocks produce only + ``post_call_failure_hook`` and no failure logging event.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + recorded: dict[str, Any] = {} + + class _StatusRecorder(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded["status"] = (kwargs.get("standard_logging_object") or {}).get("status") + + monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_guardrail_block_failure_cb", + function_id="test_guardrail_block_failure_cb", + ) + request_data = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "test_guardrail_block_failure_cb", + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {}, + } + proxy_logging.alert_types = [] + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert recorded["status"] == "failure" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ec5b994f147..68e37da8e3a 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -479,6 +479,89 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ assert logging_obj._deferred_stream_complete_args is None +@pytest.mark.asyncio +async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + On /chat/completions streams the CSW shape parks + ``(assembled ModelResponse, cache_hit)`` as the deferred args. A guardrail + that raises ``GuardrailRaisedException`` at end of stream must NOT dispatch + that deferred success logging - the request is logged via the failure path + instead, with the consumed usage carried over so the failure row bills + correctly. + """ + from litellm.exceptions import GuardrailRaisedException + from litellm.types.utils import Usage + + events: List[Any] = [] + request_data: Dict[str, Any] = {"metadata": {}} + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_chat_stream_guardrail_block", + function_id="test_chat_stream_guardrail_block", + ) + logging_obj.optional_params = {} + logging_obj.litellm_params = {} + logging_obj.standard_built_in_tools_params = None + + async def _dispatch_deferred_logging(*args): + events.append("success_dispatched") + + logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging + request_data["litellm_logging_obj"] = logging_obj + + assembled = litellm.ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + async def _upstream(): + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]} + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]} + logging_obj._deferred_stream_complete_args = (assembled, False) + + class _BlockingGuardrail(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + raise GuardrailRaisedException( + guardrail_name="g", message="blocked", blocked_content=True + ) + + monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()]) + + with pytest.raises(GuardrailRaisedException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=_upstream(), + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "callback_cleared": logging_obj._on_deferred_stream_complete is None, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens, + "response_cost_positive": logging_obj.model_call_details["response_cost"] > 0, + } + assert snapshot == { + "events": [], + "callback_cleared": True, + "args_cleared": True, + "combined_usage_total_tokens": 8, + "response_cost_positive": True, + } + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # --------------------------------------------------------------------------- From 0949f24eef0d7a3fd07b4c78f64aa0bdcee2b12e Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 01:25:36 +0000 Subject: [PATCH 2/5] refactor(proxy): tighten deferred stream logging discard docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2a62704c6da..dee4c875c1a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3742,17 +3742,12 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None: - """Discard the deferred stream-complete dispatch when the stream ends in - a failure (e.g. an end-of-stream guardrail block raising out of the - callback chain). The deferred dispatch is the success logging path — - firing it here would record the blocked request as a success callback - and a ``status=success`` spend row before the outer generator's - ``post_call_failure_hook`` writes the failure row. The CSW shape parks - ``(assembled ModelResponse, cache_hit)``; record its partial usage so - the failure row bills what the stream consumed instead of zero. The - native /v1/messages and responses shapes park ``(coroutine,)`` and - still need the flush (no success row is produced without it), so they - keep the existing fire behaviour. + """Drop the parked success dispatch when the stream ends in an exception. + + The CSW shape parks ``(assembled ModelResponse, cache_hit)``: its usage is + carried onto the logging object so the failure row bills what the stream + consumed. The native /v1/messages and responses shapes park a logging + coroutine with no recoverable usage, so they keep firing as before. """ logging_obj: Final = request_data.get("litellm_logging_obj") if logging_obj is None: From ec799686a4156daf4ac20fd954ec59885fc6eaf8 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 01:30:46 +0000 Subject: [PATCH 3/5] style(proxy): ruff format utils.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index dee4c875c1a..ba84a603dd2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3010,9 +3010,9 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance( - original_exception, (HTTPException, ProxyException, GuardrailRaisedException) - ) or (error_type == ProxyErrorTypes.auth_error) + return isinstance(original_exception, (HTTPException, ProxyException, GuardrailRaisedException)) or ( + error_type == ProxyErrorTypes.auth_error + ) async def _handle_logging_proxy_only_error( self, @@ -3755,9 +3755,7 @@ class ProxyLogging: _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( logging_obj, "_on_deferred_stream_complete", None ) - _args: Final[tuple[object, ...] | None] = getattr( - logging_obj, "_deferred_stream_complete_args", None - ) + _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) if _deferred_cb is None or _args is None: return assembled: Final = _args[0] From 7095373dd542c47fbf563e2d515633f96d1c55f3 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 20:10:18 +0000 Subject: [PATCH 4/5] fix(proxy): only discard parked stream logging for errors the failure path logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 + litellm/proxy/utils.py | 51 ++++----- .../test_post_call_failure_hook.py | 14 ++- .../proxy_logging/test_streaming_hooks.py | 106 +++++++++++++----- 4 files changed, 112 insertions(+), 61 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4ddb9ce5b8e..a7ad774b02d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -639,6 +639,8 @@ class Logging(LiteLLMLoggingBaseClass): self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None + self._on_deferred_stream_complete: Callable[..., Awaitable[None]] | None = None + self._deferred_stream_complete_args: tuple[object, ...] | None = None def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ba84a603dd2..a12bb56f8f8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -905,6 +905,9 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) + + def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields @@ -3010,9 +3013,7 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance(original_exception, (HTTPException, ProxyException, GuardrailRaisedException)) or ( - error_type == ProxyErrorTypes.auth_error - ) + return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error) async def _handle_logging_proxy_only_error( self, @@ -3568,8 +3569,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3643,8 +3645,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise # Fire deferred logging AFTER all guardrail end-of-stream blocks @@ -3741,35 +3744,29 @@ class ProxyLogging: asyncio.create_task(_deferred_cb(*_args)) @staticmethod - def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None: - """Drop the parked success dispatch when the stream ends in an exception. - - The CSW shape parks ``(assembled ModelResponse, cache_hit)``: its usage is - carried onto the logging object so the failure row bills what the stream - consumed. The native /v1/messages and responses shapes park a logging - coroutine with no recoverable usage, so they keep firing as before. + def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: + """Drop the parked success dispatch when the stream ends in an error the proxy logs + as a failure (``_PROXY_ONLY_LLM_API_ERRORS``, e.g. a post_call guardrail block) and + the CSW parked an assembled ``ModelResponse``, carrying its usage onto the logging + object so the failure row bills what the stream consumed. Returns False, leaving the + parked dispatch for the caller to flush, for any other error and for the native + /v1/messages and responses shapes that park a logging coroutine with no usage. """ logging_obj: Final = request_data.get("litellm_logging_obj") - if logging_obj is None: - return - _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( - logging_obj, "_on_deferred_stream_complete", None - ) + if not isinstance(logging_obj, Logging): + return False _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) - if _deferred_cb is None or _args is None: - return - assembled: Final = _args[0] - if not isinstance(assembled, ModelResponse): - ProxyLogging._fire_deferred_stream_logging(request_data) - return + assembled: Final = _args[0] if _args else None + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None usage: Final[Usage | None] = getattr(assembled, "usage", None) if isinstance(usage, Usage): logging_obj.record_partial_usage_for_failure( - usage, - logging_obj._response_cost_calculator(result=assembled) or 0.0, + usage, logging_obj._response_cost_calculator(result=assembled) or 0.0 ) + return True async def _arelease_max_parallel_requests_on_disconnect( self, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 49c63a91b3e..13fcccbad97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio +from datetime import datetime from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -334,15 +335,16 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( object's ``async_failure_handler`` so custom loggers see a ``failure`` status - without this, guardrail blocks produce only ``post_call_failure_hook`` and no failure logging event.""" - from datetime import datetime - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - recorded: dict[str, Any] = {} + recorded: list[object] = [] class _StatusRecorder(CustomLogger): - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - recorded["status"] = (kwargs.get("standard_logging_object") or {}).get("status") + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + standard_logging_object = kwargs.get("standard_logging_object") + recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None) monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) logging_obj = LiteLLMLoggingObj( @@ -369,4 +371,4 @@ async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( ) await asyncio.sleep(0) await asyncio.sleep(0) - assert recorded["status"] == "failure" + assert recorded == ["failure"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 68e37da8e3a..ebc831b4102 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -18,12 +19,15 @@ import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import Usage @pytest.fixture(autouse=True) @@ -479,37 +483,25 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ assert logging_obj._deferred_stream_complete_args is None -@pytest.mark.asyncio -async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( - proxy_logging, make_user_api_key_auth, monkeypatch -): - """ - On /chat/completions streams the CSW shape parks - ``(assembled ModelResponse, cache_hit)`` as the deferred args. A guardrail - that raises ``GuardrailRaisedException`` at end of stream must NOT dispatch - that deferred success logging - the request is logged via the failure path - instead, with the consumed usage carried over so the failure row bills - correctly. - """ - from litellm.exceptions import GuardrailRaisedException - from litellm.types.utils import Usage - - events: List[Any] = [] - request_data: Dict[str, Any] = {"metadata": {}} +def _armed_chat_stream( + test_name: str, request_data: dict[str, object], events: list[str] +) -> tuple[LiteLLMLoggingObj, AsyncIterator[dict[str, object]]]: + """A /chat/completions stream whose CSW shape parks ``(assembled ModelResponse, cache_hit)`` + at upstream exhaustion, with the deferred dispatch recording into ``events``.""" logging_obj = LiteLLMLoggingObj( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, call_type="acompletion", start_time=datetime.now(), - litellm_call_id="test_chat_stream_guardrail_block", - function_id="test_chat_stream_guardrail_block", + litellm_call_id=test_name, + function_id=test_name, ) logging_obj.optional_params = {} logging_obj.litellm_params = {} logging_obj.standard_built_in_tools_params = None - async def _dispatch_deferred_logging(*args): + async def _dispatch_deferred_logging(*args: object) -> None: events.append("success_dispatched") logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging @@ -521,24 +513,48 @@ async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_suc usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), ) - async def _upstream(): + async def _upstream() -> AsyncIterator[dict[str, object]]: yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]} yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]} logging_obj._deferred_stream_complete_args = (assembled, False) - class _BlockingGuardrail(CustomLogger): - async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + return logging_obj, _upstream() + + +def _raising_at_end_of_stream(error: Exception) -> CustomLogger: + class _EndOfStreamRaiser(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: async for chunk in response: yield chunk - raise GuardrailRaisedException( - guardrail_name="g", message="blocked", blocked_content=True - ) + raise error - monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()]) + return _EndOfStreamRaiser() + + +@pytest.mark.asyncio +async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + A guardrail that raises ``GuardrailRaisedException`` at end of a + /chat/completions stream must NOT dispatch the parked success logging: + the request is logged via the failure path instead, with the consumed + usage carried over so the failure row bills correctly. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_guardrail_block", request_data, events) + monkeypatch.setattr( + litellm, + "callbacks", + [_raising_at_end_of_stream(GuardrailRaisedException(guardrail_name="g", message="blocked"))], + ) with pytest.raises(GuardrailRaisedException): async for _ in proxy_logging.async_post_call_streaming_iterator_hook( - response=_upstream(), + response=upstream, user_api_key_dict=make_user_api_key_auth(), request_data=request_data, ): @@ -562,6 +578,40 @@ async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_suc } +@pytest.mark.asyncio +async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes_success_logging( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + ``post_call_failure_hook`` only routes proxy-level errors (HTTPException, + ProxyException, GuardrailRaisedException) through failure logging. A + callback that dies with any other exception after the stream completed + must keep flushing the parked success dispatch, or the request ends with + no terminal log at all. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_generic_callback_error", request_data, events) + monkeypatch.setattr(litellm, "callbacks", [_raising_at_end_of_stream(RuntimeError("callback crashed"))]) + + with pytest.raises(RuntimeError): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "failure_usage_recorded": "combined_usage_object" in logging_obj.model_call_details, + } + assert snapshot == {"events": ["success_dispatched"], "args_cleared": True, "failure_usage_recorded": False} + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # --------------------------------------------------------------------------- From 46bd3d40d7abb7f44db19fb8d81b0d9871a314f1 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 20:23:50 +0000 Subject: [PATCH 5/5] refactor(logging): bill an assembled stream on the failure log via a public Logging method Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 8 ++++++-- litellm/proxy/utils.py | 16 ++++------------ .../proxy_logging/test_post_call_failure_hook.py | 1 - 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a7ad774b02d..abac624d5ec 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -639,8 +639,6 @@ class Logging(LiteLLMLoggingBaseClass): self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None - self._on_deferred_stream_complete: Callable[..., Awaitable[None]] | None = None - self._deferred_stream_complete_args: tuple[object, ...] | None = None def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" @@ -1993,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: + """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" + usage: Final = getattr(assembled, "usage", None) + if isinstance(usage, Usage): + self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0) + async def dispatch_failure_handlers( self, exception: Exception, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a12bb56f8f8..80cf6ba3c8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3745,13 +3745,9 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: - """Drop the parked success dispatch when the stream ends in an error the proxy logs - as a failure (``_PROXY_ONLY_LLM_API_ERRORS``, e.g. a post_call guardrail block) and - the CSW parked an assembled ``ModelResponse``, carrying its usage onto the logging - object so the failure row bills what the stream consumed. Returns False, leaving the - parked dispatch for the caller to flush, for any other error and for the native - /v1/messages and responses shapes that park a logging coroutine with no usage. - """ + """Drop the parked success dispatch for an assembled chat stream that ends in an error + ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + Returns False when the parked dispatch should still be flushed by the caller.""" logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return False @@ -3761,11 +3757,7 @@ class ProxyLogging: return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - usage: Final[Usage | None] = getattr(assembled, "usage", None) - if isinstance(usage, Usage): - logging_obj.record_partial_usage_for_failure( - usage, logging_obj._response_cost_calculator(result=assembled) or 0.0 - ) + logging_obj.record_assembled_response_for_failure(assembled) return True async def _arelease_max_parallel_requests_on_disconnect( diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 13fcccbad97..d7a6124dd97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -5,7 +5,6 @@ from __future__ import annotations import asyncio from datetime import datetime -from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest