mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
parent
fb603fdb6b
commit
97f88f5513
2 changed files with 146 additions and 5 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue