From 97f88f5513e03b01435b9ebc6d5cfa8da050a2b6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 9 Sep 2026 20:00:41 -0700 Subject: [PATCH] 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 --- .../model_management_endpoints.py | 28 +++- .../test_model_management_endpoints.py | 123 ++++++++++++++++++ 2 files changed, 146 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..3ca749855ee 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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: 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 c02f886fc31..f490e245082 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 @@ -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"""