scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

fix(scim)!: delete a group through the team delete's bulk path instead of one member at a time (#47544)

Browse source 
* fix(scim): write a group's roster in one locked transaction instead of one member at a time

POST, PUT and PATCH /scim/v2/Groups resolve every member in one query and write the roster, the
membership rows and the users' team tags in one transaction under the team's advisory lock, so a
500-member push runs a fixed number of statements instead of thousands. PUT now answers with the
roster it just wrote

* test(scim): assert the roster the IdP gets back and boot the bulk roster cell with two workers

* fix(scim): skip the team lock for a rename-only PATCH and lock member rows in id order before the roster write

A PATCH naming no members opened the roster transaction and took the team advisory lock for nothing. A SIGKILLed worker then left its orphaned query engine holding that lock until the 60 s transaction timeout, and /team/delete, which takes the same lock, timed out behind it. An empty delta is now answered from a plain team read, with no transaction and no lock

Two groups trading members in parallel locked user rows in opposite orders, and one request died with a Postgres deadlock. Every user row a push touches is now locked with SELECT ... ORDER BY user_id FOR UPDATE before the first write, so the two requests serialize

The bulk roster cell gets a 120 s client through a new owned_proxy_process client_timeout, so an unfixed push fails on the statement count instead of the harness clock, requires positive counts on both pushes, and gains a parallel roster-trade test that answered 500 before the lock

* fix(scim): evict the cached team on a group rename and read a no-op roster from the writer

A rename-only PUT or PATCH wrote the new alias and left the team cached under
its id and old alias on every worker, so keys kept logging spend against the
old team alias until the TTL expired. The group write now evicts the cached
team (and broadcasts the eviction) right after the alias or metadata update.

sync_team_roster reads the team from the writer when the plan names nobody,
so a replica cannot hand back the pre-rename row; a plan already matching the
roster returns inside the lock without writing, evicting, or emitting the
members metric; the response is built from the written row, so lastModified
reflects the roster write; and detaching a removed member bumps the user's
updated_at like every other user write.

* fix(scim): bump detached members' updated_at in UTC

The detach UPDATE wrote updated_at = CURRENT_TIMESTAMP, which Postgres casts
into Prisma's naive timestamp column in the session timezone, so on a database
whose session is not UTC the removed member's row landed hours away from the
UTC instants Prisma itself writes. Write NOW() AT TIME ZONE 'UTC' instead, and
have the bulk roster cell assert the bumped row sits between two reads of the
database's own UTC clock taken around the replace.

* fix(scim): delete a group through the team delete's bulk path instead of one member at a time

* fix(scim): chunk the group write's member lookups under the bind cap

Every Prisma "in" filter the group write issues bound one parameter per
member with nothing bounding the list, which check_unbounded_in_lists
reports as 11 new findings. The single-field lookups now go through
litellm.repositories.chunked_in (caller_transaction inside the roster
transaction, per_chunk_ok for the role recompute), and the one
cross-field OR read slices its values at IN_LIST_CHUNK_SIZE and records
the bound. A push of IN_LIST_CHUNK_SIZE + 1 members now reads in two
slices, and a member named only in the last slice still resolves.

* fix(scim): demote a deleted admin group's members before the team delete so a failed delete stays retryable

* fix(scim): recompute member roles again after the team delete so a concurrent group write cannot regrant admin

* fix(proxy): send bulk cache evictions in slices so the broadcaster drops none

Veria pointed out that the invalidation publisher caps its pending tasks at
1,024 and drops the excess with a warning, and that a group push now detaches
every removed member's keys in one call instead of one member at a time, so a
removal of more than roughly 340 members (three cache keys each) silently left
stale key objects on the other replicas. The merge base never hit the cap
because each member was a separate request.

Bulk evictions (team, user, organization, and session key deletes, plus the
roster settle) now publish in slices of 256 and wait, bounded at two seconds,
for the backlog already handed to the event loop before each slice. A Redis
that never answers costs a bulk caller at most that wait per slice instead of
wedging the request, and single-key callers are untouched. Regression tests
push 1,500 keys and a 500-member removal through a slow publisher and assert
every eviction reaches Redis; both fail against the previous tip.

* fix(proxy): clear this worker's cache before pacing bulk eviction broadcasts

Greptile pointed out that pacing the broadcasts in slices also paced the
local eviction: a bulk delete waited for the publish backlog before dropping
the next slice from this worker's own memory, so on a Redis that accepts
connections but never answers a key deleted from the database kept
authenticating on the handling worker for up to two seconds per slice. The
merge base cleared every key locally at once.

Bulk evictions now drop every affected entry from this worker's memory first,
through the cache's own partition lookup since key objects live apart from
the other auth objects, and only then publish in paced slices. The roster
settle does the same for the deleted tokens, the JWT mapping keys and every
changed member's user and membership entries before any step that can wait.
Regression tests run both paths against a publisher that never answers and
check that nothing is still cached on this worker before the first bound
elapses; both fail with the local eviction removed.

* fix(scim): write recomputed member roles only where the membership still backs them

* fix(team): report the users a team delete detaches so SCIM demotes late-added members

* test(scim): audit cells for the bulk group delete

Move the SCIM group helpers into tests/integration/_support/scim.py and add the delete cells: a member's other
group survives, the admin group's members are demoted and the audit row names the SCIM caller, a member a
concurrent PUT adds while the delete waits on the team lock is demoted too, unknown, repeated, non-admin and
hostile ids, an empty group, the peer proxy, a 20-delete burst across both proxies, and a delete after a worker
was SIGKILLed

* fix(proxy): drop the Redis copies of bulk-evicted auth objects before pacing the broadcasts

* docs(proxy): drop the module docstring that repeated the PR description

* fix(team): sweep deleted team references only under the lock so a late-added member is demoted

The unlocked sweep that ran before the team's advisory lock stripped a member a
concurrent group write added after the roster was read, so neither the roster
snapshot nor the locked sweep's RETURNING named them and the SCIM recompute
left them proxy_admin after the admin group was gone. The locked sweep does the
same strip and reports whom it stripped, so it is the only sweep now.

The integration cell holds the delete on a member's user row so a PATCH lands
between the roster read and the lock, then checks the late member is demoted.
Both waiter queries in the SCIM support module return their count as text so
the first poll parses.

* chore(team): carry the unbounded-IN baseline entries to delete_validated_teams

The delete pass moved from delete_team and the unlocked sweep into delete_validated_teams, so the checker's baseline names the function the two team_id IN lists now live in

* test(scim): expect the master key alias in the audit row and send post-kill requests on a fresh connection

* fix(scim): bound bulk cache evictions to one Redis round trip and close lock-order gaps

Bulk key evictions now drop entries from memory, delete them from Redis in one
chunked DEL, and broadcast every key in one pipelined publish, so a large roster
removal no longer overflows the publisher's pending cap or costs a Redis round
trip per member. A push waits for in-flight membership loads before evicting.

User rows are locked in byte order to match the spend writer, a bulk user delete
takes every team lock before its first user write, and a POST whose roster write
fails deletes the empty team so the IdP retry creates it again.

* test(team): observe delete broadcasts at the Redis layer

* fix(scim): lock user rows in id order on every multi-row writer and create unprovisioned members concurrently

* fix(team): chunk the email lookup behind the member row lock

* fix(scim): create a member deleted mid-push again instead of failing the group write

* fix(scim): refuse strict-mode mid-push deletes with 400 and check license seats once per batch

A member deleted mid-push is created again only when scim_upsert_user allows it,
under the verb's own created_via. New group members are checked against the
license seat limit as one batch before any is created. Auth cache publish drops
log the key count instead of the key.

* fix(scim): lock member rows while recomputing admin roles after a group push

* fix(users): re-read teams under the row locks in bulk user delete so a racing group push keeps no deleted member

* fix(users): take /user/delete and SCIM DELETE /Users users off their current teams under the row locks so a racing group push keeps no deleted member

* fix(scim): keep retrying a push while deletes race it, lock rows in member bulk delete, evict removed keys from Redis first

A group push gave up after one retry when SCIM user deletes kept landing between attempts, answering 500 in 9 of 10 race rounds. It now retries while each attempt finds newly deleted members, and gives up only when a member it created again goes missing a second time.

The team member bulk delete now locks the removed members' user rows in the shared order and takes the team off their teams arrays in one array_remove statement, so it no longer writes back a teams array read before a concurrent push added another team.

A push that removes members now drops their deleted team keys from memory and Redis together with the members' other cache entries before the team eviction, so a removed key cannot refill from Redis mid-cleanup.

* fix(user-delete): keep the team alias label on the team members gauge after removals

* test(integration): count relay statements over plaintext so an SSL-enabled Postgres does not hide them

* fix(user-delete): skip a team deleted while the user delete waited on its lock instead of answering 404

* fix(scim): answer 404 and roll back when the team is deleted during a roster write

* test(integration): audit cells for SCIM bulk roster writes, user delete locking and the eviction broadcast

* test(scim): cover deleting a SCIM group through the team delete's bulk path

* fix(scim): evict a deleted group's member caches in bulk and drop cached roles after a recompute

* fix(scim): answer 404 to a group delete that lost a race to another delete

* fix(scim): chunk the locked team lookup and delete in the team delete transaction

* fix(scim): keep service-account keys and team models on a group delete and write its records once under the lock

* fix(scim): demote a deleted group's members inside the delete transaction

The admin-group role write now runs on the team delete's transaction, over the
roster read under the team lock, so a failed delete leaves roles untouched and
a member added while the delete waited is demoted with the rest. The delete
transaction also waits for a free connection as long as it may run, so a burst
of deletes queues instead of failing after Prisma's two-second default. Every
member in the 5-vs-500 statement test now holds a key.

* fix(audit): write deleted keys' audit rows in one insert

A delete that removes many keys wrote one audit insert task per key, and those tasks held pool connections the delete's follow-up writes waited on. The race test now waits on any lock wait behind the delete.
This commit is contained in:
Mateo Wang 2026-10-21 00:59:23 +00:00 • committed by GitHub
parent 3c28f4bf1e
commit e333ccabdb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1744 additions and 460 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -254,39 +254,45 @@ class KeyManagementEventHooks:
litellm_changed_by: str | None = None,
) -> None:
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
create_audit_logs_for_updates,
get_audit_log_changed_by,
is_audit_logging_enabled,
track_audit_task,
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
if not is_audit_logging_enabled():
return
for key in keys_being_deleted:
key_row = LiteLLM_VerificationToken.model_validate(key, from_attributes=True)
if key_row.token is None:
continue
_key_row = key_row.model_dump_json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key_row.token,
action="deleted",
updated_values="{}",
before_value=_key_row,
)
)
changed_by: Final = get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
key_rows: Final = (
LiteLLM_VerificationToken.model_validate(key, from_attributes=True) for key in keys_being_deleted
)
audit_rows: Final = tuple(
LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=changed_by,
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key_row.token,
action="deleted",
updated_values="{}",
before_value=key_row.model_dump_json(exclude_none=True),
)
for key_row in key_rows
if key_row.token is not None
)
if not audit_rows:
return
track_audit_task(
asyncio.create_task(
create_audit_logs_for_updates(audit_rows), name=f"{len(audit_rows)} deleted key audit logs"
)
)
@staticmethod
async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: str | None = None):

