ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(proxy): drop deleted team BYOK model name from team.models (#29820)

Browse Source 
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
This commit is contained in:
yuneng-jiang 2026-06-05 18:35:50 -07:00 committed by GitHub
parent e53bd7cbd1
commit 53cf3d8416
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 356 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

135
litellm/proxy/management_endpoints/model_management_endpoints.py
View File

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

251
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View File

@ -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."""