scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

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

Browse source 
* 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.

* fix(cli): address review feedback on EXPERIMENTAL_UI_LOGIN gate and e2e test

Restore EXPERIMENTAL_UI_LOGIN=false as an explicit opt-out: operators who set it to false keep the old boundary; unset (new default) and true both attempt NaCl decryption, which fails closed for non-blob tokens.

In the e2e test: replace the silent Redis fallback with pytest.skip so a missing Redis instance is explicit rather than silently degrading to a directly-minted token. Write the seeded flow back as JSON (proxy reads it via json.loads on cache fetch) instead of Python repr, and build the updated flow immutably.

* fix(key-management): cap CLI session token delegation budget to team ceiling

A CLI session token intentionally carries max_budget=None to avoid a per-session LLM spend cap. The key-generation delegation check (GHSA-q775-qw9r-2r4g) previously skipped non-admin callers with max_budget=None, treating them as having unlimited delegation authority. This allowed any internal user with a lite login session to mint virtual keys with arbitrary budgets.

Adds is_session_token=True to UserAPIKeyAuth for CLI session tokens and uses the caller's team budget as the delegation ceiling in that case, so the effective limit is min(requested_budget, team.max_budget) rather than unbounded.

* chore: regenerate dashboard OpenAPI types

The is_session_token field added to UserAPIKeyAuth cascades to the
dashboard schema. Regenerate types from the updated OpenAPI spec.

* fix(key-management): block personal key budget delegation from CLI session tokens

When team_table is None (personal key, no team_id in request), the personal key
has no team-budget enforcement at request time. A session token therefore cannot
delegate any explicit max_budget for a personal key -- that would open a budget
bypass path. Block the request with a clear 400 directing the caller to use a
team_id instead.

* test(auth): add unit coverage for non-admin CLI session token production path

* fix(type-check): use model_validate in _return_user_api_key_auth_obj to fix reportArgumentType gate

UserAPIKeyAuth(**user_api_key_kwargs) spread triggers a basedpyright
reportArgumentType error for each named field in UserAPIKeyAuth because
the dict's inferred value type (str | Span | LitellmUserRoles | Unknown)
is not assignable to each field's specific type. Adding is_session_token:
bool introduced +2 more such errors, breaching the gate cap.

model_validate accepts an untyped dict without per-field argument checking,
which eliminates the +2 new errors and also ratchets down the pre-existing
333 errors at those call sites. basedpyright-code-budget.json is updated
to reflect the new lower baseline (1814, down from 1934).

* fix(type-check): ratchet down reportArgumentType baseline only

The previous lint-budget-update captured all baselines from the local
environment, raising many ceilings vs the merge-base and failing the
non-gating budget_ratchet_check. Restore staging's values for every
rule and only lower reportArgumentType (1934 -> 1814) to reflect the
reduction from switching to model_validate in _return_user_api_key_auth_obj.

* fix(auth): set max_budget on CLI session token to enforce max_ui_session_budget

CLI session tokens were missing max_budget, so _virtual_key_max_budget_check
had no per-session ceiling to enforce. Operators relying on max_ui_session_budget
could be bypassed for the full token lifetime. Mirrors the existing UI token path.

* revert(auth): remove max_ui_session_budget from CLI session token

max_ui_session_budget defaults to $0.25 and is sized for the UI chat
pane (10-min sessions). CLI sessions are 24-hour tokens for real work;
capping them at that ceiling would throttle users under their actual
user/team budget. Budget enforcement for CLI sessions is via the shared
user and team counters as originally intended.

* fix(auth): cap CLI session at max_ui_session_budget only when user and team have no budget

When neither the user nor their team has a budget configured, CLI sessions
were fully uncapped. The poll endpoint now looks up the real user and team
objects from DB; if both have no max_budget, it passes litellm.max_ui_session_budget
as the token's per-key ceiling. Users or teams that already have a budget
configured are unaffected and continue to rely on the shared counters.

* fix(auth): fix black formatting and update test mock for cli_poll_key budget lookup

The get_user_object and get_team_object async calls in cli_poll_key were
not mocked in the existing test, causing MagicMock await errors. Patch
both functions at the auth_checks module level. Also apply black formatting
to ui_sso.py which CI rejected.

