mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
Merge remote-tracking branch 'origin/litellm_lite_auto_sign_in' into litellm_lite_browser_sign_in
This commit is contained in:
mateo-berri
commit
ae6d3e3bce
13 changed files with 1100 additions and 16 deletions
|
|
@ -49,6 +49,22 @@ lite config unset hidden_commands # list everything again
|
|||
|
||||
Hidden commands drop out of both `lite --help` and the interactive shell's "Available commands" block, and stay runnable so existing scripts keep working
|
||||
|
||||
## Automatic Browser Sign-In for Agents
|
||||
|
||||
When a command that calls the proxy (`lite keys`, `lite models`, `lite model-groups`, `lite users`, `lite teams`, `lite credentials`) gets a 401 and you did not pass a key with `--api-key` or `LITELLM_PROXY_API_KEY`, `lite` signs you in through your browser and then runs the same command again with the new key. That covers a proxy you never logged in to and a stored login the proxy no longer accepts. It needs no terminal, so an AI agent that runs `lite` through its shell tool finishes the command on its own once you approve the sign-in in the browser:
|
||||
|
||||
```bash
|
||||
lite --base-url https://your-proxy.example.com keys generate --models gpt-5.4-mini --key-alias experiment-1
|
||||
```
|
||||
|
||||
The sign-in is the browser flow of `lite login --pkce`, so you pick your team on the consent page, the credential is stored where `lite login` stores it, and its refresh token renews the key on its own from then on. Every line about the sign-in goes to stderr, so stdout carries only the command's own output, such as the JSON from `lite keys generate`. If another `lite` process is already signing in, the command waits for it to finish and then uses the login it stored, so parallel commands open one browser tab. When that sign-in ends without a login the proxy accepts, or is still waiting after 6 minutes, the command fails instead of opening a second tab. If another `lite` process stored a newer login for the same proxy in the meantime, that key is tried first and no browser opens. `lite` keeps one stored login, so when that login is for a different proxy, the command fails and tells you to run `lite --base-url <proxy> login --pkce` to switch, instead of replacing and revoking the other proxy's login on its own
|
||||
|
||||
Nothing signs in ahead of time. A proxy that accepts the request without a key never opens a browser, and neither does `--help`. The proxy also answers 401 when a valid key lacks the role a command needs, such as a non-admin key running `lite credentials list`, so before it signs in, `lite` checks a rejected stored login against `/key/info`. When the proxy accepts the login there, the command fails with an error that says the login likely lacks the role, and no browser opens. A key you passed yourself is never swapped out or checked again, and a 401 for it fails with an error that names `--api-key` and `LITELLM_PROXY_API_KEY`. When the browser sign-in cannot finish (the machine has no web browser or only a text-mode one such as `w3m` or `lynx`, the shell runs over SSH with no display and `BROWSER` is unset, the browser fails to open, the proxy does not serve `/.well-known/litellm-cli-auth`, you deny the request, or nobody approves it within 5 minutes), the command fails and tells you to run `lite login --pkce` from a terminal on a machine with a browser, `lite login` on one without, or set `LITELLM_PROXY_API_KEY`. `lite` prints that 5-minute wait on stderr when the sign-in starts, so an agent whose shell tool stops commands sooner knows to run the command again with a longer timeout
|
||||
|
||||
When the `CI` environment variable is set to anything but empty, `0`, or `false`, `lite` never opens a browser and the command fails at once with the same message, without checking the stored login against the proxy first. Set `CI=1` for a script or cron job that should fail instead of waiting for a sign-in
|
||||
|
||||
`lite chat`, `lite encryption migrate`, `lite http request`, `lite keys import`, and `lite teams assign-key` do not sign in on their own. `lite chat` and `lite teams assign-key` are interactive, `lite encryption migrate` is a one-time admin migration you run with an admin key, `lite keys import` also talks to a second proxy, and `lite http request` cannot tell a 401 from the proxy apart from a provider's 401 that the proxy passes through
|
||||
|
||||
## Global Options
|
||||
|
||||
- `--version`, `-v`: Print the LiteLLM Proxy client and server version and exit.
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ from .claude_settings import (
|
|||
settings_file_owners,
|
||||
)
|
||||
from .pkce_login import (
|
||||
LOGIN_TIMEOUT_SECONDS,
|
||||
Http,
|
||||
PkceFailure,
|
||||
RevocationUnavailable,
|
||||
|
|
@ -849,14 +850,14 @@ def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: Secr
|
|||
show_commands()
|
||||
|
||||
|
||||
def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave:
|
||||
def replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave:
|
||||
previous: Final = load_token(vault=vault)
|
||||
stored: Final = save_token(record, vault=vault)
|
||||
if previous is None or isinstance(stored, CredentialNotSaved):
|
||||
return stored
|
||||
revocation: Final = revoke_stored_credential(previous, http)
|
||||
if revocation is not None:
|
||||
click.echo(
|
||||
_warn(
|
||||
f"Could not revoke the previous login's refresh token on the proxy ({revocation.reason}); "
|
||||
"it expires on its own."
|
||||
)
|
||||
|
|
@ -869,7 +870,7 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None:
|
|||
if isinstance(credential, PkceFailure):
|
||||
click.echo(f"Authentication failed: {credential.reason}")
|
||||
return
|
||||
stored: Final = _replace_stored_token(pkce_token_record(base_url, credential), http, vault)
|
||||
stored: Final = replace_stored_token(pkce_token_record(base_url, credential), http, vault)
|
||||
_finish_login(base_url, credential.access_token, config_claude, stored)
|
||||
|
||||
|
||||
|
|
@ -948,7 +949,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
|
|||
|
||||
# Save token data. base_url is stored so we can verify origin
|
||||
# before reusing the key on a subsequent CLI invocation.
|
||||
stored: Final = _replace_stored_token(
|
||||
stored: Final = replace_stored_token(
|
||||
{
|
||||
"base_url": base_url.rstrip("/"),
|
||||
"key": api_key,
|
||||
|
|
@ -1089,7 +1090,11 @@ def whoami(ctx: click.Context):
|
|||
token_data: Final = load_token(vault=vault)
|
||||
|
||||
if not token_data:
|
||||
click.echo("Not authenticated. Run 'lite login' to authenticate.")
|
||||
click.echo(
|
||||
"Not authenticated. Run the command you need, such as 'lite keys generate', and lite opens your browser "
|
||||
f"to sign in and waits up to {LOGIN_TIMEOUT_SECONDS // 60} minutes for you to approve it. "
|
||||
"To sign in first from a terminal, run 'lite login --pkce'."
|
||||
)
|
||||
return
|
||||
|
||||
key_readable: Final = token_data.get("key") is not None
|
||||
|
|
|
|||
329
litellm/proxy/client/cli/commands/auto_sign_in.py
Normal file
329
litellm/proxy/client/cli/commands/auto_sign_in.py
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
import click
|
||||
import requests
|
||||
from filelock import FileLock, Timeout
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.litellm_core_utils.cli_token_utils import get_cli_token_file_path, load_cli_token
|
||||
|
||||
from ...exceptions import UnauthorizedError
|
||||
from .auth import context_secret_vault, get_stored_api_key, replace_stored_token, storage_notice
|
||||
from .pkce_login import LOGIN_TIMEOUT_SECONDS, PkceCredential, PkceFailure, pkce_token_record, run_pkce_login
|
||||
|
||||
|
||||
class BrowserSignIn(Protocol):
|
||||
def unavailable_because(self) -> PkceFailure | None: ...
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure: ...
|
||||
|
||||
|
||||
class SignInLock(Protocol):
|
||||
def acquire(self, timeout: float) -> object: ...
|
||||
|
||||
def release(self) -> object: ...
|
||||
|
||||
|
||||
_TEXT_MODE_BROWSERS: Final = frozenset({"www-browser", "links", "elinks", "lynx", "w3m"})
|
||||
_SSH_SESSION_VARIABLES: Final = ("SSH_CONNECTION", "SSH_TTY")
|
||||
_DISPLAY_VARIABLES: Final = ("DISPLAY", "WAYLAND_DISPLAY")
|
||||
_SIGNED_IN_LOGIN_WITHOUT_KEY_ROW: Final = HTTPStatus.NOT_FOUND
|
||||
_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS: Final = LOGIN_TIMEOUT_SECONDS + 60
|
||||
|
||||
|
||||
def _runs_in_ci(environ: Mapping[str, str]) -> bool:
|
||||
return environ.get("CI", "").strip().lower() not in {"", "0", "false"}
|
||||
|
||||
|
||||
def _ssh_session_without_a_display(environ: Mapping[str, str]) -> bool:
|
||||
over_ssh: Final = any(environ.get(name) for name in _SSH_SESSION_VARIABLES)
|
||||
return over_ssh and not any(environ.get(name) for name in _DISPLAY_VARIABLES)
|
||||
|
||||
|
||||
_OPEN_URL_IN_BROWSER: Final = "import sys, webbrowser; sys.exit(0 if webbrowser.get().open(sys.argv[1]) else 1)"
|
||||
|
||||
|
||||
def open_browser_detached(url: str) -> PkceFailure | None:
|
||||
launcher: Final = subprocess.run(
|
||||
(sys.executable, "-I", "-c", _OPEN_URL_IN_BROWSER, url),
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
check=False,
|
||||
)
|
||||
return None if launcher.returncode == 0 else PkceFailure("the web browser did not open")
|
||||
|
||||
|
||||
def _live_environ() -> Mapping[str, str]:
|
||||
return os.environ
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PkceBrowserSignIn:
|
||||
find_browser: Callable[[], webbrowser.BaseBrowser] = webbrowser.get
|
||||
environ: Mapping[str, str] = field(default_factory=_live_environ)
|
||||
open_browser: Callable[[str], object] = open_browser_detached
|
||||
|
||||
def unavailable_because(self) -> PkceFailure | None:
|
||||
if _runs_in_ci(self.environ):
|
||||
return PkceFailure("the CI environment variable is set, so lite does not open a browser")
|
||||
if "BROWSER" not in self.environ and _ssh_session_without_a_display(self.environ):
|
||||
return PkceFailure(
|
||||
"this is an SSH session without a display, so a browser would open where you cannot see it"
|
||||
)
|
||||
try:
|
||||
browser: Final = self.find_browser()
|
||||
except webbrowser.Error:
|
||||
return PkceFailure("no web browser is available on this machine")
|
||||
if Path(browser.name).name in _TEXT_MODE_BROWSERS:
|
||||
return PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
return None
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
echo(
|
||||
f"lite is opening your browser to sign in and waits up to {LOGIN_TIMEOUT_SECONDS // 60} minutes "
|
||||
"for you to approve it. "
|
||||
"A shell or agent that stops commands sooner should run this one with a longer timeout."
|
||||
)
|
||||
return run_pkce_login(base_url, requests.Session(), open_browser=self.open_browser, echo=echo)
|
||||
|
||||
|
||||
class _SignInContext(TypedDict):
|
||||
base_url: ReadOnly[str]
|
||||
api_key: ReadOnly[NotRequired[str | None]]
|
||||
api_key_from_token_file: ReadOnly[NotRequired[bool]]
|
||||
browser_sign_in: ReadOnly[NotRequired[BrowserSignIn]]
|
||||
sign_in_lock: ReadOnly[NotRequired[Callable[[], SignInLock]]]
|
||||
|
||||
|
||||
class _SignInContextView(Protocol):
|
||||
@property
|
||||
def obj(self) -> _SignInContext | None: ...
|
||||
|
||||
|
||||
class _ApiKeySlot(Protocol):
|
||||
def __setitem__(self, key: Literal["api_key"], value: str, /) -> None: ...
|
||||
|
||||
|
||||
class _ApiKeySlotView(Protocol):
|
||||
@property
|
||||
def obj(self) -> _ApiKeySlot: ...
|
||||
|
||||
|
||||
def _sign_in_context(view: _SignInContextView) -> _SignInContext | None:
|
||||
return view.obj
|
||||
|
||||
|
||||
def _api_key_slot(view: _ApiKeySlotView) -> _ApiKeySlot:
|
||||
return view.obj
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Completed:
|
||||
value: object
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Unauthorized:
|
||||
error: UnauthorizedError
|
||||
|
||||
|
||||
def _attempt(run: Callable[[click.Context], object], ctx: click.Context) -> _Completed | _Unauthorized:
|
||||
try:
|
||||
return _Completed(run(ctx))
|
||||
except UnauthorizedError as error:
|
||||
return _Unauthorized(error)
|
||||
|
||||
|
||||
def _attempt_with_key(
|
||||
run: Callable[[click.Context], object], ctx: click.Context, api_key: str
|
||||
) -> _Completed | _Unauthorized:
|
||||
slot: Final = _api_key_slot(ctx)
|
||||
slot["api_key"] = api_key
|
||||
return _attempt(run, ctx)
|
||||
|
||||
|
||||
def _echo_to_stderr(message: str) -> None:
|
||||
click.echo(message, err=True)
|
||||
|
||||
|
||||
def _attempt_with_newer_stored_key(
|
||||
run: Callable[[click.Context], object], ctx: click.Context, base_url: str, rejected_key: str | None
|
||||
) -> _Completed | _Unauthorized | None:
|
||||
stored_key: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
|
||||
if stored_key is None or stored_key == rejected_key:
|
||||
return None
|
||||
return _attempt_with_key(run, ctx, stored_key)
|
||||
|
||||
|
||||
def _sign_in_failure(base_url: str, failure: PkceFailure) -> str:
|
||||
if failure.proxy_predates_pkce:
|
||||
return (
|
||||
f"The proxy at {base_url} has no browser sign-in for lite, so it is likely older than v1.99.0. "
|
||||
"Upgrade it, run `lite login` if it uses SSO, or set LITELLM_PROXY_API_KEY."
|
||||
)
|
||||
return (
|
||||
f"Could not sign in to {base_url} through the browser: {failure.reason}. "
|
||||
"Run `lite login --pkce` from a terminal on a machine with a browser, `lite login` on one without, "
|
||||
"or set LITELLM_PROXY_API_KEY."
|
||||
)
|
||||
|
||||
|
||||
def _sign_in(ctx: click.Context, base_url: str, sign_in: BrowserSignIn) -> str:
|
||||
_echo_to_stderr(f"The proxy at {base_url} needs a sign-in.")
|
||||
credential: Final = sign_in(base_url, _echo_to_stderr)
|
||||
if isinstance(credential, PkceFailure):
|
||||
raise click.ClickException(_sign_in_failure(base_url, credential))
|
||||
stored: Final = replace_stored_token(
|
||||
pkce_token_record(base_url, credential), requests.Session(), context_secret_vault(ctx)
|
||||
)
|
||||
_echo_to_stderr(storage_notice(stored))
|
||||
_echo_to_stderr(f"Signed in. Continuing `{ctx.command_path}`.")
|
||||
return credential.access_token
|
||||
|
||||
|
||||
def _proxy_accepts(base_url: str, api_key: str) -> bool:
|
||||
try:
|
||||
response: Final = requests.get(
|
||||
f"{base_url.rstrip('/')}/key/info", headers={"Authorization": f"Bearer {api_key}"}, timeout=10
|
||||
)
|
||||
except requests.RequestException:
|
||||
return False
|
||||
return response.ok or response.status_code == _SIGNED_IN_LOGIN_WITHOUT_KEY_ROW
|
||||
|
||||
|
||||
def _sign_in_lock() -> SignInLock:
|
||||
return FileLock(Path(get_cli_token_file_path()).with_name("sign-in.lock"))
|
||||
|
||||
|
||||
def _try_to_lock(lock: SignInLock, timeout: float) -> Literal["held", "busy", "unlockable"]:
|
||||
try:
|
||||
lock.acquire(timeout=timeout)
|
||||
except Timeout:
|
||||
return "busy"
|
||||
except OSError:
|
||||
return "unlockable"
|
||||
return "held"
|
||||
|
||||
|
||||
def _lock_sign_in(lock: SignInLock) -> Literal["held", "waited", "unlockable"]:
|
||||
first_try: Final = _try_to_lock(lock, timeout=0)
|
||||
if first_try != "busy":
|
||||
return first_try
|
||||
_echo_to_stderr(
|
||||
"Another lite command is waiting for a browser sign-in, so this one waits for it instead of opening a second one."
|
||||
)
|
||||
match _try_to_lock(lock, timeout=_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS):
|
||||
case "held":
|
||||
return "waited"
|
||||
case "unlockable":
|
||||
return "unlockable"
|
||||
case "busy":
|
||||
raise click.ClickException(
|
||||
"Another lite command has been waiting for a browser sign-in for over "
|
||||
f"{_WAIT_FOR_ANOTHER_SIGN_IN_SECONDS // 60} minutes, so this one stopped waiting. "
|
||||
"Finish or stop that sign-in, then run this command again."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _sign_in_lock_held(lock: SignInLock) -> Generator[bool]:
|
||||
outcome: Final = _lock_sign_in(lock)
|
||||
if outcome == "unlockable":
|
||||
yield False
|
||||
return
|
||||
try:
|
||||
yield outcome == "waited"
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
|
||||
def _login_for_another_proxy(ctx: click.Context, base_url: str) -> str | None:
|
||||
login: Final = load_cli_token(vault=context_secret_vault(ctx))
|
||||
if login is None or not login.base_url or login.base_url == base_url.rstrip("/"):
|
||||
return None
|
||||
return login.base_url
|
||||
|
||||
|
||||
def _rerun_after_sign_in(run: Callable[[click.Context], object], ctx: click.Context, context: _SignInContext) -> object:
|
||||
base_url: Final = context["base_url"]
|
||||
if not context.get("api_key_from_token_file", False):
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} refused `{ctx.command_path}` with HTTP 401 for the API key from --api-key "
|
||||
"or LITELLM_PROXY_API_KEY. Check that the key is valid and has the role this command needs, "
|
||||
"or unset both to sign in through your browser instead."
|
||||
)
|
||||
sign_in: Final = context.get("browser_sign_in") or PkceBrowserSignIn()
|
||||
unavailable: Final = sign_in.unavailable_because()
|
||||
if unavailable is not None:
|
||||
_echo_to_stderr(f"The proxy at {base_url} needs a sign-in.")
|
||||
raise click.ClickException(_sign_in_failure(base_url, unavailable))
|
||||
rejected_key: Final = context.get("api_key")
|
||||
if rejected_key is not None and _proxy_accepts(base_url, rejected_key):
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} accepts your stored login but refused `{ctx.command_path}` with HTTP 401, "
|
||||
"so it likely lacks the role this command needs. Set LITELLM_PROXY_API_KEY to a key that has it."
|
||||
)
|
||||
with _sign_in_lock_held(context.get("sign_in_lock", _sign_in_lock)()) as waited_for_another_sign_in:
|
||||
return _rerun_once_signed_in(run, ctx, base_url, sign_in, rejected_key, waited_for_another_sign_in)
|
||||
|
||||
|
||||
def _rerun_once_signed_in(
|
||||
run: Callable[[click.Context], object],
|
||||
ctx: click.Context,
|
||||
base_url: str,
|
||||
sign_in: BrowserSignIn,
|
||||
rejected_key: str | None,
|
||||
waited_for_another_sign_in: bool,
|
||||
) -> object:
|
||||
retried: Final = _attempt_with_newer_stored_key(run, ctx, base_url, rejected_key)
|
||||
if isinstance(retried, _Completed):
|
||||
return retried.value
|
||||
other_proxy: Final = _login_for_another_proxy(ctx, base_url)
|
||||
if other_proxy is not None:
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} needs a sign-in, but lite keeps one stored login and it is for {other_proxy}, "
|
||||
f"so lite did not replace it on its own. Run `lite --base-url {base_url} login --pkce` to switch proxies, "
|
||||
"or set LITELLM_PROXY_API_KEY."
|
||||
)
|
||||
if waited_for_another_sign_in:
|
||||
raise click.ClickException(
|
||||
"The browser sign-in another lite command started ended while this one waited, "
|
||||
f"and the proxy at {base_url} still refuses `{ctx.command_path}`, so lite did not open a second one. "
|
||||
"Run the command again to sign in."
|
||||
)
|
||||
match _attempt_with_key(run, ctx, _sign_in(ctx, base_url, sign_in)):
|
||||
case _Completed(value=value):
|
||||
return value
|
||||
case _Unauthorized():
|
||||
raise click.ClickException(
|
||||
f"The proxy at {base_url} refused `{ctx.command_path}` with HTTP 401 again after the browser sign-in, "
|
||||
"so the signed-in user likely lacks the role this command needs. "
|
||||
"Set LITELLM_PROXY_API_KEY to a key that has it."
|
||||
)
|
||||
|
||||
|
||||
class SignInOnUnauthorizedCommand(click.Command):
|
||||
def invoke(self, ctx: click.Context) -> object:
|
||||
run: Final = super().invoke
|
||||
match _attempt(run, ctx):
|
||||
case _Completed(value=value):
|
||||
return value
|
||||
case _Unauthorized(error=error):
|
||||
context: Final = _sign_in_context(ctx)
|
||||
if context is None:
|
||||
raise error
|
||||
return _rerun_after_sign_in(run, ctx, context)
|
||||
|
||||
|
||||
class SignInOnUnauthorizedGroup(click.Group):
|
||||
command_class = SignInOnUnauthorizedCommand
|
||||
|
|
@ -10,6 +10,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
|
||||
from ...credentials import CredentialsManagementClient
|
||||
from ._cli_context import cli_context_values
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
|
||||
class _CredentialInfo(TypedDict):
|
||||
|
|
@ -39,7 +40,7 @@ def _print_json(data: object) -> None:
|
|||
rich.print_json(data=data)
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.group(cls=SignInOnUnauthorizedGroup)
|
||||
def credentials():
|
||||
"""Manage credentials for the LiteLLM proxy server"""
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from rich.table import Table
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from ...keys import KeysManagementClient
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
|
||||
class _CliContext(TypedDict):
|
||||
|
|
@ -37,7 +38,7 @@ def _cli_context(view: _CliContextView) -> _CliContext:
|
|||
return view.obj
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.group(cls=SignInOnUnauthorizedGroup)
|
||||
def keys():
|
||||
"""Manage API keys for the LiteLLM proxy server"""
|
||||
|
||||
|
|
@ -352,7 +353,7 @@ def _import_keys_to_destination(
|
|||
return imported_count, failed_count
|
||||
|
||||
|
||||
@keys.command(name="import")
|
||||
@keys.command(name="import", cls=click.Command)
|
||||
@click.option(
|
||||
"--source-base-url",
|
||||
required=True,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import rich.table
|
|||
|
||||
from ... import Client
|
||||
from ._cli_context import cli_context_values
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
|
||||
def create_client(ctx: click.Context) -> Client:
|
||||
|
|
@ -19,7 +20,7 @@ def _rendered_field(group: Mapping[str, object], key: str, default: str) -> str:
|
|||
return str(group.get(key, default))
|
||||
|
||||
|
||||
@click.group(name="model-groups")
|
||||
@click.group(name="model-groups", cls=SignInOnUnauthorizedGroup)
|
||||
def model_groups() -> None:
|
||||
"""Inspect model groups your key can access on the proxy"""
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,9 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
|
||||
# local imports
|
||||
from ... import Client
|
||||
from ...exceptions import UnauthorizedError
|
||||
from ._cli_context import cli_context_values
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from rich.console import JustifyMethod
|
||||
|
|
@ -103,7 +105,7 @@ def create_client(ctx: click.Context) -> Client:
|
|||
return Client(base_url=context["base_url"], api_key=context["api_key"])
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.group(cls=SignInOnUnauthorizedGroup)
|
||||
def models() -> None:
|
||||
"""Manage models on your LiteLLM proxy server"""
|
||||
|
||||
|
|
@ -449,6 +451,13 @@ def import_models(
|
|||
model_params=model_info_obj.model_params,
|
||||
model_info=model_info_obj.model_info,
|
||||
)
|
||||
except UnauthorizedError as error:
|
||||
if added_models:
|
||||
raise click.ClickException(
|
||||
f"The proxy rejected the API key after {len(added_models)} of {len(filtered_model_list)} "
|
||||
"models were sent, so the import stopped there."
|
||||
) from error
|
||||
raise
|
||||
except Exception:
|
||||
pass # For summary, ignore errors
|
||||
added_models.append(model_info_obj)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import webbrowser
|
|||
from base64 import urlsafe_b64encode
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from http import HTTPStatus
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol
|
||||
|
|
@ -68,6 +69,7 @@ class _TokenResponse(LiteLLMBaseModel):
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class PkceFailure:
|
||||
reason: str
|
||||
proxy_predates_pkce: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -208,7 +210,8 @@ def discover_cli_auth(base_url: str, http: Http) -> CliAuthContract | PkceFailur
|
|||
return PkceFailure(f"could not reach {url}: {exc}")
|
||||
if response.status_code != 200:
|
||||
return PkceFailure(
|
||||
f"{url} answered {response.status_code}; this proxy version does not support `lite login --pkce`"
|
||||
f"{url} answered {response.status_code}; this proxy version does not support `lite login --pkce`",
|
||||
proxy_predates_pkce=response.status_code == HTTPStatus.NOT_FOUND,
|
||||
)
|
||||
try:
|
||||
contract: Final = CliAuthContract.model_validate(response.json())
|
||||
|
|
|
|||
|
|
@ -10,8 +10,10 @@ from rich.table import Table
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.proxy.client import Client
|
||||
from litellm.proxy.client.exceptions import UnauthorizedError
|
||||
|
||||
from ._cli_context import cli_context_values
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
|
||||
class _TeamRow(TypedDict):
|
||||
|
|
@ -29,7 +31,7 @@ class _ErrorBodyView(TypedDict):
|
|||
body: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.group(cls=SignInOnUnauthorizedGroup)
|
||||
def teams():
|
||||
"""Manage teams and team assignments"""
|
||||
|
||||
|
|
@ -94,6 +96,8 @@ def list(ctx: click.Context):
|
|||
# Use list() for simpler response structure (returns array directly)
|
||||
teams: Final = client.teams.list()
|
||||
display_teams_table(teams)
|
||||
except UnauthorizedError:
|
||||
raise
|
||||
except requests.exceptions.HTTPError as e:
|
||||
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
|
||||
error_body: Final[_ErrorBodyView] = {"body": e.response.json()}
|
||||
|
|
@ -119,6 +123,8 @@ def available(ctx: click.Context):
|
|||
display_teams_table(teams)
|
||||
else:
|
||||
click.echo("No available teams to join.")
|
||||
except UnauthorizedError:
|
||||
raise
|
||||
except requests.exceptions.HTTPError as e:
|
||||
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
|
||||
error_body: Final[_ErrorBodyView] = {"body": e.response.json()}
|
||||
|
|
@ -128,7 +134,7 @@ def available(ctx: click.Context):
|
|||
raise click.Abort()
|
||||
|
||||
|
||||
@teams.command()
|
||||
@teams.command(cls=click.Command)
|
||||
@click.option("--team-id", type=str, help="Team ID to assign the key to")
|
||||
@click.pass_context
|
||||
def assign_key(ctx: click.Context, team_id: str | None):
|
||||
|
|
|
|||
|
|
@ -5,9 +5,10 @@ import rich
|
|||
|
||||
from ... import UsersManagementClient
|
||||
from ._cli_context import cli_context_values
|
||||
from .auto_sign_in import SignInOnUnauthorizedGroup
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.group(cls=SignInOnUnauthorizedGroup)
|
||||
def users():
|
||||
"""Manage users on your LiteLLM proxy server"""
|
||||
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ button {{ flex: 1; padding: 10px; border-radius: 6px; font-size: 15px; cursor: p
|
|||
<div class="container">
|
||||
<h1>Authorize CLI access</h1>
|
||||
<p>A command-line client at <code>{escape(client_origin)}</code> wants to call LiteLLM as <strong>{escape(user_id)}</strong>.</p>
|
||||
<p>Approving issues it a personal credential that expires within {CLI_JWT_EXPIRATION_HOURS} hours. <code>lite logout</code> stops it from being renewed. Only approve if you started this sign-in yourself.</p>
|
||||
<p>Approving issues it a personal credential that expires within {CLI_JWT_EXPIRATION_HOURS} hours. <code>lite logout</code> stops it from being renewed. Only approve if you started this sign-in, yourself or through an agent you are running.</p>
|
||||
<form method="post" action="{escape(complete_url)}">
|
||||
<input type="hidden" name="flow" value="{escape(flow_handle)}">
|
||||
{_team_field(teams)}
|
||||
|
|
|
|||
|
|
@ -695,7 +695,8 @@ class TestWhoamiCommand:
|
|||
|
||||
assert result.exit_code == 0
|
||||
assert "Not authenticated" in result.output
|
||||
assert "Run 'lite login'" in result.output
|
||||
assert "lite opens your browser to sign in and waits up to 5 minutes" in result.output
|
||||
assert "run 'lite login --pkce'" in result.output
|
||||
|
||||
def test_whoami_old_token(self):
|
||||
"""Test whoami with old token showing warning"""
|
||||
|
|
|
|||
711
tests/unit/proxy/client/cli/test_auto_sign_in.py
Normal file
711
tests/unit/proxy/client/cli/test_auto_sign_in.py
Normal file
|
|
@ -0,0 +1,711 @@
|
|||
import json
|
||||
import webbrowser
|
||||
from collections.abc import Callable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from click.testing import CliRunner, Result
|
||||
from filelock import Timeout
|
||||
from requests_mock import Mocker
|
||||
|
||||
from litellm.proxy.client.cli import cli
|
||||
from litellm.proxy.client.cli.commands.auth import CliTokenData, get_stored_api_key, save_token
|
||||
from litellm.proxy.client.cli.commands.auto_sign_in import PkceBrowserSignIn, open_browser_detached
|
||||
from litellm.proxy.client.cli.commands.pkce_login import PkceCredential, PkceFailure, pkce_token_record
|
||||
from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault
|
||||
|
||||
BASE_URL: Final = "https://proxy.example.com"
|
||||
SIGNED_IN_KEY: Final = "sk-signed-in"
|
||||
GENERATED: Final = {"key": "sk-generated", "key_alias": "experiment-1"}
|
||||
REJECTED: Final = {"error": {"message": "Authentication Error, No api key passed in.", "type": "auth_error"}}
|
||||
SIGNED_IN_AT: Final = 1_790_000_000.0
|
||||
EXPIRES_AT: Final = 4_102_444_800.0
|
||||
|
||||
|
||||
class FakeBrowserSignIn:
|
||||
def __init__(self, outcome: PkceCredential | PkceFailure) -> None:
|
||||
self.outcome = outcome
|
||||
self.calls: list[str] = []
|
||||
|
||||
def unavailable_because(self) -> PkceFailure | None:
|
||||
return None
|
||||
|
||||
def __call__(self, base_url: str, echo: Callable[[str], None]) -> PkceCredential | PkceFailure:
|
||||
self.calls.append(base_url)
|
||||
echo(f"Opening browser to: {base_url}/v1/cli/authorize?state=fake")
|
||||
return self.outcome
|
||||
|
||||
|
||||
def _credential(access_token: str = SIGNED_IN_KEY) -> PkceCredential:
|
||||
return PkceCredential(
|
||||
access_token=access_token,
|
||||
refresh_token="refresh-token",
|
||||
expires_at=EXPIRES_AT,
|
||||
client_id="cli-client",
|
||||
token_endpoint=f"{BASE_URL}/v1/cli/token",
|
||||
revocation_endpoint=f"{BASE_URL}/v1/cli/revoke",
|
||||
resource=BASE_URL,
|
||||
user_id="agent-user",
|
||||
team_id=None,
|
||||
)
|
||||
|
||||
|
||||
def _classic_login(key: str) -> CliTokenData:
|
||||
return {
|
||||
"base_url": BASE_URL,
|
||||
"key": key,
|
||||
"user_id": "agent-user",
|
||||
"user_email": "unknown",
|
||||
"user_role": "cli",
|
||||
"auth_header_name": "Authorization",
|
||||
"jwt_token": "",
|
||||
"timestamp": SIGNED_IN_AT,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_key_in_the_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False)
|
||||
monkeypatch.delenv("LITELLM_PROXY_URL", raising=False)
|
||||
|
||||
|
||||
def _lite(args: Sequence[str], vault: object, sign_in: FakeBrowserSignIn, global_args: Sequence[str] = ()) -> Result:
|
||||
return CliRunner().invoke(
|
||||
cli,
|
||||
[*global_args, "--base-url", BASE_URL, *args],
|
||||
obj={"secret_vault": vault, "browser_sign_in": sign_in},
|
||||
)
|
||||
|
||||
|
||||
def _bearers(route) -> list[str | None]:
|
||||
return [request.headers.get("Authorization") for request in route.request_history]
|
||||
|
||||
|
||||
def _key_check(requests_mock, status_code: int):
|
||||
return requests_mock.get(f"{BASE_URL}/key/info", status_code=status_code, json={"info": {}})
|
||||
|
||||
|
||||
def test_a_401_without_a_login_signs_in_once_stores_the_key_and_finishes_the_command(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
vault: Final = secret_vault_factory()
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(
|
||||
["keys", "generate", "--models", "gpt-5.4-mini", "--key-alias", "experiment-1"], vault, sign_in
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
assert get_stored_api_key(BASE_URL, vault=vault) == SIGNED_IN_KEY
|
||||
assert f"The proxy at {BASE_URL} needs a sign-in" in result.stderr
|
||||
assert f"Opening browser to: {BASE_URL}/v1/cli/authorize" in result.stderr
|
||||
assert "Credential stored in your OS keychain." in result.stderr
|
||||
assert "Continuing `cli keys generate`" in result.stderr
|
||||
assert SIGNED_IN_KEY not in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key_check",
|
||||
[{"status_code": 401}, {"status_code": 503}, {"exc": requests.exceptions.ConnectionError}],
|
||||
ids=["rejected", "unavailable", "unreachable"],
|
||||
)
|
||||
def test_a_stored_login_the_proxy_rejects_is_replaced_by_one_sign_in_with_only_json_on_stdout(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], key_check: dict[str, object]
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(pkce_token_record(BASE_URL, _credential("sk-stale")), vault=vault)
|
||||
requests_mock.get(f"{BASE_URL}/key/info", **key_check)
|
||||
revoke: Final = requests_mock.post(f"{BASE_URL}/v1/cli/revoke", status_code=500)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == ["Bearer sk-stale", f"Bearer {SIGNED_IN_KEY}"]
|
||||
assert revoke.call_count == 1
|
||||
assert "Could not revoke the previous login's refresh token" in result.stderr
|
||||
assert get_stored_api_key(BASE_URL, vault=vault) == SIGNED_IN_KEY
|
||||
|
||||
|
||||
def test_a_login_for_another_proxy_is_never_replaced_by_an_automatic_sign_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
other_proxy: Final = "https://other.example.com"
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token({**_classic_login("sk-other"), "base_url": other_proxy}, vault=vault)
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert f"lite keeps one stored login and it is for {other_proxy}" in result.stderr
|
||||
assert f"Run `lite --base-url {BASE_URL} login --pkce`" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == [None]
|
||||
assert get_stored_api_key(other_proxy, vault=vault) == "sk-other"
|
||||
|
||||
|
||||
def test_the_browser_sign_in_fails_at_once_on_a_machine_without_a_browser(requests_mock: Mocker) -> None:
|
||||
def no_browser() -> webbrowser.BaseBrowser:
|
||||
raise webbrowser.Error("could not locate runnable browser")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=no_browser, environ={}).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("no web browser is available on this machine")
|
||||
assert requests_mock.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ci", ["true", "1", "True"])
|
||||
def test_the_browser_sign_in_never_looks_for_a_browser_in_ci(requests_mock: Mocker, ci: str) -> None:
|
||||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser while CI is set")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(find_browser=browser_lookup, environ={"CI": ci}).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("the CI environment variable is set, so lite does not open a browser")
|
||||
assert requests_mock.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("name", "environ"),
|
||||
[
|
||||
("www-browser", {"TERM": "xterm"}),
|
||||
("links", {"TERM": "xterm"}),
|
||||
("elinks", {"TERM": "xterm"}),
|
||||
("lynx", {"TERM": "xterm"}),
|
||||
("w3m", {"TERM": "xterm"}),
|
||||
("lynx", {"BROWSER": "lynx"}),
|
||||
("/usr/bin/w3m", {"BROWSER": "/usr/bin/w3m"}),
|
||||
],
|
||||
)
|
||||
def test_a_text_mode_browser_fails_at_once_even_when_browser_picks_it(
|
||||
requests_mock: Mocker, name: str, environ: dict[str, str]
|
||||
) -> None:
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser(name), environ=environ
|
||||
).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure("the web browser is a text-mode one, which lite cannot show while a command waits")
|
||||
assert requests_mock.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ssh_variable", ["SSH_CONNECTION", "SSH_TTY"])
|
||||
def test_an_ssh_session_without_a_display_never_opens_a_browser_on_the_remote_screen(
|
||||
requests_mock: Mocker, ssh_variable: str
|
||||
) -> None:
|
||||
def browser_lookup() -> webbrowser.BaseBrowser:
|
||||
raise AssertionError("looked for a browser in an SSH session without a display")
|
||||
|
||||
outcome: Final = PkceBrowserSignIn(
|
||||
find_browser=browser_lookup, environ={ssh_variable: "203.0.113.7 52144 10.0.0.5 22"}
|
||||
).unavailable_because()
|
||||
|
||||
assert outcome == PkceFailure(
|
||||
"this is an SSH session without a display, so a browser would open where you cannot see it"
|
||||
)
|
||||
assert requests_mock.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"environ",
|
||||
[
|
||||
{"BROWSER": "firefox"},
|
||||
{"CI": "false", "BROWSER": "firefox"},
|
||||
{"CI": "0", "BROWSER": "firefox"},
|
||||
{"SSH_CONNECTION": "203.0.113.7 52144 10.0.0.5 22", "BROWSER": "firefox"},
|
||||
{"SSH_CONNECTION": "203.0.113.7 52144 10.0.0.5 22", "DISPLAY": "localhost:10.0", "BROWSER": "firefox"},
|
||||
],
|
||||
ids=["browser-picks-it", "ci-false", "ci-zero", "ssh-with-a-browser-forwarder", "ssh-with-x11-forwarding"],
|
||||
)
|
||||
def test_a_browser_the_user_picked_starts_the_sign_in_and_names_how_long_it_waits(
|
||||
requests_mock: Mocker, environ: dict[str, str]
|
||||
) -> None:
|
||||
discovery: Final = requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", status_code=404)
|
||||
echoed: list[str] = []
|
||||
sign_in: Final = PkceBrowserSignIn(find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ=environ)
|
||||
|
||||
assert sign_in.unavailable_because() is None
|
||||
outcome: Final = sign_in(BASE_URL, echoed.append)
|
||||
|
||||
assert isinstance(outcome, PkceFailure)
|
||||
assert discovery.call_count == 1
|
||||
assert echoed == [
|
||||
"lite is opening your browser to sign in and waits up to 5 minutes for you to approve it. "
|
||||
"A shell or agent that stops commands sooner should run this one with a longer timeout."
|
||||
]
|
||||
|
||||
|
||||
|
||||
def _browser_script(path: Path, ran: Path, exit_code: int) -> Path:
|
||||
path.write_text(f'#!/bin/sh\n: > "{ran}"\nexit {exit_code}\n')
|
||||
path.chmod(0o755)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("picked_exit_code", "expected"),
|
||||
[(0, None), (1, PkceFailure("the web browser did not open"))],
|
||||
ids=["picked-browser-opens", "picked-browser-fails"],
|
||||
)
|
||||
def test_the_browser_launcher_opens_only_the_browser_lite_picked_and_never_a_hidden_console_one(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, picked_exit_code: int, expected: PkceFailure | None
|
||||
) -> None:
|
||||
picked_ran: Final = tmp_path / "picked-browser-ran"
|
||||
console_ran: Final = tmp_path / "console-browser-ran"
|
||||
picked: Final = _browser_script(tmp_path / "picked-browser", picked_ran, picked_exit_code)
|
||||
_browser_script(tmp_path / "www-browser", console_ran, 0)
|
||||
monkeypatch.setenv("PATH", str(tmp_path))
|
||||
monkeypatch.setenv("TERM", "xterm")
|
||||
monkeypatch.setenv("BROWSER", str(picked))
|
||||
monkeypatch.delenv("DISPLAY", raising=False)
|
||||
monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
|
||||
|
||||
assert open_browser_detached(f"{BASE_URL}/v1/cli/authorize?state=abc") == expected
|
||||
assert picked_ran.exists()
|
||||
assert not console_ran.exists()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("command", "method", "path", "body"),
|
||||
[
|
||||
(["keys", "list", "--format", "json"], "GET", "/key/list", {"keys": []}),
|
||||
(["models", "list", "--format", "json"], "GET", "/models", {"data": []}),
|
||||
(["model-groups", "list", "--format", "json"], "GET", "/model_group/info", {"data": []}),
|
||||
(["users", "list"], "GET", "/user/list", {"users": []}),
|
||||
(["teams", "list"], "GET", "/team/list", []),
|
||||
(["teams", "available"], "GET", "/team/available", []),
|
||||
(["credentials", "list", "--format", "json"], "GET", "/credentials", {"credentials": []}),
|
||||
],
|
||||
)
|
||||
def test_every_proxy_command_group_signs_in_on_a_401_and_reruns_with_the_new_key(
|
||||
requests_mock: Mocker,
|
||||
secret_vault_factory: type[FakeSecretVault],
|
||||
command: list[str],
|
||||
method: str,
|
||||
path: str,
|
||||
body: object,
|
||||
) -> None:
|
||||
route: Final = requests_mock.register_uri(
|
||||
method, f"{BASE_URL}{path}", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": body}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(command, secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
def test_models_import_signs_in_instead_of_reporting_models_it_never_added(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], tmp_path: Path
|
||||
) -> None:
|
||||
models_file: Final = tmp_path / "models.yaml"
|
||||
models_file.write_text(
|
||||
"model_list:\n - model_name: gpt-5.4-mini\n litellm_params:\n model: openai/gpt-5.4-mini\n"
|
||||
)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/model/new", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": {}}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["models", "import", str(models_file)], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
def test_a_401_partway_through_models_import_stops_without_signing_in_or_adding_models_twice(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], tmp_path: Path
|
||||
) -> None:
|
||||
models_file: Final = tmp_path / "models.yaml"
|
||||
models_file.write_text(
|
||||
"model_list:\n"
|
||||
" - model_name: gpt-5.4-mini\n litellm_params:\n model: openai/gpt-5.4-mini\n"
|
||||
" - model_name: gpt-5.4\n litellm_params:\n model: openai/gpt-5.4\n"
|
||||
)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/model/new", [{"status_code": 200, "json": {}}, {"status_code": 401, "json": REJECTED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["models", "import", str(models_file)], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "rejected the API key after 1 of 2 models were sent" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert [request.json()["model_name"] for request in route.request_history] == ["gpt-5.4-mini", "gpt-5.4"]
|
||||
|
||||
|
||||
def test_a_401_from_the_keys_import_source_never_signs_in_to_the_destination(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
source: Final = "https://source.example.com"
|
||||
requests_mock.get(f"{source}/key/list", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(
|
||||
["keys", "import", "--source-base-url", source, "--source-api-key", "sk-source"],
|
||||
secret_vault_factory(),
|
||||
sign_in,
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert sign_in.calls == []
|
||||
assert {request.netloc for request in requests_mock.request_history} == {"source.example.com"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("global_args", "env"),
|
||||
[(["--api-key", "sk-explicit"], {}), ([], {"LITELLM_PROXY_API_KEY": "sk-explicit"})],
|
||||
ids=["flag", "env"],
|
||||
)
|
||||
def test_a_rejected_explicit_key_fails_cleanly_without_opening_a_browser(
|
||||
requests_mock: Mocker,
|
||||
secret_vault_factory: type[FakeSecretVault],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
global_args: list[str],
|
||||
env: dict[str, str],
|
||||
) -> None:
|
||||
for name, value in env.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
vault: Final = secret_vault_factory()
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in, global_args)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "with HTTP 401 for the API key from --api-key or LITELLM_PROXY_API_KEY" in result.stderr
|
||||
assert "has the role this command needs" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == ["Bearer sk-explicit"]
|
||||
assert [request.path for request in requests_mock.request_history] == ["/key/generate"]
|
||||
assert vault.writes == []
|
||||
|
||||
|
||||
def test_a_failed_sign_in_names_the_reason_and_points_at_lite_login(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
vault: Final = secret_vault_factory()
|
||||
sign_in: Final = FakeBrowserSignIn(PkceFailure("timed out waiting for the browser sign-in to finish"))
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "timed out waiting for the browser sign-in to finish" in result.stderr
|
||||
assert "Run `lite login --pkce` from a terminal on a machine with a browser" in result.stderr
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert route.call_count == 1
|
||||
assert get_stored_api_key(BASE_URL, vault=vault) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("discovery_status", "advice"),
|
||||
[
|
||||
(404, "so it is likely older than v1.99.0. Upgrade it, run `lite login` if it uses SSO"),
|
||||
(503, "Run `lite login --pkce` from a terminal on a machine with a browser"),
|
||||
],
|
||||
ids=["proxy_without_browser_sign_in", "proxy_erroring"],
|
||||
)
|
||||
def test_a_proxy_without_browser_sign_in_is_named_as_too_old_instead_of_pointing_at_lite_login_pkce(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], discovery_status: int, advice: str
|
||||
) -> None:
|
||||
requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
requests_mock.get(f"{BASE_URL}/.well-known/litellm-cli-auth", status_code=discovery_status)
|
||||
|
||||
sign_in_with_a_picked_browser: Final = PkceBrowserSignIn(
|
||||
find_browser=lambda: webbrowser.GenericBrowser("firefox"), environ={"BROWSER": "firefox"}
|
||||
)
|
||||
|
||||
result: Final = CliRunner().invoke(
|
||||
cli,
|
||||
["--base-url", BASE_URL, "keys", "generate"],
|
||||
obj={"secret_vault": secret_vault_factory(), "browser_sign_in": sign_in_with_a_picked_browser},
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert advice in result.stderr
|
||||
|
||||
|
||||
def test_a_signed_in_key_the_proxy_also_rejects_is_not_followed_by_another_sign_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "refused `cli keys generate` with HTTP 401 again after the browser sign-in" in result.stderr
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
def _rejected_while_a_sibling_signs_in(vault):
|
||||
def respond(request, context):
|
||||
save_token(_classic_login("sk-sibling"), vault=vault)
|
||||
context.status_code = 401
|
||||
return REJECTED
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def test_a_key_a_sibling_process_stored_meanwhile_is_used_before_any_browser(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(_classic_login("sk-stale"), vault=vault)
|
||||
_key_check(requests_mock, 401)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate",
|
||||
[{"json": _rejected_while_a_sibling_signs_in(vault)}, {"status_code": 200, "json": GENERATED}],
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == ["Bearer sk-stale", "Bearer sk-sibling"]
|
||||
|
||||
|
||||
def test_a_sibling_key_the_proxy_also_rejects_falls_through_to_one_sign_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(_classic_login("sk-stale"), vault=vault)
|
||||
_key_check(requests_mock, 401)
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate",
|
||||
[
|
||||
{"json": _rejected_while_a_sibling_signs_in(vault)},
|
||||
{"status_code": 401, "json": REJECTED},
|
||||
{"status_code": 200, "json": GENERATED},
|
||||
],
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == ["Bearer sk-stale", "Bearer sk-sibling", f"Bearer {SIGNED_IN_KEY}"]
|
||||
assert get_stored_api_key(BASE_URL, vault=vault) == SIGNED_IN_KEY
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key_check_status", [200, 404], ids=["virtual_key", "browser_login_without_a_key_row"])
|
||||
def test_a_stored_login_the_proxy_accepts_but_a_command_refuses_reports_a_missing_role_without_a_browser(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], key_check_status: int
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(_classic_login("sk-member"), vault=vault)
|
||||
requests_mock.get(f"{BASE_URL}/v1/models", status_code=200, json={"data": []})
|
||||
route: Final = requests_mock.get(f"{BASE_URL}/credentials", status_code=401, json=REJECTED)
|
||||
key_check: Final = _key_check(requests_mock, key_check_status)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
writes_before: Final = list(vault.writes)
|
||||
|
||||
result: Final = _lite(["credentials", "list"], vault, sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "accepts your stored login but refused `cli credentials list` with HTTP 401" in result.stderr
|
||||
assert "Set LITELLM_PROXY_API_KEY to a key that has it" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
assert _bearers(key_check) == ["Bearer sk-member"]
|
||||
assert vault.writes == writes_before
|
||||
assert "sk-member" not in result.output
|
||||
|
||||
|
||||
class AnotherCommandHoldsTheSignInLock:
|
||||
def __init__(self, when_it_lets_go: Callable[[], None], lets_go: bool = True) -> None:
|
||||
self.when_it_lets_go = when_it_lets_go
|
||||
self.lets_go = lets_go
|
||||
self.attempts = 0
|
||||
self.held = False
|
||||
|
||||
def __call__(self) -> "AnotherCommandHoldsTheSignInLock":
|
||||
return self
|
||||
|
||||
def acquire(self, timeout: float) -> object:
|
||||
self.attempts += 1
|
||||
if self.attempts == 1 or not self.lets_go:
|
||||
raise Timeout("sign-in.lock")
|
||||
self.when_it_lets_go()
|
||||
self.held = True
|
||||
return self
|
||||
|
||||
def release(self) -> object:
|
||||
self.held = False
|
||||
return None
|
||||
|
||||
|
||||
class UnwritableLoginDirectory:
|
||||
def __call__(self) -> "UnwritableLoginDirectory":
|
||||
return self
|
||||
|
||||
def acquire(self, timeout: float) -> object:
|
||||
raise PermissionError(13, "Permission denied", "/home/agent/.litellm")
|
||||
|
||||
def release(self) -> object:
|
||||
raise AssertionError("released a sign-in lock it never held")
|
||||
|
||||
|
||||
def _lite_with_lock(vault: object, sign_in: FakeBrowserSignIn, lock: object) -> Result:
|
||||
return CliRunner().invoke(
|
||||
cli,
|
||||
["--base-url", BASE_URL, "keys", "generate"],
|
||||
obj={"secret_vault": vault, "browser_sign_in": sign_in, "sign_in_lock": lock},
|
||||
)
|
||||
|
||||
|
||||
def test_a_command_that_waits_for_another_commands_sign_in_uses_its_login_instead_of_a_second_browser(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(
|
||||
lambda: save_token(_classic_login("sk-sibling"), vault=vault)
|
||||
)
|
||||
|
||||
result: Final = _lite_with_lock(vault, sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert "Another lite command is waiting for a browser sign-in" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert _bearers(route) == [None, "Bearer sk-sibling"]
|
||||
assert other_command.held is False
|
||||
|
||||
|
||||
def test_a_command_that_waited_on_a_sign_in_that_ended_without_a_login_opens_no_second_browser(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(lambda: None)
|
||||
|
||||
result: Final = _lite_with_lock(secret_vault_factory(), sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "so lite did not open a second one. Run the command again to sign in." in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
assert other_command.held is False
|
||||
|
||||
|
||||
def test_a_command_stops_waiting_for_a_sign_in_another_command_never_finishes(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
other_command: Final = AnotherCommandHoldsTheSignInLock(lambda: None, lets_go=False)
|
||||
|
||||
result: Final = _lite_with_lock(secret_vault_factory(), sign_in, other_command)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert "Another lite command has been waiting for a browser sign-in for over 6 minutes" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
def test_a_login_directory_lite_cannot_write_still_signs_in_and_finishes_the_command(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(
|
||||
f"{BASE_URL}/key/generate", [{"status_code": 401, "json": REJECTED}, {"status_code": 200, "json": GENERATED}]
|
||||
)
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite_with_lock(secret_vault_factory(), sign_in, UnwritableLoginDirectory())
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == GENERATED
|
||||
assert sign_in.calls == [BASE_URL]
|
||||
assert _bearers(route) == [None, f"Bearer {SIGNED_IN_KEY}"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("environ", "reason"),
|
||||
[
|
||||
({"CI": "1"}, "the CI environment variable is set"),
|
||||
({"SSH_CONNECTION": "203.0.113.7 52144 10.0.0.5 22"}, "this is an SSH session without a display"),
|
||||
],
|
||||
ids=["ci", "ssh_without_a_display"],
|
||||
)
|
||||
def test_a_rejected_stored_login_fails_at_once_without_a_key_check_or_lock_where_no_browser_can_open(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], environ: dict[str, str], reason: str
|
||||
) -> None:
|
||||
vault: Final = secret_vault_factory()
|
||||
save_token(_classic_login("sk-stale"), vault=vault)
|
||||
key_check: Final = _key_check(requests_mock, 401)
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=401, json=REJECTED)
|
||||
lock: Final = AnotherCommandHoldsTheSignInLock(lambda: None)
|
||||
|
||||
result: Final = CliRunner().invoke(
|
||||
cli,
|
||||
["--base-url", BASE_URL, "keys", "generate"],
|
||||
obj={"secret_vault": vault, "browser_sign_in": PkceBrowserSignIn(environ=environ), "sign_in_lock": lock},
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert isinstance(result.exception, SystemExit)
|
||||
assert reason in result.stderr
|
||||
assert "Run `lite login --pkce` from a terminal on a machine with a browser" in result.stderr
|
||||
assert _bearers(route) == ["Bearer sk-stale"]
|
||||
assert key_check.call_count == 0
|
||||
assert lock.attempts == 0
|
||||
|
||||
|
||||
def test_help_never_signs_in(requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault]) -> None:
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate", "--help"], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Generate a new API key" in result.stdout
|
||||
assert sign_in.calls == []
|
||||
assert requests_mock.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [403, 500])
|
||||
def test_an_error_other_than_401_never_signs_in(
|
||||
requests_mock: Mocker, secret_vault_factory: type[FakeSecretVault], status: int
|
||||
) -> None:
|
||||
route: Final = requests_mock.post(f"{BASE_URL}/key/generate", status_code=status, json={"error": "refused"})
|
||||
sign_in: Final = FakeBrowserSignIn(_credential())
|
||||
|
||||
result: Final = _lite(["keys", "generate"], secret_vault_factory(), sign_in)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert f"Error: HTTP {status}" in result.stderr
|
||||
assert sign_in.calls == []
|
||||
assert route.call_count == 1
|
||||
Loading…
Add table
Reference in a new issue