From 17fce8a20b71d6d1b2fbdd63363f6f2a89060e4b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 31 Aug 2026 15:04:09 -0400 Subject: [PATCH] fix(proxy): accept both deferred-stream arg shapes in native branch The proxy's `_arm_deferred_stream_dispatch` assigns a 1-arg native callback for anthropic_messages/aresponses streaming responses that aren't CustomStreamWrapper or LiteLLMCompletionStreamingIterator. But /v1/messages via the completion bridge (e.g. Together AI) returns a bare byte generator whose inner CustomStreamWrapper writes the 2-arg (assembled_response, cache_hit) shape at end-of-stream, so the native closure crashes with a TypeError inside the deferred fire. The native and bridge writer shapes are indistinguishable at arm time (the inner CSW handle is hidden inside the AnthropicStreamWrapper), so the callback discriminates on arity: 1-arg enqueues on the rooted logging worker as before; 2-arg dispatches success handlers directly. Both native Anthropic passthrough and the Together AI bridge now flow cleanly through the same code path. Also improves the `_init_custom_logger_compatible_class` failure log so it names the failing integration; the generic message was masking real S3/GCS init failures on stage (premium license, missing credentials, etc.). --- litellm/litellm_core_utils/litellm_logging.py | 6 +++- litellm/proxy/common_request_processing.py | 31 ++++++++++++++++--- .../test_deferred_guardrail_logging.py | 28 +++++++++++++++++ 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 97c4d038734..f87c24dbdf9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4719,7 +4719,11 @@ def _init_custom_logger_compatible_class( return newrelic_logger return None except Exception as e: - verbose_logger.exception("[Non-Blocking Error] Error initializing custom logger: %s", e) + verbose_logger.exception( + "Failed to initialize custom logger for integration %r: %s. Callback will not be registered.", + logging_integration, + e, + ) return None return None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ed66bb65ced..f3b942009c6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,7 +4,7 @@ import json import logging import math import traceback -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 @@ -3138,10 +3138,31 @@ 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: + # 1-arg (logging_coroutine,) is what the native passthrough/responses + # iterators write; 2-arg (assembled_response, cache_hit) is what the + # /v1/messages completion bridge writes via its inner CustomStreamWrapper. + # The proxy sees a bare byte generator in both cases, so arity dispatch + # here is the only place we can tell them apart. + if len(args) == 1 and asyncio.iscoroutine(args[0]): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=args[0]) + return + if len(args) == 2: + assembled_response, cache_hit = args + 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, + ) + return + verbose_proxy_logger.error( + "Deferred stream logging received unexpected args shape (len=%d); dropping to avoid double-logging", + len(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 8fde4cc9d5e..53a9d32baf2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1415,6 +1415,34 @@ class TestArmDeferredStreamDispatch: mock_enqueue.assert_called_once_with(async_coroutine=coro) coro.close() + @pytest.mark.asyncio + async def test_native_stream_closure_accepts_bridge_two_arg_shape(self): + """/v1/messages via the completion bridge returns a bare byte generator, + so arming lands in the native branch. But the inner CustomStreamWrapper + writes the 2-arg (assembled_response, cache_hit) shape at end-of-stream, + which the pre-fix 1-arg closure blew up on.""" + 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, + ) + closure = logging_obj._on_deferred_stream_complete + assert closure is not None + + assembled = object() + await closure(assembled, False) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper