From ecce7cdd9c8eae08931b4a89b1af9ba13cf7201b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:29:50 -0700 Subject: [PATCH] fix(proxy_cli): import proxy_server once on script-style boot (#42584) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_cli.py | 41 ++++------------ tests/test_litellm/proxy/test_proxy_cli.py | 57 ++++++++++++++++------ 2 files changed, 52 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 40140974198..ac69f2e3894 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -33,21 +33,20 @@ else: FastAPI = Any -def _deprioritize_script_dir_in_sys_path() -> None: +def _drop_script_dir_from_sys_path() -> None: """Stop ``litellm/proxy`` modules from shadowing installed packages. Running this file as a script puts its own directory at ``sys.path[0]``, so ``import a2a`` resolves to ``litellm/proxy/a2a`` instead of the ``a2a`` SDK - and A2A agent calls fail. The entry is moved to the end rather than dropped, - because the sibling-import fallbacks in this module (``from proxy_server - import ...``) still need it. No-op under the ``litellm`` console script. + and ``proxy_server`` resolves to a second copy of + ``litellm.proxy.proxy_server``. No-op under the ``litellm`` console script. """ script_dir: Final = os.path.dirname(os.path.abspath(__file__)) if sys.path and os.path.abspath(sys.path[0]) == script_dir: - sys.path.append(sys.path.pop(0)) + sys.path.pop(0) -_deprioritize_script_dir_in_sys_path() +_drop_script_dir_from_sys_path() sys.path.append(os.getcwd()) config_filename: Final = "litellm.secrets" @@ -881,7 +880,7 @@ class ProxyInitializationHelpers: default=False, help="Use prisma db push instead of prisma migrate for database schema updates", ) -@click.option("--local", is_flag=True, default=False, help="for local debugging") +@click.option("--local", is_flag=True, default=False, help="no-op, kept for backwards compatibility") @click.option( "--skip_server_startup", is_flag=True, @@ -1058,35 +1057,15 @@ def run_server( return args: Final = locals() - if local: - from proxy_server import ( + try: + from litellm.proxy.proxy_server import ( KeyManagementSettings, ProxyConfig, app, save_worker_config, ) - else: - try: - from .proxy_server import ( - KeyManagementSettings, - ProxyConfig, - app, - save_worker_config, - ) - except ModuleNotFoundError as e: - raise ModuleNotFoundError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`") - except ImportError as e: - if "litellm[proxy]" in str(e): - # user is missing a proxy dependency, ask them to pip install litellm[proxy] - raise e - else: - # this is just a local/relative import error, user git cloned litellm - from proxy_server import ( - KeyManagementSettings, - ProxyConfig, - app, - save_worker_config, - ) + except ModuleNotFoundError as e: + raise ModuleNotFoundError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`") from e if version is True: ProxyInitializationHelpers._echo_litellm_version() return diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index d2d71d6df05..b84efda3308 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -10,6 +10,8 @@ import pytest import builtins +import runpy +import sys import types import urllib.parse as urlparse @@ -18,6 +20,7 @@ import yaml from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string +from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server @@ -636,6 +639,36 @@ class TestProxyInitializationHelpers: ), f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() + @patch("uvicorn.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_script_boot_imports_the_package_proxy_server( + self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run + ): + package_proxy_server = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + sibling_proxy_server = types.ModuleType("proxy_server") + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + {"proxy_server": sibling_proxy_server, "litellm.proxy.proxy_server": package_proxy_server}, + ), + patch.object(sys, "argv", ["proxy_cli.py", "--skip_server_startup"]), + patch.object(sys, "path", list(sys.path)), + pytest.raises(SystemExit) as exit_info, + ): + runpy.run_path(proxy_cli.__file__, run_name="__main__") + + assert exit_info.value.code == 0 + package_proxy_server.save_worker_config.assert_called_once() + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") @@ -1945,13 +1978,18 @@ class TestProxyInitializationHelpers: mock_proxy_config_instance.get_config = mock_get_config mock_proxy_config.return_value = mock_proxy_config_instance - mock_proxy_server_module = MagicMock(app=mock_app) + mock_proxy_server_module = MagicMock( + app=mock_app, + ProxyConfig=mock_proxy_config, + KeyManagementSettings=mock_key_mgmt, + save_worker_config=mock_save_worker_config, + ) # Only remove DATABASE_URL and DIRECT_URL to prevent the database setup # code path from running. Do NOT use clear=True as it removes PATH, HOME, # etc., which causes imports inside run_server to break in CI (the real - # litellm.proxy.proxy_server import at line 820 of proxy_cli.py has heavy - # side effects that fail without a proper environment). + # litellm.proxy.proxy_server import has heavy side effects that fail + # without a proper environment). env_overrides = { "DATABASE_URL": "", "DIRECT_URL": "", @@ -1967,18 +2005,7 @@ class TestProxyInitializationHelpers: with ( patch.dict( "sys.modules", - { - "proxy_server": MagicMock( - app=mock_app, - ProxyConfig=mock_proxy_config, - KeyManagementSettings=mock_key_mgmt, - save_worker_config=mock_save_worker_config, - ), - # Also mock litellm.proxy.proxy_server to prevent the real - # import at line 820 of proxy_cli.py which has heavy side - # effects (FastAPI app init, logging setup, etc.) - "litellm.proxy.proxy_server": mock_proxy_server_module, - }, + {"litellm.proxy.proxy_server": mock_proxy_server_module}, ), patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"