116
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -11,7 +11,7 @@ from dataclasses import dataclass
from functools import partial, reduce
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload
from fastapi import (
APIRouter,
@ -53,7 +53,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
delete_cache_team_object,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
broadcast_evictions,
evict_and_broadcast,
evict_local,
evict_shared,
)
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
safe_get_request_headers,
@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.scim.scim_transformations import (
ScimTransformations,
)
from litellm.proxy.management_endpoints.team_endpoints import (
delete_validated_team_rows,
new_team,
team_member_add,
team_member_delete,
@ -104,6 +110,7 @@ from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.management_endpoints.scim_v2 import *
if TYPE_CHECKING:
from prisma import Prisma
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
@ -317,11 +324,15 @@ async def _check_team_exists(team_id: str) -> LiteLLM_TeamTable:
team: Final = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
if not team:
raise HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"})
raise _group_not_found(team_id)
return team
def _group_not_found(team_id: str) -> HTTPException:
return HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"})
def _extract_scim_user_data(user: SCIMUser) -> ScimUserData:
"""Extract common data from SCIMUser object."""
user_email = None
@ -482,10 +493,7 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It
Recompute and persist each user's global proxy role from their resulting team
membership. No-op unless scim_admin_group is configured, so a SCIM group write
that drops a member from the admin group demotes them just like the user
endpoints do, and the role is left untouched when the feature is off. The users'
rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the
same lock a roster write takes, so a concurrent group push cannot change a
member's teams between the read and the role write.
endpoints do, and the role is left untouched when the feature is off.
"""
admin_group: Final = await _get_scim_admin_group()
if admin_group is None:
@ -494,28 +502,46 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It
ids: Final = tuple(dict.fromkeys(user_ids))
if not ids:
return
default_role: Final = _default_scim_user_role()
async with prisma_client.tx(timeout=ROSTER_TX_TIMEOUT) as tx:
await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(ids))
rows: Final = await find_many_in(user_tx_db(tx), "user_id", ids)
team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows)))
teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids)
alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams})
resolved: Final = tuple(
(
row.user_id,
_resolve_scim_user_role(
[SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams],
admin_group,
default_role,
),
)
for row in rows
await _write_scim_member_roles_tx(tx, ids, admin_group=admin_group)
from litellm.proxy.proxy_server import user_api_key_cache
evict_local(cache_keys=ids, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=ids, cache=user_api_key_cache)
await broadcast_evictions(cache_keys=ids)
async def _write_scim_member_roles_tx(tx: "Prisma", user_ids: Sequence[str], *, admin_group: str) -> None:
"""
Write each user's role resolved from the teams they are in, on ``tx``. The users'
rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the
same lock a roster write takes, so a concurrent group push cannot change a
member's teams between the read and the role write.
"""
if not user_ids:
return
default_role: Final = _default_scim_user_role()
await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(user_ids))
rows: Final = await find_many_in(user_tx_db(tx), "user_id", user_ids)
team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows)))
teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids)
alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams})
resolved: Final = tuple(
(
row.user_id,
_resolve_scim_user_role(
[SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams],
admin_group,
default_role,
),
)
for row in rows
)
for role in dict.fromkeys(role for _, role in resolved):
await tx.execute_raw(
_SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role)
)
for role in dict.fromkeys(role for _, role in resolved):
await tx.execute_raw(
_SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role)
)
class _ResolvedUserMember(NamedTuple):
@ -3051,7 +3077,8 @@ async def update_group(
dependencies=[Depends(user_api_key_auth)],
)
async def delete_group(
group_id: str = Path(..., title="Group ID"),
group_id: Annotated[str, Path(title="Group ID")],
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Delete a group according to SCIM v2 protocol
@ -3064,23 +3091,22 @@ async def delete_group(
existing_team: Final = await _check_team_exists(group_id)
raise_if_teams_hold_ptu_shares((group_id,), llm_router)
member_ids: Final = await _get_team_member_user_ids_from_team(existing_team)
# For each member, remove this team from their teams list
for member_id in member_ids:
user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id})
if user:
current_teams = user.teams or []
if group_id in current_teams:
new_teams = [t for t in current_teams if t != group_id]
await _table(UserRepository(prisma_client)).update(
where={"user_id": member_id}, data={"teams": new_teams}
)
await _recompute_scim_member_roles(prisma_client, member_ids)
# Delete team
await _table(TeamRepository(prisma_client)).delete(where={"team_id": group_id})
admin_group: Final = await _get_scim_admin_group()
deleted: Final = await delete_validated_team_rows(
teams=(LiteLLM_TeamTable.model_validate(existing_team.model_dump()),),
prisma_client=prisma_client,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=user_api_key_dict.api_key,
user_id=user_api_key_dict.user_id,
),
litellm_changed_by=None,
before_commit=(
None if admin_group is None else partial(_write_scim_member_roles_tx, admin_group=admin_group)
),
)
if group_id not in deleted.team_ids:
raise _group_not_found(group_id)
return Response(status_code=204)

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

@ -14,7 +14,7 @@ import copy
import json
import math
import traceback
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@ -34,7 +34,7 @@ from typing import (
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, Field, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
@ -102,6 +102,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
OrganizationNotFoundError,
_cache_team_object, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
allowed_route_check_inside_route,
await_inflight_team_membership_load,
cache_team_object,
can_org_access_model,
delete_cache_key_objects,
@ -112,6 +113,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_team_object,
get_user_object,
invalidate_team_member_spend_state,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.auth_utils import (
enforce_batch_limits_are_admin_only,
@ -119,12 +121,19 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
broadcast_evictions,
evict_and_broadcast,
evict_local,
evict_shared,
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.db.db_span import db_span
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import (
@ -207,7 +216,7 @@ from litellm.proxy.management_helpers.utils import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.chunked_in import delete_many_in, find_many_in
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
@ -417,6 +426,23 @@ class _TeamDeleteTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teammembership(self) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": ...
@property
def litellm_deletedteamtable(self) -> "TableActions[prisma_models.LiteLLM_DeletedTeamTable]": ...
@property
def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
@property
def litellm_jwtkeymapping(self) -> "TableActions[prisma_models.LiteLLM_JWTKeyMapping]": ...
def _team_delete_tx(tx: "Prisma") -> _TeamDeleteTx:
return tx # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do
_LOCK_DELETED_TEAM_MEMBERS_SQL: Final = """
SELECT user_id, teams FROM "LiteLLM_UserTable" WHERE teams && $1::text[] ORDER BY user_id COLLATE "C" FOR UPDATE
"""
_STRIP_DELETED_TEAMS_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable"
@ -430,6 +456,22 @@ WHERE user_id IN (
)
"""
class _DetachedMember(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
user_id: str
teams: tuple[str, ...]
_DetachedMembers: Final = TypeAdapter(tuple[_DetachedMember, ...])
@dataclass(frozen=True, slots=True)
class DeletedTeams:
team_ids: tuple[str, ...]
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
@ -4526,18 +4568,7 @@ async def delete_team(
}'
```
"""
from litellm.proxy.management_helpers.audit_logs import (
get_audit_log_changed_by,
is_audit_logging_enabled,
)
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
litellm_proxy_admin_name,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.proxy_server import llm_router, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
@ -4567,47 +4598,32 @@ async def delete_team(
raise_if_teams_hold_ptu_shares(data.team_ids, llm_router)
await _persist_deleted_team_records(
await delete_validated_teams(
teams=team_rows,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
return deleted_teams
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if is_audit_logging_enabled():
# make an audit log for each team deleted
for team_id in data.team_ids:
team_row: LiteLLM_TeamTable | None = await prisma_client.get_data(
team_id=team_id, table_name="team", query_type="find_unique"
)
if team_row is None:
continue
async def delete_validated_teams(
teams: Sequence[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> DeletedTeams:
"""
What ``/team/delete`` does once the teams exist and the caller may delete them: every key
and BYOK model of the teams, then ``delete_validated_team_rows``.
"""
from litellm.proxy.proxy_server import (
proxy_logging_obj,
user_api_key_cache,
)
_team_row = team_row.json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=team_id,
action="deleted",
updated_values="{}",
before_value=_team_row,
)
)
)
# End of Audit logging
team_ids: Final = [team.team_id for team in teams]
## DELETE ASSOCIATED KEYS
# Fetch keys before deletion to persist them
@ -4615,7 +4631,7 @@ async def delete_team(
persist_deleted_verification_tokens,
)
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}})
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": team_ids}})
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(key.token for key in keys_to_delete),
prisma_client=prisma_client,
@ -4629,7 +4645,7 @@ async def delete_team(
litellm_changed_by=litellm_changed_by,
)
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
await prisma_client.delete_data(team_id_list=team_ids, table_name="key")
if keys_to_delete:
KeyManagementEventHooks.create_key_deleted_audit_logs(
@ -4654,18 +4670,41 @@ async def delete_team(
from litellm.proxy.proxy_server import llm_router
await delete_team_models(
team_ids=data.team_ids,
team_ids=team_ids,
prisma_client=prisma_client,
llm_router=llm_router,
)
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
member_ids_per_team: Final = await _resolve_deleted_team_member_user_ids(
teams=team_rows,
return await delete_validated_team_rows(
teams=teams,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
async def delete_validated_team_rows(
teams: Sequence[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
before_commit: Callable[["Prisma", tuple[str, ...]], Awaitable[None]] | None = None,
) -> DeletedTeams:
"""
Under every team's advisory lock, delete the teams still there with the keys their rostered
members hold on them, their membership rows and user references, and write their deleted-team
records; then evict caches and write audit rows. Keys of users off the roster and the teams'
models stay. ``team_ids`` names only the teams this call deleted. ``before_commit`` runs on the
delete's transaction with the members of the rosters read under the lock, so its writes commit
or roll back with the delete.
"""
from litellm.proxy.proxy_server import (
proxy_logging_obj,
user_api_key_cache,
)
team_ids: Final = [team.team_id for team in teams]
## DELETE TEAMS
# Both the delete and the reconcile sweep run under every team's advisory lock
# (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes),
@ -4674,14 +4713,28 @@ async def delete_team(
# the lock before this transaction starts, in which case this sweep reaches what it wrote,
# or is still waiting on the lock, in which case its own re-read happens after this commits
# and sees the row gone before it writes anything.
delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
for team_id in sorted(data.team_ids):
tx_budget: Final = timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)
async with prisma_client.tx(max_wait=tx_budget, timeout=tx_budget) as tx:
for team_id in sorted(team_ids):
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
await tx.litellm_teamtable.delete_many(where=delete_filter)
await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx)
team_tx: Final = _team_delete_tx(tx)
locked: Final = await _locked_teams_tx(team_ids=team_ids, tx=team_tx)
member_keys, jwt_mapping_cache_keys = await _delete_roster_member_keys_tx(
locked=locked,
tx=tx,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _delete_locked_teams_tx(
locked=locked, tx=team_tx, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by
)
detached_per_team: Final = await _sweep_deleted_team_references_tx(team_ids=team_ids, tx=team_tx)
if before_commit is not None:
await before_commit(tx, _locked_roster_user_ids(locked))
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
rostered_per_team: Final = await _resolve_deleted_team_member_user_ids(teams=locked, prisma_client=prisma_client)
member_ids_per_team: Final = _members_per_deleted_team(rostered_per_team, detached_per_team)
# Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and
# `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a
@ -4689,7 +4742,7 @@ async def delete_team(
# invalidated anything. Nothing fallible runs between the delete and this, or a failure there
# would strand the deleted team in cache.
await _invalidate_deleted_team_cache(
teams=team_rows,
teams=locked,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -4697,44 +4750,167 @@ async def delete_team(
member_ids_per_team=member_ids_per_team,
user_api_key_cache=user_api_key_cache,
)
await _invalidate_deleted_key_cache(
keys=member_keys,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
for deleted_team in team_rows:
_schedule_deleted_team_audit_rows(
teams=locked, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by
)
if member_keys:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=member_keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
for deleted_team in locked:
_emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()}))
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
return deleted_teams
return DeletedTeams(team_ids=tuple(team.team_id for team in locked))
async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None:
def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]:
return tuple(dict.fromkeys(chain.from_iterable(_roster_user_ids(team) for team in locked)))
def _members_per_deleted_team(
rostered: Sequence[tuple[str, Sequence[str]]], detached: Sequence[tuple[str, Sequence[str]]]
) -> tuple[tuple[str, tuple[str, ...]], ...]:
detached_by_team: Final = MappingProxyType({team_id: user_ids for team_id, user_ids in detached})
return tuple(
(team_id, tuple(dict.fromkeys((*user_ids, *detached_by_team[team_id])))) for team_id, user_ids in rostered
)
def _roster_user_ids(team: LiteLLM_TeamTable) -> tuple[str, ...]:
return tuple(member.user_id for member in team.members_with_roles if member.user_id is not None)
def _is_roster_member_key(
key: "prisma_models.LiteLLM_VerificationToken", roster_by_team: Mapping[str, frozenset[str]]
) -> bool:
return key.team_id is not None and key.user_id in roster_by_team.get(key.team_id, frozenset())
async def _locked_teams_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> tuple[LiteLLM_TeamTable, ...]:
rows: Final = await find_many_in(tx.litellm_teamtable, "team_id", team_ids)
return tuple(LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows)
async def _delete_roster_member_keys_tx(
locked: Sequence[LiteLLM_TeamTable],
tx: "Prisma",
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> tuple[tuple["prisma_models.LiteLLM_VerificationToken", ...], tuple[str, ...]]:
from litellm.proxy.management_endpoints.key_management_endpoints import (
persist_deleted_verification_tokens,
)
team_tx: Final = _team_delete_tx(tx)
roster_by_team: Final = MappingProxyType({team.team_id: frozenset(_roster_user_ids(team)) for team in locked})
team_keys: Final = await find_many_in(team_tx.litellm_verificationtoken, "team_id", tuple(roster_by_team))
member_keys: Final = tuple(key for key in team_keys if _is_roster_member_key(key, roster_by_team))
if not member_keys:
return (), ()
tokens: Final = tuple(key.token for key in member_keys)
jwt_mappings: Final = await find_many_in(team_tx.litellm_jwtkeymapping, "token", tokens)
await persist_deleted_verification_tokens(
keys=member_keys, # pyright: ignore[reportArgumentType] # generated row model carries the same columns as LiteLLM_VerificationToken
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
tx=tx,
)
await delete_many_in(team_tx.litellm_verificationtoken, "token", tokens, atomicity="caller_transaction")
return member_keys, tuple(
jwt_key_mapping_cache_key(mapping.jwt_claim_name, mapping.jwt_claim_value, mapping.jwt_issuer)
for mapping in jwt_mappings
)
async def _delete_locked_teams_tx(
locked: Sequence[LiteLLM_TeamTable],
tx: _TeamDeleteTx,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
await _persist_deleted_team_records(
teams=locked,
deleted_teams=tx.litellm_deletedteamtable,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await delete_many_in(
tx.litellm_teamtable,
"team_id",
tuple(team.team_id for team in locked),
atomicity="caller_transaction",
)
def _schedule_deleted_team_audit_rows(
teams: Sequence[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
from litellm.proxy.management_helpers.audit_logs import (
get_audit_log_changed_by,
is_audit_logging_enabled,
)
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
litellm_proxy_admin_name,
)
if not is_audit_logging_enabled():
return
for team in teams:
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=team.team_id,
action="deleted",
updated_values="{}",
before_value=team.model_dump_json(exclude_none=True),
)
)
)
async def _sweep_deleted_team_references_tx(
team_ids: Sequence[str], tx: _TeamDeleteTx
) -> tuple[tuple[str, tuple[str, ...]], ...]:
"""
Strip the deleted team ids from every user row and team-membership row that still references them.
The per-member `team_member_delete` pass above only reaches users listed in the team's
`members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps
surfacing the team on `/user/info` after the team is gone.
#36839 closed the route that created that drift, by resolving member removal off the roster
entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill
rows that already drifted, which is the state this was reported against, so the sweep still has
to run on delete.
One statement filters the array in place rather than read-filter-write: rewriting the whole
array from a snapshot read outside a transaction drops any team a concurrent `/team/member_add`
appended in between. It locks the user rows in the order `USER_ROWS_LOCK_SQL` takes them, so it
never waits on a SCIM group push in a cycle.
Strip the deleted team ids from every user and team-membership row still referencing them, on
the transaction holding the teams' advisory locks. User rows are locked first in
``USER_ROWS_LOCK_SQL`` order so the sweep never deadlocks with a SCIM group push, and one
statement filters each array in place so a team a concurrent ``/team/member_add`` appended is
kept. Returns, per team, the users read under the lock as referencing it.
"""
_ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids))
_ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)}))
async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None:
"""Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds
every id's advisory lock and deletes the team rows, so it commits or rolls back with them."""
locked: Final = _DetachedMembers.validate_python(await tx.query_raw(_LOCK_DELETED_TEAM_MEMBERS_SQL, list(team_ids)))
_ = await tx.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids))
membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}}
_ = await tx.litellm_teammembership.delete_many(where=membership_filter)
return tuple(
(team_id, tuple(member.user_id for member in locked if team_id in member.teams)) for team_id in team_ids
)
async def _invalidate_deleted_key_cache(
@ -4791,17 +4967,20 @@ async def _evict_deleted_team_member_cache(
member_user_ids: Sequence[str],
user_api_key_cache: UserApiKeyCache,
) -> None:
await evict_and_broadcast(cache_keys=tuple(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_id,
user_api_key_cache=user_api_key_cache,
)
for user_id in member_user_ids
)
_ = await asyncio.gather(
*(await_inflight_team_membership_load(user_id=user_id, team_id=team_id) for user_id in member_user_ids)
)
cache_keys: Final = tuple(member_cache_keys(team_id, member_user_ids))
evict_local(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=cache_keys, cache=user_api_key_cache)
await broadcast_evictions(cache_keys=cache_keys)
def member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]:
for user_id in user_ids:
yield user_id
yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id)
yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
async def _resolve_deleted_team_member_user_ids(
@ -4832,7 +5011,7 @@ async def _deleted_team_member_user_ids(team: LiteLLM_TeamTable, prisma_client:
def _transform_teams_to_deleted_records(
teams: list[LiteLLM_TeamTable],
teams: Sequence[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> list[dict[str, object]]:
@ -4881,17 +5060,17 @@ def _transform_teams_to_deleted_records(
async def _save_deleted_team_records(
records: list[dict[str, object]],
prisma_client: PrismaClient,
deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]",
) -> None:
"""Save deleted team records to the database."""
if not records:
return
await _deleted_team_db(prisma_client).create_many(data=records)
await deleted_teams.create_many(data=records)
async def _persist_deleted_team_records(
teams: list[LiteLLM_TeamTable],
prisma_client: PrismaClient,
teams: Sequence[LiteLLM_TeamTable],
deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]",
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> None:
@ -4903,7 +5082,7 @@ async def _persist_deleted_team_records(
)
await _save_deleted_team_records(
records=records,
prisma_client=prisma_client,
deleted_teams=deleted_teams,
)

61
litellm/proxy/management_helpers/audit_logs.py
View file

@ -4,6 +4,7 @@ Functions to create audit logs for LiteLLM Proxy
import asyncio
import json
from collections.abc import Awaitable, Sequence
from datetime import datetime, timezone
from typing import Final, TypeVar
@ -258,12 +259,19 @@ async def create_object_audit_log(
)
async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None:
async def _insert_audit_logs(insert: Awaitable[object], description: str) -> None:
try:
await repository.table.create(data={**request_data.model_dump(exclude_none=True)})
await insert
except Exception as e:
# [Non-Blocking Exception. Do not allow blocking LLM API call]
verbose_proxy_logger.error("Failed Creating audit log %s", e)
verbose_proxy_logger.error("Failed Creating %s %s", description, e)
def _json_encoded(request_data: LiteLLM_AuditLogs) -> LiteLLM_AuditLogs:
values: Final = (("updated_values", request_data.updated_values), ("before_value", request_data.before_value))
return request_data.model_copy(
update={field: json.dumps(value) for field, value in values if isinstance(value, dict)}
)
async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
@ -279,15 +287,10 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
return
verbose_proxy_logger.debug("creating audit log for %s", request_data)
if isinstance(request_data.updated_values, dict):
request_data.updated_values = json.dumps(request_data.updated_values)
if isinstance(request_data.before_value, dict):
request_data.before_value = json.dumps(request_data.before_value)
encoded: Final = _json_encoded(request_data)
# Dispatch to external audit log callbacks regardless of DB availability
await _dispatch_audit_log_to_callbacks(request_data)
await _dispatch_audit_log_to_callbacks(encoded)
if prisma_client is None:
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
@ -295,7 +298,41 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
await track_audit_task(
asyncio.create_task(
_insert_audit_log_row(AuditLogRepository(prisma_client), request_data),
name=_audit_task_name(request_data, "the database"),
_insert_audit_logs(
AuditLogRepository(prisma_client).table.create(data={**encoded.model_dump(exclude_none=True)}),
"audit log",
),
name=_audit_task_name(encoded, "the database"),
)
)
async def create_audit_logs_for_updates(rows: Sequence[LiteLLM_AuditLogs]) -> None:
"""Create audit logs for many objects with one insert."""
if not rows or not is_audit_logging_enabled():
return
from litellm.proxy.proxy_server import premium_user, prisma_client
if premium_user is not True:
return
encoded: Final = tuple(_json_encoded(row) for row in rows)
for row in encoded:
await _dispatch_audit_log_to_callbacks(row)
if prisma_client is None:
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
return
await track_audit_task(
asyncio.create_task(
_insert_audit_logs(
AuditLogRepository(prisma_client).table.create_many(
data=tuple(row.model_dump(exclude_none=True) for row in encoded)
),
f"{len(encoded)} audit logs",
),
name=f"{len(encoded)} audit logs to the database",
)
)

14
litellm/proxy/management_helpers/team_roster_sync.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import timedelta
from types import MappingProxyType
@ -23,8 +23,6 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
)
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -33,6 +31,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.team_endpoints import (
_emit_team_members_metric, # pyright: ignore[reportPrivateUsage] # same gauge /team/member_add moves
_schedule_team_membership_audit_log, # pyright: ignore[reportPrivateUsage] # same audit row /team/member_add writes
member_cache_keys,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
DETACH_TEAM_FROM_USERS_SQL,
@ -277,7 +276,7 @@ async def _settle_caches(
evicted: Final = (
*deleted_tokens,
*removal.jwt_mapping_cache_keys,
*_member_cache_keys(team.team_id, changed_user_ids),
*member_cache_keys(team.team_id, changed_user_ids),
)
evict_local(cache_keys=evicted, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=evicted, cache=user_api_key_cache)
@ -292,13 +291,6 @@ async def _settle_caches(
await broadcast_evictions(cache_keys=evicted)
def _member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]:
for user_id in user_ids:
yield user_id
yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id)
yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
async def sync_team_roster(
prisma_client: PrismaClient,
team_id: str,

9
litellm/proxy/utils.py
View file

@ -4832,6 +4832,7 @@ def jsonify_object(data: dict) -> dict:
_deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000)
_DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60
_PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5)
_PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2)
async def _lookup_deprecated_key(
@ -5238,13 +5239,17 @@ class PrismaClient:
return self.db.read_target
return self.db
def tx(self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT) -> "TransactionManager":
def tx(
self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT, max_wait: timedelta = _PRISMA_DEFAULT_TX_MAX_WAIT
) -> "TransactionManager":
"""Open an interactive transaction on the writer.
Callers go through this instead of reaching into ``self.db`` so writer
selection and read-replica routing stay encapsulated in the wrapper.
"""
return cast("TransactionManager", self.db.tx(timeout=timeout)) # cast-ok: untyped __getattr__ delegate
return cast( # cast-ok: untyped __getattr__ delegate
"TransactionManager", self.db.tx(max_wait=max_wait, timeout=timeout)
)
def get_request_status(self, payload: dict | SpendLogsPayload) -> Literal["success", "failure"]:
"""

4
tests/code_coverage_tests/unbounded_in_baseline.txt
View file

@ -80,14 +80,12 @@ litellm/proxy/management_endpoints/team_endpoints.py _get_keys_count_by_team pri
litellm/proxy/management_endpoints/team_endpoints.py _hydrate_member_user_details prisma user_id.in `sorted(user_ids)` 0
litellm/proxy/management_endpoints/team_endpoints.py _resolve_existing_member_user_ids prisma user_id.in `sorted(requested_user_ids)` 0
litellm/proxy/management_endpoints/team_endpoints.py resolve_team_daily_activity_scope prisma team_id.in `list(team_ids_list)` 0
litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references prisma team_id.in `tuple(team_ids)` 0
litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references_tx prisma team_id.in `tuple(team_ids)` 0
litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(addressed_user_ids)` 0
litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 0
litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 1
litellm/proxy/management_endpoints/team_endpoints.py _team_user_spend_sql raw-sql sl.team_id.IN `IN ({team_placeholders})` 0
litellm/proxy/management_endpoints/team_endpoints.py delete_team prisma team_id.in `data.team_ids` 0
litellm/proxy/management_endpoints/team_endpoints.py delete_team prisma team_id.in `data.team_ids` 1
litellm/proxy/management_endpoints/team_endpoints.py delete_validated_teams prisma team_id.in `team_ids` 0
litellm/proxy/management_endpoints/team_endpoints.py get_all_team_memberships prisma team_id.in `team_ids` 0
litellm/proxy/management_endpoints/team_endpoints.py list_available_teams prisma team_id.in `available_teams` 0
litellm/proxy/management_endpoints/tool_management_endpoints.py get_tool_spend prisma tool_name.in `[row.tool_name for row in top_tools]` 0

207
tests/integration/_support/scim.py Normal file
View file

@ -0,0 +1,207 @@
"""SCIM v2 group helpers for the management cells: seeded users, group pushes, the team lock, and what a group
leaves behind in Postgres."""
import os
import uuid
from collections.abc import Generator, Mapping, Sequence
from contextlib import contextmanager
from typing import Final
import httpx
import psycopg
from integration._support.client import (
JSON_OBJECT,
Gateway,
Scenario,
delete_key_if_present,
object_value,
string_value,
)
from integration._support.database import read_rows, write_rows
from pydantic import JsonValue
GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group"
PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp"
SEED_USERS_SQL: Final = """
INSERT INTO "LiteLLM_UserTable" (user_id, user_role, teams, models)
SELECT %s || '-' || lpad(n::text, 3, '0'), 'internal_user', '{}'::text[], '{}'::text[]
FROM generate_series(1, %s::int) AS n
"""
TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))"
WAITERS_ON_HELD_LOCK_SQL: Final = """
SELECT count(*)::text AS waiting
FROM pg_locks waiter
JOIN pg_stat_activity session ON session.pid = waiter.pid
WHERE waiter.locktype = 'advisory'
AND NOT waiter.granted
AND session.wait_event_type = 'Lock'
AND (waiter.classid, waiter.objid, waiter.objsubid) IN (
SELECT held.classid, held.objid, held.objsubid
FROM pg_locks held
WHERE held.locktype = 'advisory' AND held.granted AND held.pid = %s::int
)
"""
TAKE_USER_ROW_LOCK_SQL: Final = 'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s FOR UPDATE'
WAITERS_ON_ROWS_LOCKED_BY_SQL: Final = """
SELECT count(*)::text AS waiting
FROM pg_locks waiter
WHERE waiter.locktype = 'transactionid'
AND NOT waiter.granted
AND waiter.transactionid IN (
SELECT held.transactionid
FROM pg_locks held
WHERE held.locktype = 'transactionid' AND held.granted AND held.pid = %s::int
)
"""
def seed_users(scenario: Scenario, count: int) -> tuple[str, ...]:
prefix: Final = f"integration-scim-{uuid.uuid4().hex}"
write_rows(SEED_USERS_SQL, (prefix, str(count)))
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (f"{prefix}-%",))
return tuple(f"{prefix}-{index:03d}" for index in range(1, count + 1))
def members(users: Sequence[str]) -> list[JsonValue]:
return [{"value": user} for user in users]
def group_name() -> str:
return f"integration-{uuid.uuid4().hex}"
def create_group(candidate: Gateway, users: Sequence[str], *, display_name: str | None = None) -> httpx.Response:
return candidate.request(
"POST",
"/scim/v2/Groups",
{"schemas": [GROUP_SCHEMA], "displayName": display_name or group_name(), "members": members(users)},
)
def replace_group(
candidate: Gateway, team: str, users: Sequence[str], *, display_name: str | None = None
) -> httpx.Response:
return candidate.request(
"PUT",
f"/scim/v2/Groups/{team}",
{"schemas": [GROUP_SCHEMA], "id": team, "displayName": display_name or group_name(), "members": members(users)},
)
def add_to_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response:
return candidate.request(
"PATCH",
f"/scim/v2/Groups/{team}",
{"schemas": [PATCH_OP_SCHEMA], "Operations": [{"op": "add", "path": "members", "value": members(users)}]},
)
def body(response: httpx.Response) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_json(response.content)
def member_ids(group: Mapping[str, JsonValue]) -> frozenset[str]:
listed: Final = group.get("members") or []
assert isinstance(listed, list), listed
return frozenset(string_value(object_value(member)["value"]) for member in listed)
def delete_group(candidate: Gateway, team: str, *, key: str | None = None) -> httpx.Response:
return candidate.request("DELETE", f"/scim/v2/Groups/{team}", key=key)
def delete_team_if_present(scenario: Scenario, team: str) -> None:
if read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,)):
scenario.delete_team(team)
def created_team(scenario: Scenario, created: httpx.Response) -> str:
assert created.status_code == 201, created.text
team: Final = string_value(body(created)["id"])
scenario.cleanups.callback(delete_team_if_present, scenario, team)
return team
def team_key(candidate: Gateway, team: str, user: str) -> str:
created: Final = candidate.post(
"/key/generate", {"team_id": team, "user_id": user, "key_alias": f"integration-{uuid.uuid4().hex}"}
)
return string_value(created["key"])
def service_account_key(candidate: Gateway, scenario: Scenario, team: str) -> str:
created: Final = candidate.post("/key/generate", {"team_id": team, "key_alias": f"integration-{uuid.uuid4().hex}"})
key: Final = string_value(created["key"])
scenario.cleanups.callback(delete_key_if_present, candidate, key)
return key
def membership_user_ids(team: str) -> frozenset[str]:
rows: Final = read_rows('SELECT user_id FROM "LiteLLM_TeamMembership" WHERE team_id = %s', (team,))
return frozenset(string_value(row["user_id"]) for row in rows)
def keys_of(team: str) -> frozenset[str]:
rows: Final = read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,))
return frozenset(string_value(row["token"]) for row in rows)
def users_referencing(team: str) -> frozenset[str]:
rows: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE %s = ANY(teams)', (team,))
return frozenset(string_value(row["user_id"]) for row in rows)
def user_role(candidate: Gateway, user: str) -> str:
return string_value(object_value(candidate.get("/user/info", {"user_id": user})["user_info"])["user_role"])
def groups_of(candidate: Gateway, user: str) -> frozenset[str]:
groups: Final = candidate.get(f"/scim/v2/Users/{user}").get("groups") or []
assert isinstance(groups, list), groups
return frozenset(string_value(object_value(group)["value"]) for group in groups)
def assert_landed(candidate: Gateway, team: str, users: Sequence[str]) -> None:
assert member_ids(candidate.get(f"/scim/v2/Groups/{team}")) == frozenset(users)
assert membership_user_ids(team) == frozenset(users)
assert users_referencing(team) == frozenset(users)
def assert_gone(candidate: Gateway, team: str, users: Sequence[str], key: str) -> None:
assert candidate.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
assert team not in groups_of(candidate, users[0])
assert membership_user_ids(team) == frozenset()
assert users_referencing(team) == frozenset()
assert keys_of(team) == frozenset()
assert candidate.request("GET", "/v1/models", key=key).status_code == 401
@contextmanager
def held_team_lock(team: str) -> Generator[int]:
"""Hold the team's advisory lock from a session of our own and yield its backend pid; leaving the block
commits, which releases the lock to whoever is queued on it."""
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(TAKE_TEAM_LOCK_SQL, (team,))
yield holder.info.backend_pid
@contextmanager
def held_user_row(user: str) -> Generator[int]:
"""Hold a row lock on the user's row from a session of our own and yield its backend pid, so a role write
that reaches the row waits until leaving the block commits."""
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(TAKE_USER_ROW_LOCK_SQL, (user,))
yield holder.info.backend_pid
def waiters_on_lock_held_by(backend_pid: int) -> int:
return _waiting(WAITERS_ON_HELD_LOCK_SQL, backend_pid)
def waiters_on_rows_locked_by(backend_pid: int) -> int:
return _waiting(WAITERS_ON_ROWS_LOCKED_BY_SQL, backend_pid)
def _waiting(query: str, backend_pid: int) -> int:
rows: Final = read_rows(query, (str(backend_pid),))
return int(string_value(rows[0]["waiting"]))

272
tests/integration/management/test_scim_group_bulk_roster.py
View file

@ -3,7 +3,10 @@
POST /scim/v2/Groups, PUT and PATCH /scim/v2/Groups/{id} used to run /team/member_add once per member, so a
500-member push was thousands of statements and outlasted the edge in front of the proxy. One locked
transaction now resolves every member, writes the roster and the membership rows and tags the users, so a
push of 500 members runs as many statements as a push of 5.
push of 500 members runs as many statements as a push of 5. DELETE /scim/v2/Groups/{id} used to detach the
members one at a time and leave their membership rows and keys behind; it now detaches them in one locked pass,
so deleting 500 members who each hold a key costs as many statements as deleting 5 and leaves no membership rows
or member keys.
"""
import os
@ -13,6 +16,7 @@ from collections.abc import Callable, Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from hashlib import sha256
from pathlib import Path
from typing import Final, LiteralString
@ -22,7 +26,6 @@ import psycopg
import pytest
from integration._support.budget import case_timeout_seconds
from integration._support.client import (
JSON_OBJECT,
Gateway,
Scenario,
delete_key_if_present,
@ -32,14 +35,30 @@ from integration._support.client import (
)
from integration._support.database import read_rows, write_rows
from integration._support.database_relay import StatementCountingRelay, statement_counting_relay
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.scim import (
GROUP_SCHEMA,
PATCH_OP_SCHEMA,
add_to_group,
assert_gone,
assert_landed,
body,
create_group,
created_team,
delete_group,
member_ids,
members,
membership_user_ids,
replace_group,
seed_users,
team_key,
users_referencing,
)
from psycopg import sql
from pydantic import JsonValue
pytestmark: Final = pytest.mark.shares_worker_state
GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group"
PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp"
SMALL_PUSH: Final = 5
LARGE_PUSH: Final = 500
# Background work on an idle proxy (the cron-job leader poll, spend flushes) lands a few statements inside a
@ -48,7 +67,9 @@ NOISE_ALLOWANCE: Final = 100
UNFIXED_PUSH_SECONDS: Final = 120
SWAP_GROUP: Final = 200
SWAP_ROUNDS: Final = 10
OWNED_PROXY_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 2 * UNFIXED_PUSH_SECONDS)
CHAOS_GROUP: Final = 100
KEY_WRITERS: Final = 16
ROLLBACK_SECONDS: Final = 90
TEAM_ROW_SQL: Final = 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = %s'
LOCK_USER_ROW_SQL: Final = 'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s FOR UPDATE'
@ -56,16 +77,12 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t
SESSIONS_BLOCKED_BY_SQL: Final = (
"SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))"
)
BACKENDS_BLOCKED_BY_SQL: Final = "SELECT pid FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))"
TRANSACTIONS_OPEN_SINCE_SQL: Final = """
SELECT count(*)::int AS open FROM pg_stat_activity
WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid()
AND xact_start < %s::timestamptz
"""
SEED_USERS_SQL: Final = """
INSERT INTO "LiteLLM_UserTable" (user_id, user_role, teams, models)
SELECT %s || '-' || lpad(n::text, 3, '0'), 'internal_user', '{}'::text[], '{}'::text[]
FROM generate_series(1, %s::int) AS n
"""
@dataclass(frozen=True, slots=True)
@ -82,71 +99,6 @@ def _measured(relay: StatementCountingRelay, send: Callable[[], httpx.Response])
return Measured(response=response, statements=relay.statements - before, seconds=time.perf_counter() - started)
def _seed_users(scenario: Scenario, count: int) -> tuple[str, ...]:
prefix: Final = f"integration-scim-{uuid.uuid4().hex}"
write_rows(SEED_USERS_SQL, (prefix, str(count)))
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (f"{prefix}-%",))
return tuple(f"{prefix}-{index:03d}" for index in range(1, count + 1))
def _members(users: Sequence[str]) -> list[JsonValue]:
return [{"value": user} for user in users]
def _create_group(candidate: Gateway, users: Sequence[str]) -> httpx.Response:
return candidate.request(
"POST",
"/scim/v2/Groups",
{"schemas": [GROUP_SCHEMA], "displayName": f"integration-{uuid.uuid4().hex}", "members": _members(users)},
)
def _replace_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response:
return candidate.request(
"PUT",
f"/scim/v2/Groups/{team}",
{
"schemas": [GROUP_SCHEMA],
"id": team,
"displayName": f"integration-{uuid.uuid4().hex}",
"members": _members(users),
},
)
def _add_to_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response:
return candidate.request(
"PATCH",
f"/scim/v2/Groups/{team}",
{
"schemas": [PATCH_OP_SCHEMA],
"Operations": [{"op": "add", "path": "members", "value": _members(users)}],
},
)
def _body(response: httpx.Response) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_json(response.content)
def _member_ids(body: Mapping[str, JsonValue]) -> frozenset[str]:
members: Final = body.get("members") or []
assert isinstance(members, list), members
return frozenset(string_value(object_value(member)["value"]) for member in members)
def _created_team(scenario: Scenario, created: httpx.Response) -> str:
assert created.status_code == 201, created.text
team: Final = string_value(_body(created)["id"])
scenario.cleanups.callback(scenario.delete_team, team)
return team
def _membership_user_ids(team: str) -> frozenset[str]:
rows: Final = read_rows('SELECT user_id FROM "LiteLLM_TeamMembership" WHERE team_id = %s', (team,))
return frozenset(string_value(row["user_id"]) for row in rows)
_UTC_MICROSECONDS: Final = "YYYY-MM-DD HH24:MI:SS.US"
@ -163,18 +115,7 @@ def _database_utc_now() -> str:
return string_value(rows[0]["utc_now"])
def _users_referencing(team: str) -> frozenset[str]:
rows: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE %s = ANY(teams)', (team,))
return frozenset(string_value(row["user_id"]) for row in rows)
def _assert_landed(candidate: Gateway, team: str, users: Sequence[str]) -> None:
assert _member_ids(candidate.get(f"/scim/v2/Groups/{team}")) == frozenset(users)
assert _membership_user_ids(team) == frozenset(users)
assert _users_referencing(team) == frozenset(users)
@pytest.mark.timeout(300) # owned proxy boot plus three 500-member pushes
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5(
gateway: Gateway, tmp_path: Path, record_property: Callable[[str, object], None]
) -> None:
@ -191,39 +132,50 @@ def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5(
) as owned,
):
candidate: Final = owned.gateway
small_a: Final = _seed_users(scenario, SMALL_PUSH)
small_b: Final = _seed_users(scenario, SMALL_PUSH)
large: Final = _seed_users(scenario, LARGE_PUSH)
small_a: Final = seed_users(scenario, SMALL_PUSH)
small_b: Final = seed_users(scenario, SMALL_PUSH)
large: Final = seed_users(scenario, LARGE_PUSH)
created_small: Final = _measured(relay, lambda: _create_group(candidate, small_a))
small_team: Final = _created_team(scenario, created_small.response)
created_large: Final = _measured(relay, lambda: _create_group(candidate, large))
large_team: Final = _created_team(scenario, created_large.response)
assert _member_ids(_body(created_large.response)) == frozenset(large)
_assert_landed(candidate, small_team, small_a)
_assert_landed(candidate, large_team, large)
created_small: Final = _measured(relay, lambda: create_group(candidate, small_a))
small_team: Final = created_team(scenario, created_small.response)
created_large: Final = _measured(relay, lambda: create_group(candidate, large))
large_team: Final = created_team(scenario, created_large.response)
assert member_ids(body(created_large.response)) == frozenset(large)
assert_landed(candidate, small_team, small_a)
assert_landed(candidate, large_team, large)
detach_window_open: Final = _database_utc_now()
replaced_small: Final = _measured(relay, lambda: _replace_group(candidate, small_team, small_b))
replaced_large: Final = _measured(relay, lambda: _replace_group(candidate, large_team, small_a))
replaced_small: Final = _measured(relay, lambda: replace_group(candidate, small_team, small_b))
replaced_large: Final = _measured(relay, lambda: replace_group(candidate, large_team, small_a))
assert replaced_small.response.status_code == 200, replaced_small.response.text
assert replaced_large.response.status_code == 200, replaced_large.response.text
assert _member_ids(_body(replaced_small.response)) == frozenset(small_b)
_assert_landed(candidate, small_team, small_b)
_assert_landed(candidate, large_team, small_a)
assert member_ids(body(replaced_small.response)) == frozenset(small_b)
assert_landed(candidate, small_team, small_b)
assert_landed(candidate, large_team, small_a)
assert detach_window_open <= _updated_at(large[0]) <= _database_utc_now()
patched_small: Final = _measured(relay, lambda: _add_to_group(candidate, small_team, small_a))
patched_large: Final = _measured(relay, lambda: _add_to_group(candidate, large_team, large))
patched_small: Final = _measured(relay, lambda: add_to_group(candidate, small_team, small_a))
patched_large: Final = _measured(relay, lambda: add_to_group(candidate, large_team, large))
assert patched_small.response.status_code == 200, patched_small.response.text
assert patched_large.response.status_code == 200, patched_large.response.text
_assert_landed(candidate, small_team, [*small_a, *small_b])
_assert_landed(candidate, large_team, [*small_a, *large])
assert_landed(candidate, small_team, [*small_a, *small_b])
assert_landed(candidate, large_team, [*small_a, *large])
with ThreadPoolExecutor(max_workers=KEY_WRITERS) as pool:
small_keys: Final = tuple(pool.map(partial(team_key, candidate, small_team), (*small_a, *small_b)))
large_keys: Final = tuple(pool.map(partial(team_key, candidate, large_team), (*small_a, *large)))
deleted_small: Final = _measured(relay, lambda: delete_group(candidate, small_team))
deleted_large: Final = _measured(relay, lambda: delete_group(candidate, large_team))
assert deleted_small.response.status_code == 204, deleted_small.response.text
assert deleted_large.response.status_code == 204, deleted_large.response.text
assert_gone(candidate, small_team, [*small_a, *small_b], small_keys[0])
assert_gone(candidate, large_team, [*small_a, *large], large_keys[-1])
pushes: Final = (
("POST", created_small, created_large),
("PUT", replaced_small, replaced_large),
("PATCH", patched_small, patched_large),
("DELETE", deleted_small, deleted_large),
)
for verb, small, big in pushes:
record_property(
@ -238,13 +190,13 @@ def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5(
def _trade_rosters(candidate: Gateway, teams: tuple[str, str], rosters: tuple[Sequence[str], Sequence[str]]) -> None:
with ThreadPoolExecutor(max_workers=2) as pool:
left: Final = pool.submit(_replace_group, candidate, teams[0], rosters[1])
right: Final = pool.submit(_replace_group, candidate, teams[1], rosters[0])
left: Final = pool.submit(replace_group, candidate, teams[0], rosters[1])
right: Final = pool.submit(replace_group, candidate, teams[1], rosters[0])
for response in (left.result(), right.result()):
assert response.status_code == 200, response.text
@pytest.mark.timeout(300)
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_two_groups_trading_rosters_in_parallel_both_land(gateway: Gateway, tmp_path: Path) -> None:
with (
gateway.scenario() as scenario,
@ -258,15 +210,18 @@ def test_two_groups_trading_rosters_in_parallel_both_land(gateway: Gateway, tmp_
) as owned,
):
candidate: Final = owned.gateway
left_users: Final = _seed_users(scenario, SWAP_GROUP)
right_users: Final = _seed_users(scenario, SWAP_GROUP)
left: Final = _created_team(scenario, _create_group(candidate, left_users))
right: Final = _created_team(scenario, _create_group(candidate, right_users))
for round_index in range(SWAP_ROUNDS):
held: Final = (left_users, right_users) if round_index % 2 == 0 else (right_users, left_users)
left_users: Final = seed_users(scenario, SWAP_GROUP)
right_users: Final = seed_users(scenario, SWAP_GROUP)
left: Final = created_team(scenario, create_group(candidate, left_users))
right: Final = created_team(scenario, create_group(candidate, right_users))
rounds: Final = tuple(
(left_users, right_users) if round_index % 2 == 0 else (right_users, left_users)
for round_index in range(SWAP_ROUNDS)
)
for held in rounds:
_trade_rosters(candidate, (left, right), held)
_assert_landed(candidate, left, left_users)
_assert_landed(candidate, right, right_users)
assert_landed(candidate, left, left_users)
assert_landed(candidate, right, right_users)
def _int_value(rows: Sequence[Mapping[str, JsonValue]], column: str) -> int:
@ -286,6 +241,10 @@ def _sessions_blocked_by(backend_pid: int) -> int:
return _int_value(read_rows(SESSIONS_BLOCKED_BY_SQL, (str(backend_pid),)), "waiting")
def _backend_blocked_by(backend_pid: int) -> int:
return _int_value(read_rows(BACKENDS_BLOCKED_BY_SQL, (str(backend_pid),)), "pid")
def _database_now() -> str:
return string_value(read_rows("SELECT now()::text AS now", ())[0]["now"])
@ -302,36 +261,25 @@ def _roster_ids(team: str) -> tuple[str, ...]:
return tuple(string_value(object_value(member)["user_id"]) for member in roster)
def _delete_team_if_present(candidate: Gateway, team: str) -> None:
if read_rows(TEAM_ROW_SQL, (team,)):
_ = candidate.post("/team/delete", {"team_ids": [team]})
def _group_the_test_deletes(candidate: Gateway, scenario: Scenario) -> str:
created: Final = _create_group(candidate, ())
assert created.status_code == 201, created.text
team: Final = string_value(_body(created)["id"])
scenario.cleanups.callback(_delete_team_if_present, candidate, team)
return team
def test_a_group_deleted_while_its_push_waits_on_a_member_row_answers_404_and_writes_nothing(
def test_a_group_deleted_while_its_push_waits_on_a_member_row_waits_for_the_push_then_removes_every_trace(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool:
users: Final = _seed_users(scenario, SMALL_PUSH)
team: Final = _group_the_test_deletes(gateway, scenario)
with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=2) as pool:
users: Final = seed_users(scenario, SMALL_PUSH)
team: Final = created_team(scenario, create_group(gateway, ()))
with _holding_row_lock(LOCK_USER_ROW_SQL, min(users)) as holder_pid:
push: Final = pool.submit(_replace_group, gateway, team, users)
push: Final = pool.submit(replace_group, gateway, team, users)
_ = eventually(lambda: _sessions_blocked_by(holder_pid), lambda waiting: waiting >= 1, seconds=30)
deleted: Final = gateway.request("DELETE", f"/scim/v2/Groups/{team}")
assert deleted.status_code == 204, deleted.text
push_pid: Final = _backend_blocked_by(holder_pid)
deleted: Final = pool.submit(delete_group, gateway, team)
_ = eventually(lambda: _sessions_blocked_by(push_pid), lambda waiting: waiting >= 1, seconds=30)
assert not push.done(), push.result().text
pushed: Final = push.result()
assert pushed.status_code == 404, pushed.text
assert not deleted.done(), deleted.result().text
assert push.result().status_code == 200, push.result().text
assert deleted.result().status_code == 204, deleted.result().text
assert read_rows(TEAM_ROW_SQL, (team,)) == []
assert _membership_user_ids(team) == frozenset()
assert _users_referencing(team) == frozenset()
assert membership_user_ids(team) == frozenset()
assert users_referencing(team) == frozenset()
@contextmanager
@ -361,17 +309,17 @@ def test_a_push_whose_membership_write_fails_answers_5xx_keeps_the_old_roster_an
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
kept: Final = _seed_users(scenario, SMALL_PUSH)
arriving: Final = _seed_users(scenario, SMALL_PUSH)
team: Final = _created_team(scenario, _create_group(gateway, kept))
kept: Final = seed_users(scenario, SMALL_PUSH)
arriving: Final = seed_users(scenario, SMALL_PUSH)
team: Final = created_team(scenario, create_group(gateway, kept))
replacement: Final = (*kept[:2], *arriving)
with _refusing_memberships_for(max(arriving)):
refused: Final = _replace_group(gateway, team, replacement)
refused: Final = replace_group(gateway, team, replacement)
assert 500 <= refused.status_code < 600, refused.text
_assert_landed(gateway, team, kept)
retried: Final = _replace_group(gateway, team, replacement)
assert_landed(gateway, team, kept)
retried: Final = replace_group(gateway, team, replacement)
assert retried.status_code == 200, retried.text
_assert_landed(gateway, team, replacement)
assert_landed(gateway, team, replacement)
def _five_kilobyte_member(scenario: Scenario) -> str:
@ -394,8 +342,8 @@ def test_a_push_with_an_odd_member_value_never_answers_500_and_keeps_the_roster_
gateway: Gateway, shape: str
) -> None:
with gateway.scenario() as scenario:
kept: Final = _seed_users(scenario, SMALL_PUSH)
team: Final = _created_team(scenario, _create_group(gateway, kept))
kept: Final = seed_users(scenario, SMALL_PUSH)
team: Final = created_team(scenario, create_group(gateway, kept))
odd: Final = ODD_MEMBER_VALUES[shape](scenario, kept)
pushed: Final = gateway.request(
"PUT",
@ -404,14 +352,14 @@ def test_a_push_with_an_odd_member_value_never_answers_500_and_keeps_the_roster_
"schemas": [GROUP_SCHEMA],
"id": team,
"displayName": f"integration-{uuid.uuid4().hex}",
"members": [*_members(kept), {"value": odd}],
"members": [*members(kept), {"value": odd}],
},
)
assert pushed.status_code < 500, pushed.text
roster: Final = _roster_ids(team)
assert len(roster) == len(frozenset(roster)), roster
assert frozenset(kept) <= frozenset(roster), roster
_assert_landed(gateway, team, roster)
assert_landed(gateway, team, roster)
if pushed.status_code >= 400:
assert frozenset(roster) == frozenset(kept), (pushed.text, roster)
@ -443,15 +391,15 @@ PATCH_SHAPES: Final[Mapping[str, PatchShape]] = {
@pytest.mark.parametrize("shape", tuple(PATCH_SHAPES))
def test_patch_member_operation_shapes_land_the_expected_roster(gateway: Gateway, shape: str) -> None:
with gateway.scenario() as scenario:
users: Final = _seed_users(scenario, 4)
team: Final = _created_team(scenario, _create_group(gateway, users[:3]))
users: Final = seed_users(scenario, 4)
team: Final = created_team(scenario, create_group(gateway, users[:3]))
patched: Final = gateway.request(
"PATCH",
f"/scim/v2/Groups/{team}",
{"schemas": [PATCH_OP_SCHEMA], "Operations": PATCH_SHAPES[shape].operations(users)},
)
assert patched.status_code == 200, patched.text
_assert_landed(gateway, team, PATCH_SHAPES[shape].expected(users))
assert_landed(gateway, team, PATCH_SHAPES[shape].expected(users))
def _team_key_for(scenario: Scenario, team: str, user_id: str) -> str:
@ -462,7 +410,7 @@ def _team_key_for(scenario: Scenario, team: str, user_id: str) -> str:
def _push_outcome(candidate: Gateway, team: str, users: Sequence[str]) -> str:
try:
return str(_replace_group(candidate, team, users).status_code)
return str(replace_group(candidate, team, users).status_code)
except httpx.TransportError as error:
return type(error).__name__
@ -484,9 +432,9 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne
ThreadPoolExecutor(max_workers=1) as pool,
):
candidate: Final = owned.gateway
leaving: Final = _seed_users(scenario, CHAOS_GROUP)
arriving: Final = _seed_users(scenario, CHAOS_GROUP)
team: Final = _created_team(scenario, _create_group(candidate, leaving))
leaving: Final = seed_users(scenario, CHAOS_GROUP)
arriving: Final = seed_users(scenario, CHAOS_GROUP)
team: Final = created_team(scenario, create_group(candidate, leaving))
pinned_key: Final = _team_key_for(scenario, team, max(leaving))
with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid:
push: Final = pool.submit(_push_outcome, candidate, team, arriving)
@ -508,11 +456,11 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne
{"leaving": len(settled & frozenset(leaving)), "arriving": len(settled & frozenset(arriving))},
)
assert settled in (frozenset(leaving), frozenset(arriving)), sorted(settled)
_assert_landed(candidate, team, sorted(settled))
assert_landed(candidate, team, sorted(settled))
retried: Final = _replace_group(candidate, team, arriving)
retried: Final = replace_group(candidate, team, arriving)
assert retried.status_code == 200, retried.text
_assert_landed(candidate, team, arriving)
assert_landed(candidate, team, arriving)
assert (
read_rows(
'SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s',

431
tests/integration/management/test_scim_group_delete.py Normal file
View file

@ -0,0 +1,431 @@
"""DELETE /scim/v2/Groups/{id} deletes the group's team the way /team/delete does.
The members' group entries, their membership rows and the keys the members hold on the team go in one locked
pass whatever the roster size, and the configured admin group's members are demoted in that same pass, so they stay
admins until the group is gone. A member a concurrent group write added while the delete waited on the team lock is
demoted with the rest, and a group write that arrives once the delete holds the lock finds the group gone. The group
is 404 after, on every worker and on the peer proxy, and an unrelated key keeps serving.
A service-account key and a model on the team stay, so re-creating the group with the same id brings them back.
"""
import os
import re
import uuid
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from functools import partial
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import GATEWAY_LIMITS, JSON_OBJECT, Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.scim import (
GROUP_SCHEMA,
add_to_group,
assert_gone,
assert_landed,
body,
create_group,
created_team,
delete_group,
groups_of,
held_team_lock,
held_user_row,
keys_of,
member_ids,
membership_user_ids,
replace_group,
seed_users,
service_account_key,
team_key,
user_role,
waiters_on_lock_held_by,
waiters_on_rows_locked_by,
)
from pydantic import JsonValue, TypeAdapter
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
pytestmark: Final = pytest.mark.shares_worker_state
PROXY_ADMIN: Final = "proxy_admin"
DEMOTED_ROLE: Final = "internal_user_viewer"
MASTER_KEY_USER: Final = "default_user_id"
BURST_GROUPS: Final = 20
OWNED_PROXY_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
DELETED_TEAM_RECORDS_SQL: Final = 'SELECT team_id FROM "LiteLLM_DeletedTeamTable" WHERE team_id = %s'
DELETED_KEY_RECORDS_SQL: Final = 'SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s'
TEAM_MODEL_IDS_SQL: Final = """SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_info->>'team_id' = %s"""
BACKENDS_BLOCKED_BY_SQL: Final = "SELECT pid FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))"
USER_ROLE_SQL: Final = 'SELECT user_role FROM "LiteLLM_UserTable" WHERE user_id = %s'
DELETED_AUDIT_ROWS_SQL: Final = """
SELECT table_name, action, object_id, changed_by, changed_by_api_key
FROM "LiteLLM_AuditLog"
WHERE object_id IN (%s, %s) AND action = 'deleted'
ORDER BY table_name
"""
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_PIDS: Final = TypeAdapter(tuple[int, ...])
def _hashed(key: str) -> str:
return sha256(key.encode()).hexdigest()
def _admin_group_config(tmp_path: Path, admin_alias: str) -> Path:
config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
settings: Final = {**object_value(config["litellm_settings"]), "scim_admin_group": admin_alias}
path: Final = tmp_path / "scim-admin-group.yaml"
path.write_text(yaml.safe_dump({**config, "litellm_settings": settings}))
return path
def _roles_of(candidate: Gateway, users: tuple[str, ...]) -> frozenset[str]:
return frozenset(user_role(candidate, user) for user in users)
def _backends_blocked_by(holder: int) -> tuple[int, ...]:
return _PIDS.validate_python([row["pid"] for row in read_rows(BACKENDS_BLOCKED_BY_SQL, (str(holder),))])
def _stored_roles(users: tuple[str, ...]) -> frozenset[str]:
return frozenset(string_value(read_rows(USER_ROLE_SQL, (user,))[0]["user_role"]) for user in users)
def _models_status(candidate: Gateway, key: str) -> int:
return candidate.request("GET", "/v1/models", key=key).status_code
def _team_model_ids(team: str) -> frozenset[str]:
return frozenset(string_value(row["model_id"]) for row in read_rows(TEAM_MODEL_IDS_SQL, (team,)))
def _assert_revoked_everywhere(proxies: tuple[Gateway, ...], key: str) -> None:
for candidate in proxies:
eventually(partial(_models_status, candidate, key), lambda code: code == 401, seconds=30)
def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue]]:
return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key))
def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]:
text: Final = log.read_text()
return _PIDS.validate_python(_STARTED_WORKER.findall(text)), text.count("Application startup complete.")
@contextmanager
def _fresh_connection(candidate: Gateway) -> Iterator[Gateway]:
with httpx.Client(
base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS
) as client:
yield Gateway(client, candidate.key, candidate.upstream_url)
def test_deleting_one_group_leaves_the_members_other_group_intact(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
users: Final = seed_users(scenario, 3)
doomed: Final = created_team(scenario, create_group(gateway, users))
kept: Final = created_team(scenario, create_group(gateway, users))
kept_key: Final = team_key(gateway, kept, users[0])
deleted: Final = delete_group(gateway, doomed)
assert deleted.status_code == 204, deleted.text
assert gateway.request("GET", f"/scim/v2/Groups/{doomed}").status_code == 404
assert_landed(gateway, kept, users)
assert groups_of(gateway, users[0]) == frozenset({kept})
assert keys_of(kept) == frozenset({_hashed(kept_key)})
assert _models_status(gateway, kept_key) == 200
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_deleting_the_admin_group_demotes_its_members_and_writes_an_audit_row_naming_the_scim_caller(
gateway: Gateway, tmp_path: Path
) -> None:
admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}"
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
config=_admin_group_config(tmp_path, admin_alias),
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=2,
) as owned,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
users: Final = seed_users(scenario, 3)
team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias))
assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN})
key: Final = team_key(candidate, team, users[0])
deleted: Final = delete_group(candidate, team)
assert deleted.status_code == 204, deleted.text
assert_gone(candidate, team, users, key)
assert _roles_of(candidate, users) == frozenset({DEMOTED_ROLE})
rows: Final = eventually(lambda: _deleted_audit_rows(team, _hashed(key)), lambda found: len(found) >= 2, 30)
caller: Final = {"changed_by": MASTER_KEY_USER, "changed_by_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS}
assert rows == [
{"table_name": "LiteLLM_TeamTable", "action": "deleted", "object_id": team, **caller},
{"table_name": "LiteLLM_VerificationToken", "action": "deleted", "object_id": _hashed(key), **caller},
]
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_a_member_added_while_the_admin_group_delete_waits_on_the_team_lock_is_demoted_too(
gateway: Gateway, tmp_path: Path
) -> None:
admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}"
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"]},
config=_admin_group_config(tmp_path, admin_alias),
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=2,
client_timeout=60,
) as owned,
owned.gateway.scenario() as scenario,
ThreadPoolExecutor(max_workers=2) as pool,
):
candidate: Final = owned.gateway
users: Final = seed_users(scenario, 5)
(late,) = seed_users(scenario, 1)
team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias))
assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN})
key: Final = team_key(candidate, team, users[0])
with held_team_lock(team) as holder:
put: Final = pool.submit(replace_group, candidate, team, (*users, late), display_name=admin_alias)
eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 1, seconds=30)
delete: Final = pool.submit(delete_group, candidate, team)
eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 2, seconds=30)
assert _stored_roles(users) == frozenset({PROXY_ADMIN})
replaced: Final = put.result()
deleted: Final = delete.result()
assert replaced.status_code == 200, replaced.text
assert member_ids(body(replaced)) == frozenset((*users, late))
assert deleted.status_code == 204, deleted.text
assert_gone(candidate, team, (*users, late), key)
assert _roles_of(candidate, (*users, late)) == frozenset({DEMOTED_ROLE})
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_a_member_added_while_the_admin_group_delete_holds_the_team_lock_is_refused_and_never_made_admin(
gateway: Gateway, tmp_path: Path
) -> None:
admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}"
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"]},
config=_admin_group_config(tmp_path, admin_alias),
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=2,
client_timeout=60,
) as owned,
owned.gateway.scenario() as scenario,
ThreadPoolExecutor(max_workers=2) as pool,
):
candidate: Final = owned.gateway
users: Final = seed_users(scenario, 5)
(late,) = seed_users(scenario, 1)
team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias))
assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN})
key: Final = team_key(candidate, team, users[0])
with held_user_row(users[0]) as holder:
delete: Final = pool.submit(delete_group, candidate, team)
eventually(lambda: waiters_on_rows_locked_by(holder), lambda waiting: waiting == 1, seconds=30)
(deleting,) = _backends_blocked_by(holder)
added: Final = pool.submit(add_to_group, candidate, team, (late,))
eventually(lambda: _backends_blocked_by(deleting), lambda blocked: len(blocked) == 1, seconds=30)
assert _stored_roles(users) == frozenset({PROXY_ADMIN})
deleted: Final = delete.result()
assert deleted.status_code == 204, deleted.text
assert added.result().status_code == 404, added.result().text
assert_gone(candidate, team, users, key)
assert _roles_of(candidate, users) == frozenset({DEMOTED_ROLE})
assert user_role(candidate, late) != PROXY_ADMIN
def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_service_account_key_and_model_back(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
users: Final = seed_users(scenario, 2)
team: Final = created_team(scenario, create_group(gateway, users))
member_key: Final = team_key(gateway, team, users[0])
service_key: Final = service_account_key(gateway, scenario, team)
scenario.model(model_info={"team_id": team})
team_models: Final = _team_model_ids(team)
assert team_models != frozenset()
deleted: Final = delete_group(gateway, team)
assert deleted.status_code == 204, deleted.text
assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
assert membership_user_ids(team) == frozenset()
assert keys_of(team) == frozenset({_hashed(service_key)})
assert _team_model_ids(team) == team_models
_assert_revoked_everywhere((gateway,), member_key)
recreated: Final = gateway.request(
"POST",
"/scim/v2/Groups",
{"schemas": [GROUP_SCHEMA], "externalId": team, "displayName": f"integration-{uuid.uuid4().hex}"},
)
assert recreated.status_code == 201, recreated.text
assert string_value(body(recreated)["id"]) == team
eventually(partial(_models_status, gateway, service_key), lambda code: code == 200, seconds=90)
assert _models_status(gateway, member_key) == 401
def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None:
missing: Final = f"missing-{uuid.uuid4().hex}"
response: Final = delete_group(gateway, missing)
assert response.status_code == 404, response.text
assert "error" in body(response), response.text
assert missing in response.text
def test_deleting_a_group_twice_is_404_the_second_time(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
users: Final = seed_users(scenario, 2)
team: Final = created_team(scenario, create_group(gateway, users))
first: Final = delete_group(gateway, team)
assert first.status_code == 204, first.text
second: Final = delete_group(gateway, team)
assert second.status_code == 404, second.text
assert "error" in body(second), second.text
assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
assert membership_user_ids(team) == frozenset()
def test_two_deletes_racing_on_one_group_answer_204_then_404_and_record_the_team_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=2) as pool:
users: Final = seed_users(scenario, 2)
team: Final = created_team(scenario, create_group(gateway, users))
key: Final = team_key(gateway, team, users[0])
with held_team_lock(team) as holder:
deletes: Final = (pool.submit(delete_group, gateway, team), pool.submit(delete_group, gateway, team))
eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 2, seconds=30)
assert sorted(delete.result().status_code for delete in deletes) == [204, 404]
assert len(read_rows(DELETED_TEAM_RECORDS_SQL, (team,))) == 1
assert len(read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),))) == 1
assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
def test_deleting_a_group_with_a_non_admin_key_is_refused(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
users: Final = seed_users(scenario, 2)
team: Final = created_team(scenario, create_group(gateway, users))
outsider: Final = scenario.key()
response: Final = delete_group(gateway, team, key=outsider)
assert response.status_code in (401, 403), response.text
assert "error" in body(response), response.text
assert_landed(gateway, team, users)
@pytest.mark.parametrize("group_id", ("123", "x" * 5000, "a%2Fb"), ids=("int_like", "five_kb", "encoded_slash"))
def test_deleting_a_group_with_a_hostile_id_is_404(gateway: Gateway, group_id: str) -> None:
response: Final = delete_group(gateway, group_id)
assert response.status_code == 404, response.text
assert isinstance(response.json(), dict), response.text
assert gateway.request("GET", "/health/liveliness").status_code == 200
def test_deleting_an_empty_group_is_204_then_404(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
team: Final = created_team(scenario, create_group(gateway, ()))
assert member_ids(gateway.get(f"/scim/v2/Groups/{team}")) == frozenset()
deleted: Final = delete_group(gateway, team)
assert deleted.status_code == 204, deleted.text
assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
assert membership_user_ids(team) == frozenset()
def test_group_delete_on_one_proxy_revokes_the_team_key_and_the_group_on_the_peer(
gateway: Gateway, peer: Gateway
) -> None:
with gateway.scenario() as scenario:
users: Final = seed_users(scenario, 3)
team: Final = created_team(scenario, create_group(gateway, users))
key: Final = team_key(gateway, team, users[0])
assert _models_status(peer, key) == 200
assert member_ids(peer.get(f"/scim/v2/Groups/{team}")) == frozenset(users)
deleted: Final = delete_group(gateway, team)
assert deleted.status_code == 204, deleted.text
assert peer.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
_assert_revoked_everywhere((peer, gateway), key)
assert_gone(peer, team, users, key)
def test_twenty_group_deletes_split_across_two_proxies_all_land(gateway: Gateway, peer: Gateway) -> None:
proxies: Final = (gateway, peer)
with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=10) as pool:
bystander: Final = scenario.key()
rosters: Final = tuple(seed_users(scenario, 3) for _ in range(BURST_GROUPS))
teams: Final = tuple(created_team(scenario, create_group(gateway, users)) for users in rosters)
keys: Final = tuple(team_key(gateway, team, users[0]) for team, users in zip(teams, rosters))
assigned: Final = tuple(proxies[index % 2] for index in range(BURST_GROUPS))
responses: Final = tuple(pool.map(delete_group, assigned, teams))
for response in responses:
assert response.status_code == 204, response.text
for team, users, key in zip(teams, rosters, keys):
_assert_revoked_everywhere(proxies, key)
assert_gone(gateway, team, users, key)
assert peer.request("GET", f"/scim/v2/Groups/{team}").status_code == 404
assert _models_status(gateway, bystander) == 200
assert _models_status(peer, bystander) == 200
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
def test_group_delete_after_a_worker_was_killed_still_completes(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(
gateway, tmp_path, {}, remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2
) as owned,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
workers, _ = eventually(
lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 2 and found[1] == 2, seconds=120
)
first_users: Final = seed_users(scenario, 3)
second_users: Final = seed_users(scenario, 3)
first: Final = created_team(scenario, create_group(candidate, first_users))
second: Final = created_team(scenario, create_group(candidate, second_users))
first_key: Final = team_key(candidate, first, first_users[0])
second_key: Final = team_key(candidate, second, second_users[0])
with _fresh_connection(candidate) as survivor:
with owned.replacing_worker(workers[0]) as kill:
kill.freeze()
kill()
deleted_first: Final = delete_group(survivor, first)
assert deleted_first.status_code == 204, deleted_first.text
assert_gone(survivor, first, first_users, first_key)
deleted_second: Final = delete_group(survivor, second)
assert deleted_second.status_code == 204, deleted_second.text
assert_gone(survivor, second, second_users, second_key)

40
tests/unit/proxy/hooks/test_key_management_event_hooks.py
View file

@ -6,6 +6,7 @@ Validates that email and secret manager operations are independent and non-block
import asyncio
import json
from collections.abc import Sequence
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -715,14 +716,14 @@ async def test_key_deleted_hook_writes_audit_log_for_alias_deletion():
captured: Final[list[LiteLLM_AuditLogs]] = []
async def capture_audit_log(request_data: LiteLLM_AuditLogs) -> None:
captured.append(request_data)
async def capture_audit_logs(rows: Sequence[LiteLLM_AuditLogs]) -> None:
captured.extend(rows)
with (
patch("litellm.store_audit_logs", True),
patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=capture_audit_log,
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=capture_audit_logs,
),
patch.object(
KeyManagementEventHooks,
@ -747,3 +748,34 @@ async def test_key_deleted_hook_writes_audit_log_for_alias_deletion():
assert audit_row.object_id == "hashed"
assert audit_row.table_name == LitellmTableNames.KEY_TABLE_NAME
assert audit_row.changed_by == "admin"
@pytest.mark.asyncio
async def test_deleting_many_keys_writes_every_audit_row_in_one_insert():
from litellm.proxy._types import LiteLLM_VerificationToken, LitellmTableNames, UserAPIKeyAuth
from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
tokens: Final = tuple(f"hashed-{index}" for index in range(50))
prisma_client: Final = MagicMock()
prisma_client.db.litellm_auditlog.create = AsyncMock()
prisma_client.db.litellm_auditlog.create_many = AsyncMock(return_value=len(tokens))
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
):
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=[LiteLLM_VerificationToken(token=token) for token in tokens],
user_api_key_dict=UserAPIKeyAuth(user_id="admin", token="callertok"),
)
await drain_audit_tasks()
prisma_client.db.litellm_auditlog.create.assert_not_called()
prisma_client.db.litellm_auditlog.create_many.assert_awaited_once()
rows: Final = prisma_client.db.litellm_auditlog.create_many.await_args.kwargs["data"]
assert [row["object_id"] for row in rows] == list(tokens)
assert [json.loads(row["before_value"])["token"] for row in rows] == list(tokens)
assert {(row["table_name"], row["action"], row["changed_by"]) for row in rows} == {
(LitellmTableNames.KEY_TABLE_NAME, "deleted", "admin")
}

194
tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
View file

@ -3,7 +3,7 @@ import asyncio
import json
import logging
import time
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import Final, assert_never
@ -58,6 +58,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
update_user,
user_api_key_auth,
)
from litellm.proxy.management_endpoints.team_endpoints import DeletedTeams
from litellm.proxy.management_helpers.bulk_user_creation import BULK_NEW_USER_CONCURRENCY
from litellm.proxy.management_helpers.team_roster_sync import (
MembersMissing,
@ -2831,16 +2832,24 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch)
assert call_args[1]["data"]["user_role"] == LitellmUserRoles.PROXY_ADMIN
def _scim_admin_prisma(mocker, *, user_teams):
"""Prisma double and its one transaction, in which the user resolves to user_teams
and teams expose an alias equal to their id. Only the transaction answers reads,
so a role written outside it fails the role-recompute helper tests."""
user = mocker.MagicMock()
user.user_id = "member-1"
user.teams = user_teams
def _member_row(mocker, user_id: str, teams: Sequence[str]) -> MagicMock:
row = mocker.MagicMock()
row.user_id = user_id
row.teams = list(teams)
return row
def _scim_admin_prisma(mocker, *, user_teams, teammates: Sequence[tuple[str, Sequence[str]]] = ()):
"""Prisma double and its one transaction, in which the user resolves to user_teams,
``teammates`` are further (user_id, teams) rows, and teams expose an alias equal to
their id. Only the transaction answers reads, so a role written outside it fails the
role-recompute helper tests."""
user = _member_row(mocker, "member-1", user_teams)
members = (user, *(_member_row(mocker, user_id, teams) for user_id, teams in teammates))
def _team_find_unique(where):
team = mocker.MagicMock()
team.team_id = where["team_id"]
team.team_alias = where["team_id"]
return team
@ -2848,7 +2857,7 @@ def _scim_admin_prisma(mocker, *, user_teams):
return tuple(_team_find_unique({"team_id": team_id}) for team_id in where["team_id"]["in"])
def _users_find_many(where):
return (user,) if user.user_id in where["user_id"]["in"] else ()
return tuple(member for member in members if member.user_id in where["user_id"]["in"])
tx = mocker.MagicMock()
tx.query_raw = AsyncMock(return_value=[])
@ -2904,6 +2913,29 @@ async def test_recompute_scim_member_roles_grants_when_in_admin_group(mocker, mo
tx.execute_raw.assert_awaited_once_with(ANY, LitellmUserRoles.PROXY_ADMIN, ["member-1"])
@pytest.mark.asyncio
async def test_recompute_scim_member_roles_stops_serving_a_demoted_members_cached_admin_role(mocker, monkeypatch):
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import proxy_config
async def mock_get_config():
return {"litellm_settings": {"scim_admin_group": "litellm-admins"}}
monkeypatch.setattr(proxy_config, "get_config", mock_get_config)
monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False)
cache = UserApiKeyCache()
await cache.async_set_cache(
key="member-1", value={"user_id": "member-1", "user_role": LitellmUserRoles.PROXY_ADMIN.value}
)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
prisma, tx = _scim_admin_prisma(mocker, user_teams=["engineering"])
await _recompute_scim_member_roles(prisma, ["member-1"])
tx.execute_raw.assert_awaited_once_with(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1"])
assert await cache.async_get_cache(key="member-1") is None
@pytest.mark.asyncio
async def test_recompute_scim_member_roles_noop_when_admin_group_unset(mocker, monkeypatch):
"""With scim_admin_group unset the recompute helper must not touch any role,
@ -3043,44 +3075,128 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker):
assert set(recompute_mock.call_args[0][1]) == {"user1"}
@pytest.mark.asyncio
async def test_delete_group_recomputes_roles_for_members(mocker):
"""DELETE /Groups must recompute the global role for the team's members, so
deleting the admin group demotes everyone who was only admin through it."""
from litellm.proxy._types import Member
_SCIM_CALLER: Final = UserAPIKeyAuth(api_key="hashed-scim-token", user_id="idp-service-user")
_ADMIN_GROUP: Final = "litellm-admins"
_DEMOTED: Final = call(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1"])
existing_team = mocker.MagicMock()
existing_team.members_with_roles = [
Member(user_id="user1", role="user"),
Member(user_id="user2", role="user"),
]
member = mocker.MagicMock()
member.teams = ["test-team-123"]
def _group_delete_mocks(
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
team: LiteLLM_TeamTable | None,
*,
teammates: Sequence[tuple[str, Sequence[str]]] = (),
) -> tuple[MagicMock, MagicMock, AsyncMock]:
"""Prisma and its transaction, whose users are left in no team once the delete detached
them and whose group lookup answers ``team``, plus the patched team delete, which hands
its transaction and the roster it read under the lock, member-1 and every teammate, to
the caller's write."""
from litellm.proxy.proxy_server import proxy_config
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
mock_prisma_client.db.litellm_teamtable.delete = AsyncMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.update = AsyncMock()
async def mock_get_config():
return {"litellm_settings": {"scim_admin_group": _ADMIN_GROUP}}
monkeypatch.setattr(proxy_config, "get_config", mock_get_config)
monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False)
prisma, tx = _scim_admin_prisma(mocker, user_teams=[], teammates=teammates)
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client),
AsyncMock(return_value=prisma),
)
recompute_mock = mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles",
AsyncMock(),
roster: Final = ("member-1", *(user_id for user_id, _ in teammates))
async def delete_rows(
*, before_commit: Callable[[object, tuple[str, ...]], Awaitable[None]] | None = None, **_: object
) -> DeletedTeams:
if before_commit is not None:
await before_commit(tx, roster)
return DeletedTeams(team_ids=(_ADMIN_GROUP,))
delete_team_mock: Final = mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2.delete_validated_team_rows",
AsyncMock(side_effect=delete_rows),
)
return prisma, tx, delete_team_mock
await delete_group(group_id="test-team-123")
recompute_mock.assert_awaited_once()
assert list(recompute_mock.call_args[0][1]) == ["user1", "user2"]
def _admin_group_team() -> LiteLLM_TeamTable:
return LiteLLM_TeamTable(team_id=_ADMIN_GROUP, members_with_roles=[Member(user_id="member-1", role="user")])
@pytest.mark.asyncio
async def test_delete_group_demotes_a_member_who_was_admin_only_through_the_deleted_group(
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
):
"""DELETE /Groups answers 204 and writes the non-admin role for a member who was
admin only through the deleted group, once, inside the team delete."""
_, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team())
response: Final = await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER)
assert response.status_code == 204
assert tx.execute_raw.await_args_list == [_DEMOTED]
@pytest.mark.asyncio
async def test_delete_group_demotes_every_member_on_the_roster_the_team_delete_read_under_the_lock(
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
):
"""A member a concurrent PUT added to the admin group after this delete read the
group is on the roster the team delete read under the lock, so the one role write
covers that late member too."""
_, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team(), teammates=(("member-2", ()),))
response: Final = await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER)
assert response.status_code == 204
assert tx.execute_raw.await_args_list == [
call(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1", "member-2"]),
]
@pytest.mark.asyncio
async def test_delete_group_fails_when_the_role_writes_fail(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""The role writes run on the team delete's transaction, so when they fail the
request fails with them and the identity provider's retry finds the group."""
_, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team())
tx.execute_raw.side_effect = RuntimeError("database unavailable")
with pytest.raises(ProxyException) as raised:
await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER)
assert raised.value.code == "500"
@pytest.mark.asyncio
async def test_delete_group_answers_404_when_a_concurrent_delete_removed_the_team_first(
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
):
"""A concurrent delete that removed the group while this one waited on the team lock
leaves the team delete nothing to delete or report, so this request answers 404, the
same as a group that never existed."""
_, _, delete_team_mock = _group_delete_mocks(mocker, monkeypatch, _admin_group_team())
delete_team_mock.side_effect = None
delete_team_mock.return_value = DeletedTeams(team_ids=())
with pytest.raises(ProxyException) as raised:
await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER)
assert raised.value.code == "404"
@pytest.mark.asyncio
async def test_delete_group_of_an_unknown_group_answers_404_without_deleting_or_touching_roles(
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch
):
_, tx, delete_team_mock = _group_delete_mocks(mocker, monkeypatch, None)
with pytest.raises(ProxyException) as raised:
await delete_group(group_id="missing-team", user_api_key_dict=_SCIM_CALLER)
assert raised.value.code == "404"
delete_team_mock.assert_not_awaited()
tx.execute_raw.assert_not_called()
@pytest.mark.asyncio
@ -3121,12 +3237,10 @@ async def test_delete_group_is_refused_while_the_team_holds_a_ptu_share(mocker,
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", shared_router)
with pytest.raises(ProxyException) as refused:
await delete_group(group_id="test-team-123")
await delete_group(group_id="test-team-123", user_api_key_dict=_SCIM_CALLER)
assert refused.value.code == "400"
assert "test-team-123" in str(refused.value.message) and "shared-ptu" in str(refused.value.message)
mock_prisma_client.db.litellm_teamtable.delete.assert_not_called()
mock_prisma_client.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio

6
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -5123,11 +5123,11 @@ async def test_delete_user_writes_deleted_audit_log_for_user_keys(mocker):
mocker.patch("litellm.store_audit_logs", True)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
async def _capture(rows):
captured.extend(rows)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=_capture,
)

