mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
parent
496f5b9859
commit
5c40ec618e
8 changed files with 337 additions and 10 deletions
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue