From 129c4a703b7224965b7d0c6ff402a4db2b85a635 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 20 Sep 2026 01:21:01 -0700 Subject: [PATCH] fix(proxy): only warn about the deprecated flag when it came from the CLI USE_V2_MIGRATION_RESOLVER=true is a supported way to select v2, but click sets the same parameter from that env var, so the deprecation notice fired for environment-based config that is not deprecated. The notice now keys off click's parameter source. Also drops an em dash from the notice, and moves the resolver decision under mock-free tests by making it take the env value as an argument. --- litellm/proxy/proxy_cli.py | 25 ++++++--- tests/test_litellm/proxy/test_proxy_cli.py | 63 ++++++++++++++-------- 2 files changed, 58 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 14d0331c0ff..78885461724 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final import click import httpx +from click.core import ParameterSource from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict @@ -181,12 +182,21 @@ def append_query_params(url: str | None, params: dict) -> str: return modified_url -def resolve_v2_migration_resolver(*, use_legacy_flag: bool) -> bool: +def resolve_v2_migration_resolver(*, use_legacy_flag: bool, env_value: str | None) -> bool: from litellm_proxy_extras.utils import str_to_bool if use_legacy_flag: return False - return bool(str_to_bool(os.getenv("USE_V2_MIGRATION_RESOLVER", "true"))) + if env_value is None: + return True + return bool(str_to_bool(env_value)) + + +def deprecated_v2_flag_passed_on_cli() -> bool: + ctx: Final = click.get_current_context(silent=True) + if ctx is None: + return False + return ctx.get_parameter_source("use_v2_migration_resolver") is ParameterSource.COMMANDLINE class ProxyInitializationHelpers: @@ -1368,14 +1378,15 @@ def run_server( check_prisma_schema_diff(db_url=None) else: use_v2_resolver: Final = resolve_v2_migration_resolver( - use_legacy_flag=use_legacy_migration_resolver + use_legacy_flag=use_legacy_migration_resolver, + env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"), ) - if use_v2_migration_resolver and use_v2_resolver: + if deprecated_v2_flag_passed_on_cli() and use_v2_resolver: print( "\033[1;33mLiteLLM Proxy: --use_v2_migration_resolver is " - "deprecated and has no effect \u2014 the v2 migration resolver " - "is now the default. You can safely remove it. To opt back " - "into the legacy v1 resolver, pass " + "deprecated and has no effect, because the v2 migration " + "resolver is now the default. You can safely remove it. To " + "opt back into the legacy v1 resolver, pass " "--use_legacy_migration_resolver.\033[0m" ) if not use_v2_resolver: diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index d2835142194..a38470d1fdf 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2203,6 +2203,7 @@ class TestRunServerDbSetup: mock_setup_database, mock_atexit_register, mock_subprocess_run, + capsys, ): """USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver. @@ -2248,44 +2249,58 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_once_with( use_migrate=True, use_v2_resolver=True ) + assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @pytest.mark.parametrize( - "argv_extra, env_extra, expected_v2", + "use_legacy_flag, env_value, expected", [ - ([], {}, True), - ([], {"USE_V2_MIGRATION_RESOLVER": "false"}, False), - (["--use_legacy_migration_resolver"], {}, False), - ( - ["--use_legacy_migration_resolver"], - {"USE_V2_MIGRATION_RESOLVER": "true"}, - False, - ), - (["--use_v2_migration_resolver"], {}, True), + (False, None, True), + (False, "true", True), + (False, "false", False), + (True, None, False), + (True, "true", False), ], ids=[ - "default-is-v2", - "env-false-opts-out", - "legacy-flag-opts-out", + "unset-env-defaults-to-v2", + "env-true-selects-v2", + "env-false-selects-v1", + "legacy-flag-selects-v1", "legacy-flag-beats-env-true", - "deprecated-v2-flag-still-accepted", ], ) + def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected): + from litellm.proxy.proxy_cli import resolve_v2_migration_resolver + + assert ( + resolve_v2_migration_resolver( + use_legacy_flag=use_legacy_flag, env_value=env_value + ) + is expected + ) + + def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self): + from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli + + assert deprecated_v2_flag_passed_on_cli() is False + @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_migration_resolver_selection( + def test_legacy_resolver_flag_reaches_database_setup( self, mock_should_update_schema, mock_check_schema_diff, mock_setup_database, mock_atexit_register, mock_subprocess_run, - argv_extra, - env_extra, - expected_v2, ): + """--use_legacy_migration_resolver must reach the database setup call. + + The resolver decision itself is covered mock-free above; this is the + one wiring check that the flag is threaded through run_server. + """ from litellm.proxy.proxy_cli import run_server mock_subprocess_run.return_value = MagicMock(returncode=0) @@ -2302,11 +2317,9 @@ class TestRunServerDbSetup: clean_env = { k: v for k, v in os.environ.items() - if k - not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") } clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" - clean_env.update(env_extra) with ( patch.dict(os.environ, clean_env, clear=True), @@ -2319,12 +2332,16 @@ class TestRunServerDbSetup: ), ): run_server.main( - ["--local", "--skip_server_startup", *argv_extra], + [ + "--local", + "--skip_server_startup", + "--use_legacy_migration_resolver", + ], standalone_mode=False, ) mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=expected_v2 + use_migrate=True, use_v2_resolver=False )