From 53cf3d8416fe801a579b4fd59cf16e1d2255a3d7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 5 Jun 2026 18:35:50 -0700 Subject: [PATCH] fix(proxy): drop deleted team BYOK model name from team.models (#29820) Deleting a team-scoped BYOK model left its public name in team.models, so /models with a team key kept listing the now-deleted "ghost" model. delete_model stripped team.models using only litellm_modeltable alias lookups, but models added via /model/new with a team_id never create an alias row; their public name lives only in team.models and model_info.team_public_model_name, so it was never removed. The team cache was also left stale because the delete path skipped _refresh_cached_team. The cleanup now keys off team_public_model_name (falling back to alias keys), runs after the deployment row is deleted, and strips a public name only when no remaining team deployment still backs it, so a load-balanced replica is not revoked and concurrent deletes cannot leave a ghost. The updated team row is refreshed in cache so /models reflects the change immediately --- .../model_management_endpoints.py | 135 +++++++--- .../test_model_management_endpoints.py | 251 ++++++++++++++++++ 2 files changed, 356 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 404ed4491e..e4ecda3fe3 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,7 +13,7 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from typing import Dict, List, Literal, Optional, Tuple, Union, cast +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field @@ -39,6 +39,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( + _refresh_cached_team, team_model_add, team_model_delete, ) @@ -592,6 +593,98 @@ async def _get_team_deployments( return result +async def _get_team_public_model_names( + team_id: str, + prisma_client: PrismaClient, +) -> Set[str]: + """ + Public model names currently backed by a deployment in the team. + + Called on delete (after the deployment row is removed) so a public name that is + load-balanced across several deployments stays in team.models while a replica + still serves it. + """ + deployments = await _get_team_deployments(team_id, prisma_client) + public_names: Set[str] = set() + for row in deployments: + model_info = row.model_info + if isinstance(model_info, str): + try: + model_info = json.loads(model_info) + except (TypeError, ValueError): + continue + if isinstance(model_info, dict): + public_name = model_info.get("team_public_model_name") + if public_name: + public_names.add(public_name) + return public_names + + +async def _remove_unbacked_team_models( + model_params: Deployment, + prisma_client: PrismaClient, + user_api_key_cache: Any, + proxy_logging_obj: Any, +) -> None: + """ + Strip a deleted team model's public name(s) from team.models and refresh the cache. + + Must be called after the deployment row is deleted: a public name is removed only + when no remaining team deployment still backs it, so a load-balanced replica isn't + revoked while siblings serve it, and concurrent deletes can't leave a ghost. + """ + team_id = model_params.model_info.team_id + if team_id is None: + return + + # BYOK models carry an internal `model_name_{team_id}_{uuid}` name that can never + # be a team alias value, so skip the full litellm_modeltable scan for them. + removed_model_aliases: List[Tuple[str, str]] = [] + if not model_params.model_name.startswith(f"model_name_{team_id}_"): + removed_model_aliases = await delete_team_model_alias( + public_model_name=model_params.model_name, + prisma_client=prisma_client, + ) + names_to_remove = { + alias + for alias_team_id, alias in removed_model_aliases + if alias_team_id == team_id + } + if model_params.model_info.team_public_model_name is not None: + names_to_remove.add(model_params.model_info.team_public_model_name) + + if names_to_remove: + names_to_remove -= await _get_team_public_model_names( + team_id=team_id, prisma_client=prisma_client + ) + + if not names_to_remove: + return + + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + if existing_team_row is None: + return + + updated_team_row = await prisma_client.db.litellm_teamtable.update( + where={"team_id": team_id}, + data={ + "models": [ + model + for model in existing_team_row.models + if model not in names_to_remove + ] + }, + include={"object_permission": True}, # type: ignore + ) + await _refresh_cached_team( + team_row=updated_team_row, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _update_existing_team_model_assignment( team_id: str, public_model_name: str, @@ -831,7 +924,9 @@ async def delete_model( llm_router, premium_user, prisma_client, + proxy_logging_obj, store_model_in_db, + user_api_key_cache, ) if prisma_client is None: @@ -859,35 +954,6 @@ async def delete_model( premium_user=premium_user, ) - # delete team model alias - if model_params.model_info.team_id is not None: - removed_model_aliases = await delete_team_model_alias( - public_model_name=model_params.model_name, - prisma_client=prisma_client, - ) - - valid_team_model_aliases = [ - model - for team_id, model in removed_model_aliases - if team_id == model_params.model_info.team_id - ] - - ## UPDATE TEAM TO NOT LIST MODEL ## - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": model_params.model_info.team_id} - ) - if existing_team_row is not None: - existing_team_row.models = [ - model - for model in existing_team_row.models - if model not in valid_team_model_aliases - ] - - await prisma_client.db.litellm_teamtable.update( - where={"team_id": model_params.model_info.team_id}, - data={"models": existing_team_row.models}, - ) - # update DB if store_model_in_db is True: """ @@ -909,6 +975,15 @@ async def delete_model( if llm_router is not None: llm_router.delete_deployment(id=model_info.id) + # Runs after the row delete so the sibling check sees post-delete state. + if model_params.model_info.team_id is not None: + await _remove_unbacked_team_models( + model_params=model_params, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## CREATE AUDIT LOG ## asyncio.create_task( create_object_audit_log( 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 f16074a049..a9bb2b09a1 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 @@ -1664,6 +1664,257 @@ class TestAddAndDeleteModelLifecycle: assert str(exc_info.value.code) == "400" +class TestDeleteTeamBYOKModelGhost: + """Regression for issue #22594. + + A team BYOK model (added via /model/new with model_info.team_id) stores its + public name only in team.models and model_info.team_public_model_name -- it + never creates a litellm_modeltable alias row. delete_model used to strip + team.models using alias lookups alone, so the public name lingered forever + and showed up as a 'ghost' in /models. It also skipped the team cache + refresh, so even a corrected DB write would lag behind the cache TTL. + """ + + @pytest.mark.asyncio + async def test_delete_strips_public_name_and_refreshes_cache(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + + team_id = "team-byok-ghost" + model_id = "byok-model-123" + public_name = "my-team-gpt" + kept_name = "kept-team-model" + + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_abc-uuid", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": public_name, + }, + created_by="admin", + updated_by="admin", + ) + + def _team(models): + return LiteLLM_TeamTable( + team_id=team_id, + team_alias="byok-team", + members_with_roles=[Member(user_id="admin", role="admin")], + models=models, + ) + + team_row = _team([public_name, kept_name]) + updated_team_row = _team([kept_name]) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + # After the row delete no team deployment remains -> nothing backs the public name. + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=updated_team_row + ) + # Team BYOK models have no alias row; delete_team_model_alias finds nothing. + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + + admin_user = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()) as mock_refresh, + ): + result = await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=admin_user, + ) + + assert "deleted successfully" in result["message"] + + mock_prisma.db.litellm_teamtable.update.assert_awaited_once() + update_kwargs = mock_prisma.db.litellm_teamtable.update.await_args.kwargs + assert public_name not in update_kwargs["data"]["models"] + assert kept_name in update_kwargs["data"]["models"] + assert update_kwargs["include"] == {"object_permission": True} + + mock_refresh.assert_awaited_once() + assert mock_refresh.await_args.kwargs["team_row"] is updated_team_row + # BYOK internal name can't be an alias value -> the alias-table scan is skipped. + mock_prisma.db.litellm_modeltable.find_many.assert_not_awaited() + + @pytest.mark.asyncio + async def test_delete_non_internal_team_model_still_scans_aliases(self): + """A team model whose name is not the BYOK internal shape must still run the + alias cleanup (delete_team_model_alias), preserving legacy behavior.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + + team_id = "team-legacy" + model_id = "legacy-model-1" + public_name = "legacy-public" + + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=public_name, # not the model_name_{team_id}_ internal shape + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": public_name, + }, + created_by="admin", + updated_by="admin", + ) + team_row = LiteLLM_TeamTable( + team_id=team_id, + team_alias="legacy-team", + members_with_roles=[Member(user_id="admin", role="admin")], + models=[public_name, "kept"], + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_modeltable = AsyncMock() + # No alias row matches -> delete_team_model_alias returns nothing, but it still ran. + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + + admin_user = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + result = await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=admin_user, + ) + + assert "deleted successfully" in result["message"] + # Non-internal name -> the alias-table scan runs. + mock_prisma.db.litellm_modeltable.find_many.assert_awaited() + + @pytest.mark.asyncio + async def test_delete_keeps_public_name_when_sibling_backs_it(self): + """A public name load-balanced across two team deployments must stay in + team.models when one replica is deleted but a sibling still backs it.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + + team_id = "team-lb" + deleted_id = "replica-1" + sibling_id = "replica-2" + public_name = "lb-gpt" + + def _row(model_id): + return LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_{model_id}", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": public_name, + }, + created_by="admin", + updated_by="admin", + ) + + deleted_row = _row(deleted_id) + sibling_row = _row(sibling_id) + team_row = LiteLLM_TeamTable( + team_id=team_id, + team_alias="lb-team", + members_with_roles=[Member(user_id="admin", role="admin")], + models=[public_name], + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=deleted_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock( + return_value=deleted_row + ) + # After the deleted replica's row is gone, the sibling still backs the public name. + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[sibling_row] + ) + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + + admin_user = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()) as mock_refresh, + ): + result = await delete_model_endpoint( + model_info=ModelInfoDelete(id=deleted_id), + user_api_key_dict=admin_user, + ) + + assert "deleted successfully" in result["message"] + # The public name is still backed by the sibling, so team.models is untouched. + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() + mock_refresh.assert_not_awaited() + + class TestGetTeamDeployments: """Tests for _get_team_deployments which filters by model_name prefix + Python-side team_id check."""