mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
refactor(model-management): drop redundant team lookup on model delete
Move the orphaned-team handling into can_user_make_model_call behind an allow_missing_team flag instead of pre-checking team existence in delete_model. The endpoint no longer issues its own litellm_teamtable lookup, so deleting a model whose team still exists hits the team table once instead of twice. The auth behavior is unchanged: a proxy admin can delete a model whose team was deleted, any other caller gets a 403, and add/update/health keep the strict "team must exist" validation.
This commit is contained in:
parent
b1b131659b
commit
d80dfe6ca2
2 changed files with 92 additions and 27 deletions
|
|
@ -857,6 +857,7 @@ class ModelManagementAuthChecks:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
premium_user: bool,
|
||||
allow_missing_team: bool = False,
|
||||
) -> Literal[True]:
|
||||
## Check team model auth
|
||||
if (
|
||||
|
|
@ -867,6 +868,18 @@ class ModelManagementAuthChecks:
|
|||
where={"team_id": model_params.model_info.team_id}
|
||||
)
|
||||
if team_obj_row is None:
|
||||
# The team was deleted. Callers that opt in (e.g. model deletion) may
|
||||
# act on the orphaned model, but only as a proxy admin -- without the
|
||||
# team there is no team-admin membership left to verify.
|
||||
if allow_missing_team:
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "Only a proxy admin can delete a model whose team has been deleted."
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
|
|
@ -947,32 +960,14 @@ async def delete_model(
|
|||
)
|
||||
|
||||
model_params = Deployment(**model_in_db.model_dump())
|
||||
|
||||
team_id = model_params.model_info.team_id if model_params.model_info else None
|
||||
team_exists = team_id is None or (
|
||||
await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
is not None
|
||||
await ModelManagementAuthChecks.can_user_make_model_call(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
allow_missing_team=True,
|
||||
)
|
||||
|
||||
if team_exists:
|
||||
await ModelManagementAuthChecks.can_user_make_model_call(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
elif user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
# The model's team was deleted; without it team-admin membership can't be
|
||||
# verified, so only a proxy admin may delete the orphaned model.
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "Only a proxy admin can delete a model whose team has been deleted."
|
||||
},
|
||||
)
|
||||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1915,15 +1915,16 @@ class TestDeleteTeamBYOKModelGhost:
|
|||
mock_refresh.assert_not_awaited()
|
||||
|
||||
|
||||
class TestDeleteOrphanedTeamModel:
|
||||
"""Deleting a team BYOK model whose team was deleted.
|
||||
class TestDeleteModelTeamAuth:
|
||||
"""Team auth on the /model/delete path.
|
||||
|
||||
A model added via /model/new with model_info.team_id is orphaned once its
|
||||
team is deleted: can_user_make_model_call looked the team up and raised
|
||||
'Team id=... does not exist in db' before the delete could run, so the model
|
||||
was undeletable from the Models + Endpoints page. Without the team, team-admin
|
||||
membership can't be verified, so a proxy admin (and only a proxy admin) may
|
||||
delete the orphan; a missing team must never let a non-admin through.
|
||||
delete the orphan; a missing team must never let a non-admin through. The team
|
||||
is also looked up exactly once -- the auth check must not add a second query.
|
||||
"""
|
||||
|
||||
def _orphaned_model_mocks(self, team_id, model_id):
|
||||
|
|
@ -2028,6 +2029,75 @@ class TestDeleteOrphanedTeamModel:
|
|||
assert str(exc_info.value.code) == "403"
|
||||
mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_team_delete_looks_up_team_once(self):
|
||||
"""The auth check must not add a redundant team query on the live-team path."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelInfoDelete,
|
||||
delete_model as delete_model_endpoint,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ProxyException
|
||||
|
||||
team_id = "live-team-1"
|
||||
model_id = "live-byok-1"
|
||||
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": "live-gpt",
|
||||
},
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
)
|
||||
team_row = LiteLLM_TeamTable(
|
||||
team_id=team_id,
|
||||
team_alias="live-team",
|
||||
members_with_roles=[Member(user_id="admin", role="admin")],
|
||||
models=["live-gpt"],
|
||||
)
|
||||
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_modeltable = AsyncMock()
|
||||
mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# A team member who is not the team admin: rejected before the delete runs,
|
||||
# so the only team lookup is the single one inside the auth check.
|
||||
non_admin = UserAPIKeyAuth(
|
||||
user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
)
|
||||
|
||||
_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()),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await delete_model_endpoint(
|
||||
model_info=ModelInfoDelete(id=model_id),
|
||||
user_api_key_dict=non_admin,
|
||||
)
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert mock_prisma.db.litellm_teamtable.find_unique.await_count == 1
|
||||
mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited()
|
||||
|
||||
|
||||
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