From b479aa0014a609e6923de42308d4c204161e7a60 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 22 Mar 2026 06:56:49 +0530 Subject: [PATCH] fix: req changes by greptile --- .../key_management_endpoints.py | 3 +++ .../management_endpoints/team_endpoints.py | 20 ++++++++++++++++--- litellm/proxy/management_helpers/utils.py | 10 ++++++---- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c3b957d14a2..d493d2bf7c8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -667,6 +667,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 team_pool=team_pool, ) + # When effective_models is non-empty, enforce per-member restrictions. + # When empty (team.models=[] = "allow all" and no overrides configured), + # skip validation — any model is permitted, matching runtime auth behavior. if effective_models: # if 'all-team-models' was requested, restrict it to the effective models if "all-team-models" in (data.models or []): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 64c372c6eaa..3e7fd9b20ec 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2541,6 +2541,7 @@ async def team_member_update( # noqa: PLR0915 # (it was just written to the DB by _upsert_budget_and_membership). # Only fetch from DB when models weren't provided (e.g., role-only update). stored_models = data.models if data.models is not None else [] + _tm_row = None if data.models is None: from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.proxy_server import user_api_key_cache @@ -2557,11 +2558,24 @@ async def team_member_update( # noqa: PLR0915 ) stored_models = (_tm_row.models or []) if _tm_row is not None else [] else: - # Models were explicitly changed — invalidate cache so key-gen sees fresh data + # Models were explicitly changed — invalidate cache and fetch fresh membership + from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.proxy_server import user_api_key_cache _cache_key = f"team_membership:{received_user_id}:{data.team_id}" await user_api_key_cache.async_delete_cache(key=_cache_key) + _tm_row = await get_team_membership( + user_id=received_user_id, + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + # Resolve tpm/rpm: prefer request values, then budget table (authoritative), + # then existing JSON blob (may be stale for pre-deployment members). + _budget = _tm_row.litellm_budget_table if (_tm_row and hasattr(_tm_row, "litellm_budget_table") and _tm_row.litellm_budget_table) else None + resolved_tpm = data.tpm_limit if data.tpm_limit is not None else (getattr(_budget, "tpm_limit", None) if _budget else None) + resolved_rpm = data.rpm_limit if data.rpm_limit is not None else (getattr(_budget, "rpm_limit", None) if _budget else None) if data.role is not None or data.models is not None: team_members: List[Member] = [] @@ -2573,8 +2587,8 @@ async def team_member_update( # noqa: PLR0915 role=data.role or member.role, user_email=data.user_email or member.user_email, models=stored_models, - tpm_limit=data.tpm_limit if data.tpm_limit is not None else getattr(member, "tpm_limit", None), - rpm_limit=data.rpm_limit if data.rpm_limit is not None else getattr(member, "rpm_limit", None), + tpm_limit=resolved_tpm, + rpm_limit=resolved_rpm, ) ) else: diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 1580422c550..f92087555e3 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -234,8 +234,10 @@ async def add_new_member( # noqa: PLR0915 else: _budget_id = default_team_budget_id + # Create membership row when budget, or models is explicitly provided (even []). + # models=None means "not specified" (skip), models=[] means "explicitly no overrides". if ( - (_budget_id or new_member.models) + (_budget_id or new_member.models is not None) and returned_user is not None and returned_user.user_id is not None ): @@ -245,9 +247,9 @@ async def add_new_member( # noqa: PLR0915 } if _budget_id: membership_create_data["budget_id"] = _budget_id - if new_member.models: - # Defense-in-depth: validate member models are within team models - if team_models: + if new_member.models is not None: + # Defense-in-depth: validate non-empty member models are within team models + if new_member.models and team_models: disallowed = set(new_member.models) - set(team_models) if disallowed: raise HTTPException(