From a5ed93e8b74bc1251be0cdc519004dc448d3abfd Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sun, 1 Mar 2026 16:44:45 +0530 Subject: [PATCH] fix: budget disconnect bug and skip overrides for non-configured teams 1. Fix _upsert_budget_and_membership: when only models is passed (no budget params), the early-return was disconnecting the existing budget. Now models-only updates leave the budget untouched. 2. Fix override logic running unconditionally for all teams when feature flag is enabled. Teams that haven't configured default_models or member models now skip the override path entirely, preserving existing team_object.models / access_group behavior. Only teams with actual overrides configured enter the effective models computation. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 43 ++++++++++--------- .../management_endpoints/common_utils.py | 12 ++++-- .../key_management_endpoints.py | 9 +++- 3 files changed, 39 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a8955f46d15..2bcb6dcb59c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2582,28 +2582,31 @@ async def can_team_access_model( """ models_to_check: List[str] = team_object.models if team_object else [] if LITELLM_TEAM_MODEL_OVERRIDES and valid_token: - # Compute effective models: team defaults + per-user overrides - effective_models = compute_effective_team_models( - team_default_models=valid_token.team_default_models, - team_member_models=valid_token.team_member_models, - ) - - if effective_models: - models_to_check = effective_models - elif team_object and team_object.models: - # Fallback: effective models empty but team has models configured. - # Graceful degradation prevents misconfiguration cliff when feature - # is enabled without populating default_models. - models_to_check = team_object.models - else: - # Both effective models and team.models are empty — deny access - raise ProxyException( - message=f"Team not allowed to access model. No models available for user in this team. Model={model}.", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, + # Only apply override logic when overrides are actually configured. + # Teams that haven't set default_models or member models continue + # to use the original team_object.models / access_group path unchanged. + if valid_token.team_default_models or valid_token.team_member_models: + effective_models = compute_effective_team_models( + team_default_models=valid_token.team_default_models, + team_member_models=valid_token.team_member_models, ) + if effective_models: + models_to_check = effective_models + elif team_object and team_object.models: + # Fallback: effective models empty but team has models configured. + # Graceful degradation prevents misconfiguration cliff when feature + # is enabled without populating default_models. + models_to_check = team_object.models + else: + # Both effective models and team.models are empty — deny access + raise ProxyException( + message=f"Team not allowed to access model. No models available for user in this team. Model={model}.", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + try: return _can_object_call_model( model=model, diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 03043c5e79d..29b5b31a486 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -331,14 +331,18 @@ async def _upsert_budget_and_membership( If any of these values exist, a budget is updated or created and linked to the team membership. """ if max_budget is None and tpm_limit is None and rpm_limit is None: - # disconnect the budget since all limits are None - update_data: Dict[str, Any] = {"litellm_budget_table": {"disconnect": True}} if models is not None: - update_data["models"] = models + # Only updating models — do not touch the existing budget + await tx.litellm_teammembership.update( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + data={"models": models}, + ) + return + # No budget params and no models — disconnect the budget await tx.litellm_teammembership.update( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - data=update_data, + data={"litellm_budget_table": {"disconnect": True}}, ) return diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f7ccf0fc202..fc7e67e03e5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -917,11 +917,18 @@ async def _validate_key_models_against_effective_team_models( if membership: member_models = membership.models or [] + team_default_models = team_table.default_models or [] + + # Skip override validation when neither default_models nor member models + # are configured. Teams not using overrides continue with existing behavior. + if not team_default_models and not member_models: + return + # 2. Compute effective models from litellm.proxy.auth.auth_checks import compute_effective_team_models effective_models = compute_effective_team_models( - team_default_models=team_table.default_models or [], + team_default_models=team_default_models, team_member_models=member_models, )