From ee0c0cc7e8f5834a2b64bba442e0a10723f0d254 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Fri, 7 Aug 2026 22:18:59 -0700
Subject: [PATCH] feat(proxy): let USE_V2_MIGRATION_RESOLVER select the v2
 migration resolver

--use_v2_migration_resolver had no env var, and the Helm migrations job runs
prisma_migration.py, which calls run_server with a fixed argv. There was no
seam to pass the flag, so a Helm install could not reach the v2 resolver at
all. Reading it from the environment makes the existing flag configurable
from a deployment.
---
 litellm/proxy/proxy_cli.py                 |  1 +
 tests/test_litellm/proxy/test_proxy_cli.py | 58 ++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index ab159e84b6a..6a0b3c6bfb2 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -916,6 +916,7 @@ class ProxyInitializationHelpers:
         "path that can cause schema thrashing during rolling deploys where two "
         "LiteLLM versions contend for the same DB. Default is the v1 resolver."
     ),
+    envvar="USE_V2_MIGRATION_RESOLVER",
 )
 @click.option(
     "--reload",
diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py
index 2ccaa0df440..20d17b5a510 100644
--- a/tests/test_litellm/proxy/test_proxy_cli.py
+++ b/tests/test_litellm/proxy/test_proxy_cli.py
@@ -1925,6 +1925,64 @@ class TestRunServerDbSetup:
             assert exc_info.value.code == 1
             mock_setup_database.assert_not_called()
 
+    @patch("subprocess.run")
+    @patch("atexit.register")
+    @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
+    @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
+    @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
+    def test_v2_migration_resolver_opts_in_via_env_var(
+        self,
+        mock_should_update_schema,
+        mock_check_schema_diff,
+        mock_setup_database,
+        mock_atexit_register,
+        mock_subprocess_run,
+    ):
+        """USE_V2_MIGRATION_RESOLVER must select the v2 resolver.
+
+        The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`,
+        which calls run_server with a fixed argv, so a deployment has no way to
+        pass --use_v2_migration_resolver and an env var is the only route in.
+        """
+        from litellm.proxy.proxy_cli import run_server
+
+        mock_subprocess_run.return_value = MagicMock(returncode=0)
+        mock_should_update_schema.return_value = True
+        mock_setup_database.return_value = True
+
+        mock_proxy_module = MagicMock(
+            app=MagicMock(),
+            ProxyConfig=MagicMock(),
+            KeyManagementSettings=MagicMock(),
+            save_worker_config=MagicMock(),
+        )
+
+        clean_env = {
+            k: v
+            for k, v in os.environ.items()
+            if k not in ("DATABASE_URL", "DIRECT_URL")
+        }
+        clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
+        clean_env["USE_V2_MIGRATION_RESOLVER"] = "true"
+
+        with (
+            patch.dict(os.environ, clean_env, clear=True),
+            patch.dict(
+                "sys.modules",
+                {
+                    "proxy_server": mock_proxy_module,
+                    "litellm.proxy.proxy_server": mock_proxy_module,
+                },
+            ),
+        ):
+            run_server.main(
+                ["--local", "--skip_server_startup"], standalone_mode=False
+            )
+
+        mock_setup_database.assert_called_once_with(
+            use_migrate=True, use_v2_resolver=True
+        )
+
 
 # --- Module-level helpers for worker startup hook tests ---