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
This commit is contained in:
Sameer Kankute
2026-03-23 17:33:07 +05:30
parent 298df75066
commit 8aa58bdcaa
2 changed files with 83 additions and 2 deletions
+15
View File
@@ -1296,6 +1296,10 @@ def _update_model_if_team_alias_exists(
"gpt-4o": "gpt-4o-team-1"
}
- requested_model = "gpt-4o-team-1"
Note: model_aliases for team models are deprecated. This function only applies
to legacy non-team-scoped aliases. Team-scoped deployments use team_public_model_name
and are resolved via map_team_model in route_llm_request.
"""
_model = data.get("model")
if (
@@ -1303,6 +1307,17 @@ def _update_model_if_team_alias_exists(
and user_api_key_dict.team_model_aliases
and _model in user_api_key_dict.team_model_aliases
):
from litellm.proxy.proxy_server import llm_router
# Skip alias rewrite if this model resolves to team-specific deployments
# (team models use team_public_model_name, not model_aliases)
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
data["model"] = user_api_key_dict.team_model_aliases[_model]
return
@@ -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"""