This commit is contained in:
Mateo Wang 2026-06-27 12:00:51 -07:00 • committed by GitHub
commit ebaf4bc16e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 387 additions and 22 deletions

View file

@ -31,6 +31,17 @@ def load_cli_token() -> Optional[dict]:
return None
def get_stored_base_url() -> Optional[str]:
"""Return the proxy base URL recorded by the last `lite login`, if any."""
token_data = load_cli_token()
if not token_data:
return None
base_url = token_data.get("base_url")
if isinstance(base_url, str) and base_url:
return base_url
return None
def get_litellm_gateway_api_key(
expected_base_url: Optional[str] = None,
) -> Optional[str]:

View file

@ -213,20 +213,25 @@ def _is_interactive() -> bool:
def _resolve_api_key(ctx: click.Context) -> str:
base_url = ctx.obj["base_url"]
api_key = ctx.obj.get("api_key")
if api_key:
return api_key
if not _is_interactive():
raise click.ClickException(
message = (
"No LiteLLM key found. Set LITELLM_PROXY_API_KEY (or pass --api-key) for "
"non-interactive use, or run `lite login` from a terminal."
)
if ctx.obj.get("base_url_is_default"):
message += (
f" If your proxy is not at {ctx.obj['base_url']}, also set "
"LITELLM_PROXY_URL (or pass --base-url)."
)
raise click.ClickException(message)
click.echo("No LiteLLM credentials found; starting login...")
ctx.invoke(login)
api_key = get_stored_api_key(expected_base_url=base_url)
api_key = get_stored_api_key(expected_base_url=ctx.obj["base_url"])
if not api_key:
raise click.ClickException(
"Login did not produce an API key; cannot start the agent."
@ -240,9 +245,9 @@ _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy."
def _launch(
ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool
) -> None:
base_url = ctx.obj["base_url"]
started_interactive = _is_interactive()
api_key = _resolve_api_key(ctx)
base_url = ctx.obj["base_url"]
display_name, _ = agent_profile(binary)
click.echo(

View file

@ -538,6 +538,36 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
return None
def _can_prompt_for_input() -> bool:
return sys.stdin.isatty()
def _normalize_base_url_input(value: str) -> str:
cleaned = value.strip().rstrip("/")
if not cleaned.startswith(("http://", "https://")):
raise click.UsageError("The proxy URL must start with http:// or https://")
return cleaned
def _resolve_login_base_url(ctx: click.Context) -> str:
"""Return the proxy URL to log in against, prompting when it was never chosen.
The prompt only fires when the URL is the hardcoded localhost fallback (no
flag, no LITELLM_PROXY_URL, no URL stored by a previous login) and stdin is
a terminal, so configured and scripted invocations never see it.
"""
base_url = ctx.obj["base_url"]
if not ctx.obj.get("base_url_is_default") or not _can_prompt_for_input():
return base_url
entered = click.prompt(
"LiteLLM proxy URL", default=base_url, value_proc=_normalize_base_url_input
)
ctx.obj["base_url"] = entered
ctx.obj["base_url_is_default"] = False
return entered
@click.command(name="login")
@click.pass_context
def login(ctx: click.Context):
@ -545,10 +575,18 @@ def login(ctx: click.Context):
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
from litellm.proxy.client.cli.interface import show_commands
base_url = ctx.obj["base_url"]
base_url = _resolve_login_base_url(ctx)
try:
cli_sso_flow = _start_cli_sso_flow(base_url=base_url)
except (requests.RequestException, ValueError) as e:
raise click.ClickException(
f"Could not start an SSO login with the LiteLLM proxy at {base_url}: {e}. "
"Check that the URL points at your LiteLLM proxy, or change it with "
"--base-url or LITELLM_PROXY_URL."
)
try:
key_id = cli_sso_flow["login_id"]
poll_secret = cli_sso_flow["poll_secret"]
user_code = cli_sso_flow["user_code"]

View file

@ -1,10 +1,13 @@
# stdlib imports
import os
from typing import Optional
# third party imports
import click
from click import ParameterSource
from litellm._version import version as litellm_version
from litellm.litellm_core_utils.cli_token_utils import get_stored_base_url
from litellm.proxy.client.health import HealthManagementClient
from .commands.agents import agent_commands
@ -20,6 +23,8 @@ from .commands.teams import teams
from .commands.users import users
from .interface import interactive_shell
DEFAULT_BASE_URL = "http://localhost:4000"
def print_version(base_url: str, api_key: Optional[str]):
"""Print CLI and server version info."""
@ -37,6 +42,25 @@ def print_version(base_url: str, api_key: Optional[str]):
click.echo(f"Could not retrieve server version: {e}")
def _eager_version_base_url(ctx: click.Context) -> str:
"""Resolve the base URL for `--version`, which can fire before --base-url is parsed."""
return (
ctx.params.get("base_url")
or os.environ.get("LITELLM_PROXY_URL")
or get_stored_base_url()
or DEFAULT_BASE_URL
)
def _print_version_callback(
ctx: click.Context, param: click.Parameter, value: bool
) -> None:
if not value or ctx.resilient_parsing:
return
print_version(_eager_version_base_url(ctx), ctx.params.get("api_key"))
ctx.exit()
@click.group(invoke_without_command=True)
@click.option(
"--version",
@ -45,23 +69,13 @@ def print_version(base_url: str, api_key: Optional[str]):
is_eager=True,
expose_value=False,
help="Show the LiteLLM Proxy CLI and server version and exit.",
callback=lambda ctx, param, value: (
(
print_version(
ctx.params.get("base_url") or "http://localhost:4000",
ctx.params.get("api_key"),
)
or ctx.exit()
)
if value and not ctx.resilient_parsing
else None
),
callback=_print_version_callback,
)
@click.option(
"--base-url",
envvar="LITELLM_PROXY_URL",
show_envvar=True,
default="http://localhost:4000",
default=DEFAULT_BASE_URL,
help="Base URL of the LiteLLM proxy server",
)
@click.option(
@ -75,12 +89,22 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None:
"""LiteLLM Proxy CLI - Manage your LiteLLM proxy server"""
ctx.ensure_object(dict)
base_url_is_default = (
ctx.get_parameter_source("base_url") == ParameterSource.DEFAULT
)
if base_url_is_default:
stored_base_url = get_stored_base_url()
if stored_base_url:
base_url = stored_base_url
base_url_is_default = False
# If no API key provided via flag or environment variable, try to load from saved token.
# Pass base_url so we only use the stored key when it was issued for this server.
if api_key is None:
api_key = get_stored_api_key(expected_base_url=base_url)
ctx.obj["base_url"] = base_url
ctx.obj["base_url_is_default"] = base_url_is_default
ctx.obj["api_key"] = api_key
# If no subcommand was invoked, start interactive mode

View file

@ -3,14 +3,13 @@ Unit tests for CLI token utilities
"""
import json
import os
import tempfile
from pathlib import Path
from unittest.mock import mock_open, patch
import pytest
from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
from litellm.litellm_core_utils.cli_token_utils import (
get_litellm_gateway_api_key,
get_stored_base_url,
)
class TestCLITokenUtils:
@ -87,3 +86,35 @@ class TestCLITokenUtils:
result = get_litellm_gateway_api_key()
assert result is None
class TestGetStoredBaseUrl:
"""Test reading the proxy base URL stored by `lite login`"""
def test_returns_stored_base_url(self):
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value={"key": "sk-test", "base_url": "https://llm.acme.com"},
):
assert get_stored_base_url() == "https://llm.acme.com"
def test_returns_none_when_no_token_file(self):
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value=None,
):
assert get_stored_base_url() is None
def test_returns_none_when_base_url_missing(self):
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value={"key": "sk-old-token"},
):
assert get_stored_base_url() is None
def test_returns_none_when_base_url_empty(self):
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value={"key": "sk-test", "base_url": ""},
):
assert get_stored_base_url() is None

View file

@ -393,6 +393,65 @@ class TestAgentCommands:
assert "LITELLM_PROXY_API_KEY" in result.output
mock_run.assert_not_called()
def test_non_interactive_without_key_mentions_url_knobs_when_url_defaulted(self):
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=False),
patch(f"{AGENTS_MODULE}.run_agent") as mock_run,
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={
"base_url": "http://localhost:4000",
"base_url_is_default": True,
"api_key": None,
},
)
assert result.exit_code != 0
assert "LITELLM_PROXY_API_KEY" in result.output
assert "LITELLM_PROXY_URL" in result.output
assert "http://localhost:4000" in result.output
mock_run.assert_not_called()
def test_base_url_chosen_during_login_is_used_for_launch(self):
captured = {}
@click.command()
@click.pass_context
def fake_login(ctx):
ctx.obj["base_url"] = "https://llm.acme.com"
ctx.obj["base_url_is_default"] = False
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
patch(f"{AGENTS_MODULE}.login", fake_login),
patch(
f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login"
) as mock_get,
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda base_url, api_key, command, **k: captured.update(
base_url=base_url
),
),
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={
"base_url": "http://localhost:4000",
"base_url_is_default": True,
"api_key": None,
},
)
assert result.exit_code == 0, result.output
assert captured["base_url"] == "https://llm.acme.com"
assert "routing Claude Code through proxy at https://llm.acme.com" in (
result.output
)
mock_get.assert_called_once_with(expected_base_url="https://llm.acme.com")
def test_interactive_without_key_logs_in_then_launches(self):
captured = {}

View file

@ -463,6 +463,114 @@ class TestLoginCommand:
assert "❌ Authentication failed: Invalid value" in result.output
class TestLoginBaseUrlPrompt:
"""Login prompts for the proxy URL when it was never chosen by the user"""
AUTH = "litellm.proxy.client.cli.commands.auth"
def setup_method(self):
self.runner = CliRunner()
def _ready_response(self):
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "ready",
"key": "test.jwt",
"user_id": "test-user",
"team_id": None,
"teams": [],
}
return mock_response
def _invoke_login(self, obj, input=None, can_prompt=True, post_side_effect=None):
post_kwargs = (
{"side_effect": post_side_effect}
if post_side_effect
else {"return_value": _mock_cli_sso_start_response()}
)
with (
patch(f"{self.AUTH}._can_prompt_for_input", return_value=can_prompt),
patch("webbrowser.open") as mock_browser,
patch("requests.post", **post_kwargs),
patch("requests.get", return_value=self._ready_response()),
patch(f"{self.AUTH}.save_token") as mock_save,
patch("litellm.proxy.client.cli.interface.show_commands"),
):
result = self.runner.invoke(login, obj=obj, input=input)
return result, mock_browser, mock_save
def test_prompts_and_uses_entered_url(self):
obj = {"base_url": "http://localhost:4000", "base_url_is_default": True}
result, mock_browser, mock_save = self._invoke_login(
obj, input="https://llm.acme.com\n"
)
assert result.exit_code == 0, result.output
assert "LiteLLM proxy URL [http://localhost:4000]:" in result.output
assert "https://llm.acme.com/sso/key/generate" in mock_browser.call_args[0][0]
assert mock_save.call_args[0][0]["base_url"] == "https://llm.acme.com"
assert obj["base_url"] == "https://llm.acme.com"
assert obj["base_url_is_default"] is False
def test_prompt_accepting_default_keeps_localhost(self):
obj = {"base_url": "http://localhost:4000", "base_url_is_default": True}
result, mock_browser, _ = self._invoke_login(obj, input="\n")
assert result.exit_code == 0, result.output
assert "http://localhost:4000/sso/key/generate" in mock_browser.call_args[0][0]
def test_entered_url_trailing_slash_is_stripped(self):
obj = {"base_url": "http://localhost:4000", "base_url_is_default": True}
result, mock_browser, _ = self._invoke_login(
obj, input="https://llm.acme.com/\n"
)
assert result.exit_code == 0, result.output
assert obj["base_url"] == "https://llm.acme.com"
assert "https://llm.acme.com/sso/key/generate" in mock_browser.call_args[0][0]
def test_url_without_scheme_reprompts(self):
obj = {"base_url": "http://localhost:4000", "base_url_is_default": True}
result, mock_browser, _ = self._invoke_login(
obj, input="llm.acme.com\nhttps://llm.acme.com\n"
)
assert result.exit_code == 0, result.output
assert "must start with http:// or https://" in result.output
assert "https://llm.acme.com/sso/key/generate" in mock_browser.call_args[0][0]
def test_no_prompt_when_url_was_configured(self):
obj = {"base_url": "https://llm.acme.com", "base_url_is_default": False}
result, mock_browser, _ = self._invoke_login(obj)
assert result.exit_code == 0, result.output
assert "LiteLLM proxy URL" not in result.output
assert "https://llm.acme.com/sso/key/generate" in mock_browser.call_args[0][0]
def test_no_prompt_when_not_interactive(self):
obj = {"base_url": "http://localhost:4000", "base_url_is_default": True}
result, mock_browser, _ = self._invoke_login(obj, can_prompt=False)
assert result.exit_code == 0, result.output
assert "LiteLLM proxy URL" not in result.output
assert "http://localhost:4000/sso/key/generate" in mock_browser.call_args[0][0]
def test_unreachable_proxy_names_url_and_override_knobs(self):
import requests
obj = {"base_url": "http://localhost:4000", "base_url_is_default": False}
result, _, mock_save = self._invoke_login(
obj, post_side_effect=requests.ConnectionError("connection refused")
)
assert result.exit_code != 0
assert "http://localhost:4000" in result.output
assert "--base-url" in result.output
assert "LITELLM_PROXY_URL" in result.output
mock_save.assert_not_called()
class TestLogoutCommand:
"""Test logout CLI command"""

View file

@ -50,3 +50,92 @@ def test_cli_version_command(cli_runner):
assert f"LiteLLM Proxy CLI Version: {litellm_version}" in result.output
assert "LiteLLM Proxy Server URL: http://localhost:4000" in result.output
assert "LiteLLM Proxy Server Version: 1.2.3" in result.output
class TestBaseUrlResolution:
"""The base URL stored by `lite login` becomes the default for later runs"""
def _server_url_line(self, cli_runner, args, stored_base_url, monkeypatch):
monkeypatch.delenv("LITELLM_PROXY_URL", raising=False)
with (
patch(
"litellm.proxy.client.health.HealthManagementClient.get_server_version",
return_value="1.2.3",
),
patch(
"litellm.proxy.client.cli.main.get_stored_base_url",
return_value=stored_base_url,
),
patch(
"litellm.proxy.client.cli.main.get_stored_api_key",
return_value=None,
),
):
result = cli_runner.invoke(cli, args)
assert result.exit_code == 0, result.output
return result.output
def test_stored_base_url_used_when_nothing_specified(self, cli_runner, monkeypatch):
output = self._server_url_line(
cli_runner, ["version"], "https://llm.acme.com", monkeypatch
)
assert "LiteLLM Proxy Server URL: https://llm.acme.com" in output
def test_flag_wins_over_stored_base_url(self, cli_runner, monkeypatch):
output = self._server_url_line(
cli_runner,
["--base-url", "http://flag:1234", "version"],
"https://llm.acme.com",
monkeypatch,
)
assert "LiteLLM Proxy Server URL: http://flag:1234" in output
def test_env_var_wins_over_stored_base_url(self, cli_runner, monkeypatch):
monkeypatch.setenv("LITELLM_PROXY_URL", "http://env:1234")
with (
patch(
"litellm.proxy.client.health.HealthManagementClient.get_server_version",
return_value="1.2.3",
),
patch(
"litellm.proxy.client.cli.main.get_stored_base_url",
return_value="https://llm.acme.com",
),
patch(
"litellm.proxy.client.cli.main.get_stored_api_key",
return_value=None,
),
):
result = cli_runner.invoke(cli, ["version"])
assert result.exit_code == 0, result.output
assert "LiteLLM Proxy Server URL: http://env:1234" in result.output
def test_falls_back_to_localhost_without_stored_url(self, cli_runner, monkeypatch):
output = self._server_url_line(cli_runner, ["version"], None, monkeypatch)
assert "LiteLLM Proxy Server URL: http://localhost:4000" in output
def test_version_flag_uses_stored_base_url(self, cli_runner, monkeypatch):
output = self._server_url_line(
cli_runner, ["--version"], "https://llm.acme.com", monkeypatch
)
assert "LiteLLM Proxy Server URL: https://llm.acme.com" in output
def test_version_flag_env_var_wins_over_stored(self, cli_runner, monkeypatch):
monkeypatch.setenv("LITELLM_PROXY_URL", "http://env:1234")
with (
patch(
"litellm.proxy.client.health.HealthManagementClient.get_server_version",
return_value="1.2.3",
),
patch(
"litellm.proxy.client.cli.main.get_stored_base_url",
return_value="https://llm.acme.com",
),
):
result = cli_runner.invoke(cli, ["--version"])
assert result.exit_code == 0, result.output
assert "LiteLLM Proxy Server URL: http://env:1234" in result.output
def test_version_flag_falls_back_to_localhost(self, cli_runner, monkeypatch):
output = self._server_url_line(cli_runner, ["--version"], None, monkeypatch)
assert "LiteLLM Proxy Server URL: http://localhost:4000" in output