fix(proxy_cli): import canonical proxy_server module in script-mode fallback

Running the proxy as a script (python litellm/proxy/proxy_cli.py) made the
relative-import fallback in run_server load proxy_server a second time as a
bare top-level module. That shadow copy won the litellm.callbacks
registration, while uvicorn served the canonical litellm.proxy.proxy_server
module, so budget/spend hooks incremented in-memory counters on one instance
and auth checks read them from another. Without Redis, key-level and
end-user model budgets were never enforced on such runs. Import the
canonical module in the fallback so only one module instance exists.
This commit is contained in:
ryan-crabbe-berri 2026-08-01 16:55:54 -07:00
parent b1fd20f4cd
commit 4a3819a602
2 changed files with 34 additions and 2 deletions

View file

@ -1013,8 +1013,7 @@ def run_server(
# 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 (
from litellm.proxy.proxy_server import (
KeyManagementSettings,
ProxyConfig,
app,

View file

@ -2277,3 +2277,36 @@ class TestPostgresStatementTimeoutOptions:
standalone_mode=False,
)
return {k: os.environ[k] for k in ("DATABASE_URL", "DIRECT_URL") if k in os.environ}
@pytest.mark.xdist_group("proxy_cli")
class TestScriptModeImportsCanonicalProxyServer:
def test_no_shadow_proxy_server_module_in_script_mode(self):
"""Script-mode runs (`python litellm/proxy/proxy_cli.py`) hit the relative-import
fallback in run_server. The old bare `from proxy_server import ...` loaded a second
full copy of proxy_server as a top-level module; that shadow copy won the
litellm.callbacks registration while uvicorn served the canonical module, so
in-memory budget counters were incremented on one instance and read on another
and budgets were never enforced without Redis."""
import subprocess
repo_root = Path(__file__).resolve().parents[3]
driver = (
"import sys, runpy\n"
"sys.path.insert(0, 'litellm/proxy')\n"
"sys.argv = ['proxy_cli.py', '--version']\n"
"try:\n"
" runpy.run_path('litellm/proxy/proxy_cli.py', run_name='__main__')\n"
"except SystemExit:\n"
" pass\n"
"assert 'litellm.proxy.proxy_server' in sys.modules, 'canonical module not loaded'\n"
"assert 'proxy_server' not in sys.modules, 'shadow top-level proxy_server module loaded'\n"
)
result = subprocess.run(
[sys.executable, "-c", driver],
cwd=repo_root,
capture_output=True,
text=True,
timeout=120,
)
assert result.returncode == 0, f"stdout={result.stdout}\nstderr={result.stderr}"