fix(cli): prompt for the proxy URL at login instead of silently assuming localhost

A first-time user running lite claude with no LITELLM_PROXY_URL set was sent
into an SSO login against http://localhost:4000 that failed with a raw
connection error, and nothing told them the URL was the problem or how to set
it. Login now prompts for the proxy URL when it was never chosen (no flag, no
env var, no previous login), remembers it via the URL already stored in
token.json so later runs default to it, and turns an unreachable proxy into a
clear error naming the URL and the --base-url/LITELLM_PROXY_URL overrides.
This commit is contained in:
mateo-berri 2026-06-10 18:36:57 -07:00
parent 496f5b9859
commit 5c40ec618e
8 changed files with 337 additions and 10 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

@ -3,8 +3,10 @@ from typing import Optional
# third party imports
import click
from click.core 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
@ -75,12 +77,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,66 @@ 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