mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
scope cascade cleanup to global models only
Team-scoped model deletions should not cascade to other teams' model lists. A team admin could set a deployment name matching a shared model and then delete it to remove that model from unrelated teams. Restrict the cascade to models with no team_id (global deployments).
This commit is contained in:
parent
42086bdbd4
commit
13480a9c3a
2 changed files with 81 additions and 18 deletions
|
|
@ -808,25 +808,27 @@ async def delete_model(
|
|||
data={"models": existing_team_row.models},
|
||||
)
|
||||
|
||||
# Remove deleted model from all teams that reference it by name
|
||||
try:
|
||||
affected_teams = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"models": {"has": model_params.model_name}}
|
||||
)
|
||||
for team in affected_teams:
|
||||
updated_models = [
|
||||
m for m in team.models if m != model_params.model_name
|
||||
]
|
||||
await prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": team.team_id},
|
||||
data={"models": updated_models},
|
||||
# Remove deleted global model from all teams that reference it by name.
|
||||
# Only for non-team-scoped models; team-scoped cleanup is handled above.
|
||||
if model_params.model_info.team_id is None:
|
||||
try:
|
||||
affected_teams = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"models": {"has": model_params.model_name}}
|
||||
)
|
||||
for team in affected_teams:
|
||||
updated_models = [
|
||||
m for m in team.models if m != model_params.model_name
|
||||
]
|
||||
await prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": team.team_id},
|
||||
data={"models": updated_models},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to remove model %s from teams: %s",
|
||||
model_params.model_name,
|
||||
e,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to remove model %s from teams: %s",
|
||||
model_params.model_name,
|
||||
e,
|
||||
)
|
||||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
|
|
|
|||
|
|
@ -1439,6 +1439,67 @@ class TestDeleteModelCascadeTeams:
|
|||
)
|
||||
assert team_b_call.kwargs["data"]["models"] == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_scoped_delete_does_not_cascade(self):
|
||||
"""A team-scoped model deletion should NOT cascade to other teams."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
delete_model as delete_model_endpoint,
|
||||
ModelInfoDelete,
|
||||
)
|
||||
|
||||
model_id = "team-scoped-model"
|
||||
admin_user = UserAPIKeyAuth(
|
||||
user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
team_row = MagicMock()
|
||||
team_row.team_id = "team-owner"
|
||||
team_row.models = ["gpt-4"]
|
||||
|
||||
db_row = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="gpt-4",
|
||||
litellm_params={"model": "openai/gpt-4"},
|
||||
model_info={"id": model_id, "team_id": "team-owner"},
|
||||
created_by="test-admin",
|
||||
updated_by="test-admin",
|
||||
)
|
||||
|
||||
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_teamtable = AsyncMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_teamtable.update = AsyncMock()
|
||||
|
||||
_PS = "litellm.proxy.proxy_server"
|
||||
_MM = "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}.proxy_logging_obj", MagicMock()),
|
||||
patch(f"{_PS}.general_settings", {}),
|
||||
patch(f"{_PS}.premium_user", True),
|
||||
patch(f"{_PS}.llm_router", MagicMock()),
|
||||
patch(f"{_MM}.delete_team_model_alias", AsyncMock(return_value=[])),
|
||||
patch(
|
||||
f"{_MM}.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
AsyncMock(return_value=True),
|
||||
),
|
||||
):
|
||||
await delete_model_endpoint(
|
||||
model_info=ModelInfoDelete(id=model_id),
|
||||
user_api_key_dict=admin_user,
|
||||
)
|
||||
|
||||
# The cascade find_many should NOT be called for team-scoped models
|
||||
mock_prisma.db.litellm_teamtable.find_many.assert_not_called()
|
||||
|
||||
|
||||
class TestGetTeamDeployments:
|
||||
"""Tests for _get_team_deployments which filters by model_name prefix + Python-side team_id check."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue