From c122c6c18afbf10a308da199212b26caa9a09b79 Mon Sep 17 00:00:00 2001 From: wufg2002 <285603640@qq.com> Date: Thu, 19 Mar 2026 20:05:44 +0800 Subject: [PATCH] fix: populate team_member_rpm/tpm_limit for Virtual Key auth flow 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 --- litellm/proxy/auth/user_api_key_auth.py | 34 +++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index eb6a5bdb994..44072136ab1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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: