scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix: populate team_member_rpm/tpm_limit for Virtual Key auth flow

Browse source 
JWT auth flow already populates team_member_rpm_limit and
team_member_tpm_limit via JWTAuthManager.auth_builder(), but the
Virtual Key auth flow never queried team membership, so per-member
rate limits were silently ignored for Virtual Key users.

This fix queries team membership after token validation in the Virtual
Key path and applies the same safe_get_team_member_rpm/tpm_limit()
logic that JWT auth uses.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
wufg2002 2026-03-19 20:05:44 +08:00
parent 3093ef844e
commit c122c6c18a
1 changed files with 34 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/auth/user_api_key_auth.py
View file

@ -39,6 +39,7 @@ from litellm.proxy.auth.auth_checks import (
get_jwt_key_mapping_object,
get_key_object,
get_project_object,
get_team_membership,
get_team_object,
get_user_object,
is_valid_fallback_model,
@ -1168,6 +1169,39 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
valid_token
) # updating it here, allows all downstream reporting / checks to use the updated budget
# Populate team_member_rpm_limit / team_member_tpm_limit for Virtual Key auth.
# JWT auth flow already gets this from JWTAuthManager.auth_builder(),
# but Virtual Key flow never queries team membership. Fix that here.
if (
valid_token.user_id is not None
and valid_token.team_id is not None
and valid_token.team_member_rpm_limit is None
and valid_token.team_member_tpm_limit is None
):
try:
_team_membership = await get_team_membership(
user_id=valid_token.user_id,
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if _team_membership is not None:
valid_token.team_member_rpm_limit = (
_team_membership.safe_get_team_member_rpm_limit()
)
valid_token.team_member_tpm_limit = (
_team_membership.safe_get_team_member_tpm_limit()
)
except Exception as e:
verbose_logger.debug(
"user_api_key_auth: Unable to get team membership for "
"user_id={}, team_id={}. Error={}".format(
valid_token.user_id, valid_token.team_id, str(e)
)
)
user_obj: Optional[LiteLLM_UserTable] = None
valid_token_dict: dict = {}
if valid_token is not None: