fix(key_generate): scope session-token team-key budget exemption to caller-supplied team_id (#29641)
Some checks failed
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Security / security (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled

#29612 exempts UI/CLI session tokens from the key budget ceiling when they
create a team key, keyed on data.team_id. That value is read after the
default_key_generate_params loop can populate team_id, so on deployments that
set default_key_generate_params.team_id a request the caller did not scope to a
team is treated as a team key and skips the ceiling. Capture _requested_team_id
before defaults run and key the exemption off it, mirroring how
_requested_max_budget is already captured. Requests the caller did not scope to a
team keep the ceiling.
This commit is contained in:
yuneng-jiang 2026-06-03 19:11:53 -07:00 • committed by GitHub
parent c7f1bcfd0d
commit e9417603a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 51 additions and 4 deletions

View file

@ -691,10 +691,12 @@ async def _common_key_generation_helper( # noqa: PLR0915
prisma_client=prisma_client,
)
# Capture the caller-supplied max_budget before any defaults or upperbound
# params can fill it, so the ceiling check only fires when the caller
# explicitly requested a budget.
# Capture caller-supplied max_budget and team_id before any defaults or
# upperbound params can fill them, so the ceiling check and its team-key
# exemption key off what the caller explicitly requested, not a value that
# default_key_generate_params injected.
_requested_max_budget = data.max_budget
_requested_team_id = data.team_id
# check if user set default key/generate params on config.yaml
if litellm.default_key_generate_params is not None:
@ -728,7 +730,7 @@ async def _common_key_generation_helper( # noqa: PLR0915
# at request time. Personal keys keep the ceiling; nothing else bounds them.
is_ui_session_team_key = (
user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID
and data.team_id is not None
and _requested_team_id is not None
)
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value

View file

@ -11623,3 +11623,48 @@ async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
@pytest.mark.asyncio
async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption():
"""
Security regression for GHSA-q775: the team-key exemption must key off the
team_id the CALLER supplied, not one injected by default_key_generate_params.
With default_key_generate_params.team_id set, a UI session token's personal-key
request (no team_id) would otherwise have team_id auto-filled before the ceiling
check, flipping is_ui_session_team_key to True and bypassing the ceiling. The
request must still be rejected. Mirrors how _requested_max_budget is captured
before defaults run.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
data = GenerateKeyRequest(max_budget=500)
assert data.team_id is None
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-ui-session",
user_id="user-1",
team_id=UI_SESSION_TOKEN_TEAM_ID,
max_budget=0.25,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", False),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
patch("litellm.default_key_generate_params", {"team_id": "injected-team"}),
):
with pytest.raises((HTTPException, ProxyException)) as exc_info:
await _common_key_generation_helper(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
team_table=None,
)
err = exc_info.value
code = getattr(err, "status_code", None) or getattr(err, "code", None)
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
assert str(code) == "400"
assert "cannot exceed" in msg.lower()