From c4ea0e93c80772d1615e36eb1052b2179465885f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 23 Apr 2026 14:48:02 -0700 Subject: [PATCH 1/2] fix: drain logging worker in test_router_caching_ttl to remove flake The mocked async_increment_cache_pipeline is invoked from Router's deployment_callback_on_success, registered as an async success callback. Those callbacks are enqueued to GLOBAL_LOGGING_WORKER and run on a background task, so the mock may not have been called yet when the test asserts on it. Flush the worker before asserting. --- tests/local_testing/test_tpm_rpm_routing_v2.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/local_testing/test_tpm_rpm_routing_v2.py b/tests/local_testing/test_tpm_rpm_routing_v2.py index 9de5625c63..211af56642 100644 --- a/tests/local_testing/test_tpm_rpm_routing_v2.py +++ b/tests/local_testing/test_tpm_rpm_routing_v2.py @@ -547,6 +547,8 @@ async def test_router_caching_ttl(): assert router.cache.redis_cache is not None + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + increment_cache_kwargs = {} with patch.object( router.cache, @@ -555,6 +557,10 @@ async def test_router_caching_ttl(): ) as mock_client: await router.acompletion(model=model, messages=messages) + # Async success callbacks are dispatched to GLOBAL_LOGGING_WORKER's + # background queue; drain it before asserting the mock was invoked. + await GLOBAL_LOGGING_WORKER.flush() + # mock_client.assert_called_once() print(f"mock_client.call_args.kwargs: {mock_client.call_args.kwargs}") print(f"mock_client.call_args.args: {mock_client.call_args.args}") From c14a73fa59138eaf86c268c1b1879f57d9c7689c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 23 Apr 2026 15:06:33 -0700 Subject: [PATCH 2/2] fix: make LoggingWorker.flush() wait for in-flight callbacks The previous `while not self._queue.empty(): await self._queue.join()` pattern skipped the join entirely when the worker had already dequeued a task but not yet called task_done(). asyncio.Queue.join() tracks _unfinished_tasks (incremented by put, decremented by task_done), not queue depth, so it already handles that case on its own. --- litellm/litellm_core_utils/logging_worker.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 7f00c47c1f..3db3700ee0 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -370,11 +370,17 @@ class LoggingWorker: self._running_tasks.clear() async def flush(self) -> None: - """Flush the logging queue.""" + """Flush the logging queue. + + Waits until every enqueued task has completed. ``queue.join()`` blocks + on the queue's unfinished-task counter (decremented by ``task_done()``), + so it correctly handles items that have been dequeued but whose + callback hasn't finished yet — ``queue.empty()`` would return True in + that window and cause us to skip the wait. + """ if self._queue is None: return - while not self._queue.empty(): - await self._queue.join() + await self._queue.join() async def clear_queue(self): """