fix(teams): release a moved team model's public name from its former team

Moving a team model to another team registered the public name on the
destination but never removed it from the source team's models list, so the
source team kept a grant to a name nothing of its own served and a later
proxy-wide model of that name was callable by its keys. patch_model now runs
the same cleanup a delete runs once the move has been reloaded: the name is
kept only while a sibling row or a live model group of that exact name still
backs it, and a legacy alias entry is scrubbed the same way.

Resolves LIT-7445
This commit is contained in:
Tin Chi Lo 2026-09-09 20:00:41 -07:00
parent fb603fdb6b
commit 97f88f5513
2 changed files with 146 additions and 5 deletions

View file

@ -835,7 +835,9 @@ async def patch_model(
llm_router,
premium_user,
prisma_client,
proxy_logging_obj,
store_model_in_db,
user_api_key_cache,
)
try:
@ -953,6 +955,19 @@ async def patch_model(
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
reload_outcome: Final = await clear_cache()
source_team_id: Final = db_model.model_info.team_id if db_model.model_info else None
if (
source_team_id is not None
and patch_data.model_info is not None
and patch_data.model_info.team_id not in (None, source_team_id)
):
await _remove_unbacked_team_models(
model_params=db_model,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
)
## CREATE AUDIT LOG ##
asyncio.create_task(
@ -1507,11 +1522,12 @@ async def _remove_unbacked_team_models(
llm_router: Router | None = None,
) -> None:
"""
Strip a deleted team model's public name(s) from team.models and refresh the cache.
Strip a team model's public name(s) from team.models once the row no longer backs them, 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.
Runs after the row is deleted, or after a move to another team has been reloaded into the
router: 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, concurrent deletes can't
leave a ghost, and a moved row's former team does not keep a grant nothing of its own serves.
Legacy team models (created before team_public_model_name existed) store a
``{public_name: "model_name_{team_id}_{uuid}"}`` entry in the team's model_aliases,
@ -1523,7 +1539,9 @@ async def _remove_unbacked_team_models(
A public name that still resolves to a live router deployment (e.g. a gateway-level
model group shared with the team) is kept in team.models, so deleting a per-team
duplicate does not revoke the team's access to the shared deployment.
duplicate does not revoke the team's access to the shared deployment. Team rows are
indexed under their internal ``model_name_{team_id}_{uuid}`` key, so a row that was just
moved does not keep its public name alive for the team it left.
"""
team_id: Final = model_params.model_info.team_id
if team_id is None:

View file

@ -15,6 +15,7 @@ from litellm.proxy._types import (
LiteLLM_ModelTable,
LiteLLM_ProxyModelTable,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LitellmUserRoles,
Member,
ReconcileOutcome,
@ -1393,6 +1394,128 @@ class TestTeamModelUpdate:
# update_team (model_aliases write) must NOT be called in the new implementation
mock_update_team.assert_not_called()
# Moving a row between teams registers its public name on the destination; the source must
# give the name up too, unless a sibling row there still backs it, or it keeps a grant to a
# name nothing of its own serves. A legacy row lists its name through a team alias rather
# than team_public_model_name; the alias is scrubbed and its key released the same way.
@pytest.mark.asyncio
@pytest.mark.parametrize(
"moved_row, source_backing, released",
[
("named", "nothing", True),
("named", "sibling row with the same public name", False),
("legacy alias", "nothing", True),
],
)
async def test_team_move_releases_the_public_name_from_the_source_team(self, moved_row, source_backing, released):
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.model_management_endpoints import patch_model
named = moved_row == "named"
moved = LiteLLM_ProxyModelTable(
model_id="m-moved",
model_name="model_name_team_a_uuid1",
litellm_params={"model": "azure/gpt-4o-mini"},
model_info={"id": "m-moved", "team_id": "team_a", **({"team_public_model_name": "shared-name"} if named else {})},
created_by="admin",
updated_by="admin",
)
sibling = MagicMock()
sibling.model_name = "model_name_team_a_uuid2"
sibling.model_info = {"team_id": "team_a", "team_public_model_name": "shared-name"}
alias_row = MagicMock()
alias_row.id = 1
alias_row.model_aliases = {"shared-name": "model_name_team_a_uuid1"}
alias_row.team.team_id = "team_a"
async def team_update(where, data, include=None):
return LiteLLM_TeamTable(team_id=where["team_id"], models=data.get("models", ["shared-name"]))
prisma = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=moved)
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=moved)
prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
return_value=[sibling] if source_backing.startswith("sibling") else []
)
prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[] if named else [alias_row])
prisma.db.litellm_modeltable.update = AsyncMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(
return_value=LiteLLM_TeamTable(team_id="team_a", models=["shared-name"])
)
prisma.db.litellm_teamtable.update = AsyncMock(side_effect=team_update)
prisma.db.execute_raw = AsyncMock()
cache = UserApiKeyCache()
router = MagicMock(**{"get_model_ids.return_value": ["m-moved"], "model_name_to_deployment_indices": {}})
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: same
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same
patch("litellm.proxy.proxy_server.proxy_logging_obj", None), # test-quality-ok: same
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: same
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: same
patch( # test-quality-ok: the reload needs a live router; the team list write is what is under test
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
await patch_model(
model_id="m-moved",
patch_data=updateDeployment(model_info=ModelInfo(team_id="team_b")),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
source_team = await cache.async_get_cache(key="team_id:team_a", model_type=LiteLLM_TeamTableCachedObj)
assert (source_team.models if source_team is not None else ["shared-name"]) == ([] if released else ["shared-name"])
assert prisma.db.litellm_modeltable.update.await_count == (0 if named else 1)
# The cleanup keeps a name only while a deployment of exactly that name is live; a wildcard
# or alias that would merely answer a call to it is not a grant the team should keep.
@pytest.mark.asyncio
async def test_cleanup_ignores_wildcards_that_would_match_the_released_name(self):
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.model_management_endpoints import _remove_unbacked_team_models
released = Deployment(
model_name="model_name_team_a_uuid1",
litellm_params=LiteLLM_Params(model="anthropic/claude-haiku-4-5"),
model_info=ModelInfo(id="m-moved", team_id="team_a", team_public_model_name="anthropic/team-name"),
)
llm_router = Router(
model_list=[
{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*", "api_key": "k"}},
{
"model_name": "model_name_team_b_uuid1",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"},
"model_info": {"id": "m-moved", "team_id": "team_b", "team_public_model_name": "anthropic/team-name"},
},
],
model_group_alias={"anthropic/team-name": "anthropic/*"},
)
async def update(where, data, include=None):
return LiteLLM_TeamTable(team_id="team_a", models=data["models"])
prisma = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(
return_value=LiteLLM_TeamTable(team_id="team_a", models=["anthropic/team-name"])
)
prisma.db.litellm_teamtable.update = AsyncMock(side_effect=update)
prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[])
cache = UserApiKeyCache()
await _remove_unbacked_team_models(
model_params=released,
prisma_client=prisma,
user_api_key_cache=cache,
proxy_logging_obj=None,
llm_router=llm_router,
)
cached_team = await cache.async_get_cache(key="team_id:team_a", model_type=LiteLLM_TeamTableCachedObj)
assert cached_team is not None and cached_team.models == []
@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"""