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.
This commit is contained in:
Yuneng Jiang 2026-09-20 01:21:01 -07:00
parent 9aec964bac
commit 129c4a703b
No known key found for this signature in database
2 changed files with 58 additions and 30 deletions

View file

@ -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:

View file

@ -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
)