fix: drain shutdown callbacks and persist pending tag spend

litellm/litellm_core_utils/logging_worker.py

@@ -473,8 +473,12 @@ class LoggingWorker:
         self.start()
         self.enqueue(async_coroutine)
 
-    async def stop(self) -> None:
+    async def stop(self, *, discard_pending: bool = False) -> None:
         """Stop the logging worker and clean up resources."""
+        if discard_pending and self._queue is not None:
+            for task in self._drain_pending(self._queue):
+                task["coroutine"].close()
+                self._queue.task_done()
         if self._timeout_summary_task is not None:
             self._timeout_summary_task.cancel()
             self._timeout_summary_task = None
@@ -496,6 +500,11 @@ class LoggingWorker:
         # Wait for cancellation to settle; ignore errors raised during shutdown.
         await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
 
+        if discard_pending and self._queue is not None:
+            for task in self._unstarted_dequeued_tasks():
+                task["coroutine"].close()
+                self._untrack_dequeued(task)
+                self._queue.task_done()
         self._worker_task = None
         # Drop references to completed tasks so we can restart cleanly.
         self._running_tasks.clear()

litellm/proxy/proxy_server.py

@@ -1268,14 +1268,30 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
 
 
 async def _flush_spend_on_shutdown() -> None:
+    from litellm.proxy.utils import update_daily_tag_spend
+
     await flush_spend_counters_on_shutdown()
+    if prisma_client is not None:
+        await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
     await _flush_spend_logs_queue_on_shutdown()
 
 
 async def _flush_spend_around_logging_drain_on_shutdown(
     flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown,
 ) -> None:
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
     await flush_spend()
+    verbose_proxy_logger.info("Draining success callbacks before the final spend flush")
+    try:
+        await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=GracefulShutdownManager.get_timeout())
+    except TimeoutError:
+        verbose_proxy_logger.warning("Timed out draining success callbacks on shutdown; accounting may be incomplete")
+    finally:
+        try:
+            await asyncio.wait_for(GLOBAL_LOGGING_WORKER.stop(discard_pending=True), timeout=5.0)
+        except TimeoutError:
+            verbose_proxy_logger.warning("Timed out cancelling success callbacks on shutdown")
     await drain_logging_tasks()
     await flush_spend()
 

litellm/utils.py

@@ -1526,14 +1526,12 @@ def _schedule_async_success_logging(
 
     def _enqueue_async_logging() -> None:
         logging_obj.mark_async_success_log_enqueued()
-        asyncio.create_task(
-            _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,
-            )
+        _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,
         )
 
     if not logging_obj.defer_async_logging:
