mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
feat: add lite auth print-token for Claude Code apiKeyHelper support (#32846)
* feat: add silent CLI token refresh for apiKeyHelper support
lite auth print-token prints a valid proxy credential for use as Claude
Code's apiKeyHelper, transparently refreshing it first if the cached JWT
is stale. This unblocks MDM-managed apiKeyHelper deployments (managed
via `lite auth print-token`) that need silent mid-session credential
rotation without restarting the client.
Refresh capability is backed by a virtual key minted with an empty model
list and cli_refresh metadata, kept strictly separate from the actual
(short-lived, real-model-scoped) call credential -- so a leak of the
credential that flows through every LLM request and subprocess env var
can't also self-renew. The refresh flow is single-use: /sso/cli/refresh
mints a fresh JWT + refresh token pair and blocks the presented refresh
token immediately, so a replay can't mint a second pair from it.
Server: /sso/cli/refresh (rotate) and /sso/cli/logout (revoke) endpoints.
lite login now also stores a refresh token; lite logout revokes it
server-side instead of only clearing the local file.
* fix: allow non-admin users to hit CLI refresh routes; resolve apiKeyHelper base_url from token.json
Found via a live end-to-end test against a real proxy + real Claude Code
session: /sso/cli/refresh and /sso/cli/logout were unreachable for any
non-proxy-admin caller, since Depends(user_api_key_auth) pulls in a
route-RBAC gate that 403s any route not on an explicit allowlist. That
made the feature unusable for actual end users, who authenticate as
internal_user. Add both routes to internal_user_routes; the handlers
already do their own fine-grained check (metadata.cli_refresh) same as
/key/block does today.
Also: `lite auth print-token` required an explicit --base-url/
LITELLM_PROXY_URL matching the stored token's origin, defaulting to
localhost:4000 otherwise. But apiKeyHelper is configured bare (no
flags), so this always mismatched a real deployment. Track whether
--base-url was explicitly passed (via click's ParameterSource) and, if
not, resolve the server from token.json directly instead of the CLI
default.
* test: mock refresh-token minting in test_cli_poll_key_tolerates_missing_user_row
Landed on litellm_internal_staging after this branch's refresh-key minting
change; needs the same mock as the other cli_poll_key tests since minting
now runs unconditionally whenever a JWT is generated.
* fix(ci): update test_cli_auth.py for refresh_token contract, regenerate schema.d.ts
_poll_for_authentication now always includes "refresh_token" in its
returned dict, and _handle_team_selection_during_polling returns a dict
instead of a bare JWT string -- test_cli_auth.py predates this branch's
refresh-token work and still asserted the old shapes.
schema.d.ts regenerated via `npm run gen:api` to pick up the new
/sso/cli/refresh and /sso/cli/logout routes (plus unrelated drift from
other PRs merged since it was last generated).
* fix(ci): apply CI's own schema.d.ts diff (enterprise routes I can't generate locally)
Local `npm run gen:api` only sees OSS routes -- this machine's
litellm_enterprise editable install points at a now-deleted temp
directory, so it silently drops enterprise-only routes from the spec.
Applied the exact diff CI's own generation produced instead of
re-running the generator locally.
* fix: close refresh-token race, fail closed on DB down, fix logout base_url
Addresses Greptile review findings on the CLI refresh-token PR:
- cli_refresh_token minted a new JWT + refresh token BEFORE blocking the
presented one. Two concurrent requests bearing the same refresh token
could both pass auth and both mint fresh pairs, yielding four live
credentials from one consumed token. Now the presented token is
consumed atomically first via update_many (only succeeding if it flips
blocked from False/None to True); the loser gets count=0 and is
rejected before anything is minted.
- When prisma_client is None, refresh silently returned a new JWT
without ever being able to mark the presented token consumed, leaving
it valid indefinitely. Now fails closed with a 500 instead.
- `lite logout` sent its revocation POST to ctx.obj["base_url"], which
defaults to localhost:4000 when --base-url isn't passed -- the same
bug print_token had before the base_url_explicit fix, just missed
here. Now resolves the same way: trust the stored token's origin
unless the caller explicitly overrode --base-url.
* fix(ci): satisfy ruff format and narrow token_data type in logout
* fix(security): never trust refresh-token metadata for authorization
Addresses a real privilege-escalation path Veria flagged: cli_refresh_token
read team_id, team_alias, and max_budget straight off the presented
token's own metadata and used them to authorize the new JWT. Since any
authenticated user can self-mint a virtual key with arbitrary metadata
via the ordinary /key/generate endpoint, a self-forged key with
{"cli_refresh": true, "team_id": "<any-team>", "max_budget": 999999999}
would sail through _require_cli_refresh_token's only check
(metadata.cli_refresh == True) and get a JWT scoped to a team the
caller never belonged to, with a budget it never had -- full
cross-team / budget bypass, and a removed team member could keep
refreshing team-scoped sessions indefinitely.
Metadata's team_id is now treated as an untrusted UX hint only: honored
solely if the CALLER (identified by the authenticated key's own
user_id, not client input) is a current member per a fresh
get_user_object lookup. team_alias and max_budget are never read back
from metadata at all -- team_alias comes from a live get_team_object
lookup and max_budget is recomputed with the exact same capping logic
the initial SSO login poll uses. _mint_cli_refresh_token no longer
accepts or stores team_alias/max_budget, only the team_id hint.
Added regression tests proving: a forged/stale team_id is dropped
(falls back to no team, not silently honored), and a forged max_budget
in metadata never reaches the issued JWT.
* fix(ci): catch HTTPException specifically instead of bare Exception (BLE001)
* fix: un-consume refresh token if minting the replacement fails
Greptile flagged a real reliability gap: cli_refresh_token blocks the
presented token atomically, then does several more DB calls before
returning a replacement (user lookup, team lookup, JWT mint, new
refresh-key mint). Since this endpoint exists specifically for fully
unattended apiKeyHelper operation, a single transient failure in that
window (DB hiccup, etc.) permanently stranded the user: their old
token was already dead and no new one was issued, with no recovery
path short of a full interactive browser re-login.
Wrap that window in try/except; on any failure, best-effort revert the
consumed token back to usable (blocked=False) before re-raising, so a
retry can succeed. Standard compensating-action pattern since
generate_key_helper_fn doesn't take an injectable transaction, so
wrapping the whole thing in a real DB transaction isn't practical here.
* fix(security): refresh key had unrestricted model access, not none
Critical bug: _mint_cli_refresh_token used models=[] intending "no LLM
access", but that's backwards in this codebase. Per
_check_model_access_helper: `len(filtered_models) == 0 and len(models)
== 0` -> all_model_access = True. An empty models list on a key with no
team_id means UNRESTRICTED access to every model, not zero access. The
CLI refresh token -- meant to be usable for nothing but silently
exchanging itself for a new JWT -- was actually a fully unrestricted
API key for its entire 90-day lifetime, completely undermining the
whole point of keeping it separate from the short-lived call
credential.
Fixed with two independent layers: allowed_routes hard-restricts the
key to exactly /sso/cli/refresh and /sso/cli/logout (the real enforced
boundary, checked in the shared user_api_key_auth dependency for every
route); models is set to an unmatchable sentinel string as
defense-in-depth in case any code path only consults the models field.
Added an end-to-end regression test that exercises the actual
model-access-control function against a key shaped like the minted
refresh token, rather than only asserting on what arguments were passed
to the key-generation call -- the latter kind of test is exactly what
let the original bug ship, since asserting `models == []` is equally
consistent with "no access" and "unrestricted access" without checking
what the access-control code actually does with that shape.
Also: the compensating-rollback added for reliability un-blocked a
consumed refresh token even when the underlying user no longer exists.
That's a permanent, intentional rejection, not a transient failure --
un-blocking it would let a stale refresh token become valid again for a
different account if the user_id is ever reused/re-registered. Moved
the user-existence check outside the rollback-on-failure block so it
stays permanently blocked.
* refactor: rotate CLI refresh tokens via regenerate_key_fn instead of hand-rolled consume/rollback
The refresh token is already a plain litellm virtual key, so rotation can
delegate to the same atomic DB update /key/regenerate uses instead of a
bespoke update_many + compensating-rollback dance. This makes silent CLI
refresh an Enterprise feature, same as regular key regeneration.
* refactor: replace CLI stateless JWT + refresh-key pair with one self-rotating virtual key
The CLI previously minted two credentials on login: a stateless self-signed
JWT for LLM calls, and a separate DB-backed refresh-only key (scoped away
from ever calling an LLM) just to authorize minting a new JWT. Collapse
this into a single real virtual key, used directly as the LLM bearer token
and re-presented to /sso/cli/refresh to rotate its own secret in place.
This also means the CLI session key now shows up in the Admin UI's Keys
page and can be revoked/regenerated like any other key, rather than being
an invisible, unmanageable stateless token.
* refactor: drop silent CLI refresh, key just expires and requires re-login
/sso/cli/refresh only ever benefited Enterprise deployments (regenerate_key_fn's
gate), while everyone else already fell through to "re-run lite login" on
failure. Cut the endpoint, the rotation logic, and the client-side refresh
path entirely; print-token now just prints the cached key until it hits its
LITELLM_CLI_JWT_EXPIRATION_HOURS duration, then fails fast telling the user
to log in again. Session key itself is unaffected: still a real, revocable
virtual key visible in the Keys UI, `lite logout` still revokes it directly.
* fix(ci): regenerate schema.d.ts after removing /sso/cli/refresh route
* revert: go back to stateless JWT, keep only lite auth print-token
The virtual-key redesign (revocable, Keys-UI-visible credential) wasn't
needed just to support print-token, and cost real server-side surface
(a mint path, a logout-revoke endpoint, migrated tests/docs) for a property
this repo doesn't need yet. Reverting cli_poll_key/_types.py/schema.d.ts
back to the original stateless-JWT design; the only durable addition from
this whole effort is `lite auth print-token` (reads the cached credential,
prints it while fresh, fails with a clear message once it's past
LITELLM_CLI_JWT_EXPIRATION_HOURS) plus the base_url_explicit plumbing it
needs. `lite logout` goes back to clearing the local file only, since a
stateless JWT can't be revoked server-side.
* refactor: move CLI token freshness check to cli_token_utils, drop unnecessary renames
Addresses review: the freshness check is a pure token-shape/timestamp
util, not command logic, so it belongs alongside the other SDK-level
CLI token helpers (load_cli_token, get_litellm_gateway_api_key) rather
than in commands/auth.py. Also reverted a few incidental jwt_token/
session_key variable and string renames that weren't load-bearing.
This commit is contained in:
parent
b21c4ce865
commit
36aec560db
6 changed files with 218 additions and 105 deletions
|
|
@ -7,6 +7,7 @@ This module has no dependencies on proxy code and can be safely imported at the
|
|||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
|
@ -68,3 +69,17 @@ def get_litellm_gateway_api_key(
|
|||
if stored_url != expected_base_url.rstrip("/"):
|
||||
return None
|
||||
return token_data["key"]
|
||||
|
||||
|
||||
def is_cli_token_fresh(token_data: dict, buffer_hours: float = 0.1) -> bool:
|
||||
"""Check whether a cached CLI token (as stored in token.json) is still
|
||||
within its expiration window. Used by `lite auth print-token` to fail
|
||||
fast, without a network round trip, once the cached token is past
|
||||
`LITELLM_CLI_JWT_EXPIRATION_HOURS`."""
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
|
||||
timestamp = token_data.get("timestamp")
|
||||
if not isinstance(timestamp, (int, float)):
|
||||
return False
|
||||
age_hours = (time.time() - timestamp) / 3600
|
||||
return age_hours < (CLI_JWT_EXPIRATION_HOURS - buffer_hours)
|
||||
|
|
|
|||
|
|
@ -336,11 +336,12 @@ sequenceDiagram
|
|||
|
||||
### Authentication Commands
|
||||
|
||||
The CLI provides three authentication commands:
|
||||
The CLI provides these authentication commands:
|
||||
|
||||
- **`lite login`** - Start SSO authentication flow
|
||||
- **`lite logout`** - Clear stored authentication token
|
||||
- **`lite whoami`** - Show current authentication status
|
||||
- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); fails once the token has expired
|
||||
|
||||
### Authentication Flow Steps
|
||||
|
||||
|
|
@ -376,7 +377,7 @@ Authentication tokens are stored in `~/.litellm/token.json` with restricted file
|
|||
}
|
||||
```
|
||||
|
||||
The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`.
|
||||
The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while fresh and fails once it expires -- there is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`.
|
||||
|
||||
### Usage
|
||||
|
||||
|
|
|
|||
|
|
@ -469,7 +469,7 @@ To pin the model, pass the agent's own model flag (for example `lite claude --mo
|
|||
|
||||
The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit.
|
||||
|
||||
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
|
||||
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
|
||||
|
||||
## Environment Variables
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from rich.console import Console
|
|||
from rich.table import Table
|
||||
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
||||
|
||||
|
||||
# Token storage utilities
|
||||
|
|
@ -593,6 +594,44 @@ def logout():
|
|||
click.echo("✅ Logged out successfully. Authentication token cleared.")
|
||||
|
||||
|
||||
@click.command(name="print-token")
|
||||
@click.pass_context
|
||||
def print_token(ctx: click.Context):
|
||||
"""Print a valid API token for this proxy.
|
||||
|
||||
Designed to be used as Claude Code's `apiKeyHelper`
|
||||
(https://docs.claude.com/en/docs/claude-code/settings): stdout must
|
||||
contain only the token, so all diagnostics go to stderr. The token
|
||||
expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); once
|
||||
expired, run `lite login` again.
|
||||
"""
|
||||
token_data = load_token()
|
||||
if not token_data:
|
||||
click.echo("Not authenticated. Run 'lite login'.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
# apiKeyHelper is invoked bare (no --base-url), so unless the caller
|
||||
# explicitly pointed us at a server, trust whichever one `lite login`
|
||||
# actually issued this token for -- that's the whole point of not
|
||||
# needing a wrapper command.
|
||||
if ctx.obj.get("base_url_explicit"):
|
||||
base_url = ctx.obj["base_url"]
|
||||
if token_data.get("base_url") != base_url.rstrip("/"):
|
||||
click.echo("Not authenticated for this server. Run 'lite login'.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
if not is_cli_token_fresh(token_data):
|
||||
click.echo("Token expired. Run 'lite login' again.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
api_key = token_data.get("key")
|
||||
if not api_key:
|
||||
click.echo("No token available. Run 'lite login'.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
click.echo(api_key)
|
||||
|
||||
|
||||
@click.command(name="whoami")
|
||||
def whoami():
|
||||
"""Show current authentication status"""
|
||||
|
|
@ -616,8 +655,16 @@ def whoami():
|
|||
click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
|
||||
|
||||
|
||||
@click.group(name="auth")
|
||||
def auth_group():
|
||||
"""Manage CLI authentication (apiKeyHelper support, etc.)"""
|
||||
|
||||
|
||||
auth_group.add_command(print_token)
|
||||
|
||||
|
||||
# Export functions for use by other CLI commands
|
||||
__all__ = ["login", "logout", "whoami", "prompt_team_selection"]
|
||||
__all__ = ["login", "logout", "print_token", "auth_group", "whoami", "prompt_team_selection"]
|
||||
|
||||
# Export individual commands instead of grouping them
|
||||
# login, logout, and whoami will be added as top-level commands
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from litellm._version import version as litellm_version
|
|||
from litellm.proxy.client.health import HealthManagementClient
|
||||
|
||||
from .commands.agents import agent_commands
|
||||
from .commands.auth import get_stored_api_key, login, logout, whoami
|
||||
from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami
|
||||
from .commands.chat import chat
|
||||
from .commands.credentials import credentials
|
||||
from .commands.encryption import encryption
|
||||
|
|
@ -87,6 +87,12 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None:
|
|||
|
||||
ctx.obj["base_url"] = base_url
|
||||
ctx.obj["api_key"] = api_key
|
||||
# `--base-url` defaults to localhost:4000 for local dev convenience, but
|
||||
# apiKeyHelper is invoked bare (no flags) -- commands that must work
|
||||
# unattended (print-token) need to tell "user didn't say" apart from
|
||||
# "user said localhost:4000 on purpose" so they can fall back to
|
||||
# whatever server the stored token was actually issued for.
|
||||
ctx.obj["base_url_explicit"] = ctx.get_parameter_source("base_url") != click.core.ParameterSource.DEFAULT
|
||||
|
||||
# If no subcommand was invoked, start interactive mode
|
||||
if ctx.invoked_subcommand is None:
|
||||
|
|
@ -104,6 +110,8 @@ def version(ctx: click.Context):
|
|||
cli.add_command(login)
|
||||
cli.add_command(logout)
|
||||
cli.add_command(whoami)
|
||||
# Add the auth command group (e.g. `lite auth print-token`, used as Claude Code's apiKeyHelper)
|
||||
cli.add_command(auth_group, name="auth")
|
||||
# Add the models command group
|
||||
cli.add_command(models)
|
||||
# Add the credentials command group
|
||||
|
|
|
|||
|
|
@ -5,13 +5,12 @@ import time
|
|||
from pathlib import Path
|
||||
from unittest.mock import Mock, mock_open, patch
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.proxy.client.cli.commands.auth import (
|
||||
clear_token,
|
||||
get_stored_api_key,
|
||||
|
|
@ -19,6 +18,7 @@ from litellm.proxy.client.cli.commands.auth import (
|
|||
load_token,
|
||||
login,
|
||||
logout,
|
||||
print_token,
|
||||
save_token,
|
||||
whoami,
|
||||
)
|
||||
|
|
@ -78,12 +78,9 @@ class TestTokenUtilities:
|
|||
|
||||
with (
|
||||
patch("builtins.open", mock_open()) as mock_file,
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.chmod") as mock_chmod,
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
save_token(token_data)
|
||||
|
|
@ -93,9 +90,7 @@ class TestTokenUtilities:
|
|||
mock_chmod.assert_called_once_with("/test/path/token.json", 0o600)
|
||||
|
||||
# Verify JSON content was written correctly
|
||||
written_content = "".join(
|
||||
call[0][0] for call in mock_file().write.call_args_list
|
||||
)
|
||||
written_content = "".join(call[0][0] for call in mock_file().write.call_args_list)
|
||||
parsed_content = json.loads(written_content)
|
||||
assert parsed_content == token_data
|
||||
|
||||
|
|
@ -109,12 +104,9 @@ class TestTokenUtilities:
|
|||
|
||||
with (
|
||||
patch("builtins.open", mock_open(read_data=json.dumps(token_data))),
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=True),
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
result = load_token()
|
||||
|
|
@ -124,12 +116,9 @@ class TestTokenUtilities:
|
|||
def test_load_token_file_not_exists(self):
|
||||
"""Test loading token when file doesn't exist"""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=False),
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
result = load_token()
|
||||
|
|
@ -140,12 +129,9 @@ class TestTokenUtilities:
|
|||
"""Test loading token with invalid JSON"""
|
||||
with (
|
||||
patch("builtins.open", mock_open(read_data="invalid json")),
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=True),
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
result = load_token()
|
||||
|
|
@ -156,12 +142,9 @@ class TestTokenUtilities:
|
|||
"""Test loading token with IO error"""
|
||||
with (
|
||||
patch("builtins.open", side_effect=IOError("Permission denied")),
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=True),
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
result = load_token()
|
||||
|
|
@ -171,13 +154,10 @@ class TestTokenUtilities:
|
|||
def test_clear_token_file_exists(self):
|
||||
"""Test clearing token when file exists"""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("os.remove") as mock_remove,
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
clear_token()
|
||||
|
|
@ -187,13 +167,10 @@ class TestTokenUtilities:
|
|||
def test_clear_token_file_not_exists(self):
|
||||
"""Test clearing token when file doesn't exist"""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.get_token_file_path"
|
||||
) as mock_path,
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=False),
|
||||
patch("os.remove") as mock_remove,
|
||||
):
|
||||
|
||||
mock_path.return_value = "/test/path/token.json"
|
||||
|
||||
clear_token()
|
||||
|
|
@ -238,10 +215,7 @@ class TestTokenUtilities:
|
|||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com")
|
||||
== "sk-prod"
|
||||
)
|
||||
assert get_stored_api_key(expected_base_url="https://real-proxy.com") == "sk-prod"
|
||||
|
||||
def test_get_stored_api_key_base_url_match_trailing_slash(self):
|
||||
"""Trailing slash on expected_base_url is normalised before comparison"""
|
||||
|
|
@ -250,10 +224,7 @@ class TestTokenUtilities:
|
|||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com/")
|
||||
== "sk-prod"
|
||||
)
|
||||
assert get_stored_api_key(expected_base_url="https://real-proxy.com/") == "sk-prod"
|
||||
|
||||
def test_get_stored_api_key_base_url_mismatch(self):
|
||||
"""Stored key is NOT returned when expected_base_url differs from stored origin"""
|
||||
|
|
@ -271,9 +242,7 @@ class TestTokenUtilities:
|
|||
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
|
||||
return_value=token_data,
|
||||
):
|
||||
assert (
|
||||
get_stored_api_key(expected_base_url="https://real-proxy.com") is None
|
||||
)
|
||||
assert get_stored_api_key(expected_base_url="https://real-proxy.com") is None
|
||||
|
||||
|
||||
class TestLoginCommand:
|
||||
|
|
@ -307,11 +276,8 @@ class TestLoginCommand:
|
|||
) as mock_post,
|
||||
patch("requests.get", return_value=mock_response) as mock_get,
|
||||
patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save,
|
||||
patch(
|
||||
"litellm.proxy.client.cli.interface.show_commands"
|
||||
) as mock_show_commands,
|
||||
patch("litellm.proxy.client.cli.interface.show_commands") as mock_show_commands,
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -326,9 +292,7 @@ class TestLoginCommand:
|
|||
assert "Verification code: ABCD-EFGH" in result.output
|
||||
mock_post.assert_called_once()
|
||||
mock_get.assert_called()
|
||||
assert mock_get.call_args.kwargs["headers"] == {
|
||||
"x-litellm-cli-poll-secret": "poll-secret"
|
||||
}
|
||||
assert mock_get.call_args.kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"}
|
||||
|
||||
# Verify JWT was saved
|
||||
mock_save.assert_called_once()
|
||||
|
|
@ -355,7 +319,6 @@ class TestLoginCommand:
|
|||
patch("requests.get", return_value=mock_response),
|
||||
patch("time.sleep"),
|
||||
):
|
||||
|
||||
# Mock time.sleep to avoid actual delays in tests
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
|
|
@ -377,7 +340,6 @@ class TestLoginCommand:
|
|||
patch("requests.get", return_value=mock_response),
|
||||
patch("time.sleep"),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -399,7 +361,6 @@ class TestLoginCommand:
|
|||
),
|
||||
patch("time.sleep"),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -415,7 +376,6 @@ class TestLoginCommand:
|
|||
patch("requests.post", return_value=_mock_cli_sso_start_response()),
|
||||
patch("requests.get", side_effect=KeyboardInterrupt),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -440,7 +400,6 @@ class TestLoginCommand:
|
|||
patch("requests.get", return_value=mock_response),
|
||||
patch("time.sleep"),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -456,7 +415,6 @@ class TestLoginCommand:
|
|||
patch("requests.post", return_value=_mock_cli_sso_start_response()),
|
||||
patch("requests.get", side_effect=ValueError("Invalid value")),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -496,9 +454,7 @@ class TestWhoamiCommand:
|
|||
"timestamp": time.time() - 3600, # 1 hour ago
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data
|
||||
):
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data):
|
||||
result = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -510,9 +466,7 @@ class TestWhoamiCommand:
|
|||
|
||||
def test_whoami_not_authenticated(self):
|
||||
"""Test whoami when user is not authenticated"""
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token", return_value=None
|
||||
):
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=None):
|
||||
result = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -528,9 +482,7 @@ class TestWhoamiCommand:
|
|||
"timestamp": time.time() - (25 * 3600), # 25 hours ago
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data
|
||||
):
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data):
|
||||
result = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -540,21 +492,16 @@ class TestWhoamiCommand:
|
|||
def test_whoami_missing_fields(self):
|
||||
"""Test whoami with token missing some fields"""
|
||||
token_data = {
|
||||
"timestamp": time.time()
|
||||
- 3600
|
||||
"timestamp": time.time() - 3600
|
||||
# Missing user_email, user_id, user_role
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data
|
||||
):
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data):
|
||||
result = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "✅ Authenticated" in result.output
|
||||
assert (
|
||||
"Unknown" in result.output
|
||||
) # Should show "Unknown" for missing fields
|
||||
assert "Unknown" in result.output # Should show "Unknown" for missing fields
|
||||
|
||||
def test_whoami_no_timestamp(self):
|
||||
"""Test whoami with token missing timestamp"""
|
||||
|
|
@ -572,7 +519,6 @@ class TestWhoamiCommand:
|
|||
),
|
||||
patch("time.time", return_value=1000),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(whoami)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -625,20 +571,13 @@ class TestCLIKeyRegenerationFlow:
|
|||
patch("webbrowser.open") as mock_browser,
|
||||
patch(
|
||||
"requests.post",
|
||||
return_value=_mock_cli_sso_start_response(
|
||||
login_id="cli-session-uuid-456"
|
||||
),
|
||||
return_value=_mock_cli_sso_start_response(login_id="cli-session-uuid-456"),
|
||||
),
|
||||
patch(
|
||||
"requests.get", side_effect=[mock_first_response, mock_second_response]
|
||||
) as mock_get,
|
||||
patch("requests.get", side_effect=[mock_first_response, mock_second_response]) as mock_get,
|
||||
patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save,
|
||||
patch(
|
||||
"litellm.proxy.client.cli.interface.show_commands"
|
||||
) as mock_show_commands,
|
||||
patch("litellm.proxy.client.cli.interface.show_commands") as mock_show_commands,
|
||||
patch("click.prompt", return_value="2"),
|
||||
): # User selects index 2
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -659,9 +598,7 @@ class TestCLIKeyRegenerationFlow:
|
|||
first_poll_url = mock_get.call_args_list[0][0][0]
|
||||
assert "cli-session-uuid-456" in first_poll_url
|
||||
assert "team_id=" not in first_poll_url
|
||||
assert mock_get.call_args_list[0].kwargs["headers"] == {
|
||||
"x-litellm-cli-poll-secret": "poll-secret"
|
||||
}
|
||||
assert mock_get.call_args_list[0].kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"}
|
||||
|
||||
# Second poll should include team_id=team-beta
|
||||
second_poll_url = mock_get.call_args_list[1][0][0]
|
||||
|
|
@ -670,10 +607,7 @@ class TestCLIKeyRegenerationFlow:
|
|||
# Verify JWT was saved
|
||||
mock_save.assert_called_once()
|
||||
saved_data = mock_save.call_args[0][0]
|
||||
assert (
|
||||
saved_data["key"]
|
||||
== "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt"
|
||||
)
|
||||
assert saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt"
|
||||
assert saved_data["user_id"] == "test-user-456"
|
||||
|
||||
mock_show_commands.assert_called_once()
|
||||
|
|
@ -698,15 +632,12 @@ class TestCLIKeyRegenerationFlow:
|
|||
patch("webbrowser.open") as mock_browser,
|
||||
patch(
|
||||
"requests.post",
|
||||
return_value=_mock_cli_sso_start_response(
|
||||
login_id="cli-session-uuid-solo"
|
||||
),
|
||||
return_value=_mock_cli_sso_start_response(login_id="cli-session-uuid-solo"),
|
||||
),
|
||||
patch("requests.get", return_value=mock_response),
|
||||
patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save,
|
||||
patch("litellm.proxy.client.cli.interface.show_commands"),
|
||||
):
|
||||
|
||||
result = self.runner.invoke(login, obj=mock_context.obj)
|
||||
|
||||
assert result.exit_code == 0
|
||||
|
|
@ -722,7 +653,118 @@ class TestCLIKeyRegenerationFlow:
|
|||
# Verify JWT was saved
|
||||
mock_save.assert_called_once()
|
||||
saved_data = mock_save.call_args[0][0]
|
||||
assert (
|
||||
saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt"
|
||||
)
|
||||
assert saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt"
|
||||
assert saved_data["user_id"] == "test-user-solo"
|
||||
|
||||
|
||||
class TestPrintTokenCommand:
|
||||
"""Test `lite auth print-token`, used as Claude Code's apiKeyHelper.
|
||||
|
||||
stdout must contain *only* the token -- Claude Code treats stdout
|
||||
verbatim as the bearer token, so any diagnostic text on stdout would
|
||||
corrupt authentication.
|
||||
|
||||
apiKeyHelper is configured as a bare command (managed-settings.json sets
|
||||
just `"apiKeyHelper": "lite auth print-token"`, no --base-url flag) --
|
||||
so in the common case ctx.obj has no explicit base_url at all, and the
|
||||
command must resolve the server from whatever `lite login` stored in
|
||||
token.json, not from a CLI default. `--base-url`/`LITELLM_PROXY_URL`
|
||||
only matters when a caller explicitly overrides it (tracked via
|
||||
ctx.obj["base_url_explicit"], set by the `cli` group from
|
||||
click's ParameterSource).
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
self.runner = CliRunner()
|
||||
|
||||
def test_no_stored_token_fails_cleanly(self):
|
||||
with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=None):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "Not authenticated" in result.output
|
||||
|
||||
def test_bare_invocation_resolves_server_from_stored_token(self):
|
||||
"""The apiKeyHelper's real invocation shape: no --base-url given at
|
||||
all. Must use token.json's own base_url, not a hardcoded default."""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
"base_url": "https://litellm-proxy.corp.com",
|
||||
"key": "sk-prod-fresh",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
),
|
||||
patch("requests.post") as mock_post,
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert result.output.strip() == "sk-prod-fresh"
|
||||
mock_post.assert_not_called()
|
||||
|
||||
def test_explicit_base_url_mismatch_fails_cleanly(self):
|
||||
"""When the caller *does* explicitly pass --base-url, a token issued
|
||||
for a different server must never be printed."""
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
"base_url": "https://other-server.com",
|
||||
"key": "sk-should-not-print",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
):
|
||||
result = self.runner.invoke(
|
||||
print_token,
|
||||
obj={"base_url": "http://localhost:4000", "base_url_explicit": True},
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "sk-should-not-print" not in result.output
|
||||
|
||||
def test_fresh_cached_key_printed_without_network_call(self):
|
||||
"""A recently-issued key should be printed straight from cache -- no
|
||||
refresh call on every single invocation (apiKeyHelper gets called
|
||||
frequently)."""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
"base_url": "http://localhost:4000",
|
||||
"key": "sk-cached-fresh",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
),
|
||||
patch("requests.post") as mock_post,
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert result.output.strip() == "sk-cached-fresh"
|
||||
mock_post.assert_not_called()
|
||||
|
||||
def test_stale_key_fails_fast_without_network_call(self):
|
||||
"""There is no silent refresh: an expired cached key must fail
|
||||
loudly (stderr, nonzero exit) telling the user to `lite login`
|
||||
again, rather than making a network call or printing a dead key
|
||||
that will just 401 Claude Code."""
|
||||
old_timestamp = time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
"base_url": "http://localhost:4000",
|
||||
"key": "sk-stale-key",
|
||||
"timestamp": old_timestamp,
|
||||
},
|
||||
),
|
||||
patch("requests.post") as mock_post,
|
||||
):
|
||||
result = self.runner.invoke(print_token, obj={})
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "sk-stale-key" not in result.output
|
||||
assert "lite login" in result.output
|
||||
mock_post.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue