From 67cee9d82460a1bf5b89e90ae43afb13a66cc2e4 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 17:16:14 -0400 Subject: [PATCH] fix(proxy): accept both deferred logging arg shapes on bridged /v1/messages streams --- litellm/proxy/common_request_processing.py | 33 +++++++-- .../proxy/test_common_request_processing.py | 73 +++++++++++++++++++ 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fc83c1ddeed..b3b804de1b3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3123,10 +3123,15 @@ class ProxyBaseLLMRequestProcessing: 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. + - Everything else on anthropic_messages/aresponses gets a fallback + closure that dispatches on the stored args shape, because both + producers are plain async generators the arming site cannot tell + apart: native iterators store a ready-made (logging_coroutine,) + to enqueue, while bridged /v1/messages (AnthropicStreamWrapper's + SSE generator) shares its inner CustomStreamWrapper's logging_obj + and stores (assembled_response, cache_hit). - Raw async generators from passthrough routes bypass all three and + Raw async generators from passthrough routes bypass all of these and would orphan the closure, so they are not armed here. The router wraps iterators that cannot carry _hidden_params in @@ -3182,10 +3187,24 @@ class ProxyBaseLLMRequestProcessing: 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 (Coroutine() as 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.warning( + "deferred stream logging dropped: unexpected args shape %s", tuple(map(type, args)) + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 6d6aad22ca3..b791f14ac95 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7719,3 +7719,76 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()] assert len(records) == 1 assert (records[0].exc_info is not None) is expect_traceback + + +class _RecordingDeferredLoggingObj: + def __init__(self) -> None: + self.dispatched: list[tuple[object, object]] = [] + self._on_deferred_stream_complete: Callable | None = None + self._deferred_stream_complete_args: tuple | None = None + + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: + self.dispatched.append((result, cache_hit)) + + +class TestArmDeferredStreamDispatchFallback: + """The anthropic_messages fallback closure must accept both stored args shapes. + + A bridged /v1/messages stream (non-Anthropic model behind the Anthropic + adapter) is a plain async generator, so arming falls through to the + fallback closure, but its inner CustomStreamWrapper stores + (assembled_response, cache_hit). Before the fix the 1-arg closure raised + TypeError on fire, which reached clients as a trailing SSE error frame and + dropped the request's spend log. + """ + + def _arm(self, logging_obj: _RecordingDeferredLoggingObj) -> None: + async def bridged_sse_stream() -> AsyncGenerator[bytes, None]: + yield b"event: message_stop\n\n" + + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=bridged_sse_stream(), + route_type="anthropic_messages", + user_api_key_dict=ProxyUserAPIKeyAuth(), + logging_obj=logging_obj, + ) + assert logging_obj._on_deferred_stream_complete is not None + + @pytest.mark.asyncio + async def test_bridged_messages_args_shape_dispatches_success_logging(self): + logging_obj = _RecordingDeferredLoggingObj() + self._arm(logging_obj) + + assembled_response: Final = object() + logging_obj._deferred_stream_complete_args = (assembled_response, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + + for _ in range(3): + await asyncio.sleep(0) + assert logging_obj.dispatched == [(assembled_response, False)] + assert logging_obj._on_deferred_stream_complete is None + assert logging_obj._deferred_stream_complete_args is None + + @pytest.mark.asyncio + async def test_native_coroutine_shape_still_enqueues_to_logging_worker(self): + logging_obj = _RecordingDeferredLoggingObj() + self._arm(logging_obj) + + ran: Final = asyncio.Event() + + async def native_logging_coroutine() -> None: + ran.set() + + callback = logging_obj._on_deferred_stream_complete + assert callback is not None + await callback(native_logging_coroutine()) + + await asyncio.wait_for(ran.wait(), timeout=5) + assert logging_obj.dispatched == []