fix(cli): mint per-session agent credential on lite login

The `lite login` command was producing a shared UI session token that broke agent use in three ways: a $0.25 budget cap (from max_ui_session_budget) that killed agent sessions in minutes, a fixed identity "cli-jwt-token" shared across every user preventing per-session spend attribution, and auth gated behind EXPERIMENTAL_UI_LOGIN so the token was rejected on default deployments.

This fixes all three. Each login now generates a unique cli-session-{uuid} token with no per-key budget cap (enforced via shared team/user counters instead), and the decrypt path activates for any non-sk- token without requiring EXPERIMENTAL_UI_LOGIN.
This commit is contained in:
Sameer Kankute 2026-06-23 15:37:02 +05:30
parent 8bca05d311
commit 36c2b029ab
No known key found for this signature in database
8 changed files with 491 additions and 19 deletions

View file

@ -1530,7 +1530,7 @@ LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token"
CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS = 600
CLI_JWT_TOKEN_NAME = "cli-jwt-token"
CLI_SESSION_KEY_PREFIX = "cli-session"
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
CLI_JWT_EXPIRATION_HOURS = int(
os.getenv("CLI_JWT_EXPIRATION_HOURS")

View file

@ -23,7 +23,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import LimitedSizeOrderedDict
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
CLI_JWT_TOKEN_NAME,
CLI_SESSION_KEY_PREFIX,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
@ -2432,6 +2432,7 @@ class ExperimentalUIJWTToken:
Returns:
Encrypted JWT token string
"""
import secrets
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
@ -2453,11 +2454,16 @@ class ExperimentalUIJWTToken:
# Use first team if user has teams
_team_id = user_info.teams[0] if len(user_info.teams) > 0 else None
# Unique per-login identity so concurrent sessions get isolated spend
# attribution; budget is enforced via the shared team/user counters
# (team_id/user_id), not a per-key max_budget.
session_token = f"{CLI_SESSION_KEY_PREFIX}-{secrets.token_urlsafe(16)}"
session_alias = f"{CLI_SESSION_KEY_PREFIX}-{user_info.user_id}"
valid_token = UserAPIKeyAuth(
token=CLI_JWT_TOKEN_NAME,
key_name=CLI_JWT_TOKEN_NAME,
key_alias=CLI_JWT_TOKEN_NAME,
max_budget=litellm.max_ui_session_budget,
token=session_token,
key_name=session_alias,
key_alias=session_alias,
expires=expires,
user_id=user_info.user_id,
team_id=_team_id,

View file

@ -84,7 +84,6 @@ from litellm.proxy.utils import (
normalize_route_for_root_path,
)
from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
try:
@ -1513,8 +1512,11 @@ async def _user_api_key_auth_builder(
verbose_logger.debug("api key not found in cache.")
valid_token = None
## Check UI Hash Key
if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
## Check UI/CLI Hash Key
# Embedded session tokens (lite login, UI dashboard) are encrypted
# blobs, never sk- keys. Attempt decryption only for non-sk- keys;
# decryption fails closed for anything that isn't a genuine blob.
if valid_token is None and not api_key.startswith("sk-"):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(
api_key
)

View file

@ -376,6 +376,8 @@ 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`.
### Usage
Once authenticated, the CLI will automatically use the stored token for all requests. You no longer need to specify `--api-key` for subsequent commands.

View file

@ -465,6 +465,12 @@ Options (these belong to the wrapper, so put them before the agent's own flags):
To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints.
#### About the `lite login` credential
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.
## Environment Variables
The CLI respects the following environment variables:

View file

@ -1,10 +1,17 @@
import pytest
import asyncio
import aiohttp
import json
from httpx import AsyncClient
import secrets
import uuid
from typing import Any, Optional
import aiohttp
import pytest
from httpx import AsyncClient
PROXY_BASE = "http://0.0.0.0:4000"
MASTER_HEADERS = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
CLI_SSO_MODEL = "fake-openai-endpoint"
async def make_calls_until_budget_exceeded(session, key: str, call_function, **kwargs):
"""Helper function to make API calls until budget is exceeded. Verify that the budget is exceeded error is returned."""
@ -300,17 +307,249 @@ async def generate_team_key(
async def create_team(
session,
max_budget=None,
models: Optional[list[str]] = None,
team_alias: Optional[str] = None,
):
"""Helper function to create a new team"""
url = "http://0.0.0.0:4000/team/new"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {
"max_budget": max_budget,
}
async with session.post(url, headers=headers, json=data) as response:
url = f"{PROXY_BASE}/team/new"
data: dict[str, Any] = {"max_budget": max_budget}
if models is not None:
data["models"] = models
if team_alias is not None:
data["team_alias"] = team_alias
async with session.post(url, headers=MASTER_HEADERS, json=data) as response:
return await response.json()
async def create_user(
session,
*,
user_id: str,
user_email: str,
teams: list[str],
models: list[str],
):
url = f"{PROXY_BASE}/user/new"
data = {
"user_id": user_id,
"user_email": user_email,
"teams": teams,
"models": models,
"auto_create_key": False,
}
async with session.post(url, headers=MASTER_HEADERS, json=data) as response:
return await response.json()
async def add_team_member(
session,
*,
team_id: str,
user_id: str,
user_email: str,
):
url = f"{PROXY_BASE}/team/member_add"
data = {
"team_id": team_id,
"member": [{"user_id": user_id, "user_email": user_email, "role": "user"}],
}
async with session.post(url, headers=MASTER_HEADERS, json=data) as response:
return await response.json()
def _mint_cli_sso_token(
*,
user_id: str,
user_email: str,
team_id: str,
team_alias: str,
models: list[str],
) -> str:
"""
Mint the encrypted CLI session JWT returned by /sso/cli/poll after SSO login.
Uses the same code path as ExperimentalUIJWTToken.get_cli_jwt_auth_token.
"""
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
proxy_server.master_key = "sk-1234"
user_info = LiteLLM_UserTable(
user_id=user_id,
user_email=user_email,
user_role=LitellmUserRoles.INTERNAL_USER.value,
teams=[team_id],
models=models,
)
return ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=user_info,
team_id=team_id,
team_alias=team_alias,
)
async def obtain_cli_sso_token_via_poll_flow(
session,
*,
user_id: str,
user_email: str,
team_id: str,
team_alias: str,
models: list[str],
) -> str:
"""
Obtain a CLI SSO JWT through the same HTTP flow as `litellm-proxy login`:
/sso/cli/start -> (SSO callback) -> /sso/cli/complete -> /sso/cli/poll.
When the proxy SSO session cache is not shared with the test runner (otel CI
uses an isolated in-container cache), falls back to minting the identical JWT
that /sso/cli/poll would return.
"""
async with session.post(f"{PROXY_BASE}/sso/cli/start") as resp:
resp.raise_for_status()
start = await resp.json()
login_id = start["login_id"]
poll_secret = start["poll_secret"]
user_code = start["user_code"]
browser_complete_token = secrets.token_urlsafe(32)
seeded = await _seed_cli_sso_flow_in_shared_redis(
login_id=login_id,
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=team_alias,
models=models,
browser_complete_token=browser_complete_token,
)
if not seeded:
return _mint_cli_sso_token(
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=team_alias,
models=models,
)
async with session.post(
f"{PROXY_BASE}/sso/cli/complete/{login_id}",
data={
"user_code": user_code,
"browser_complete_token": browser_complete_token,
},
headers={"Content-Type": "application/x-www-form-urlencoded"},
) as resp:
assert resp.status == 200, await resp.text()
poll_headers = {
"x-litellm-cli-poll-secret": poll_secret,
}
async with session.get(
f"{PROXY_BASE}/sso/cli/poll/{login_id}",
params={"team_id": team_id},
headers=poll_headers,
) as resp:
poll = await resp.json()
assert poll.get("status") == "ready", poll
assert "key" in poll, poll
return poll["key"]
async def _seed_cli_sso_flow_in_shared_redis(
*,
login_id: str,
user_id: str,
user_email: str,
team_id: str,
team_alias: str,
models: list[str],
browser_complete_token: str,
) -> bool:
"""Seed the CLI SSO flow in Redis when tests share the proxy's Redis instance."""
import ast
import os
try:
import redis
except ImportError:
return False
host = os.getenv("REDIS_HOST")
if not host:
return False
try:
client = redis.Redis(
host=host,
port=int(os.getenv("REDIS_PORT", "6379")),
password=os.getenv("REDIS_PASSWORD") or None,
decode_responses=True,
)
client.ping()
except Exception:
return False
from litellm.proxy.management_endpoints.ui_sso import (
_get_cli_sso_flow_cache_key,
_hash_cli_sso_secret,
)
cache_key = _get_cli_sso_flow_cache_key(login_id)
raw_flow = client.get(cache_key)
if raw_flow is None:
return False
try:
flow = ast.literal_eval(raw_flow)
except (SyntaxError, ValueError):
return False
if not isinstance(flow, dict):
return False
flow["sso_complete"] = True
flow["user_code_verified"] = False
flow["session_data"] = {
"user_id": user_id,
"user_role": "internal_user",
"models": models,
"user_email": user_email,
"teams": [team_id],
"team_details": [{"team_id": team_id, "team_alias": team_alias}],
}
flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
client.setex(cache_key, 600, str(flow))
return True
async def make_calls_until_team_budget_exceeded_cli_sso(
session,
token: str,
team_id: str,
model: str,
):
"""Like make_calls_until_budget_exceeded but asserts team budget blocked the CLI SSO token."""
MAX_CALLS = 200
call_count = 0
try:
while call_count < MAX_CALLS:
await chat_completion(session=session, key=token, model=model)
call_count += 1
await asyncio.sleep(0.1)
pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls")
except Exception as e:
error_dict = e.body
assert error_dict["code"] == "429"
assert error_dict["type"] == "budget_exceeded"
message = error_dict["message"]
assert "Budget has been exceeded!" in message
assert "Team=" in message, f"Expected team budget error, got: {message}"
assert team_id in message, f"Expected team id in error, got: {message}"
return call_count
@pytest.mark.asyncio
async def test_team_budget_enforcement():
"""
@ -342,6 +581,67 @@ async def test_team_budget_enforcement():
), "Should make at least one successful call before team budget exceeded"
@pytest.mark.asyncio
async def test_team_budget_enforcement_cli_sso_token():
"""
Team budget enforcement for CLI SSO session tokens (litellm-proxy login JWT).
1. Create team with a tiny max_budget and a user on that team
2. Obtain a CLI SSO JWT (HTTP poll flow when Redis is shared, else mint)
3. Make chat completion calls until the team budget is exceeded
4. Verify HTTP 429 budget_exceeded names the team
"""
user_id = f"cli-budget-user-{uuid.uuid4().hex[:8]}"
user_email = f"{user_id}@example.com"
team_alias = f"cli-budget-team-{uuid.uuid4().hex[:8]}"
async with aiohttp.ClientSession() as session:
team_response = await create_team(
session=session,
max_budget=0.0000000005,
models=[CLI_SSO_MODEL],
team_alias=team_alias,
)
team_id = team_response["team_id"]
await create_user(
session,
user_id=user_id,
user_email=user_email,
teams=[team_id],
models=[CLI_SSO_MODEL],
)
await add_team_member(
session,
team_id=team_id,
user_id=user_id,
user_email=user_email,
)
cli_token = await obtain_cli_sso_token_via_poll_flow(
session,
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=team_alias,
models=[CLI_SSO_MODEL],
)
assert not cli_token.startswith(
"sk-"
), "CLI SSO token must not be a virtual key"
calls_made = await make_calls_until_team_budget_exceeded_cli_sso(
session=session,
token=cli_token,
team_id=team_id,
model=CLI_SSO_MODEL,
)
assert (
calls_made > 0
), "Should make at least one successful call before team budget exceeded"
@pytest.mark.asyncio
async def test_team_and_key_budget_enforcement():
"""

View file

@ -459,7 +459,9 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values
assert token_data["user_id"] == "test_user"
assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value
assert token_data["models"] == ["gpt-3.5-turbo"]
assert token_data["max_budget"] == litellm.max_ui_session_budget
# CLI session tokens carry no per-key budget; spend is enforced via the
# shared team/user counters. The $0.25 UI session cap must not leak in.
assert token_data.get("max_budget") is None
# Verify expiration time is set to 24 hours (default)
assert "expires" in token_data
@ -504,6 +506,35 @@ def test_get_cli_jwt_auth_token_custom_expiration(
assert expires <= get_utc_datetime() + timedelta(hours=48, minutes=1)
def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values):
"""Each CLI login mints a unique token id (per-session spend isolation) while
keeping a stable, user-scoped key_alias for log grouping. A regression that
pins token back to a constant would collapse both ids and fail here."""
from litellm.constants import CLI_SESSION_KEY_PREFIX
def _decode(token: str) -> dict:
decrypted = decrypt_value_helper(
token, key="ui_hash_key", exception_type="debug"
)
assert decrypted is not None
return json.loads(decrypted)
first = _decode(
ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)
)
second = _decode(
ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)
)
assert first["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-")
assert second["token"].startswith(f"{CLI_SESSION_KEY_PREFIX}-")
assert first["token"] != second["token"]
expected_alias = f"{CLI_SESSION_KEY_PREFIX}-test_user"
assert first["key_alias"] == second["key_alias"] == expected_alias
assert first["key_name"] == second["key_name"] == expected_alias
@pytest.mark.asyncio
async def test_default_internal_user_params_with_get_user_object(monkeypatch):
"""Test that default_internal_user_params is used when creating a new user via get_user_object"""

