mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
8bca05d311
commit
36c2b029ab
8 changed files with 491 additions and 19 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue