From 46ff6bf654942117afd911909e855bc62869ac25 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 23 Mar 2026 17:33:07 +0530 Subject: [PATCH] fix(routing): prevent stale model_aliases from interfering with team routing - Skip model_aliases rewrite if model resolves to team deployments - Add test coverage for sibling-preservation branch - Update MockPrismaClient to support sibling deployment scenarios Made-with: Cursor --- litellm/proxy/litellm_pre_call_utils.py | 57 ++------------- .../test_model_management_endpoints.py | 70 ++++++++++++++++++- 2 files changed, 75 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ba9577f35d7..1a7bbe0474b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1,7 +1,6 @@ import asyncio import copy import time -from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from fastapi import Request @@ -27,7 +26,6 @@ _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) from litellm.router import Router -from litellm.secret_managers.main import get_secret_bool from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -38,11 +36,6 @@ from litellm.types.utils import ( ) service_logger_obj = ServiceLogging() # used for tracking latency on OTEL -# Bounded dedup for stale-alias warnings (FIFO eviction when over cap). -_MAX_STALE_ALIAS_WARNING_KEYS = 10_000 -_STALE_TEAM_ALIAS_WARNING_KEYS: OrderedDict[str, None] = OrderedDict() -# Cache the stale alias bypass flag at module load to avoid hot-path secret lookups -_ENABLE_TEAM_STALE_ALIAS_BYPASS: Optional[bool] = None if TYPE_CHECKING: @@ -1318,50 +1311,14 @@ def _update_model_if_team_alias_exists( # Skip alias rewrite if this model resolves to team-specific deployments # (team models use team_public_model_name, not model_aliases) - aliased_target = user_api_key_dict.team_model_aliases[_model] + if ( + llm_router + and user_api_key_dict.team_id + and llm_router.map_team_model(_model, user_api_key_dict.team_id) is not None + ): + return - # Optional bypass for stale aliases from pre-PR deployments: - # only enabled via feature flag to preserve backwards compatibility. - # Cached at module level to avoid hot-path secret lookups on every request. - global _ENABLE_TEAM_STALE_ALIAS_BYPASS - if _ENABLE_TEAM_STALE_ALIAS_BYPASS is None: - _ENABLE_TEAM_STALE_ALIAS_BYPASS = get_secret_bool( - "LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", False - ) - enable_stale_alias_bypass = _ENABLE_TEAM_STALE_ALIAS_BYPASS - # Check if the alias points to a team-scoped UUID name - # (format: "model_name_{team_id}_{uuid}") - is_stale_team_alias = aliased_target.startswith( - f"model_name_{user_api_key_dict.team_id}_" - ) - if is_stale_team_alias and llm_router: - # This is a stale alias from pre-PR deployments. - # Check if current team deployments exist for the public name. - key = (user_api_key_dict.team_id, _model) - if key in llm_router.team_model_to_deployment_indices: - if enable_stale_alias_bypass: - # Team deployments exist; skip stale alias - return - warning_key = f"{user_api_key_dict.team_id}:{_model}:{aliased_target}" - if warning_key not in _STALE_TEAM_ALIAS_WARNING_KEYS: - _STALE_TEAM_ALIAS_WARNING_KEYS[warning_key] = None - while ( - len(_STALE_TEAM_ALIAS_WARNING_KEYS) - > _MAX_STALE_ALIAS_WARNING_KEYS - ): - _STALE_TEAM_ALIAS_WARNING_KEYS.popitem(last=False) - verbose_proxy_logger.warning( - "Stale team model alias detected for model='%s', team_id='%s'. " - "New sibling deployments may be unreachable. " - "Set LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS=true to enable " - "team-scoped sibling routing.", - str(_model).replace("\n", "").replace("\r", ""), - str(user_api_key_dict.team_id) - .replace("\n", "") - .replace("\r", ""), - ) - - data["model"] = aliased_target + data["model"] = user_api_key_dict.team_model_aliases[_model] return diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index fd4f3d56b10..dcfd5847bd5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -28,9 +28,15 @@ from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment class MockPrismaClient: - def __init__(self, team_exists: bool = True, user_admin: bool = True): + def __init__( + self, + team_exists: bool = True, + user_admin: bool = True, + sibling_deployments: list = None, + ): self.team_exists = team_exists self.user_admin = user_admin + self.sibling_deployments = sibling_deployments or [] self.db = self async def find_unique(self, where): @@ -47,7 +53,7 @@ class MockPrismaClient: return None async def find_many(self, where): - return [] + return self.sibling_deployments @property def litellm_teamtable(self): @@ -742,6 +748,66 @@ class TestTeamModelUpdate: # team_model_add must be called to add public name to team's models list mock_team_model_add.assert_called_once() + @pytest.mark.asyncio + async def test_rename_preserves_old_name_when_siblings_exist(self): + """Test that renaming a deployment preserves old public name when sibling deployments still use it""" + from unittest.mock import MagicMock + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + # Create a deployment being renamed + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo( + team_id="team_123", team_public_model_name="old-public-name" + ), + ) + + # Create a sibling deployment that still uses the old public name + sibling_deployment = MagicMock() + sibling_deployment.model_name = "model_name_team_123_uuid2" + sibling_deployment.model_info = { + "team_id": "team_123", + "team_public_model_name": "old-public-name", + } + + prisma_client = MockPrismaClient( + team_exists=True, sibling_deployments=[sibling_deployment] + ) + + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_delete, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_add: + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + # team_model_delete should NOT be called because sibling exists + mock_delete.assert_not_called() + # team_model_add should be called to add new public name + mock_add.assert_called_once() + @pytest.mark.asyncio async def test_patch_model_with_team_id_validates_permissions(self): """Test PATCH with team_id runs same validation as POST for team permissions"""