From 4204f199eaff544560763b8a4587885293a42e48 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:08:25 -0700 Subject: [PATCH] fix(logging): track overflow callbacks through shutdown --- litellm/litellm_core_utils/logging_worker.py | 102 ++++++++------ litellm/proxy/proxy_server.py | 7 +- litellm/utils.py | 8 +- .../integration/spend/test_shutdown_flush.py | 11 +- .../litellm_core_utils/test_logging_worker.py | 131 ++++++++++++++++++ .../test_deferred_guardrail_logging.py | 103 ++++---------- tests/unit/proxy/test_proxy_server.py | 35 ++++- tests/unit/test_utils.py | 31 +++++ 8 files changed, 299 insertions(+), 129 deletions(-) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 8b25e9d06fe..0a7e1bbb7f2 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -6,7 +6,7 @@ import atexit import contextvars import inspect import logging -from collections.abc import Coroutine, Iterator +from collections.abc import Awaitable, Callable, Coroutine, Iterator from typing import Final from typing_extensions import TypedDict @@ -66,6 +66,8 @@ class LoggingWorker: max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE, concurrency: int = LOGGING_WORKER_CONCURRENCY, timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS, + *, + retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, ): self.timeout = timeout self.max_queue_size = max_queue_size @@ -74,6 +76,9 @@ class LoggingWorker: self._queue: asyncio.Queue[LoggingTask] | None = None self._worker_task: asyncio.Task[None] | None = None self._running_tasks: set[asyncio.Task[None]] = set() + self._overflow_tasks: set[asyncio.Task[None]] = set() + self._pending_enqueues: dict[int, LoggingTask] = {} + self._retry_sleep = retry_sleep self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks self._sem: asyncio.Semaphore | None = None self._bound_loop: asyncio.AbstractEventLoop | None = None @@ -168,6 +173,7 @@ class LoggingWorker: self._sem = None self._worker_task = None self._running_tasks.clear() + self._overflow_tasks.clear() self._timeout_summary_task = None self._timeout_burst_count = 0 self._queue = new_queue @@ -185,6 +191,10 @@ class LoggingWorker: self._sem = asyncio.Semaphore(self.concurrency) if self._worker_task is None or self._worker_task.done(): self._worker_task = asyncio.create_task(self._worker_loop()) + if not self._overflow_tasks: + self._aggressive_clear_in_progress = False + for pending in tuple(self._pending_enqueues.values()): + self._schedule_delayed_enqueue_retry(pending) async def _process_log_task( self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]" @@ -278,12 +288,24 @@ class LoggingWorker: # Capture the current context when enqueueing task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context()) + self._enqueue_task(task) + + def _enqueue_task(self, task: LoggingTask) -> None: + if self._queue is None: + return try: self._queue.put_nowait(task) except asyncio.QueueFull: - # Queue is full - handle it appropriately verbose_logger.exception("LoggingWorker queue is full") + self._pending_enqueues[id(task)] = task self._handle_queue_full(task) + else: + self._pending_enqueues.pop(id(task), None) + + def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None: + task: Final = asyncio.create_task(coroutine) + self._overflow_tasks.add(task) + task.add_done_callback(self._overflow_tasks.discard) def _should_start_aggressive_clear(self) -> bool: """ @@ -326,7 +348,7 @@ class LoggingWorker: if self._should_start_aggressive_clear(): self._mark_aggressive_clear_started() # Schedule clearing as async task so enqueue returns immediately (non-blocking) - asyncio.create_task(self._aggressively_clear_queue_async(task)) + self._start_overflow_task(self._aggressively_clear_queue_async(task)) else: # Cooldown active or clear in progress, schedule a delayed retry self._schedule_delayed_enqueue_retry(task) @@ -363,7 +385,7 @@ class LoggingWorker: delay: Final = self._calculate_retry_delay() # Schedule the retry as a background task - asyncio.create_task(self._retry_enqueue_task(task, delay)) + self._start_overflow_task(self._retry_enqueue_task(task, delay)) except RuntimeError: # No event loop, drop the task as we can't schedule a retry pass @@ -373,17 +395,8 @@ class LoggingWorker: Retry enqueueing the task after delay, preserving original context. This is called as a background task from _schedule_delayed_enqueue_retry. """ - await asyncio.sleep(delay) - - # Try to enqueue the task directly, preserving its original context - if self._queue is None: - return - - try: - self._queue.put_nowait(task) - except asyncio.QueueFull: - # Still full - handle it appropriately (clear or retry again) - self._handle_queue_full(task) + await self._retry_sleep(delay) + self._enqueue_task(task) def _extract_tasks_from_queue(self) -> list[LoggingTask]: """ @@ -423,32 +436,18 @@ class LoggingWorker: return extracted_tasks: Final = self._extract_tasks_from_queue() - uncounted: Final = None if new_task is None or self._requeue(new_task) else new_task - runs: Final = ( - *(self._process_single_task(task) for task in extracted_tasks), - *(() if uncounted is None else (self._process_single_task(uncounted, counted=False),)), - ) - if runs: - await asyncio.gather(*runs) + if new_task is not None: + self._enqueue_task(new_task) + if extracted_tasks: + await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks)) except Exception as e: verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e) finally: # Always reset the flag even if an error occurs self._aggressive_clear_in_progress = False - def _requeue(self, task: LoggingTask) -> bool: - """Put a rejected task back on the queue once the clear made room; track it as dequeued otherwise.""" - if self._queue is None: - return False - try: - self._queue.put_nowait(task) - except asyncio.QueueFull: - self._track_dequeued(task) - return False - return True - - async def _process_single_task(self, task: LoggingTask, *, counted: bool = True) -> None: - """Process a single task; ``counted`` is False for a task the full queue never accepted.""" + async def _process_single_task(self, task: LoggingTask) -> None: + """Process a task extracted from the queue and mark it done.""" if self._queue is None: return @@ -462,8 +461,7 @@ class LoggingWorker: pass finally: self._untrack_dequeued(task) - if counted: - self._queue.task_done() + self._queue.task_done() def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None: """ @@ -474,7 +472,11 @@ class LoggingWorker: @property def pending_count(self) -> int: - return (self._queue.qsize() if self._queue is not None else 0) + len(self._dequeued_tasks) + return ( + (self._queue.qsize() if self._queue is not None else 0) + + len(self._dequeued_tasks) + + len(self._pending_enqueues) + ) async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None: """Stop the logging worker and clean up resources.""" @@ -490,14 +492,21 @@ class LoggingWorker: self._timeout_summary_task = None self._emit_timeout_summary() - if self._worker_task is None and not self._running_tasks: - # No worker launched and no in-flight tasks to drain. - return + if discard_pending: + rejected: Final = tuple(self._pending_enqueues.values()) + self._pending_enqueues.clear() + for task in rejected: + task["coroutine"].close() + if rejected: + verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected)) - tasks_to_cancel: Final[list[asyncio.Task[None]]] = list(self._running_tasks) - if self._worker_task: - # Include the main worker loop so it stops fetching work. - tasks_to_cancel.append(self._worker_task) + tasks_to_cancel: Final = ( + *self._running_tasks, + *self._overflow_tasks, + *((self._worker_task,) if self._worker_task is not None else ()), + ) + if not tasks_to_cancel: + return for task in tasks_to_cancel: # Propagate cancellation to every pending task. @@ -535,6 +544,9 @@ class LoggingWorker: return self.start() await self._queue.join() + while self._overflow_tasks: + await asyncio.wait(tuple(self._overflow_tasks)) + await self._queue.join() async def clear_queue(self): """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87cea9b3787..dc85ad9e011 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1294,8 +1294,13 @@ async def _flush_spend_around_logging_drain_on_shutdown( remaining, callback_timeout, ) + + async def drain_callbacks() -> None: + await drain_tasks(callback_timeout) + await worker.flush() + try: - await asyncio.wait_for(worker.flush(), timeout=callback_timeout) + await asyncio.wait_for(drain_callbacks(), timeout=callback_timeout) except asyncio.TimeoutError: verbose_proxy_logger.warning( "Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete", diff --git a/litellm/utils.py b/litellm/utils.py index 72c37c3760e..e5632e4ea63 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1539,15 +1539,19 @@ def _schedule_async_success_logging( return def _enqueue_async_logging() -> None: + from litellm.responses.streaming_iterator import spawn_logging_task + logging_obj.mark_async_success_log_enqueued() - asyncio.create_task( + spawn_logging_task( + asyncio.get_running_loop(), _client_async_logging_helper( logging_obj=logging_obj, result=result, start_time=start_time, end_time=end_time, is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + ), + task_name="success callback enqueue", ) if not logging_obj.defer_async_logging: diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 0dda707dd9c..5c2e9cf98da 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -137,10 +137,12 @@ def test_tagged_request_survives_shutdown_before_the_tag_timer( @pytest.mark.parametrize("redis_buffer", [False, True]) +@pytest.mark.parametrize("overflow", [False, True]) def test_callback_backlog_is_accounted_before_shutdown_disconnects( - gateway: Gateway, tmp_path: Path, redis_buffer: bool + gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool ) -> None: release: Final = threading.Event() + requests: Final = 6 if overflow else 3 tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}") with ( wire_server(_provider) as wire, @@ -153,6 +155,7 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects( { "DATABASE_URL": os.environ["DATABASE_URL"], "LOGGING_WORKER_CONCURRENCY": "2", + "LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000", "SHUTDOWN_CALLBACK_URL": logger.url, "GRACEFUL_SHUTDOWN_TIMEOUT": "15", "LITELLM_LOG": "DEBUG", @@ -163,9 +166,9 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects( owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"] ) try: - for _ in range(3): + for _ in range(requests): _tagged_chat(owned.gateway, key, tags) - eventually(logger.received.qsize, lambda count: count == 2, seconds=10) + eventually(logger.received.qsize, lambda count: count >= 2, seconds=10) owned.begin_shutdown() eventually( lambda: ( @@ -177,7 +180,7 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects( ) release.set() owned.process.wait(timeout=90) - _assert_shutdown_accounting(key, tags, 3) + _assert_shutdown_accounting(key, tags, requests) finally: release.set() diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index 30ec851d54f..fccc11617bb 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -197,6 +197,137 @@ async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> N await worker.stop() +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_size", (1, 2)) +async def test_flush_waits_for_delayed_overflow_and_preserves_callback_context(queue_size: int) -> None: + retry_started: Final = asyncio.Event() + release_retry: Final = asyncio.Event() + completed: Final[asyncio.Queue[tuple[int, str]]] = asyncio.Queue() + identity: Final = contextvars.ContextVar("overflow_identity", default="outside") + + async def retry_sleep(delay: float) -> None: + retry_started.set() + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=queue_size, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + completed.put_nowait((index, identity.get())) + + worker.start() + token: Final = identity.set("request") + for index in range(4): + worker.enqueue(callback(index)) + identity.reset(token) + flushing: Final = asyncio.create_task(worker.flush()) + try: + await asyncio.wait_for(retry_started.wait(), timeout=2) + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=2) + await asyncio.sleep(0) + assert not flushing.done(), "Queue completion must not hide delayed overflow callbacks" + assert worker.pending_count > 0 + release_retry.set() + await asyncio.wait_for(flushing, timeout=2) + assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == [ + (index, "request") for index in range(4) + ] + assert worker.pending_count == 0 + finally: + release_retry.set() + await worker.stop(discard_pending=True) + await asyncio.gather(flushing, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("start_retry", (False, True)) +async def test_shutdown_discards_overflow_without_late_enqueue_or_replay(start_retry: bool) -> None: + import inspect + + retry_started: Final = asyncio.Event() + release_retry: Final = asyncio.Event() + executed: Final[asyncio.Queue[int]] = asyncio.Queue() + + async def retry_sleep(delay: float) -> None: + retry_started.set() + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + executed.put_nowait(index) + + callbacks: Final = tuple(callback(index) for index in range(4)) + worker._ensure_queue() + for callback_coro in callbacks: + worker.enqueue(callback_coro) + if start_retry: + await asyncio.wait_for(retry_started.wait(), timeout=2) + await worker.stop(discard_pending=True, timeout=1) + completed_at_stop: Final = executed.qsize() + release_retry.set() + await worker.flush() + try: + assert executed.qsize() == completed_at_stop, "Discarded overflow must not run when the worker restarts" + assert worker.pending_count == 0 + assert all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in callbacks) + finally: + await worker.stop(discard_pending=True) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_aggressive_clear_callbacks_before_returning() -> None: + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1) + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def noop() -> None: + pass + + worker._ensure_queue() + worker.enqueue(blocked()) + worker.enqueue(noop()) + worker.enqueue(noop()) + await asyncio.wait_for(started.wait(), timeout=2) + await worker.stop(discard_pending=True, timeout=1) + assert cancelled.is_set(), "Overflow execution must participate in shutdown cancellation" + assert worker.pending_count == 0 + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=1) + + +@pytest.mark.asyncio +async def test_stopped_overflow_can_be_flushed_when_not_discarded() -> None: + release_retry: Final = asyncio.Event() + completed: Final[asyncio.Queue[int]] = asyncio.Queue() + + async def retry_sleep(delay: float) -> None: + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + completed.put_nowait(index) + + worker._ensure_queue() + for index in range(4): + worker.enqueue(callback(index)) + await worker.stop() + release_retry.set() + try: + await asyncio.wait_for(worker.flush(), timeout=2) + assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == list(range(4)) + finally: + await worker.stop(discard_pending=True) + + class TestLoggingWorker: """Test cases for LoggingWorker functionality.""" diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 7e5d0f6342e..56c68170902 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -239,52 +239,37 @@ class TestHasPostCallGuardrailsForPassthrough: @pytest.mark.asyncio -async def test_deferred_flag_stores_and_executes_closure(): - """ - When defer_async_logging is True on logging_obj: - 1. wrapper_async stores a callable closure instead of calling create_task - 2. Calling the closure fires create_task - 3. Sync callbacks fire immediately (not deferred) - """ - mock_logging_obj = MagicMock() - mock_logging_obj.defer_async_logging = True - mock_logging_obj.enqueue_deferred_logging = None +@pytest.mark.parametrize("deferred", (False, True)) +async def test_success_logging_waits_for_deferred_closure_when_requested(deferred: bool) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks - await litellm.acompletion( - model="gpt-3.5-turbo", + recorder: Final = _RecordingLogger() + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="deferred-success", + function_id="deferred-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.defer_async_logging = deferred + response: Final = await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response="Hello!", - litellm_logging_obj=mock_logging_obj, + litellm_logging_obj=logging_obj, ) - - # Closure was stored - enqueue_fn = mock_logging_obj.enqueue_deferred_logging - assert callable(enqueue_fn), "Closure should be stored on logging_obj" - - # Sync callbacks fired immediately - mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() - - # Calling the closure fires create_task - created_tasks = [] - real_create_task = asyncio.create_task - - def tracking_create_task(coro): - task = real_create_task(coro) - created_tasks.append(task) - return task - - with patch("asyncio.create_task", side_effect=tracking_create_task): - enqueue_fn() - - assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task" - - for task in created_tasks: - if not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass + if deferred: + assert recorder.standard_logging_object is None + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + await drain_logging_tasks() + await GLOBAL_LOGGING_WORKER.flush() + assert recorder.standard_logging_object is not None + assert recorder.standard_logging_object["id"] == response.id @pytest.mark.asyncio @@ -401,40 +386,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6) -# --------------------------------------------------------------------------- -# 3. Non-streaming regression: without flag, create_task fires normally -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_no_flag_fires_create_task_normally(): - """Without defer_async_logging, wrapper_async calls create_task as before.""" - created_tasks = [] - real_create_task = asyncio.create_task - - def tracking_create_task(coro): - task = real_create_task(coro) - created_tasks.append(task) - return task - - with patch("asyncio.create_task", side_effect=tracking_create_task): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="Hello!", - ) - - assert len(created_tasks) >= 1 - - for task in created_tasks: - if not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass - - # --------------------------------------------------------------------------- # 4. Non-streaming: deferred success log is suppressed when guardrail raises # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 43e30d88bce..c835e1e27ea 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -409,13 +409,14 @@ async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain( monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") worker: Final = LoggingWorker() persisted: Final = asyncio.Event() + drain_budgets: Final[asyncio.Queue[float]] = asyncio.Queue() GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) async def flush_spend() -> None: persisted.set() async def drain_tasks(timeout: float) -> None: - assert timeout == 5.0 + drain_budgets.put_nowait(timeout) try: with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): @@ -424,6 +425,7 @@ async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain( ) assert "remaining_s=24.000 callback_timeout_s=14.000" in caplog.text assert persisted.is_set() + assert tuple(drain_budgets.get_nowait() for _ in range(drain_budgets.qsize())) == (14.0, 5.0) finally: await worker.stop() GracefulShutdownManager.reset() @@ -4225,3 +4227,34 @@ def test_detailed_debug_logs_request_body_only_as_received_and_as_sent( assert len(body_dumps) == 2, [message[:120] for message in body_dumps] assert body_dumps[0].startswith("Request received by LiteLLM:") assert "POST Request Sent from LiteLLM" in body_dumps[1] + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.responses.streaming_iterator import spawn_logging_task + + worker: Final = LoggingWorker() + persisted: Final[asyncio.Queue[str]] = asyncio.Queue() + accounting: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def callback() -> None: + accounting.put_nowait("completed request") + + async def producer() -> None: + await asyncio.sleep(0) + await asyncio.sleep(0) + worker.ensure_initialized_and_enqueue(callback()) + + async def flush_spend() -> None: + while not accounting.empty(): + persisted.put_nowait(accounting.get_nowait()) + + spawn_logging_task(asyncio.get_running_loop(), producer(), task_name="success enqueue") + try: + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) + assert persisted.qsize() == 1, "An enqueue producer must finish before callback draining stops" + assert persisted.get_nowait() == "completed request" + assert worker.pending_count == 0 + finally: + await worker.stop(discard_pending=True) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 436c966e5d5..7dda31ec383 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10903,3 +10903,34 @@ def test_custom_provider_named_like_a_built_in_warns_once( warnings: Final = [r for r in caplog.records if f"provider '{provider}'" in r.getMessage()] assert len(warnings) == expected_warnings assert provider in litellm._custom_providers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred", (False, True)) +async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronously(deferred: bool) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.utils import _schedule_async_success_logging + + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="acompletion", + start_time=now, + litellm_call_id="shutdown-enqueue", + function_id="shutdown-enqueue", + ) + logging_obj.defer_async_logging = deferred + _schedule_async_success_logging(logging_obj, ModelResponse(), now, now, False) + if deferred: + assert not logging_obj.async_success_log_enqueued + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + assert "has_logged_async_success" not in logging_obj.model_call_details + await drain_logging_tasks() + assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" + await GLOBAL_LOGGING_WORKER.flush() + assert logging_obj.model_call_details["has_logged_async_success"] is True