From 687a62e5612bd177a96f7aea13bb67fee76927e6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 26 Jun 2026 21:35:15 +0530 Subject: [PATCH] fix(cli): mint per-session agent credential on lite login (#31072) * 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 --- basedpyright-code-budget.json | 2 +- litellm/constants.py | 2 +- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 16 +- litellm/proxy/auth/user_api_key_auth.py | 16 +- litellm/proxy/client/README.md | 2 + litellm/proxy/client/cli/README.md | 6 + .../key_management_endpoints.py | 42 ++- litellm/proxy/management_endpoints/ui_sso.py | 46 ++- tests/otel_tests/test_e2e_budgeting.py | 285 +++++++++++++++++- .../proxy/auth/test_auth_checks.py | 56 +++- .../proxy/auth/test_user_api_key_auth.py | 203 +++++++++++++ .../test_key_management_endpoints.py | 107 +++++++ .../proxy/management_endpoints/test_ui_sso.py | 142 ++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 15 files changed, 887 insertions(+), 44 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 1af0148e452..f2b54e1f889 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -4,7 +4,7 @@ "slack": 2500 }, "reportArgumentType": { - "baseline": 1934, + "baseline": 1814, "slack": 180 }, "reportAssignmentType": { diff --git a/litellm/constants.py b/litellm/constants.py index d2f2e89eca3..d137c639778 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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") diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index be5d7a2db78..ed030c31b19 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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] = ( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 88db2a2b7ea..a1c70d78902 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e439f6a5998..2dcc7bdbd19 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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( diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index c2ce28884c7..f33367a96c2 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -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. diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 333e2029e46..f53e7db4e6b 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -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: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9f316256ac9..67f9103d4b7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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})." ) }, ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 480861fb517..dc128e88d51 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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) diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py index f61befac4fb..44542558002 100644 --- a/tests/otel_tests/test_e2e_budgeting.py +++ b/tests/otel_tests/test_e2e_budgeting.py @@ -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(): """ diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index c8dc0ea5ed6..f56e309a552 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 7219ab58799..f80a8b28cc3 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 97397fb06be..a92eaa3a5f6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 44abc7acf21..976048b9521 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b53acf930f2..cdc91371dfe 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;