fix(proxy-extras): only spend a migrate-deploy attempt when a pass made no progress

The v2 migration resolver gave `prisma migrate deploy` four attempts, and
every recovery path ended in a bare `continue`, so each one burned an attempt.
A database first brought up with `--use_prisma_db_push` has a full schema and
no migrations ledger, so the baseline spent attempt one and the first three
migrations whose objects already existed spent the rest. The proxy then exited
before binding its port, and that database could never be moved onto the
resolver.

The retry budget now counts only attempts that got nowhere. Creating the
baseline, and each migration newly marked applied, leaves the budget alone, so
a push-created database works through its pre-existing objects one pass at a
time. Timeouts, deadlock rollbacks, advisory-lock waits, and a repeat of a
recovery that already ran still spend an attempt, so a run that stops making
progress gives up exactly as before.

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -6,6 +6,7 @@ import shutil
 import subprocess
 import tempfile
 import time
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Optional
 
@@ -45,6 +46,38 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
 
 _MIGRATION_DEADLOCK_MARKER = "deadlock detected"
 
+MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
+
+
+@dataclass(frozen=True)
+class _MigrateAttemptBudget:
+    """Retries left, and the recoveries already run.
+
+    A recovery that lands something new costs nothing, so a database full of
+    objects `prisma db push` created works through them one per pass. Anything
+    that made no progress spends an attempt, so a stuck run still gives up.
+    """
+
+    attempts_left: int
+    recoveries: frozenset[str] = frozenset()
+
+    @property
+    def exhausted(self) -> bool:
+        return self.attempts_left <= 0
+
+    @property
+    def attempt_number(self) -> int:
+        return MAX_MIGRATE_DEPLOY_ATTEMPTS - self.attempts_left + 1
+
+    def spend(self) -> "_MigrateAttemptBudget":
+        return replace(self, attempts_left=self.attempts_left - 1)
+
+    def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
+        if recovery in self.recoveries:
+            return self.spend()
+        return replace(self, recoveries=self.recoveries | {recovery})
+
+
 _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
 _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
     r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
@@ -716,6 +749,9 @@ class ProxyExtrasDBManager:
         Ahead-of-HEAD state (DB has migrations newer than this build ships)
         is logged as a warning, not a fatal error — users whose DBs got into
         weird shapes from the old thrashing should still be able to start.
+
+        The retry budget only counts attempts that made no progress: see
+        _MigrateAttemptBudget.
         """
         schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
         migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
@@ -749,8 +785,9 @@ class ProxyExtrasDBManager:
         original_dir = os.getcwd()
         os.chdir(migrations_dir)
         deploy_timeout = prisma_migrate_deploy_timeout()
+        budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
         try:
-            for attempt in range(4):
+            while not budget.exhausted:
                 try:
                     result = subprocess.run(
                         [_get_prisma_command(), "migrate", "deploy"],
@@ -767,10 +804,11 @@ class ProxyExtrasDBManager:
                     logger.warning(
                         "prisma migrate deploy attempt %s timed out after %ss, retrying. "
                         "Raise %s if this database needs longer to apply its pending migrations.",
-                        attempt + 1,
+                        budget.attempt_number,
                         deploy_timeout,
                         PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
                     )
+                    budget = budget.spend()
                     time.sleep(random.randrange(5, 15))
                     continue
 
@@ -781,7 +819,14 @@ class ProxyExtrasDBManager:
                         logger.info(
                             "Schema exists but no migrations ledger — creating baseline"
                         )
-                        ProxyExtrasDBManager._create_baseline_migration(schema_path)
+                        baselined = ProxyExtrasDBManager._create_baseline_migration(
+                            schema_path
+                        )
+                        budget = (
+                            budget.after_recovery("baseline")
+                            if baselined
+                            else budget.spend()
+                        )
                         continue
 
                     if "P3009" in stderr:
@@ -818,6 +863,7 @@ class ProxyExtrasDBManager:
                                     f"intervention may be required.\n\n"
                                     f"Detail: {resolve_err}"
                                 ) from resolve_err
+                            budget = budget.after_recovery(f"resolved:{name}")
                             continue
                         if migration_match:
                             migration_name = migration_match.group(1)
@@ -831,6 +877,7 @@ class ProxyExtrasDBManager:
                                     migration_name,
                                 )
                                 ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
+                                budget = budget.spend()
                                 time.sleep(random.randrange(5, 15))
                                 continue
                         raise RuntimeError(
@@ -876,6 +923,7 @@ class ProxyExtrasDBManager:
                                     f"intervention may be required.\n\n"
                                     f"Detail: {resolve_err}"
                                 ) from resolve_err
+                            budget = budget.after_recovery(f"resolved:{name}")
                             continue
 
                         if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
@@ -888,6 +936,7 @@ class ProxyExtrasDBManager:
                             ProxyExtrasDBManager._roll_back_migration_best_effort(
                                 migration_match.group(1)
                             )
+                            budget = budget.spend()
                             time.sleep(random.randrange(5, 15))
                             continue
 
@@ -900,8 +949,9 @@ class ProxyExtrasDBManager:
                         logger.info(
                             "prisma migrate deploy attempt %s deadlocked against "
                             "a concurrent migrate deploy, retrying",
-                            attempt + 1,
+                            budget.attempt_number,
                         )
+                        budget = budget.spend()
                         time.sleep(random.randrange(5, 15))
                         continue
 
@@ -909,8 +959,9 @@ class ProxyExtrasDBManager:
                         logger.info(
                             "prisma migrate deploy attempt %s timed out waiting for "
                             "the advisory lock a concurrent migrate deploy holds, retrying",
-                            attempt + 1,
+                            budget.attempt_number,
                         )
+                        budget = budget.spend()
                         time.sleep(random.randrange(5, 15))
                         continue
 
@@ -920,9 +971,9 @@ class ProxyExtrasDBManager:
                     ) from e
 
             raise RuntimeError(
-                "Database migration failed after 4 attempts (retry loop "
-                "exhausted by timeouts, deadlock retries, or repeated "
-                "idempotent-recovery continues). Check database connectivity, "
+                f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
+                "attempts that made no progress (timeouts, deadlock retries, or a "
+                "recovery that had already run once). Check database connectivity, "
                 "load, and _prisma_migrations ledger state, and raise "
                 f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
             )

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -703,3 +703,170 @@ class TestSpendLogsPartitionDetectionMissingPsycopg:
         assert any(
             "psycopg is not installed" in record.message for record in caplog.records
         )
+
+
+_ATTEMPT_BUDGET = 4
+
+_P3005_STDERR = """Error: P3005
+
+The database schema is not empty. Read more about how to baseline an existing production database: https://pris.ly/d/migrate-baseline
+"""
+
+
+def _p3018_stderr(migration_name):
+    return f"""Error: P3018
+
+A migration failed to apply. New migrations cannot be applied before the error is recovered from.
+
+Migration name: {migration_name}
+
+Database error code: 42P07
+
+Database error:
+ERROR: relation "SomeTable" already exists
+"""
+
+
+class _MigrateDeployHarness:
+    """Drives _setup_database_v2 with a scripted sequence of
+    `prisma migrate deploy` outcomes, with every recovery command faked out so
+    nothing touches a database or the packaged migrations directory."""
+
+    def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False):
+        import subprocess as subprocess_module
+
+        import litellm_proxy_extras.utils as utils_module
+
+        self.deploy_calls = []
+        self.resolved = []
+        self.baselines = 0
+        self._outcomes = list(outcomes)
+        self._repeat_last = repeat_last
+        self._subprocess_module = subprocess_module
+
+        monkeypatch.delenv("DATABASE_URL", raising=False)
+        monkeypatch.setattr(
+            ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))
+        )
+        monkeypatch.setattr(
+            ProxyExtrasDBManager,
+            "_create_baseline_migration",
+            staticmethod(self._fake_baseline),
+        )
+        monkeypatch.setattr(
+            ProxyExtrasDBManager,
+            "_roll_back_migration",
+            staticmethod(lambda name: None),
+        )
+        monkeypatch.setattr(
+            ProxyExtrasDBManager,
+            "_resolve_specific_migration",
+            staticmethod(self.resolved.append),
+        )
+        monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run)
+        monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
+
+        self.baseline_succeeds = True
+
+    def _fake_baseline(self, *args, **kwargs):
+        self.baselines += 1
+        return self.baseline_succeeds
+
+    def _next_outcome(self):
+        if self._outcomes:
+            if self._repeat_last and len(self._outcomes) == 1:
+                return self._outcomes[0]
+            return self._outcomes.pop(0)
+        raise AssertionError("prisma migrate deploy called more times than scripted")
+
+    def _fake_run(self, cmd, **kwargs):
+        assert cmd[1:] == ["migrate", "deploy"], f"unexpected prisma command: {cmd}"
+        self.deploy_calls.append(cmd)
+        outcome = self._next_outcome()
+        if outcome == "ok":
+            return _FakeCompleted()
+        if outcome == "timeout":
+            raise self._subprocess_module.TimeoutExpired(cmd, 1)
+        raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome)
+
+    def run(self):
+        return ProxyExtrasDBManager._setup_database_v2(use_migrate=True)
+
+
+class TestMigrateDeployAttemptAccounting:
+    """A `prisma db push` database has a full schema and no ledger, so the v2
+    resolver baselines it and then works through every migration whose objects
+    already exist. Those recoveries make progress, so they must not spend the
+    retry budget, which is there to stop a run that is getting nowhere."""
+
+    def test_a_push_created_database_finishes_bootstrapping(
+        self, monkeypatch, tmp_path
+    ):
+        already_there = [
+            "20250329084805_new_cron_job_table",
+            "20250806095134_rename_alias_to_server_name_mcp_table",
+            "20260224203854_add_agent_object_permissions_table",
+            "20260301120000_fourth_table",
+            "20260302120000_fifth_table",
+            "20260303120000_sixth_table",
+        ]
+        harness = _MigrateDeployHarness(
+            monkeypatch,
+            tmp_path,
+            [_P3005_STDERR]
+            + [_p3018_stderr(name) for name in already_there]
+            + ["ok"],
+        )
+
+        assert harness.run() is True
+        assert harness.baselines == 1
+        assert harness.resolved == already_there
+        assert len(harness.deploy_calls) == len(already_there) + 2
+
+    def test_repeated_recovery_of_one_migration_still_gives_up(
+        self, monkeypatch, tmp_path
+    ):
+        harness = _MigrateDeployHarness(
+            monkeypatch,
+            tmp_path,
+            [_p3018_stderr("20250329084805_new_cron_job_table")],
+            repeat_last=True,
+        )
+
+        with pytest.raises(RuntimeError):
+            harness.run()
+        assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1
+
+    def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path):
+        harness = _MigrateDeployHarness(
+            monkeypatch, tmp_path, ["timeout"], repeat_last=True
+        )
+
+        with pytest.raises(RuntimeError):
+            harness.run()
+        assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
+
+    def test_a_baseline_that_never_lands_stops_after_the_budget(
+        self, monkeypatch, tmp_path
+    ):
+        harness = _MigrateDeployHarness(
+            monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True
+        )
+        harness.baseline_succeeds = False
+
+        with pytest.raises(RuntimeError):
+            harness.run()
+        assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
+
+    def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path):
+        harness = _MigrateDeployHarness(
+            monkeypatch,
+            tmp_path,
+            ["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"],
+            repeat_last=True,
+        )
+
+        with pytest.raises(RuntimeError):
+            harness.run()
+        assert len(harness.deploy_calls) == 1
+        assert harness.resolved == []