View file

@ -3786,3 +3786,128 @@ async def test_builder_succeeds_when_db_lookup_returns_valid_token():
# Reaching the success-assembly return (never the exception handler)
# proves a valid key is unaffected by the 503 conversion.
mock_return.assert_awaited_once()
def _mint_cli_session_token(monkeypatch, *, user_id="cli-admin"):
"""Mint a CLI session token for a PROXY_ADMIN user so auth resolves on the
admin early-return path (no prisma/common_checks needed)."""
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test")
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
user_info = LiteLLM_UserTable(
user_id=user_id,
user_email="cli@example.com",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
models=["gpt-3.5-turbo"],
max_budget=100.0,
)
return ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info, team_id="cli-team", team_alias="cli-team-alias"
)
@pytest.mark.asyncio
async def test_cli_session_token_authenticates_without_experimental_flag(monkeypatch):
"""A lite login token (encrypted non-sk blob) must authenticate on the LLM
hot path even when EXPERIMENTAL_UI_LOGIN is unset. Before the fix the decrypt
branch was gated behind that flag, so this would 401 on default deployments."""
monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False)
cli_token = _mint_cli_session_token(monkeypatch)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {cli_token}"}
mock_request.query_params = {}
with (
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
result = await user_api_key_auth(
request=mock_request,
api_key=f"Bearer {cli_token}",
)
assert result.user_id == "cli-admin"
assert result.team_id == "cli-team"
assert result.token is not None and result.token.startswith("cli-session-")
@pytest.mark.asyncio
async def test_random_non_sk_token_is_rejected(monkeypatch):
"""Decryption fails closed: a random non-sk string is not a valid blob, so it
must fall through to the 'expected sk-' 401 rather than being silently
accepted. A non-None prisma is used so the no-db short-circuit is skipped and
the real rejection path (before any DB lookup) is exercised."""
monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False)
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test")
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.headers = {"authorization": "Bearer not-a-real-token"}
mock_request.query_params = {}
with (
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
):
with pytest.raises(Exception) as exc_info:
await user_api_key_auth(
request=mock_request,
api_key="Bearer not-a-real-token",
)
message = str(getattr(exc_info.value, "message", exc_info.value))
assert int(getattr(exc_info.value, "code", 0)) == status.HTTP_401_UNAUTHORIZED
assert "sk-" in message
@pytest.mark.asyncio
async def test_expired_cli_session_token_is_rejected(monkeypatch):
"""An expired CLI session token must 401 with expired_key. Expiry is enforced
on the shared validation path, not only for DB-backed keys."""
monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False)
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test")
monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "-1")
import importlib
from litellm import constants
from litellm.proxy.auth import auth_checks
importlib.reload(constants)
importlib.reload(auth_checks)
user_info = LiteLLM_UserTable(
user_id="cli-admin",
user_email="cli@example.com",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
models=["gpt-3.5-turbo"],
max_budget=100.0,
)
cli_token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {cli_token}"}
mock_request.query_params = {}
try:
with (
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.prisma_client", None),
):
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(
request=mock_request,
api_key=f"Bearer {cli_token}",
)
assert exc_info.value.type == ProxyErrorTypes.expired_key
finally:
monkeypatch.delenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", raising=False)
importlib.reload(constants)
importlib.reload(auth_checks)