From d80dfe6ca29e03701c462f02cba62a69ef761766 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 8 Jun 2026 12:46:24 -0700 Subject: [PATCH] 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. --- .../model_management_endpoints.py | 43 +++++------ .../test_model_management_endpoints.py | 76 ++++++++++++++++++- 2 files changed, 92 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cd217e4a344..73d801c8707 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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: """ 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 d3f21cf061c..2ba604e5da6 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 @@ -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."""