mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
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:
parent
0beec45c13
commit
9f1e970a68
2 changed files with 85 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue