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)]