diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 317203cc2d5..3b0888b9c0e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3630,14 +3630,12 @@ class ProxyBaseLLMRequestProcessing: # callbacks fire regardless of the call-type heuristic in # _is_sync_litellm_request (which only recognizes a subset of # async markers stored in litellm_params). - asyncio.create_task( - _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( - _response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) + await _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( + _response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, ) except Exception as e: verbose_proxy_logger.exception( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 61c08bd54e9..185b7d4f51d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1103,6 +1103,21 @@ def _call_type_for_route(route: str | None) -> str | None: _PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) _LOG_DB_METRICS_CALL_TYPES: Final = frozenset(("get_data", "insert_data", "update_data", "delete_data")) +_DEFERRED_STREAM_CALLBACK_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: retains background tasks + + +async def _run_deferred_stream_callback(callback: Coroutine[object, object, object]) -> None: + try: + await callback + except Exception as e: + verbose_proxy_logger.exception("Error in deferred stream callback: %s", e) + + +def _retain_deferred_stream_callback(callback: Coroutine[object, object, object]) -> asyncio.Task[None]: + task: Final = asyncio.create_task(_run_deferred_stream_callback(callback)) + _DEFERRED_STREAM_CALLBACK_TASKS.add(task) + task.add_done_callback(_DEFERRED_STREAM_CALLBACK_TASKS.discard) + return task def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: @@ -4172,7 +4187,7 @@ class ProxyLogging: record_served_output_texts(logging_obj.model_call_details, served_stream_output_texts(served_chunks)) @staticmethod - def fire_deferred_stream_logging(request_data: dict) -> None: + def fire_deferred_stream_logging(request_data: dict) -> asyncio.Task[None] | None: """ Fire the deferred streaming logging callback after the full streaming pipeline (including guardrail end-of-stream blocks) has completed. @@ -4183,8 +4198,7 @@ class ProxyLogging: """ logging_obj: Final = request_data.get("litellm_logging_obj") if logging_obj is None: - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + return None _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( logging_obj, "_on_deferred_stream_complete", None @@ -4193,7 +4207,8 @@ class ProxyLogging: 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 - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_deferred_cb(*_args)) + return _retain_deferred_stream_callback(_deferred_cb(*_args)) + return None _fire_deferred_stream_logging = fire_deferred_stream_logging diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index d1369f92c9f..e4e1f6e3151 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -736,8 +736,9 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task assert guardrail_called is True, "Guardrail hook should be called" assert ( @@ -946,8 +947,9 @@ class TestDeferredStreamingClosure: # CSW stored args; now simulate what ProxyLogging does request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task assert hook_called is True, "Production closure must call guardrail hook" assert ( @@ -1316,9 +1318,12 @@ class TestFireDeferredStreamLogging: mock_logging_obj._deferred_stream_complete_args = ("test_response", True) request_data = {"litellm_logging_obj": mock_logging_obj} - ProxyLogging._fire_deferred_stream_logging(request_data) - await GLOBAL_LOGGING_WORKER.flush() + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue: + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data) + assert deferred_task is not None + await deferred_task + enqueue.assert_not_called() assert callback_called is True assert callback_args["response"] == "test_response" assert callback_args["cache_hit"] is True @@ -1326,6 +1331,34 @@ class TestFireDeferredStreamLogging: assert mock_logging_obj._on_deferred_stream_complete is None assert mock_logging_obj._deferred_stream_complete_args is None + @pytest.mark.asyncio + async def test_deferred_guardrails_wait_for_success_logging(self): + logging_started = asyncio.Event() + release_logging = asyncio.Event() + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + logging_started.set() + await release_logging.wait() + + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) + + with patch("litellm.callbacks", []): + deferred_task = asyncio.create_task( + ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + ) + await logging_started.wait() + assert deferred_task.done() is False + release_logging.set() + await deferred_task + @pytest.mark.asyncio async def test_noop_when_no_deferred_args(self): """_fire_deferred_stream_logging should be a no-op when no deferred @@ -1516,8 +1549,9 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1550,8 +1584,9 @@ class TestArmDeferredStreamDispatch: assembled = object() logging_obj._deferred_stream_complete_args = (assembled, False) - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - await GLOBAL_LOGGING_WORKER.flush() + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task assert recorded["result"] is assembled assert recorded["cache_hit"] is False @@ -1577,13 +1612,14 @@ class TestArmDeferredStreamDispatch: 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: - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + 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 + ): + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_called_once_with(async_coroutine=coro) assert recorded == {} coro.close() @@ -1610,13 +1646,14 @@ class TestArmDeferredStreamDispatch: 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}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + 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 + ): + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_not_called() assert recorded["result"] is assembled @@ -1645,10 +1682,9 @@ class TestArmDeferredStreamDispatch: ) as mock_enqueue, caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): - ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) - deferred = mock_enqueue.call_args.kwargs["async_coroutine"] - mock_enqueue.reset_mock() - await deferred + deferred_task = ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + assert deferred_task is not None + await deferred_task mock_enqueue.assert_not_called() assert recorded == {} 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 f462340a808..86429080997 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,10 +10,10 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime from typing import Any, Dict, Final, List -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException @@ -23,7 +23,6 @@ from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) @@ -180,6 +179,13 @@ async def _one_chunk() -> AsyncGenerator[object, None]: yield "chunk" +async def _wait_until(condition: Callable[[], bool]) -> None: + for _ in range(200): + if condition(): + return + await asyncio.sleep(0.01) + + class _AttributeStream: _hidden_params = {"model_id": "m-1"} model = "gpt-x" @@ -564,7 +570,7 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 2) assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)] @@ -599,7 +605,7 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 1) assert [event[0] for event in events] == ["logging_dispatched"] assert logging_obj._deferred_stream_complete_args is None @@ -723,7 +729,7 @@ async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes request_data=request_data, ): pass - await GLOBAL_LOGGING_WORKER.flush() + await _wait_until(lambda: len(events) == 1) snapshot = { "events": events, @@ -749,12 +755,10 @@ async def test_fire_deferred_stream_logging_fires_callback(): logging_obj._on_deferred_stream_complete = deferred logging_obj._deferred_stream_complete_args = ("payload",) - with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue") as enqueue: - ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) - - enqueue.assert_called_once() + deferred_task = ProxyLogging._fire_deferred_stream_logging(request_data={"litellm_logging_obj": logging_obj}) + assert deferred_task is not None assert captured == {} - await enqueue.call_args.kwargs["async_coroutine"] + await deferred_task snapshot = { "arg": captured["arg"], "callback_cleared": logging_obj._on_deferred_stream_complete is None,