From 7d6ef778ffb92141546773d35731a8309e8480cd Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Wed, 28 Oct 2026 18:29:28 -0700
Subject: [PATCH] fix: type shutdown cancellation tasks and document its clock

---
 litellm/litellm_core_utils/logging_worker.py        | 6 +++---
 litellm/proxy/db/db_spend_update_writer.py          | 1 -
 litellm/proxy/shutdown/graceful_shutdown_manager.py | 5 +++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py
index b2156216416..715b0e21dec 100644
--- a/litellm/litellm_core_utils/logging_worker.py
+++ b/litellm/litellm_core_utils/logging_worker.py
@@ -72,8 +72,8 @@ class LoggingWorker:
         self.concurrency = concurrency
         self.timeout_summary_window = timeout_summary_window
         self._queue: asyncio.Queue[LoggingTask] | None = None
-        self._worker_task: asyncio.Task | None = None
-        self._running_tasks: set[asyncio.Task] = set()
+        self._worker_task: asyncio.Task[None] | None = None
+        self._running_tasks: set[asyncio.Task[None]] = set()
         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
@@ -495,7 +495,7 @@ class LoggingWorker:
             # No worker launched and no in-flight tasks to drain.
             return
 
-        tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
+        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)
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 538779f11b4..16dbe2b4131 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1900,7 +1900,6 @@ class DBSpendUpdateWriter:
             proxy_logging_obj=proxy_logging_obj,
         )
 
-
         ################## Daily End-User Spend Update Transactions ##################
         # Aggregate all in memory daily end-user spend transactions and commit to db
         await self._flush_daily_spend_queue(
diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py
index fdc5bf07406..9aeee0453ff 100644
--- a/litellm/proxy/shutdown/graceful_shutdown_manager.py
+++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py
@@ -24,7 +24,7 @@ import os
 import signal
 import threading
 import time
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from types import FrameType
 from typing import Final
@@ -99,7 +99,7 @@ class GracefulShutdownManager:
 
     @classmethod
     @contextmanager
-    def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Iterator[None]:
+    def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Generator[None, None, None]:
         if threading.current_thread() is not threading.main_thread():
             yield
             return
@@ -144,6 +144,7 @@ class GracefulShutdownManager:
                 live ``InFlightRequestsMiddleware`` counter; injectable for tests.
             poll_interval: Seconds between counter polls.
             log_interval: Minimum seconds between ``drain_waiting`` log lines.
+            clock: Monotonic clock shared with the shutdown timestamp.
 
         Returns:
             Number of requests that drained while waiting (>= 0).