From 5c40ec618e31fff9a55f19288bbbf35352d6b505 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 18:36:57 -0700 Subject: [PATCH] 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. --- litellm/litellm_core_utils/cli_token_utils.py | 11 ++ litellm/proxy/client/cli/commands/agents.py | 13 ++- litellm/proxy/client/cli/commands/auth.py | 40 ++++++- litellm/proxy/client/cli/main.py | 12 ++ .../test_cli_token_utils.py | 41 ++++++- .../proxy/client/cli/test_agents.py | 59 ++++++++++ .../proxy/client/cli/test_auth_commands.py | 108 ++++++++++++++++++ .../proxy/client/cli/test_global_options.py | 63 ++++++++++ 8 files changed, 337 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index eb01359cdc0..6e47706eabb 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -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]: diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index f39ffb3e864..ebe65e065d0 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -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( diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b06d86d5965..19233270c72 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -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"] diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index b8c483f4b08..ce426d26c44 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index 27fc5eb4bd0..06204da7e19 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -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 diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index afd1696a89f..190d486db8f 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -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 = {} diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 4ee8b502aa2..ba716330bbd 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -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""" diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 3a19f735c1b..548a62877a7 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -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