scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(proxy): delete large teams without per-member transaction fan-out

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-24 17:23:21 +00:00
parent 793627ce83
commit 0a023a4e1a
6 changed files with 168 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -4517,25 +4517,6 @@ async def delete_team(
llm_router=llm_router,
)
# ## DELETE TEAM MEMBERSHIPS
for team_row in team_rows:
### get all team members
team_members = team_row.members_with_roles
### call team_member_delete for each team member
tasks = []
for team_member in team_members:
tasks.append(
_team_member_delete(
data=TeamMemberDeleteRequest(
team_id=team_row.team_id,
user_id=team_member.user_id,
user_email=team_member.user_email,
),
user_api_key_dict=user_api_key_dict,
)
)
await asyncio.gather(*tasks)
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
## DELETE TEAMS
@ -4565,6 +4546,10 @@ async def delete_team(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _invalidate_deleted_team_member_cache(
teams=team_rows,
user_api_key_cache=user_api_key_cache,
)
for deleted_team in team_rows:
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
@ -4641,6 +4626,31 @@ async def _invalidate_deleted_team_cache(
)
async def _invalidate_deleted_team_member_cache(
teams: Sequence[LiteLLM_TeamTable],
user_api_key_cache: UserApiKeyCache,
) -> None:
for team in teams:
await _evict_deleted_team_member_cache(team=team, user_api_key_cache=user_api_key_cache)
async def _evict_deleted_team_member_cache(team: LiteLLM_TeamTable, user_api_key_cache: UserApiKeyCache) -> None:
member_user_ids: Final = tuple(
sorted({member.user_id for member in team.members_with_roles if member.user_id is not None})
)
await evict_and_broadcast(cache_keys=member_user_ids, user_api_key_cache=user_api_key_cache)
await asyncio.gather(
*(
invalidate_team_member_spend_state(
user_id=user_id,
team_id=team.team_id,
user_api_key_cache=user_api_key_cache,
)
for user_id in member_user_ids
)
)
def _transform_teams_to_deleted_records(
teams: list[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,

1
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -35,6 +35,7 @@
- {id: mgmt.team.update.team_admin_cannot_grow_budget, module: mgmt, tier: P0, surface: api, assertions: [team_admin_cannot_grow_budget], source: "team_endpoints.py:1203", fail_before_fix: proven, rationale: "With max_budget enabled, a team admin may keep or lower its team's budget; raising or removing it is 403 and writes nothing, also under an organization's larger cap"}
- {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"}
- {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"}
- {id: mgmt.team.delete.membership_larger_than_db_pool, module: mgmt, tier: P0, surface: api, assertions: [membership_larger_than_db_pool], source: "team_endpoints.py:4362", rationale: "Deleting a team with more members than the Prisma connection pool still completes instead of exhausting the pool and answering 500", fail_before_fix: proven}
- {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"}
- {id: mgmt.team.info.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "team_endpoints.py:2244", rationale: "Metadata+members+budgets"}
- {id: mgmt.team.daily_activity.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "vendor testing strategy §9.20 / LIT-4778", rationale: "GET /team/daily/activity returns results+metadata for a valid date range"}

21
tests/e2e/management/management_client.py
View file

@ -411,6 +411,27 @@ class ManagementClient:
assert last is not None
raise AssertionError(last)
def add_team_members(self, team_id: str, members: list[TeamMemberEntry]) -> None:
"""Bulk form of /team/member_add: `member` accepts a list, so one call
seeds a whole roster the way an admin import does."""
_ = unwrap(
self.proxy.transport.post(
"/team/member_add",
headers=self.proxy.management_headers(),
json=TeamMemberAddBody(team_id=team_id, member=members),
response_type=NoBody,
)
)
def delete_team_status(self, team_id: str) -> StreamingResponse:
"""POST /team/delete judged by HTTP outcome: the raw status and body, so a
test can assert on what a caller actually sees when the delete fails."""
return self.proxy.transport.send(
"/team/delete",
headers=self.proxy.management_headers(),
json=TeamDeleteBody(team_ids=[team_id]),
)
def delete_team_member(self, team_id: str, user_id: str) -> None:
_ = unwrap(
self.proxy.transport.post(

39
tests/e2e/management/test_management_e2e.py
View file

@ -37,6 +37,7 @@ from models import (
OrgUpdateBody,
TagListEntry,
TagNewBody,
TeamMemberEntry,
TeamNewBody,
TeamUpdateBody,
UserNewBody,
@ -48,6 +49,7 @@ pytestmark = pytest.mark.e2e
REGENERATE_GRACE_PERIOD = "15s"
REGENERATE_GRACE_SECONDS = 15.0
TEAM_DELETE_POOL_OVERFLOW_MEMBERS = 250
def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T:
@ -479,6 +481,43 @@ class TestTeamRoutes:
client, rejected, "team-bound key was still accepted on chat (never rejected 401) after team deletion"
)
@pytest.mark.covers("mgmt.team.delete.membership_larger_than_db_pool")
def test_team_delete_succeeds_for_team_larger_than_db_pool(
self, client: ManagementClient, resources: ResourceManager
) -> None:
"""Customer repro: /team/delete fans one transaction per member out over a
Prisma pool of 10 connections, each queued on the team's advisory lock,
so a team bigger than the pool must still delete cleanly instead of
answering 500 P2028."""
team_id = _create_team(client, resources, f"e2e-mgmt-team-{unique_marker()}", [])
user_ids = tuple(
_create_user(
client,
resources,
UserNewBody(
user_email=f"e2e-mgmt-bulk-{i}-{unique_marker()}@example.com",
user_role="internal_user",
),
)
for i in range(TEAM_DELETE_POOL_OVERFLOW_MEMBERS)
)
client.add_team_members(team_id, [TeamMemberEntry(role="user", user_id=user_id) for user_id in user_ids])
seated = len(client.team_info(team_id).members_with_roles)
assert seated >= len(user_ids), (
f"/team/info lists {seated} members after the bulk /team/member_add, expected at least {len(user_ids)}"
)
outcome = client.delete_team_status(team_id)
assert outcome.status_code == 200, (
f"/team/delete on a {len(user_ids)}-member team must succeed, got "
f"{outcome.status_code}: {outcome.body[:500]}"
)
probe = client.team_info_status(team_id)
assert probe.status_code == 404, (
f"deleted team {team_id} still resolves: /team/info returned {probe.status_code}: {probe.body[:300]}"
)
@pytest.mark.covers("mgmt.team.member_add.persists")
def test_member_add_and_delete_persist_to_team_info(
self, client: ManagementClient, resources: ResourceManager

2
tests/e2e/models.py
View file

@ -1458,7 +1458,7 @@ class TeamInfoResponse(BaseModel):
class TeamMemberAddBody(BaseModel):
team_id: str
member: TeamMemberEntry
member: TeamMemberEntry | list[TeamMemberEntry]
class TeamMemberDeleteBody(BaseModel):

87
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -8869,10 +8869,6 @@ async def test_delete_team_persists_deleted_teams(
"litellm.proxy.proxy_server.litellm_proxy_admin_name",
"admin",
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints._team_member_delete",
AsyncMock(return_value=(team1, (), ())),
)
data = DeleteTeamRequest(team_ids=["team-1"])
@ -9015,6 +9011,83 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
assert cache_state_when_rows_deleted["doomed_still_cached"] is True
@pytest.mark.asyncio
async def test_delete_team_evicts_member_caches_with_one_transaction(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
Regression pin for LIT-8533: `delete_team` used to fan out one
`_team_member_delete` per roster entry via `asyncio.gather`, and each opened
its own `prisma_client.tx()` and queued on the team's advisory lock, so a
team larger than the Prisma pool exhausted it and the late transactions died
on P2028. Every member-side db effect is already covered by the key delete
and the locked sweep, so the only work left is evicting each member's cache
entries, which needs no transaction at all.
"""
from litellm.proxy._types import DeleteTeamRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
member_user_ids = tuple(f"member-{i}" for i in range(3))
team = LiteLLM_TeamTable(
team_id="team-doomed",
team_alias="doomed-team",
members_with_roles=[Member(user_id=user_id, role="user") for user_id in member_user_ids],
metadata={},
model_max_budget={},
model_spend={},
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0})
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.execute_raw = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_tx_cm = MagicMock()
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma_client)
fresh_cache = UserApiKeyCache()
for user_id in member_user_ids:
fresh_cache.set_cache(key=user_id, value=UserAPIKeyAuth(user_id=user_id))
fresh_cache.set_cache(key="bystander-user", value=UserAPIKeyAuth(user_id="bystander-user"))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
await delete_team(
data=DeleteTeamRequest(team_ids=["team-doomed"]),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
),
litellm_changed_by="admin-user",
)
assert mock_prisma_client.tx.call_count == 1, (
f"delete_team must run a single locked transaction for the whole delete, not one per member; "
f"prisma_client.tx() was entered {mock_prisma_client.tx.call_count} times for "
f"{len(member_user_ids)} members"
)
for user_id in member_user_ids:
assert fresh_cache.get_cache(key=user_id) is None, (
f"member {user_id}'s cached user object survived the team delete"
)
assert fresh_cache.get_cache(key="bystander-user") is not None
@pytest.mark.asyncio
async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
monkeypatch,
@ -14146,12 +14219,6 @@ async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
removals = [(team, members, members[1:]), (team, members[1:], ())]
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints._team_member_delete",
AsyncMock(side_effect=lambda **_kwargs: removals.pop(0)),
)
await delete_team(
data=DeleteTeamRequest(team_ids=["team-gone"]),
http_request=MagicMock(),