* fix(auth): skip fallback budget cap when team lookup fails for cli session token

* test(auth): pin cli session budget cap to user/team budget presence

The session_max_budget fallback in cli_poll_key only applied
max_ui_session_budget when neither the user nor the resolved team had a
budget. The existing coverage exercised only the team-lookup-failure
branch. Add two regression tests: a user with a configured budget must
not receive the fallback cap, and a session with no user and no team
budget must fall back to max_ui_session_budget. Mutating either guard
out of the branch now fails these tests.

* fix: remove CLI poll session budget cap

* revert(auth): restore CLI session fallback budget cap

Bugbot autofix (60b81fb8) removed the user/team budget lookup in
cli_poll_key and stopped passing max_budget to the session token,
making CLI sessions fully uncapped whenever neither the user nor the
team has an explicit budget.

That reintroduces the unbounded-spend bypass veria flagged as High
("CLI session budget bypass"): on deployments that rely on
max_ui_session_budget rather than per-user/team budgets, a completed
lite login could run LLM calls with no ceiling for the whole token
lifetime. The fallback only applies when no other budget bounds the
session, so users and teams with a configured budget are unaffected and
keep relying on their shared counters.

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute 2026-06-26 21:35:15 +05:30 • committed by GitHub
parent 2b496bc7f7
commit 687a62e561
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 887 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -4,7 +4,7 @@
"slack": 2500
},
"reportArgumentType": {
"baseline": 1934,
"baseline": 1814,
"slack": 180
},
"reportAssignmentType": {

2
litellm/constants.py
View file

@ -1532,7 +1532,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")

1
litellm/proxy/_types.py
View file

@ -2527,6 +2527,7 @@ class UserAPIKeyAuth(
user_spend: Optional[float] = None
user_max_budget: Optional[float] = None
request_route: Optional[str] = None
is_session_token: bool = False
budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
user: Optional[Any] = None # Expanded user object when expand=user is used
created_by_user: Optional[Any] = (

16
litellm/proxy/auth/auth_checks.py
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,
@ -2417,6 +2417,7 @@ class ExperimentalUIJWTToken:
user_info: LiteLLM_UserTable,
team_id: Optional[str] = None,
team_alias: Optional[str] = None,
max_budget: Optional[float] = None,
) -> str:
"""
Generate a JWT token for CLI authentication with configurable expiration.
@ -2432,6 +2433,7 @@ class ExperimentalUIJWTToken:
Returns:
Encrypted JWT token string
"""
import secrets
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
@ -2453,18 +2455,22 @@ class ExperimentalUIJWTToken:
# Use first team if user has teams
_team_id = user_info.teams[0] if len(user_info.teams) > 0 else None
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,
max_budget=max_budget,
user_id=user_info.user_id,
team_id=_team_id,
team_alias=team_alias,
models=user_info.models,
max_parallel_requests=None,
user_role=LitellmUserRoles(user_info.user_role),
is_session_token=True,
)
return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True))

