fix(auth): join team-member budget so rpm/tpm limits are enforced

The combined_view query joined LiteLLM_TeamMembership but only selected
tm.spend, never the team-member-budget's rpm/tpm. As a result
LiteLLM_VerificationTokenView.team_member_rpm_limit / _tpm_limit always
came back None, the rate limiter's `is not None` guard skipped the
team_member descriptor, and the configured limits were silently ignored.

Add the join through tm.budget_id and select the two limits under the
field names the view already exposes. Verified end-to-end by running
the query against a real Postgres with seeded data.

Fixes #26378
This commit is contained in:
Kiyeon Jeon 2026-04-25 18:40:35 +09:00
parent 0beec45c13
commit 9f1e970a68
2 changed files with 85 additions and 0 deletions

View file

@ -3086,6 +3086,8 @@ class PrismaClient:
t.organization_id as org_id,
p.project_alias AS project_alias,
tm.spend AS team_member_spend,
tmb.tpm_limit AS team_member_tpm_limit,
tmb.rpm_limit AS team_member_rpm_limit,
m.aliases AS team_model_aliases,
-- Added comma to separate b.* columns
b.max_budget AS litellm_budget_table_max_budget,
@ -3101,6 +3103,7 @@ class PrismaClient:
FROM "LiteLLM_VerificationToken" AS v
LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id
LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id
LEFT JOIN "LiteLLM_BudgetTable" AS tmb ON tm.budget_id = tmb.budget_id
LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id
LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id
LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id

View file

@ -0,0 +1,82 @@
"""
Verifies that the combined-view query result names align with what
LiteLLM_VerificationTokenView and UserAPIKeyAuth read for team member
rate limits.
The actual SQL query in litellm/proxy/utils.py was verified manually
against a real Postgres (see PR description). This module guards the
Python-side contract: if the SQL aliases drift away from these field
names, the rate limit will silently stop being enforced again the
exact failure mode of https://github.com/BerriAI/litellm/issues/26378.
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy._types import (
LiteLLM_VerificationTokenView,
UserAPIKeyAuth,
)
def test_view_carries_team_member_rate_limits_from_combined_view_row():
"""
Construct the view from the same kwargs the SQL row delivers.
A row with team_member_tpm_limit / team_member_rpm_limit columns
must populate the matching fields on the view.
"""
row = {
"token": "hashed-test-token",
"user_id": "user-1",
"team_id": "team-1",
"team_member_spend": 0.5,
"team_member_tpm_limit": 100,
"team_member_rpm_limit": 3,
}
view = LiteLLM_VerificationTokenView(**row)
assert view.team_member_spend == 0.5
assert view.team_member_tpm_limit == 100
assert view.team_member_rpm_limit == 3
def test_user_api_key_auth_inherits_team_member_rate_limits_from_view():
"""
The auth flow converts the view to UserAPIKeyAuth via
`UserAPIKeyAuth(**view.model_dump(exclude_none=True))`. The two
team_member fields must round-trip so the rate limiter
(parallel_request_limiter_v3) sees them.
"""
view = LiteLLM_VerificationTokenView(
token="hashed-test-token",
user_id="user-1",
team_id="team-1",
team_member_tpm_limit=100,
team_member_rpm_limit=3,
)
auth = UserAPIKeyAuth(**view.model_dump(exclude_none=True))
assert auth.team_member_tpm_limit == 100
assert auth.team_member_rpm_limit == 3
def test_view_leaves_team_member_rate_limits_none_when_missing():
"""
LEFT JOIN against a membership without a budget row produces NULLs
for these columns. The view must accept that and leave the fields
None otherwise the rate limiter would mis-fire on members with
no per-member limits configured.
"""
view = LiteLLM_VerificationTokenView(
token="hashed-test-token",
user_id="user-1",
team_id="team-1",
team_member_spend=0.0,
)
assert view.team_member_tpm_limit is None
assert view.team_member_rpm_limit is None