From a637b32528910330991f4671693e2ca539b02aa9 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 28 Jul 2026 14:33:21 -0700
Subject: [PATCH] fix(proxy): kill retirements scheduled after the shutdown
 flush synchronously

An in-flight refresh triggered by the __getattr__ stale-token fallback
can complete its rotation after flush_engine_retirements has taken its
snapshot, scheduling a retirement task that no longer has an owner and
gets abandoned at event-loop teardown. Once the flush has run,
_schedule_engine_retirement now SIGKILLs the old engine inline instead
of creating a task. The event loop is single threaded and both the
flag-set-plus-snapshot and the schedule are synchronous, so every
schedule either lands in the flush snapshot or sees the flag; no
retirement can escape both
---
 litellm/proxy/db/prisma_client.py                 |  5 +++++
 .../db/test_prisma_planned_engine_restart.py      | 15 +++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py
index 645def469a5..ca5a60e78b3 100644
--- a/litellm/proxy/db/prisma_client.py
+++ b/litellm/proxy/db/prisma_client.py
@@ -229,6 +229,7 @@ class PrismaWrapper:
         self._last_refresh_time: datetime | None = None
         self._active_drain_tracker = self._instrument_prisma_client(original_prisma)
         self._retirement_tasks: frozenset[_EngineRetirement] = frozenset()
+        self._retirements_flushed = False
 
         # Coordination for planned engine restarts (issue #29176). Every
         # `recreate_prisma_client` SIGTERMs the running query-engine on
@@ -321,6 +322,7 @@ class PrismaWrapper:
         )
 
     async def flush_engine_retirements(self) -> None:
+        self._retirements_flushed = True
         pending = self._retirement_tasks
         if not pending:
             return
@@ -333,6 +335,9 @@ class PrismaWrapper:
     def _schedule_engine_retirement(self, pid: int, tracker: _PrismaDrainTracker | None) -> None:
         if pid <= 0:
             return
+        if self._retirements_flushed:
+            self._kill_engine_process_now(pid)
+            return
         retirement_task = asyncio.create_task(self._retire_engine_when_drained(pid, tracker))
         self._retirement_tasks = self._retirement_tasks.union((_EngineRetirement(pid=pid, task=retirement_task),))
         retirement_task.add_done_callback(self._retirement_finished)
diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py
index 04ab76cd09c..2a28b1bdacb 100644
--- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py
+++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py
@@ -443,6 +443,21 @@ async def test_stop_token_refresh_flushes_pending_engine_retirement(mock_prisma_
     assert wrapper._retirement_tasks == frozenset()
 
 
+@pytest.mark.asyncio
+async def test_retirement_scheduled_after_flush_kills_engine_synchronously(mock_prisma_binary):
+    """A rotation that completes during shutdown (e.g. an in-flight __getattr__
+    refresh) must not create an unowned retirement task after the flush; the
+    old engine is killed inline instead."""
+    wrapper = _make_wrapper(engine_pid=111, iam=True)
+    await wrapper.stop_token_refresh_task()
+
+    with patch("os.kill") as mock_kill:
+        wrapper._schedule_engine_retirement(333, wrapper._active_drain_tracker)
+
+    assert mock_kill.call_args_list == [call(333, signal.SIGKILL)]
+    assert wrapper._retirement_tasks == frozenset()
+
+
 @pytest.mark.asyncio
 async def test_safe_refresh_cancellation_restores_token_and_cleans_replacement(
     mock_prisma_binary, monkeypatch