From 7095373dd542c47fbf563e2d515633f96d1c55f3 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 20:10:18 +0000 Subject: [PATCH] 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 # ---------------------------------------------------------------------------