16
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1513,8 +1513,16 @@ 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
# Attempt decryption for non-sk- tokens unless the operator has
# explicitly set EXPERIMENTAL_UI_LOGIN=false to disable it.
# Unset (None) keeps the new default of always attempting decryption;
# decryption fails closed for anything that is not a genuine blob.
if (
valid_token is None
and not api_key.startswith("sk-")
and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False
):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(
api_key
)
@ -2726,9 +2734,9 @@ async def _return_user_api_key_auth_obj(
user_api_key_kwargs.update(
user_role=LitellmUserRoles.PROXY_ADMIN,
)
return UserAPIKeyAuth(**user_api_key_kwargs)
return UserAPIKeyAuth.model_validate(user_api_key_kwargs)
else:
return UserAPIKeyAuth(**user_api_key_kwargs)
return UserAPIKeyAuth.model_validate(user_api_key_kwargs)
def get_api_key_from_custom_header(

2
litellm/proxy/client/README.md
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.

6
litellm/proxy/client/cli/README.md
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:

42
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -743,29 +743,53 @@ async def _common_key_generation_helper(
_enforce_upperbound_key_params(data, fill_defaults=True)
# Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller
# with an explicit budget cannot grant a key a higher budget than their own.
# Callers with max_budget=None (unlimited) can delegate any budget.
# A UI/CLI session token's max_budget is a per-session chat spend cap
# (max_ui_session_budget), not a delegation authority, so it is exempt only
# when creating a team key - that key's spend is bounded by the team budget
# at request time. Personal keys keep the ceiling; nothing else bounds them.
# cannot grant a key a higher budget than their own authority.
is_ui_session_team_key = (
user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID
and _requested_team_id is not None
)
# Session tokens (lite login) carry max_budget=None to avoid a per-session
# LLM spend cap, but that None must not be read as "unlimited delegation
# authority". A personal key (no team) has no team-budget enforcement at
# request time, so a session token cannot delegate any budget for one.
if (
user_api_key_dict.is_session_token
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not is_ui_session_team_key
and _requested_max_budget is not None
and team_table is None
):
raise HTTPException(
status_code=400,
detail={
"error": (
f"max_budget ({_requested_max_budget}) cannot be set without "
"specifying team_id when using a CLI session token."
)
},
)
delegation_ceiling = (
user_api_key_dict.max_budget
if user_api_key_dict.max_budget is not None
else (
team_table.max_budget
if user_api_key_dict.is_session_token and team_table is not None
else None
)
)
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not is_ui_session_team_key
and _requested_max_budget is not None
and user_api_key_dict.max_budget is not None
and _requested_max_budget > user_api_key_dict.max_budget
and delegation_ceiling is not None
and _requested_max_budget > delegation_ceiling
):
raise HTTPException(
status_code=400,
detail={
"error": (
f"max_budget ({_requested_max_budget}) cannot exceed the caller's "
f"own max_budget ({user_api_key_dict.max_budget})."
f"own max_budget ({delegation_ceiling})."
)
},
)

46
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -2243,8 +2243,12 @@ async def cli_poll_key(
key_id: The CLI login session ID
team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams.
"""
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
get_team_object,
get_user_object,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache)
@ -2320,18 +2324,46 @@ async def cli_poll_key(
None,
)
# Create user object for JWT generation
user_info = LiteLLM_UserTable(
user_id=user_id,
user_role=session_data["user_role"],
models=session_data.get("models", []),
max_budget=litellm.max_ui_session_budget,
)
# Generate CLI JWT on-demand (expiration configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS)
# Pass selected team_id to ensure JWT has correct team
user_db_obj = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
user_budget = user_db_obj.max_budget if user_db_obj is not None else None
team_budget: Optional[float] = None
team_budget_resolved = False
if team_id is not None:
try:
team_obj = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
team_budget = team_obj.max_budget
team_budget_resolved = True
except Exception:
pass
session_max_budget = (
litellm.max_ui_session_budget
if user_budget is None
and (team_id is None or (team_budget_resolved and team_budget is None))
else None
)
jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=user_info, team_id=team_id, team_alias=team_alias
user_info=user_info,
team_id=team_id,
team_alias=team_alias,
max_budget=session_max_budget,
)
# Delete cache entry (single-use)

285
tests/otel_tests/test_e2e_budgeting.py
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,216 @@ 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()
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:
pytest.skip("Shared Redis not available; skipping full poll-flow test")
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 json
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
updated_flow = {
**flow,
"sso_complete": True,
"user_code_verified": False,
"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}],
},
"browser_complete_token_hash": _hash_cli_sso_secret(browser_complete_token),
}
client.setex(cache_key, 600, json.dumps(updated_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 +548,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():
"""

56
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -459,7 +459,12 @@ 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
# is_session_token=True causes key_management_endpoints to use the team
# budget as the delegation ceiling instead of treating None as unlimited.
assert token_data.get("is_session_token") is True
# Verify expiration time is set to 24 hours (default)
assert "expires" in token_data
@ -504,6 +509,55 @@ 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
def test_get_cli_jwt_auth_token_applies_fallback_budget(valid_sso_user_defined_values):
token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
valid_sso_user_defined_values, max_budget=litellm.max_ui_session_budget
)
decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug")
assert decrypted is not None
assert json.loads(decrypted).get("max_budget") == litellm.max_ui_session_budget
def test_get_cli_jwt_auth_token_no_fallback_when_budget_provided(
valid_sso_user_defined_values,
):
token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
valid_sso_user_defined_values, max_budget=None
)
decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug")
assert decrypted is not None
assert json.loads(decrypted).get("max_budget") is None
@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"""

203
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -3786,3 +3786,206 @@ 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)
@pytest.mark.asyncio
async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypatch):
"""A CLI session token minted for a non-admin (INTERNAL_USER) must flow
through the production auth path, not the admin early-return. The key
regression: without the fix, valid_token is None after the decrypt block
(gated behind EXPERIMENTAL_UI_LOGIN), so the builder raises 401 at the
sk- guard. With the fix, valid_token is set, the sk- guard is skipped, and
_return_user_api_key_auth_obj is called with the correct identity."""
monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False)
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test")
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
user_info = LiteLLM_UserTable(
user_id="internal-user-1",
user_email="user@example.com",
user_role=LitellmUserRoles.INTERNAL_USER.value,
models=[],
)
cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info, team_id="team-abc", team_alias="my-team"
)
import litellm.proxy.proxy_server as _proxy_server_mod
from fastapi import Request
from starlette.datastructures import URL
assembled = UserAPIKeyAuth(
user_id="internal-user-1",
team_id="team-abc",
is_session_token=True,
)
attrs = _proxy_attrs_for_db_lookup()
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
with (
patch(
"litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj",
new_callable=AsyncMock,
return_value=assembled,
) as mock_assemble,
patch(
"litellm.proxy.auth.user_api_key_auth.get_user_object",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.auth.user_api_key_auth.get_team_object",
new_callable=AsyncMock,
side_effect=__import__("fastapi").HTTPException(status_code=404),
),
):
result = await _user_api_key_auth_builder(
request=request,
api_key=f"Bearer {cli_token}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
mock_assemble.assert_awaited_once()
call_kwargs = mock_assemble.call_args.kwargs
assert call_kwargs["valid_token_dict"]["user_id"] == "internal-user-1"
assert call_kwargs["valid_token_dict"]["team_id"] == "team-abc"
assert call_kwargs["valid_token_dict"]["is_session_token"] is True
assert call_kwargs["valid_token_dict"]["user_role"] == LitellmUserRoles.INTERNAL_USER
assert result.is_session_token is True

107
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -12610,3 +12610,110 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
)
assert kwargs["use_substring_matching"] is False
assert kwargs["user_id"] == "alice"
@pytest.mark.asyncio
async def test_cli_session_token_delegation_ceiling_blocked_by_team_budget():
team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
team_id="team-1",
is_session_token=True,
)
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=GenerateKeyRequest(max_budget=1000.0),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=team,
)
assert exc_info.value.status_code == 400
assert "max_budget" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_cli_session_token_delegation_allowed_within_team_budget():
team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
team_id="team-1",
is_session_token=True,
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
):
result = await _common_key_generation_helper(
data=GenerateKeyRequest(max_budget=25.0),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=team,
)
assert result is not None
@pytest.mark.asyncio
async def test_regular_unlimited_user_delegation_ceiling_not_applied():
team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
team_id="team-1",
is_session_token=False,
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
):
result = await _common_key_generation_helper(
data=GenerateKeyRequest(max_budget=1000.0),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=team,
)
assert result is not None
@pytest.mark.asyncio
async def test_cli_session_token_personal_key_with_budget_blocked():
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
team_id="team-1",
is_session_token=True,
)
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=GenerateKeyRequest(max_budget=1000.0),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None, # no team in request = personal key
)
assert exc_info.value.status_code == 400
assert "team_id" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_cli_session_token_personal_key_without_budget_allowed():
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
team_id="team-1",
is_session_token=True,
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
):
result = await _common_key_generation_helper(
data=GenerateKeyRequest(max_budget=None),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None, # no team in request = personal key, but no explicit budget
)
assert result is not None

142
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -2870,17 +2870,20 @@ class TestCLIKeyRegenerationFlow:
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
return_value=mock_jwt_token,
) as mock_get_jwt,
patch(
"litellm.proxy.auth.auth_checks.get_user_object",
new=AsyncMock(return_value=mock_user_info),
),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(side_effect=Exception("no team")),
),
):
# Mock the user lookup
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=mock_user_info
)
# Act - Second poll with team_id
result = await cli_poll_key(
key_id=session_key,
@ -2895,15 +2898,140 @@ class TestCLIKeyRegenerationFlow:
assert result["team_id"] == selected_team
assert result["teams"] == ["team-a", "team-b", "team-c"]
# Verify JWT was generated with correct team
# Verify JWT was generated with correct team and no budget cap
# (team lookup failed, but team_id is set, so fallback cap must not apply)
mock_get_jwt.assert_called_once()
jwt_call_args = mock_get_jwt.call_args
assert jwt_call_args.kwargs["team_id"] == selected_team
assert jwt_call_args.kwargs["team_alias"] == "Team B"
assert jwt_call_args.kwargs["max_budget"] is None
# Verify session was deleted after JWT generation
mock_cache.delete_cache.assert_called_once()
@pytest.mark.asyncio
async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self):
"""A user with a configured budget must not get the max_ui_session_budget fallback cap."""
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
session_data = {
"user_id": "budgeted-user",
"user_role": "internal_user",
"teams": [],
"team_details": [],
"models": ["gpt-4"],
"user_email": "budgeted@example.com",
}
mock_user_info = LiteLLM_UserTable(
user_id="budgeted-user",
user_role="internal_user",
teams=[],
models=["gpt-4"],
max_budget=100.0,
)
mock_cache = MagicMock()
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": True,
"session_data": session_data,
}
mock_jwt_token = "eyJhbGciOiJIUzI1NiJ9.budgeted.token"
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
return_value=mock_jwt_token,
) as mock_get_jwt,
patch(
"litellm.proxy.auth.auth_checks.get_user_object",
new=AsyncMock(return_value=mock_user_info),
),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(
side_effect=AssertionError("team lookup must be skipped")
),
),
):
result = await cli_poll_key(
key_id="cli-session-budgeted",
team_id=None,
x_litellm_cli_poll_secret="poll-secret",
)
assert result["status"] == "ready"
mock_get_jwt.assert_called_once()
assert mock_get_jwt.call_args.kwargs["max_budget"] is None
@pytest.mark.asyncio
async def test_cli_poll_key_caps_session_when_user_and_team_have_no_budget(self):
"""With no user and no team budget, the session falls back to max_ui_session_budget."""
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
from litellm.proxy.management_endpoints.ui_sso import (
_hash_cli_sso_secret,
cli_poll_key,
)
session_data = {
"user_id": "unbudgeted-user",
"user_role": "internal_user",
"teams": ["team-x"],
"team_details": [{"team_id": "team-x", "team_alias": "Team X"}],
"models": ["gpt-4"],
"user_email": "unbudgeted@example.com",
}
mock_user_info = LiteLLM_UserTable(
user_id="unbudgeted-user",
user_role="internal_user",
teams=["team-x"],
models=["gpt-4"],
max_budget=None,
)
mock_team = LiteLLM_TeamTableCachedObj(team_id="team-x", max_budget=None)
mock_cache = MagicMock()
mock_cache.get_cache.return_value = {
"poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
"sso_complete": True,
"user_code_verified": True,
"session_data": session_data,
}
mock_jwt_token = "eyJhbGciOiJIUzI1NiJ9.unbudgeted.token"
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
patch("litellm.proxy.proxy_server.prisma_client"),
patch(
"litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token",
return_value=mock_jwt_token,
) as mock_get_jwt,
patch(
"litellm.proxy.auth.auth_checks.get_user_object",
new=AsyncMock(return_value=mock_user_info),
),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(return_value=mock_team),
),
):
result = await cli_poll_key(
key_id="cli-session-unbudgeted",
team_id="team-x",
x_litellm_cli_poll_secret="poll-secret",
)
assert result["status"] == "ready"
mock_get_jwt.assert_called_once()
assert (
mock_get_jwt.call_args.kwargs["max_budget"] == litellm.max_ui_session_budget
)
class TestGetAppRolesFromIdToken:
"""Test the get_app_roles_from_id_token method"""

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -32300,6 +32300,11 @@ export interface components {
end_user_tpm_limit?: number | null;
/** Expires */
expires?: string | null;
/**
* Is Session Token
* @default false
*/
is_session_token: boolean;
/** Jwt Claims */
jwt_claims?: {
[key: string]: unknown;