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