fix: req changes by greptile

This commit is contained in:
Harshit28j 2026-03-22 06:56:49 +05:30
parent d8a47ebb69
commit b479aa0014
3 changed files with 26 additions and 7 deletions

View file

@ -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 []):

View file

@ -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:

View file

@ -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(