@@ -1544,13 +1542,13 @@ def _schedule_async_success_logging(
     logging_obj.enqueue_deferred_logging = _enqueue_async_logging
 
 
-async def _client_async_logging_helper(
+def _client_async_logging_helper(
     logging_obj: LiteLLMLoggingObject,
-    result,
-    start_time,
-    end_time,
+    result: object,
+    start_time: datetime.datetime,
+    end_time: datetime.datetime,
     is_completion_with_fallbacks: bool,
-):
+) -> None:
     if (
         is_completion_with_fallbacks is False
     ):  # don't log the parent event litellm.completion_with_fallbacks as a 'log_success_event', this will lead to double logging the same call - https://github.com/BerriAI/litellm/issues/7477

tests/integration/spend/test_shutdown_flush.py

@@ -1,3 +1,4 @@
+import hashlib
 import json
 import os
 import threading
@@ -26,6 +27,154 @@ MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new"
 COMMIT_DELAY_SECONDS: Final = 15
 BURST_REQUESTS: Final = 30
 
+_GATED_LOGGER: Final = """
+import os
+import httpx
+from litellm.integrations.custom_logger import CustomLogger
+
+class ShutdownLogger(CustomLogger):
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        async with httpx.AsyncClient(timeout=25, trust_env=False) as client:
+            await client.post(os.environ["SHUTDOWN_CALLBACK_URL"])
+
+logger = ShutdownLogger()
+"""
+
+
+def _tag_shutdown_config(directory: Path, upstream: str, redis_buffer: bool, *, gated: bool = False) -> Path:
+    path: Final = directory / "tag-shutdown.yaml"
+    if gated:
+        (directory / "shutdown_callbacks.py").write_text(_GATED_LOGGER)
+    path.write_text(
+        yaml.safe_dump(
+            {
+                "model_list": [
+                    {
+                        "model_name": "shutdown-model",
+                        "litellm_params": {
+                            "model": "openai/gpt-4o-mini",
+                            "api_base": upstream + "/v1",
+                            "api_key": "synthetic-key",
+                            "input_cost_per_token": 0.01,
+                            "output_cost_per_token": 0.02,
+                        },
+                    }
+                ],
+                "general_settings": {
+                    "master_key": "os.environ/LITELLM_MASTER_KEY",
+                    "database_url": "os.environ/DATABASE_URL",
+                    "proxy_batch_write_at": 3600,
+                    "use_redis_transaction_buffer": redis_buffer,
+                    "disable_model_info_refresh": True,
+                },
+                "litellm_settings": {
+                    "cache": True,
+                    "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
+                    "callbacks": ["shutdown_callbacks.logger"] if gated else [],
+                },
+            }
+        )
+    )
+    return path
+
+
+def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None:
+    response: Final = gateway.request(
+        "POST",
+        "/v1/chat/completions",
+        {
+            "model": "shutdown-model",
+            "messages": [{"role": "user", "content": uuid.uuid4().hex}],
+            "metadata": {"tags": list(tags)},
+        },
+        key=key,
+    )
+    assert response.status_code == 200, response.text
+
+
+def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None:
+    hashed: Final = hashlib.sha256(key.encode()).hexdigest()
+    for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")):
+        rows: Final = read_rows(
+            f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
+        )
+        assert rows[0]["requests"] == requests, (table, rows)
+        assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
+    key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,))
+    assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9)
+    for tag in tags:
+        tag_rows: Final = read_rows(
+            'SELECT sum(api_requests)::int AS requests, sum(spend) AS spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s',
+            (hashed, tag),
+        )
+        assert tag_rows[0]["requests"] == requests, (tag, tag_rows)
+        assert tag_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (tag, tag_rows)
+
+
+@pytest.mark.parametrize("redis_buffer", [False, True])
+def test_tagged_request_survives_shutdown_before_the_tag_timer(
+    gateway: Gateway, tmp_path: Path, redis_buffer: bool
+) -> None:
+    tags: Final = (f"shutdown-a-{uuid.uuid4().hex}", f"shutdown-b-{uuid.uuid4().hex}")
+    overrides: Final = {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG"}
+    with wire_server(_provider) as wire:
+        config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer)
+        with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned:
+            key: Final = string_value(
+                owned.gateway.post("/key/generate", {"user_id": f"shutdown-{uuid.uuid4().hex}"})["key"]
+            )
+            _tagged_chat(owned.gateway, key, tags)
+            owned.begin_shutdown()
+            owned.process.wait(timeout=90)
+        with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted:
+            _assert_shutdown_accounting(key, tags, 1)
+            restarted.begin_shutdown()
+            restarted.process.wait(timeout=90)
+        _assert_shutdown_accounting(key, tags, 1)
+
+
+@pytest.mark.parametrize("redis_buffer", [False, True])
+def test_callback_backlog_is_accounted_before_shutdown_disconnects(
+    gateway: Gateway, tmp_path: Path, redis_buffer: bool
+) -> None:
+    release: Final = threading.Event()
+    tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}")
+    with wire_server(_provider) as wire, wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger:
+        config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
+        with owned_proxy_process(
+            gateway,
+            tmp_path,
+            {
+                "DATABASE_URL": os.environ["DATABASE_URL"],
+                "LOGGING_WORKER_CONCURRENCY": "2",
+                "SHUTDOWN_CALLBACK_URL": logger.url,
+                "GRACEFUL_SHUTDOWN_TIMEOUT": "15",
+                "LITELLM_LOG": "DEBUG",
+            },
+            config=config,
+        ) as owned:
+            key: Final = string_value(
+                owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"]
+            )
+            try:
+                for _ in range(3):
+                    _tagged_chat(owned.gateway, key, tags)
+                eventually(logger.received.qsize, lambda count: count == 2, seconds=10)
+                owned.begin_shutdown()
+                eventually(
+                    lambda: (
+                        owned.process.poll() is not None
+                        or "Draining success callbacks before the final spend flush" in owned.log.read_text()
+                    ),
+                    lambda ready: ready,
+                    seconds=30,
+                )
+                release.set()
+                owned.process.wait(timeout=90)
+                _assert_shutdown_accounting(key, tags, 3)
+            finally:
+                release.set()
+
 
 def _api_requests(table: str, column: str, identity: str) -> int:
     rows: Final = read_rows(

tests/unit/litellm_core_utils/test_logging_worker.py

@@ -82,6 +82,56 @@ class _RecordCollector(logging.Handler):
         self.records.append(record)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("running", [False, True])
+async def test_stop_discards_waiting_callbacks_without_executing_them(running: bool) -> None:
+    worker: Final = LoggingWorker(concurrency=1)
+    started: Final = asyncio.Event()
+    cancelled: Final = asyncio.Event()
+    executed: Final = asyncio.Event()
+
+    async def blocked() -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            cancelled.set()
+
+    async def waiting() -> None:
+        executed.set()
+
+    if running:
+        worker.ensure_initialized_and_enqueue(blocked())
+        await started.wait()
+    else:
+        worker._ensure_queue()
+    worker.enqueue(waiting())
+    await worker.stop(discard_pending=True)
+    assert not executed.is_set()
+    assert cancelled.is_set() == running
+    assert worker._queue is not None
+    assert worker._queue.empty()
+    await asyncio.wait_for(worker._queue.join(), timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_started() -> None:
+    worker: Final = LoggingWorker(concurrency=1)
+    executed: Final = asyncio.Event()
+
+    async def callback() -> None:
+        executed.set()
+
+    worker.ensure_initialized_and_enqueue(callback())
+    await asyncio.sleep(0)
+    await worker.stop(discard_pending=True)
+    assert not executed.is_set()
+    worker._flush_on_exit()
+    assert not executed.is_set()
+    assert worker._queue is not None
+    await asyncio.wait_for(worker._queue.join(), timeout=1)
+
+
 class TestLoggingWorker:
     """Test cases for LoggingWorker functionality."""
 

tests/unit/proxy/test_proxy_server.py

@@ -260,6 +260,131 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm
 
 token = MASTER_KEY
 
+
+@pytest.mark.asyncio
+async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None:
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from litellm.responses.streaming_iterator import spawn_logging_task
+
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    first_flush: Final = asyncio.Event()
+    pending: Final[asyncio.Queue[str]] = asyncio.Queue()
+    persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
+
+    async def enqueue_spend() -> None:
+        pending.put_nowait("accounted request")
+
+    async def callback() -> None:
+        started.set()
+        await release.wait()
+        spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting")
+
+    async def flush_spend() -> None:
+        while not pending.empty():
+            persisted.put_nowait(pending.get_nowait())
+        first_flush.set()
+
+    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    await started.wait()
+    shutdown: Final = asyncio.create_task(
+        litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+    )
+    try:
+        await first_flush.wait()
+        waited_for_callback: Final = not shutdown.done()
+        release.set()
+        await shutdown
+        assert waited_for_callback, "Shutdown returned while accounting callbacks were still running"
+        assert persisted.get_nowait() == "accounted request"
+        assert pending.empty()
+    finally:
+        release.set()
+        await GLOBAL_LOGGING_WORKER.flush()
+        await GLOBAL_LOGGING_WORKER.stop()
+
+
+@pytest.mark.asyncio
+async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0")
+    started: Final = asyncio.Event()
+    cancelled: Final = asyncio.Event()
+    flushed_after_cancellation: Final = asyncio.Event()
+
+    async def callback() -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            cancelled.set()
+
+    async def flush_spend() -> None:
+        if cancelled.is_set():
+            flushed_after_cancellation.set()
+
+    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    await started.wait()
+    try:
+        with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+            await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+        assert cancelled.is_set()
+        assert flushed_after_cancellation.is_set()
+        assert "Timed out draining success callbacks" in caplog.text
+    finally:
+        await GLOBAL_LOGGING_WORKER.stop()
+
+
+@pytest.mark.asyncio
+async def test_shutdown_bounds_callback_cancellation_and_flushes_accounting(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0")
+    started: Final = asyncio.Event()
+    cancelling: Final = asyncio.Event()
+    cancelled: Final = asyncio.Event()
+    flushed: Final = asyncio.Event()
+    loop: Final = asyncio.get_running_loop()
+    clock: Final = loop.time
+
+    async def callback() -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            cancelling.set()
+            try:
+                await asyncio.Event().wait()
+            finally:
+                cancelled.set()
+
+    async def flush_spend() -> None:
+        if cancelled.is_set():
+            flushed.set()
+
+    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    await started.wait()
+    try:
+        with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+            shutdown: Final = asyncio.create_task(
+                litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+            )
+            await cancelling.wait()
+            monkeypatch.setattr(loop, "time", lambda: clock() + 6)
+            await shutdown
+        assert cancelled.is_set()
+        assert flushed.is_set()
+        assert "Timed out cancelling success callbacks" in caplog.text
+    finally:
+        monkeypatch.setattr(loop, "time", clock)
+        await GLOBAL_LOGGING_WORKER.stop()
+
+
 headers = {"Authorization": f"Bearer {token}"}
 
 example_completion_result = {

tests/unit/test_utils.py

@@ -10882,3 +10882,43 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_
     )
 
     assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("deferred", [False, True])
+async def test_success_logging_is_enqueued_before_the_worker_is_flushed(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.utils import _schedule_async_success_logging
+
+    completed: Final[asyncio.Queue[object]] = asyncio.Queue()
+
+    async def callback(kwargs: object, response: object, start: object, end: object) -> None:
+        completed.put_nowait(response)
+
+    now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
+    logging_obj: Final = Logging(
+        model="synthetic-model",
+        messages=[],
+        stream=False,
+        call_type="acompletion",
+        start_time=now.timestamp(),
+        litellm_call_id="logging-enqueue",
+        function_id="logging-enqueue",
+        dynamic_async_success_callbacks=[callback],
+    )
+    logging_obj.defer_async_logging = deferred
+    result: Final = litellm.ModelResponse(id="accounted-response")
+    try:
+        _schedule_async_success_logging(logging_obj, result, now, now, False)
+        if deferred:
+            assert completed.empty()
+            assert logging_obj.enqueue_deferred_logging is not None
+            logging_obj.enqueue_deferred_logging()
+        await GLOBAL_LOGGING_WORKER.flush()
+        assert completed.get_nowait() is result
+        assert completed.empty()
+    finally:
+        await asyncio.sleep(0)
+        await GLOBAL_LOGGING_WORKER.flush()
+        await GLOBAL_LOGGING_WORKER.stop()