391
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -29,6 +29,7 @@ from litellm.proxy._types import (
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLM_VerificationToken,
LitellmTableNames,
LitellmUserRoles,
Member,
@ -50,7 +51,9 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_PERMISSIONS,
)
from litellm.proxy.management_endpoints.team_endpoints import (
_LOCK_DELETED_TEAM_MEMBERS_SQL,
_STRIP_DELETED_TEAMS_FROM_USERS_SQL,
DeletedTeams,
GetTeamMemberPermissionsResponse,
UpdateTeamMemberPermissionsRequest,
_build_team_list_where_conditions,
@ -63,6 +66,8 @@ from litellm.proxy.management_endpoints.team_endpoints import (
_validate_team_member_reset_spend_value,
aggregated_date_range_error,
delete_team,
delete_validated_team_rows,
delete_validated_teams,
list_available_teams,
new_team,
reset_team_member_budget_fn,
@ -212,21 +217,50 @@ def _wire_member_delete_tx(prisma_client):
return tx
def _wire_team_delete_tx(prisma_client):
"""`/team/delete` deletes the team rows and runs its post-delete reference sweep under
every team's advisory lock in one transaction, so a mocked client has to hand its own
table mocks (and db-level execute_raw) back out of `tx()` for existing per-table
assertions on `prisma_client.db.*` to keep seeing those calls."""
def _wire_team_delete_tx(
prisma_client,
*,
present_team_ids=None,
locked_teams: tuple[LiteLLM_TeamTable, ...] = (),
team_keys: tuple[LiteLLM_VerificationToken, ...] = (),
):
"""`/team/delete` records, deletes and sweeps the team rows under every team's advisory
lock in one transaction, so a mocked client has to hand its own table mocks back out of
`tx()` for existing per-table assertions on `prisma_client.db.*` to keep seeing those
calls. Under the lock every requested team still exists unless `present_team_ids` names
the ones left, and reads as the matching `locked_teams` entry or a bare row. The keys the
lock finds on those teams are `team_keys`. The locked sweep reads the users it detaches
through a `query_raw` whose rows name them and their teams, so the default answers none."""
locked_by_id = {team.team_id: team for team in locked_teams}
async def teams_still_there(where: dict[str, dict[str, list[str]]]) -> list[LiteLLM_TeamTable]:
return [
locked_by_id.get(team_id, LiteLLM_TeamTable(team_id=team_id))
for team_id in where["team_id"]["in"]
if present_team_ids is None or team_id in present_team_ids
]
async def keys_on_teams(where: dict[str, dict[str, list[str]]]) -> list[LiteLLM_VerificationToken]:
return [key for key in team_keys if key.team_id in where["team_id"]["in"]]
tx = SimpleNamespace(
litellm_teamtable=prisma_client.db.litellm_teamtable,
litellm_teamtable=SimpleNamespace(
find_many=teams_still_there,
delete_many=prisma_client.db.litellm_teamtable.delete_many,
),
litellm_verificationtoken=SimpleNamespace(find_many=keys_on_teams, delete_many=AsyncMock(return_value=0)),
litellm_deletedverificationtoken=SimpleNamespace(create_many=AsyncMock()),
litellm_jwtkeymapping=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_teammembership=prisma_client.db.litellm_teammembership,
litellm_deletedteamtable=prisma_client.db.litellm_deletedteamtable,
query_raw=AsyncMock(return_value=[]),
execute_raw=prisma_client.db.execute_raw,
execute_raw=AsyncMock(return_value=0),
)
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx)
tx_cm.__aexit__ = AsyncMock(return_value=None)
prisma_client.tx = MagicMock(return_value=tx_cm)
return tx
# Mock prisma_client
@ -5489,13 +5523,13 @@ async def test_team_member_delete_writes_deleted_audit_log_for_member_keys(
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
async def _capture(rows):
captured.extend(rows)
with (
patch("litellm.store_audit_logs", True),
patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=_capture,
),
):
@ -5564,12 +5598,12 @@ async def test_delete_team_writes_deleted_audit_log_for_team_keys(
)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
async def _capture(rows):
captured.extend(rows)
monkeypatch.setattr("litellm.store_audit_logs", True)
monkeypatch.setattr(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
_capture,
)
await delete_team(
@ -9050,7 +9084,7 @@ async def test_save_deleted_team_records():
},
]
await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client)
await _save_deleted_team_records(records=records, deleted_teams=mock_prisma_client.db.litellm_deletedteamtable)
mock_create_many.assert_called_once_with(data=records)
@ -9061,7 +9095,7 @@ async def test_save_deleted_team_records_empty_list():
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client)
await _save_deleted_team_records(records=[], deleted_teams=mock_prisma_client.db.litellm_deletedteamtable)
mock_create_many.assert_not_called()
@ -9091,7 +9125,7 @@ async def test_persist_deleted_team_records():
await _persist_deleted_team_records(
teams=[team],
prisma_client=mock_prisma_client,
deleted_teams=mock_prisma_client.db.litellm_deletedteamtable,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
@ -9295,7 +9329,7 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
# The locked delete-and-sweep transaction /team/member_add serializes against, kept
# separate from mock_tx above (the BYOK-model-cleanup transaction, unrelated to this lock).
_wire_team_delete_tx(mock_prisma_client)
_wire_team_delete_tx(mock_prisma_client, locked_teams=(doomed_team,))
mock_lock_tx = mock_prisma_client.tx.return_value.__aenter__.return_value
fresh_cache = UserApiKeyCache()
@ -9325,24 +9359,23 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
litellm_changed_by="admin-user",
)
assert mock_execute_raw.await_args_list == [
call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"]),
call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"]),
assert mock_execute_raw.await_args_list == [], (
"no sweep runs outside the lock: one there strips a member a concurrent roster write added "
"after the caller read the team, before the locked sweep can report them"
)
# one pass for the membership rows, under the lock alongside the delete
assert mock_membership_delete_many.await_args_list == [call(where={"team_id": {"in": ("team-doomed",)}})]
assert mock_lock_tx.query_raw.await_args_list == [
call(TEAM_ADVISORY_LOCK_SQL, "team-doomed"),
call(_LOCK_DELETED_TEAM_MEMBERS_SQL, ["team-doomed"]),
], (
"the unlocked sweep must run once to catch pre-existing drift, and the locked sweep "
"(alongside the delete, under the same advisory lock member_add takes) must run again "
"so a member_add that wrote its reference just before losing the lock is still reaped"
)
# same two passes for the membership rows, the second under the lock alongside the delete
assert mock_membership_delete_many.await_args_list == [
call(where={"team_id": {"in": ("team-doomed",)}}),
call(where={"team_id": {"in": ("team-doomed",)}}),
]
assert mock_lock_tx.query_raw.await_args_list == [call(TEAM_ADVISORY_LOCK_SQL, "team-doomed")], (
"the advisory lock must be acquired before the team row is deleted"
"the advisory lock must be acquired before the team row is deleted, and the user sweep "
"runs under it (the same lock member_add takes) so a member_add that wrote its "
"reference just before losing the lock is still reaped"
)
assert mock_lock_tx.execute_raw.await_args_list == [call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"])]
assert fresh_cache.get_cache(key="team_id:team-doomed") is None
assert fresh_cache.get_cache(key="team_alias:doomed-team") is None
@ -9417,7 +9450,7 @@ async def test_delete_team_evicts_member_caches_with_one_transaction(
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)
_wire_team_delete_tx(mock_prisma_client, locked_teams=(team,))
fresh_cache = UserApiKeyCache()
for user_id in member_user_ids:
@ -9461,6 +9494,234 @@ async def test_delete_team_evicts_member_caches_with_one_transaction(
)
@pytest.mark.asyncio
async def test_delete_validated_team_rows_records_the_roster_the_lock_read_and_writes_its_users_before_commit(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
A roster write that lands between the caller's read of the team and the locked delete adds
a member the caller never saw, and a user row can reference the team without a roster entry.
The deleted-team record and the caller's write, which runs on the delete's transaction before
it commits, both take the roster read under the lock, so they name the late member and not the
stray user, while both lose their cached user object.
"""
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
read_by_caller = LiteLLM_TeamTable(
team_id="team-doomed",
team_alias="doomed-team",
members_with_roles=[Member(user_id="member-1", role="user")],
metadata={},
model_max_budget={},
model_spend={},
)
read_under_lock = read_by_caller.model_copy(
update={
"members_with_roles": [
Member(user_id="member-1", role="user"),
Member(user_id="late-member", role="user"),
]
}
)
mock_prisma_client = AsyncMock()
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_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)
tx = _wire_team_delete_tx(mock_prisma_client, locked_teams=(read_under_lock,))
tx.query_raw = AsyncMock(
side_effect=[
[],
[
{"user_id": "member-1", "teams": ["team-kept", "team-doomed"]},
{"user_id": "late-member", "teams": ["team-doomed"]},
{"user_id": "stray-user", "teams": ["team-doomed"]},
],
]
)
fresh_cache = UserApiKeyCache()
for user_id in ("member-1", "late-member", "stray-user", "bystander-user"):
fresh_cache.set_cache(key=user_id, value=UserAPIKeyAuth(user_id=user_id))
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")
async def write_before_commit(_tx: object, _user_ids: tuple[str, ...]) -> None:
mock_prisma_client.tx.return_value.__aexit__.assert_not_awaited()
before_commit = AsyncMock(side_effect=write_before_commit)
deleted = await delete_validated_team_rows(
teams=(read_by_caller,),
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
),
litellm_changed_by="admin-user",
before_commit=before_commit,
)
assert deleted == DeletedTeams(team_ids=("team-doomed",))
before_commit.assert_awaited_once_with(tx, ("member-1", "late-member"))
(written,) = mock_prisma_client.db.litellm_deletedteamtable.create_many.await_args_list
(record,) = written.kwargs["data"]
assert [member["user_id"] for member in json.loads(record["members_with_roles"])] == ["member-1", "late-member"]
for user_id in ("member-1", "late-member", "stray-user"):
assert fresh_cache.get_cache(key=user_id) is None, user_id
assert fresh_cache.get_cache(key="bystander-user") is not None
@pytest.mark.asyncio
async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_members_hold_on_the_team(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
The SCIM group delete removes the team the way a SCIM member removal removes each member:
the keys rostered members hold on the team are tombstoned and deleted, and a key with no
user or held by someone off the roster stays.
"""
team = LiteLLM_TeamTable(
team_id="team-doomed",
members_with_roles=[Member(user_id="member-1", role="user")],
metadata={},
model_max_budget={},
model_spend={},
)
team_keys = tuple(
LiteLLM_VerificationToken(token=token, team_id="team-doomed", user_id=user_id)
for token, user_id in (("hashed-member", "member-1"), ("hashed-service-account", None), ("hashed-outsider", "admin-made"))
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
mock_access_group_tx_cm = MagicMock()
mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock())
mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm)
tx = _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=team_keys)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
await delete_validated_team_rows(
teams=(team,),
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-scim-token"),
litellm_changed_by=None,
)
(tombstones,) = tx.litellm_deletedverificationtoken.create_many.await_args_list
assert [record["token"] for record in tombstones.kwargs["data"]] == ["hashed-member"]
(deleted_keys,) = tx.litellm_verificationtoken.delete_many.await_args_list
assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}}
@pytest.mark.asyncio
async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock(
monkeypatch: pytest.MonkeyPatch,
):
"""Two deletes racing on one team both read it before the lock, and only the one that still
finds it under the lock deletes it, so only that one writes its deleted audit row."""
audit_logger = _wire_audit_log_callback(monkeypatch)
teams = tuple(
LiteLLM_TeamTable(team_id=team_id, metadata={}, model_max_budget={}, model_spend={})
for team_id in ("team-raced", "team-kept-for-us")
)
mock_prisma = AsyncMock()
mock_prisma.delete_data = AsyncMock(return_value={"deleted_keys": 0})
mock_prisma.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_auditlog.create = 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.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma, present_team_ids={"team-kept-for-us"})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
await delete_validated_teams(
teams=teams,
prisma_client=mock_prisma,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1", api_key="sk-a"),
litellm_changed_by=None,
)
await _settle_audit_log_tasks()
assert [
(p["object_id"], p["action"])
for p in audit_logger.payloads
if p["table_name"] == LitellmTableNames.TEAM_TABLE_NAME
] == [("team-kept-for-us", "deleted")]
@pytest.mark.asyncio
async def test_delete_validated_teams_records_and_reports_only_the_teams_still_there_under_the_lock(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
A concurrent delete that removed one of the teams while this one waited on the lock
already wrote that team's deleted-team record, so this delete writes records only for
the teams it still finds under the lock and reports only those as deleted.
"""
teams = tuple(
LiteLLM_TeamTable(
team_id=team_id,
team_alias=team_id,
members_with_roles=[],
metadata={},
model_max_budget={},
model_spend={},
)
for team_id in ("team-raced", "team-kept-for-us")
)
mock_prisma_client = AsyncMock()
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_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, present_team_ids={"team-kept-for-us"})
monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
deleted = await delete_validated_teams(
teams=teams,
prisma_client=mock_prisma_client,
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
),
litellm_changed_by=None,
)
assert deleted.team_ids == ("team-kept-for-us",)
(written,) = mock_prisma_client.db.litellm_deletedteamtable.create_many.await_args_list
assert [record["team_id"] for record in written.kwargs["data"]] == ["team-kept-for-us"]
@pytest.mark.asyncio
async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
monkeypatch,
@ -9558,8 +9819,7 @@ async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
# the unlocked pre-delete sweep succeeds, the locked post-delete sweep blows up
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")])
mock_prisma_client.db.execute_raw = AsyncMock()
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
@ -9568,6 +9828,10 @@ async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves
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)
# the advisory lock is taken, then the locked post-delete user sweep blows up
mock_prisma_client.tx.return_value.__aenter__.return_value.query_raw = AsyncMock(
side_effect=[[], ConnectionError("db went away")]
)
fresh_cache = UserApiKeyCache()
cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team")
@ -9607,21 +9871,26 @@ class _BroadcastPipeline:
return self
async def execute(self) -> None:
for message in self._messages:
await self._redis.publish("", message)
await self._redis.send(tuple(self._messages))
class _BroadcastRecordingRedis:
namespace = None
def __init__(self) -> None:
def __init__(self, round_trip_seconds: float = 0.0) -> None:
self.broadcast_keys: list[str] = []
self._round_trip_seconds = round_trip_seconds
def init_pubsub_client(self) -> "_BroadcastRecordingRedis":
return self
async def send(self, messages: Sequence[str]) -> None:
if self._round_trip_seconds:
await asyncio.sleep(self._round_trip_seconds)
self.broadcast_keys.extend(json.loads(message)["cache_key"] for message in messages)
async def publish(self, channel: str, message: str) -> int:
self.broadcast_keys.append(json.loads(message)["cache_key"])
await self.send((message,))
return 1
def pipeline(self, transaction: bool = True) -> _BroadcastPipeline:
@ -9667,7 +9936,7 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(
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)
_wire_team_delete_tx(mock_prisma_client, locked_teams=(team,))
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
@ -9697,6 +9966,44 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(
assert redis.broadcast_keys == ["hashed-doomed-key", "team_id:team-doomed", "team_alias:doomed-team"]
@pytest.mark.asyncio
async def test_deleting_a_500_member_team_reaches_every_member_cache_on_other_workers(
monkeypatch: pytest.MonkeyPatch,
):
from itertools import chain
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.management_endpoints.team_endpoints import _evict_deleted_team_member_cache
redis: Final = _BroadcastRecordingRedis(round_trip_seconds=0.001)
monkeypatch.setattr(pubsub_module, "coordination_redis_cache", lambda: redis)
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
members: Final = tuple(f"member-{index:03d}" for index in range(500))
await _evict_deleted_team_member_cache(
team_id="team-large", member_user_ids=members, user_api_key_cache=UserApiKeyCache()
)
await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks
expected: Final = frozenset(
chain.from_iterable(
(
member,
team_membership_auth_cache_key(team_id="team-large", user_id=member),
team_membership_reservation_cache_key(user_id=member, team_id="team-large"),
)
for member in members
)
)
assert expected - frozenset(redis.broadcast_keys) == frozenset()
@pytest.mark.asyncio
async def test_delete_team_survives_a_failing_cache_backend(
monkeypatch,
@ -15783,6 +16090,8 @@ async def test_new_team_and_delete_team_both_drive_the_mirror(
prisma.delete_data = AsyncMock(return_value=[team_row])
prisma.db.execute_raw = AsyncMock(return_value=0)
prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0)
prisma.db.litellm_teamtable.delete_many = AsyncMock(return_value=1)
_wire_team_delete_tx(prisma, locked_teams=(team_row,))
await delete_team(
data=DeleteTeamRequest(team_ids=["team-gone"]),

14
tests/unit/proxy/management_helpers/test_bulk_user_deletion.py
View file

@ -1007,12 +1007,12 @@ async def test_bulk_delete_writes_deleted_audit_log_for_deleted_keys(mocker):
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
async def _capture(rows):
captured.extend(rows)
mocker.patch("litellm.store_audit_logs", True)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=_capture,
)
await _delete(prisma, ["u1"])
@ -1041,12 +1041,12 @@ async def test_bulk_member_delete_writes_deleted_audit_log_for_removed_team_keys
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
async def _capture(rows):
captured.extend(rows)
mocker.patch("litellm.store_audit_logs", True)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=_capture,
)
await _remove(prisma, "t1", [{"user_id": "u1"}])
@ -1076,7 +1076,7 @@ async def test_bulk_delete_skips_the_key_audit_log_when_the_tx_rolls_back(mocker
mocker.patch("litellm.store_audit_logs", True)
mock_audit_write = mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
"litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates",
new=mocker.AsyncMock(),
)
results = await _delete(prisma, ["u1", "u2"], cache=cache)