From 8366e15be070d55e3b4255762e7bb253cc5fc6f5 Mon Sep 17 00:00:00 2001 From: maciej-berri Date: Sat, 14 Nov 2026 14:32:48 -0800 Subject: [PATCH] fix(proxy): close streams before rebuilding the served output for logging (#51487) * fix(proxy): close streams before rebuilding the served output for logging The streaming iterator hook awaited _record_served_stream_output before it returned, so every streamed response stayed open after its last event while the proxy re-parsed all served chunks. That held the connection about 11 ms plus 25 ms per 1,000 output tokens, and SDK clients wait for the close. The parked deferred logging is now claimed at stream end and run in a rooted background task after the served output is recorded, so logging still sees the served text and the client sees the stream end right away. * fix(proxy): cap background served-output rebuilds and fall back to waiting past it Each background rebuild holds its served chunks until it runs, and the stream's concurrency slot is already released. Past 64 pending rebuilds a stream waits for its own rebuild as it did before, so overload keeps its backpressure and no spend logging is dropped. --- litellm/proxy/utils.py | 62 +++++++-- .../test_proxy_logging_hook_detection.py | 11 +- tests/unit/proxy/test_proxy_utils.py | 128 ++++++++++++++++++ .../proxy_logging/test_streaming_hooks.py | 8 +- 4 files changed, 193 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5d02078fdf9..ce7a041e147 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -404,6 +404,8 @@ class _TeamMembershipRow(TypedDict): team_members_with_roles: ReadOnly[Iterable[_TeamMemberRow] | None] +_ROOTED_STREAM_LOGGING_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep tasks alive +_MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS: Final = 64 _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' @@ -5335,8 +5337,17 @@ class ProxyLogging: # completed. unified_guardrail writes guardrail_information during # its end-of-stream block (inside current_response), so by the time # we reach this point the metadata is fully populated. - await ProxyLogging._record_served_stream_output(request_data, served_chunks) - ProxyLogging.fire_deferred_stream_logging(request_data) + # Rebuilding the served output only feeds logging, so it runs after the stream closes, + # unless enough rebuilds are already pending that this stream should wait like it used to. + if len(_ROOTED_STREAM_LOGGING_TASKS) >= _MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS: + await ProxyLogging._record_served_stream_output(request_data, served_chunks) + ProxyLogging.fire_deferred_stream_logging(request_data) + return + ProxyLogging._spawn_rooted( + ProxyLogging._log_served_stream( + request_data, served_chunks, ProxyLogging._take_deferred_stream_logging(request_data) + ) + ) async def async_post_call_event_stream_iterator_hook( self, @@ -5598,6 +5609,39 @@ class ProxyLogging: texts: Final = await offload_token_count(served_stream_output_texts)(served_chunks) record_served_output_texts(logging_obj.model_call_details, texts) + @staticmethod + async def _log_served_stream( + request_data: Mapping[str, object], + served_chunks: Sequence[object], + deferred_logging: Coroutine[object, object, object] | None, + ) -> None: + try: + await ProxyLogging._record_served_stream_output(request_data, served_chunks) + finally: + if deferred_logging is not None: + _ = await deferred_logging + + @staticmethod + def _spawn_rooted(work: Coroutine[object, object, None]) -> None: + task: Final = asyncio.create_task(work) + _ROOTED_STREAM_LOGGING_TASKS.add(task) + task.add_done_callback(_ROOTED_STREAM_LOGGING_TASKS.discard) + + @staticmethod + def _take_deferred_stream_logging(request_data: Mapping[str, object]) -> Coroutine[object, object, object] | None: + logging_obj: Final = request_data.get("litellm_logging_obj") + if logging_obj is None: + return None + _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 None + logging_obj._on_deferred_stream_complete = None # pyright: ignore[reportAttributeAccessIssue] # slot parked by the stream iterators + logging_obj._deferred_stream_complete_args = None # pyright: ignore[reportAttributeAccessIssue] # slot parked by the stream iterators + return _deferred_cb(*_args) + @staticmethod def fire_deferred_stream_logging(request_data: dict) -> None: """ @@ -5608,17 +5652,9 @@ class ProxyLogging: scheduling via create_task (which would race with unified_guardrail's end-of-stream block). This method retrieves and fires them. """ - 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 not None and _args is not None: - logging_obj._on_deferred_stream_complete = None - logging_obj._deferred_stream_complete_args = None - asyncio.create_task(_deferred_cb(*_args)) + deferred_logging: Final = ProxyLogging._take_deferred_stream_logging(request_data) + if deferred_logging is not None: + asyncio.create_task(deferred_logging) _fire_deferred_stream_logging = fire_deferred_stream_logging diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 7e34bf6741f..5dff6dca1a4 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -1,3 +1,5 @@ +import asyncio + import pytest from fastapi import HTTPException @@ -591,6 +593,13 @@ async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeyp monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")]) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) logging_obj = _streaming_logging_obj() + served_texts_seen_by_logging = asyncio.get_running_loop().create_future() + + async def deferred_logging() -> None: + served_texts_seen_by_logging.set_result(logging_obj.model_call_details.get(SERVED_OUTPUT_TEXTS_KEY)) + + logging_obj._on_deferred_stream_complete = deferred_logging + logging_obj._deferred_stream_complete_args = () async def fake_stream(): yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))]) @@ -606,7 +615,7 @@ async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeyp delivered_text += choice.delta.content or "" assert "zebra" not in delivered_text - assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,) + assert await asyncio.wait_for(served_texts_seen_by_logging, timeout=5) == (delivered_text,) @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 698f4787a42..52b3e9529b3 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -3768,3 +3768,131 @@ async def test_query_raw_within_statement_timeout_leaves_other_database_errors_a await query_raw_within_statement_timeout(prisma, "SELECT 1") assert exc_info.value is failure + + +@pytest.mark.asyncio +async def test_stream_ends_without_waiting_for_served_output_and_logging_still_sees_it(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY + from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop + from litellm.proxy.utils import ProxyLogging + + class PassThroughIterator(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + + monkeypatch.setattr(litellm, "callbacks", [PassThroughIterator()]) + logging_obj: Final = Logging( + model="claude-sonnet-4-5", + messages=[], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="served-output-close", + function_id="served-output-close", + ) + served_texts_seen_by_logging: Final = asyncio.get_running_loop().create_future() + + async def deferred_logging() -> None: + served_texts_seen_by_logging.set_result(logging_obj.model_call_details.get(SERVED_OUTPUT_TEXTS_KEY)) + + logging_obj._on_deferred_stream_complete = deferred_logging + logging_obj._deferred_stream_complete_args = () + events: Final = _anthropic_thinking_sse(50) + + async def upstream(): + for event in events: + yield event + + async def drain() -> list[object]: + return [ + chunk + async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1", request_route="/v1/messages"), + request_data={"model": "claude-sonnet-4-5", "litellm_logging_obj": logging_obj}, + ) + ] + + offload_limiter: Final = _count_offload_limiter_for_this_loop() + holders: Final = tuple(object() for _ in range(int(offload_limiter.total_tokens))) + for holder in holders: + await offload_limiter.acquire_on_behalf_of(holder) + + served: Final = await asyncio.wait_for(drain(), timeout=5) + for holder in holders: + offload_limiter.release_on_behalf_of(holder) + + assert served == events + assert await asyncio.wait_for(served_texts_seen_by_logging, timeout=5) == ("final answer",) + + +@pytest.mark.asyncio +async def test_streams_wait_for_their_served_output_once_background_rebuilds_hit_the_cap(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY + from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop + from litellm.proxy.utils import _MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS, ProxyLogging + + class PassThroughIterator(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + + monkeypatch.setattr(litellm, "callbacks", [PassThroughIterator()]) + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + events: Final = _anthropic_thinking_sse(1) + logging_objs: Final = tuple( + Logging( + model="claude-sonnet-4-5", + messages=[], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=f"served-output-cap-{index}", + function_id=f"served-output-cap-{index}", + ) + for index in range(_MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS + 1) + ) + + async def upstream(): + for event in events: + yield event + + async def drain(logging_obj: Logging) -> list[object]: + return [ + chunk + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1", request_route="/v1/messages"), + request_data={"model": "claude-sonnet-4-5", "litellm_logging_obj": logging_obj}, + ) + ] + + offload_limiter: Final = _count_offload_limiter_for_this_loop() + holders: Final = tuple(object() for _ in range(int(offload_limiter.total_tokens))) + for holder in holders: + await offload_limiter.acquire_on_behalf_of(holder) + + for logging_obj in logging_objs[:-1]: + assert await asyncio.wait_for(drain(logging_obj), timeout=5) == events + over_cap: Final = asyncio.ensure_future(drain(logging_objs[-1])) + await asyncio.sleep(0.2) + over_cap_waited: Final = not over_cap.done() + + for holder in holders: + offload_limiter.release_on_behalf_of(holder) + + assert over_cap_waited + assert await asyncio.wait_for(over_cap, timeout=5) == events + + async def every_stream_recorded() -> None: + while any(SERVED_OUTPUT_TEXTS_KEY not in obj.model_call_details for obj in logging_objs): + await asyncio.sleep(0.01) + + await asyncio.wait_for(every_stream_recorded(), timeout=10) diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 8f8f10565e9..1b0ad8d8fdd 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -563,8 +563,12 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea request_data=request_data, ): pass - await asyncio.sleep(0) - await asyncio.sleep(0) + + async def logging_dispatched() -> None: + while len(events) < 2: + await asyncio.sleep(0.01) + + await asyncio.wait_for(logging_dispatched(), timeout=5) assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)]