From 261807114a50e09b7fe6b48a7d8709dbf25e9c02 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:10:58 -0700 Subject: [PATCH] fix(proxy): accept both deferred stream logging arg shapes on native routes (#40869) * fix(proxy): accept both deferred stream logging arg shapes on native routes _arm_deferred_stream_dispatch armed a one-argument closure on every anthropic_messages/aresponses stream that was not a CustomStreamWrapper or a LiteLLMCompletionStreamingIterator. The bridged /v1/messages path returns a plain SSE generator that shares its inner CustomStreamWrapper logging_obj, so it stores (assembled_response, cache_hit) and _fire_deferred_stream_logging raised TypeError, dropping spend logs and callbacks and ending the stream with an error. The closure now dispatches on the stored args shape: a single coroutine is enqueued, a two-tuple runs success handlers, anything else is logged and dropped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert dropped deferred payload via caplog instead of patching the logger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 66 ++++++++-------- .../test_deferred_guardrail_logging.py | 78 ++++++++++++++++++- 2 files changed, 107 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 289fe086379..e6ed60ba177 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,7 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -3210,20 +3210,24 @@ class ProxyBaseLLMRequestProcessing: end-of-stream blocks complete, so the spend log sees guardrail_information. - Three closure shapes, matching who owns logging for the stream: + Two closure shapes, matching who owns logging for the stream: - CustomStreamWrapper (chat completions) stores (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares - its inner CustomStreamWrapper's logging_obj, so it stores the same - (assembled_response, cache_hit) shape; the closure only dispatches - success logging, matching the route's pre-existing hook surface. - - Native anthropic_messages/aresponses iterators store a single - ready-made logging coroutine to enqueue. + - Every other anthropic_messages/aresponses stream gets a closure + that dispatches on the stored args shape, because the arming site + cannot tell the producers apart: native iterators store a single + ready-made logging coroutine to enqueue, while bridged streams + (LiteLLMCompletionStreamingIterator, and the plain SSE generator + AnthropicStreamWrapper returns for bridged /v1/messages) share + their inner CustomStreamWrapper's logging_obj and so store + (assembled_response, cache_hit); for those the closure only + dispatches success logging, matching the route's pre-existing + hook surface. - Raw async generators from passthrough routes bypass all three and - would orphan the closure, so they are not armed here. + Raw async generators from passthrough routes bypass both and would + orphan the closure, so they are not armed here. The router wraps iterators that cannot carry _hidden_params in HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the @@ -3257,31 +3261,27 @@ class ProxyBaseLLMRequestProcessing: if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): return - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - - if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): - _captured_bridge_logging_obj: Final = logging_obj - - async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: - await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + _captured_native_logging_obj: Final = logging_obj + + async def _on_deferred_native_stream_complete(*args: object) -> None: + match args: + case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + case (assembled_response, cache_hit): + await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + case _: + verbose_proxy_logger.error( + "Deferred stream logging dropped: unexpected stored args shape %s", + tuple(type(arg).__name__ for arg in args), + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 11c3d2f8b20..c550a0a41d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,6 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio +import logging from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1422,7 +1423,7 @@ class TestArmDeferredStreamDispatch: async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, _ = self._dispatch_recording_logging_obj() + logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): yield b"x" @@ -1433,20 +1434,89 @@ class TestArmDeferredStreamDispatch: user_api_key_dict=MagicMock(), logging_obj=logging_obj, ) - closure = logging_obj._on_deferred_stream_complete - assert closure is not None + assert logging_obj._on_deferred_stream_complete is not None async def _logging_coroutine(): return None coro = _logging_coroutine() + logging_obj._deferred_stream_complete_args = (coro,) with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: - await closure(coro) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) mock_enqueue.assert_called_once_with(async_coroutine=coro) + assert recorded == {} coro.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_raw_generator_stream_storing_csw_arg_shape_dispatches_success(self, route_type): + """Bridged /v1/messages returns AnthropicStreamWrapper's plain SSE + generator, which shares its inner CustomStreamWrapper's logging_obj and + so stores (assembled_response, cache_hit). The closure armed for a raw + generator must accept that shape too, or _fire_deferred_stream_logging + raises TypeError and the request loses its spend log and callbacks.""" + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type=route_type, + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, True) + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded["result"] is assembled + assert recorded["cache_hit"] is True + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())]) + async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + logging_obj._deferred_stream_complete_args = stored_args + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue, + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded == {} + dropped = [r for r in caplog.records if r.getMessage().startswith("Deferred stream logging dropped")] + assert len(dropped) == 1 + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper