scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(team): serialize member_add, member_delete, and delete under the team's advisory lock (#37969)

Browse source 
* fix(proxy): make /team/member_delete's four cleanups atomic

The team roster update, the user.teams update, the team membership
delete, and the team-scoped verification token delete ran as four
sequential writes with no transaction around them, so a failure
between any two left the removal half applied. Thread a single
prisma transaction through all four writes, following the same
tx.<table> pattern /team/member_add and /team/member_update already
use, so either all four land or none do.

* fix(team): serialize member_add, member_delete, and delete under the team's advisory lock

/team/member_add validated a team exists and then wrote the user's teams array and
a membership row without holding anything across that gap, so a /team/delete could
commit its reference sweeps in between and leave a member pointing at a team id that
no longer exists. The write path already re-read members_with_roles under a row lock
before this change, but SELECT ... FOR UPDATE can deadlock with the access-group
endpoints, which lock an access group and then a team.

member_add now takes pg_advisory_xact_lock(hashtext(team_id)) before re-reading the
team and only writes if it is still there, so a delete that already committed is
visible before any write happens. delete_team takes the same lock around its own
row delete and reference sweep, so the two requests can never interleave: whichever
acquires the lock first runs to completion before the other's read can proceed.

Dropping the row lock from member_add's read also dropped the incidental protection
it gave against a concurrent member_delete, which still wrote from the snapshot it
validated against, unlocked, and could silently overwrite whatever member_add had
just committed. member_delete now takes the same advisory lock and re-reads the
roster under it before computing its own write, so it can never resurrect a member
by overwriting from stale data.

Resolves LIT-5544

* fix(team): run member writes on the advisory lock's transaction

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(team): keep member writes on the lock holder's connection after merge

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(team): keep the transactional member create an upsert on user_id

The transaction path was creating the email-identified user row outright, where the
regular client path upserts on user_id. Share one upsert helper between both member
paths so the create stays idempotent on the lock holder's connection.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(team): read member_delete's user and key rows on the lock-holding transaction

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-08-25 14:55:11 -07:00 • committed by GitHub
parent 9962f4fba8
commit 6c0c91c5ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 949 additions and 210 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -16,11 +16,12 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, JsonValue
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
@ -116,6 +117,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
AccessGroupSyncTx,
invalidate_access_group_caches,
reconcile_team_access_group_membership,
@ -134,6 +136,7 @@ from litellm.proxy.management_helpers.team_metadata_validation import (
validate_team_metadata_if_configured,
)
from litellm.proxy.management_helpers.utils import (
MemberWriteTx,
add_new_member,
management_endpoint_wrapper,
)
@ -330,11 +333,44 @@ class _TeamIdInFilter(TypedDict, total=False):
team_id: Mapping[str, Sequence[str]]
class _DeletedTeamsResult(TypedDict):
deleted_teams: ReadOnly[Sequence[str]]
class _ErrorDetail(TypedDict):
error: ReadOnly[str]
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
class _MemberDeleteTx(Protocol):
"""The tables `/team/member_delete` reads while it holds the team's advisory lock.
Reading them off the transaction keeps the whole endpoint on the one pooled connection
it already checked out: a request that has the lock but still needs another connection
can be starved by the lock waiters, which is a deadlock rather than a wait when enough
of them hold the rest of the pool."""
@property
def litellm_usertable(self) -> "_PrismaTableActions[LiteLLM_UserTable]": ...
@property
def litellm_verificationtoken(self) -> "_PrismaTableActions[LiteLLM_VerificationToken]": ...
class _TeamDeleteTx(AccessGroupSyncTx, Protocol):
async def execute_raw(self, query: str, *args: object) -> int: ...
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
@property
def litellm_teammembership(self) -> "_PrismaTableActions[LiteLLM_TeamMembership]": ...
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams)
"""
@ -2580,8 +2616,13 @@ async def _process_team_members(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
tx: MemberWriteTx | None = None,
) -> tuple[list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
"""Process and add new team members."""
"""Process and add new team members.
``tx`` is the caller's open transaction, when it has one, so the member writes run on the
connection it already holds instead of checking out a second one.
"""
updated_users: Final[list[LiteLLM_UserTable]] = []
updated_team_memberships: Final[list[LiteLLM_TeamMembership]] = []
@ -2607,6 +2648,7 @@ async def _process_team_members(
default_team_budget_id=default_team_budget_id,
allowed_models=member_allowed_models,
budget_duration=data.budget_duration,
tx=tx,
)
except Exception as e:
raise HTTPException(
@ -2629,6 +2671,7 @@ async def _process_team_members(
default_team_budget_id=default_team_budget_id,
allowed_models=member_allowed_models,
budget_duration=data.budget_duration,
tx=tx,
)
except Exception as e:
raise HTTPException(
@ -2708,65 +2751,40 @@ async def _add_team_members_to_team(
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
"""Add team members to the team.
"""Add team members to the team, under the team's advisory lock.
The members_with_roles reconciliation runs inside a transaction that locks
the team row with ``SELECT ... FOR UPDATE`` before reading the current
membership. Concurrent /team/member_add calls for the same team therefore
serialize on the row lock and each appends onto the other's committed
result, instead of both rewriting the whole JSON array from a stale
snapshot (which silently drops one member on the losing write).
The lock (``TEAM_ADVISORY_LOCK_SQL``, keyed on the team id) is taken first, and the
team is re-read under it before any write, so a delete that already committed is
visible here before this call writes anything: the user and membership writes only
happen once the re-read proves the team is still live. /team/delete takes the same
lock around its own sweep-and-delete, so the two can never interleave; whichever
acquires the lock first runs to completion before the other's re-read can proceed.
The same lock serializes this against /team/delete: the delete cannot remove
the row while the reconcile holds it, and a reconcile that finds the row
already gone cleans up after itself rather than leaving the member pointing
at a deleted team id.
"""
# Process and add new members
updated_users, updated_team_memberships = await _process_team_members(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
updated_team: Final = await _write_members_with_roles_locked(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
updated_users=updated_users,
)
if updated_team is None:
await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client)
raise HTTPException(
status_code=404,
detail={"error": f"Team={data.team_id} was deleted while this member add was running"},
)
return updated_team, updated_users, updated_team_memberships
async def _write_members_with_roles_locked(
data: TeamMemberAddRequest,
complete_team_data: LiteLLM_TeamTable,
prisma_client: PrismaClient,
updated_users: list[LiteLLM_UserTable],
) -> LiteLLM_TeamTable | None:
"""Reconcile members_with_roles under the team row lock. None when the team row is gone.
That read is at least as recent as the user and membership writes the caller
already made, so a missing row means /team/delete committed after them. Its
post-delete sweep can have run before those writes landed, which is why the
caller sweeps this team id again rather than only reporting the 404.
The user and membership writes run on this transaction too, not on a second
connection from the pool: a lock waiter that needs a connection it hasn't got yet is
a waiter that can deadlock the pool, since enough concurrent adds for one team would
hold every connection waiting on the lock while the holder waits for a free one.
"""
async with prisma_client.tx() as tx:
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
if locked_members is None:
return None
gone_detail: Final[_ErrorDetail] = {
"error": f"Team={data.team_id} was deleted while this member add was running"
}
raise HTTPException(status_code=404, detail=gone_detail)
complete_team_data.members_with_roles = locked_members
updated_users, updated_team_memberships = await _process_team_members(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
tx=tx,
)
await _update_team_members_list(
data=data,
complete_team_data=complete_team_data,
@ -2774,11 +2792,13 @@ async def _write_members_with_roles_locked(
)
_db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles]
return await tx.litellm_teamtable.update(
updated_team: Final = await tx.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)},
)
return updated_team, updated_users, updated_team_memberships
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
"""Update the Prometheus team members gauge after a membership change.
@ -3159,10 +3179,6 @@ async def team_member_add(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
# Check if updated_team is None
if updated_team is None:
raise HTTPException(status_code=404, detail={"error": f"Team with id {data.team_id} not found"})
_emit_team_members_metric(complete_team_data)
await _create_team_member_add_audit_logs(
@ -3276,45 +3292,63 @@ async def team_member_delete(
)
## DELETE MEMBER FROM TEAM
removed_team_members, new_team_members = _cleanup_members_with_roles(
existing_team_row=existing_team_row,
data=data,
)
if not removed_team_members:
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
existing_team_row.members_with_roles = new_team_members
_db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
## DELETE TEAM ID from USER ROW, IF EXISTS ##
# get user row
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
key_val: Final[Mapping[str, object]] = (
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
)
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
# Also clean up any existing team membership rows for this user and team
user_ids_to_delete: Final = removed_user_ids.union(
(data.user_id,) if data.user_id is not None else (),
(user.user_id for user in existing_user_rows if user.user_id),
)
## DELETE KEYS CREATED BY USER FOR THIS TEAM
# Fetch keys before deletion so their audit records can be persisted alongside the delete.
# An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
where={
"user_id": {"in": sorted(user_ids_to_delete)},
"team_id": data.team_id,
}
)
# All four cleanups run on one connection so a failure between them leaves
# no partial removal: either every write below lands, or none of them do.
# Everything from here on runs under the team's advisory lock, the same one
# /team/member_add and /team/delete take: without it, this endpoint's own row-level
# update lock used to be the only thing serializing it against a concurrent member_add,
# and only by accident (their SELECT ... FOR UPDATE contended for the same row lock this
# UPDATE takes). Now that member_add reads under the advisory lock instead, this has to
# take it too, and re-read the roster under it rather than off the snapshot validated
# above, or a member_add that commits in between can have its addition silently
# overwritten by this delete computing from stale data.
async with prisma_client.tx() as tx:
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
fresh_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
if fresh_members is None:
raise HTTPException(
status_code=400,
detail={"error": f"Team id={data.team_id} does not exist in db"},
)
removed_team_members, new_team_members = _cleanup_members_with_roles(
existing_team_row=LiteLLM_TeamTable(team_id=data.team_id, members_with_roles=fresh_members),
data=data,
)
if not removed_team_members:
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
existing_team_row.members_with_roles = new_team_members
_db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
## DELETE TEAM ID from USER ROW, IF EXISTS ##
# get user row
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
key_val: Final[Mapping[str, object]] = (
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
)
member_tx: Final[_MemberDeleteTx] = tx
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await member_tx.litellm_usertable.find_many(
where=key_val
)
# Also clean up any existing team membership rows for this user and team
user_ids_to_delete: Final = removed_user_ids.union(
(data.user_id,) if data.user_id is not None else (),
(user.user_id for user in existing_user_rows if user.user_id),
)
## DELETE KEYS CREATED BY USER FOR THIS TEAM
# Fetch keys before deletion so their audit records can be persisted alongside the delete.
# An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await member_tx.litellm_verificationtoken.find_many(
where={
"user_id": {"in": sorted(user_ids_to_delete)},
"team_id": data.team_id,
}
)
await tx.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_new_team_members)},
@ -4009,7 +4043,21 @@ async def delete_team(
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
## DELETE TEAMS
deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team")
# 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),
# sorted so two overlapping batch deletes always request their locks in the same order.
# A member_add mid-flight for one of these teams either finishes its write and releases
# 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() as tx:
for team_id in sorted(data.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)
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
# 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
@ -4022,12 +4070,6 @@ async def delete_team(
proxy_logging_obj=proxy_logging_obj,
)
# Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep
# and the delete would have re-appended the reference; an add still in flight sees the row
# missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and
# keeping the first one means a failure here still leaves a team the admin can retry deleting.
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
for deleted_team in team_rows:
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
@ -4056,6 +4098,16 @@ async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client:
_ = 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."""
for team_id in team_ids:
_ = await tx.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id)
membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}}
_ = await tx.litellm_teammembership.delete_many(where=membership_filter)
async def _invalidate_deleted_key_cache(
keys: Sequence[LiteLLM_VerificationToken],
user_api_key_cache: UserApiKeyCache,

10
litellm/proxy/management_helpers/access_group_team_sync.py
View file

@ -23,9 +23,11 @@ from pydantic import BaseModel, TypeAdapter
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
# hashtext collisions only cost two unrelated teams a little serialization, and the
# lock is never taken by the access-group endpoints, so it cannot join their
# access-group-then-team lock order to form a cycle.
_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
# lock is never taken by the access-group endpoints as a SELECT ... FOR UPDATE row lock,
# so it cannot join their access-group-then-team lock order to form a cycle. team_endpoints
# reuses this exact statement to serialize /team/member_add and /team/delete against each
# other and against this mirror, rather than defining a second, divergent lock on the same key.
TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
@ -138,7 +140,7 @@ async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id:
concurrent write for a different team cannot be lost the way a read-modify-write of
the whole array can, and the pair commits together or not at all.
"""
await tx.query_raw(_LOCK_TEAM_SQL, team_id)
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))

128
litellm/proxy/management_helpers/utils.py
View file

@ -34,7 +34,7 @@ from litellm.proxy._types import ( # key request types; user request types; tea
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.utils import PrismaClient
from litellm.proxy.utils import PrismaClient, jsonify_object
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.repositories.user_repository import UserRepository
@ -79,6 +79,8 @@ class _PrismaUserTable(Protocol):
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> _PrismaUserRecord | None: ...
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_PrismaUserRecord]: ...
class _PrismaTeamMembershipTable(Protocol):
"""Team membership table actions the management helpers issue."""
@ -86,6 +88,73 @@ class _PrismaTeamMembershipTable(Protocol):
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
class MemberWriteTx(Protocol):
"""Transaction surface `add_new_member` writes through when the caller owns one.
A caller already holding a transaction, and with it a pooled connection plus that
transaction's locks, passes it here so these writes reuse that connection rather than
checking out another one that lock waiters may already have drained from the pool.
"""
@property
def litellm_usertable(self) -> _PrismaUserTable: ...
@property
def litellm_budgettable(self) -> _PrismaBudgetTable: ...
@property
def litellm_teammembership(self) -> _PrismaTeamMembershipTable: ...
def _user_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaUserTable:
return tx.litellm_usertable if tx is not None else UserRepository(prisma_client).table
def _budget_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaBudgetTable:
return tx.litellm_budgettable if tx is not None else BudgetRepository(prisma_client).table
def _team_membership_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaTeamMembershipTable:
return tx.litellm_teammembership if tx is not None else TeamMembershipRepository(prisma_client).table
async def _find_users_by_email(
prisma_client: PrismaClient, tx: MemberWriteTx | None, user_email: str
) -> Sequence[_PrismaUserRecord]:
if tx is not None:
return await tx.litellm_usertable.find_many(where={"user_email": user_email})
rows: Final[Sequence[_PrismaUserRecord] | None] = await prisma_client.get_data(
key_val={"user_email": user_email},
table_name="user",
query_type="find_all",
)
return rows if rows is not None else ()
async def _upsert_user_row(
user_table: _PrismaUserTable, user_id: str, create_data: Mapping[str, object]
) -> _PrismaUserRecord | None:
"""Insert the user row if it is absent, leaving an existing row as it is.
Upserting keeps concurrent provisioning of the same new user from racing on create.
The update branch re-states user_id rather than being empty because Prisma only
compiles an upsert down to INSERT ... ON CONFLICT when the update is non-empty, and
otherwise falls back to a racy SELECT-then-INSERT.
"""
return await user_table.upsert(
where={"user_id": user_id},
data={"create": create_data, "update": {"user_id": user_id}},
)
async def _create_user_row(
prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object]
) -> _PrismaUserRecord | None:
if tx is not None:
return await _upsert_user_row(tx.litellm_usertable, str(user_data["user_id"]), jsonify_object(user_data))
return await prisma_client.insert_data(data=user_data, table_name="user")
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
user_info: Final = litellm.default_internal_user_params or {}
@ -206,6 +275,7 @@ async def _clone_team_default_budget_for_member(
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
budget_duration_override: str | None = None,
tx: MemberWriteTx | None = None,
) -> str | None:
"""
Create a new budget row that copies the values from the team's default
@ -220,7 +290,7 @@ async def _clone_team_default_budget_for_member(
member while keeping the default's other limits, so an admin can set a
member's reset cadence without discarding the team default's max_budget.
"""
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id})
if default_budget is None:
return None
@ -248,7 +318,7 @@ async def _clone_team_default_budget_for_member(
if cloned_data.get("budget_duration"):
cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"])
new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data)
new_budget: Final[_PrismaBudgetRecord] = await budget_table.create(data=cloned_data)
return new_budget.budget_id
@ -260,6 +330,7 @@ async def _resolve_member_budget_id(
allowed_models: list[str] | None,
budget_duration: str | None,
default_team_budget_id: str | None,
tx: MemberWriteTx | None = None,
) -> str | None:
"""
Resolve the budget a new team member should be linked to.
@ -279,6 +350,7 @@ async def _resolve_member_budget_id(
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
budget_duration_override=budget_duration,
tx=tx,
)
if not has_explicit_limit and budget_duration is None:
@ -295,12 +367,14 @@ async def _resolve_member_budget_id(
if budget_duration is not None:
budget_data["budget_duration"] = budget_duration
budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
response: Final = await budget_table.create(data=budget_data)
return response.budget_id
async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None:
async def _append_team_id_if_absent(
prisma_client: PrismaClient, user_id: str, team_id: str, tx: MemberWriteTx | None = None
) -> None:
"""Append team_id to a user's teams array, only if it is not already present.
The row-level filter makes the append a no-op once the team is present, so
@ -309,7 +383,7 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
number of teams a user belongs to). Teams added concurrently for a different
team id are unaffected, since each update filters on its own team id.
"""
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
user_table: Final[_PrismaUserTable] = _user_table(prisma_client, tx)
await user_table.update_many(
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
data={"teams": {"push": [team_id]}},
@ -326,6 +400,7 @@ async def add_new_member(
default_team_budget_id: str | None = None,
allowed_models: list[str] | None = None,
budget_duration: str | None = None,
tx: MemberWriteTx | None = None,
) -> tuple[LiteLLM_UserTable, LiteLLM_TeamMembership | None]:
"""
Add a new member to a team
@ -334,49 +409,41 @@ async def add_new_member(
- add team member w/ budget to team member table
Returns created/existing user + team membership w/ budget id
Callers already inside a transaction pass it as ``tx`` so every write here runs on that
connection instead of borrowing more from the pool while the caller's locks are held.
"""
returned_user: LiteLLM_UserTable | None = None
returned_team_membership: LiteLLM_TeamMembership | None = None
## ADD TEAM ID, to USER TABLE IF NEW ##
if new_member.user_id is not None:
new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id)
# Upsert ensures the user row exists atomically (no create race when the
# same new user is provisioned concurrently), seeding teams on create.
# The teams append lives in the filtered update below rather than the
# upsert's update branch so an already-existing user does not get a
# duplicate team id. The update branch still has to write something:
# Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
# is non-empty, and falls back to a racy SELECT-then-INSERT when it is
# not, so this re-states user_id as a no-op rather than being empty.
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
_returned_user: _PrismaUserRecord | None = await user_table.upsert(
where={"user_id": new_member.user_id},
data={
"create": {"teams": [team_id], **new_user_defaults},
"update": {"user_id": new_member.user_id},
},
# The teams append lives in the filtered update below rather than the upsert's
# update branch so an already-existing user does not get a duplicate team id.
_returned_user: _PrismaUserRecord | None = await _upsert_user_row(
_user_table(prisma_client, tx),
new_member.user_id,
{"teams": [team_id], **new_user_defaults},
)
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx)
if _returned_user is not None:
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif new_member.user_email is not None:
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
## user email is not unique acc. to prisma schema -> future improvement
### for now: check if it exists in db, if not - insert it
existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
key_val={"user_email": new_member.user_email},
table_name="user",
query_type="find_all",
existing_user_row: Final[Sequence[_PrismaUserRecord]] = await _find_users_by_email(
prisma_client, tx, new_member.user_email
)
if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0):
if len(existing_user_row) == 0:
new_user_defaults["teams"] = [team_id]
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
_returned_user = await _create_user_row(prisma_client, tx, new_user_defaults)
if _returned_user is not None:
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif len(existing_user_row) == 1:
user_info: Final = existing_user_row[0]
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id)
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx)
returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump())
elif len(existing_user_row) > 1:
raise HTTPException(
@ -392,10 +459,11 @@ async def add_new_member(
allowed_models=allowed_models,
budget_duration=budget_duration,
default_team_budget_id=default_team_budget_id,
tx=tx,
)
if _budget_id and returned_user is not None and returned_user.user_id is not None:
membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table
membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
_returned_team_membership: Final = await membership_table.create(
data={
"team_id": team_id,

21
litellm/repositories/team_repository.py
View file

@ -58,19 +58,22 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
return LiteLLM_TeamTable.model_validate(data)
async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None:
"""Return the team's members_with_roles, locking the row FOR UPDATE.
"""Return the team's members_with_roles. The caller must already hold
``TEAM_ADVISORY_LOCK_SQL`` for this team_id on ``tx`` before calling this.
``None`` when the team row is gone, which a caller holding the lock can
only see if a delete committed under it, as opposed to ``[]`` for a team
that simply has no members.
``None`` when the team row is gone, which is only possible under that lock if
a delete committed before this read, as opposed to ``[]`` for a team that
simply has no members.
Must be called inside a transaction so the row lock is held until
commit. This serializes concurrent membership writers on the team row
so the losing writer appends onto the winner's committed result instead
of overwriting it from a stale snapshot.
A plain read is enough here because the advisory lock, not a row lock, is what
serializes this against a concurrent writer: ``SELECT ... FOR UPDATE`` would
additionally take a row lock on ``LiteLLM_TeamTable``, and the access-group
endpoints lock an access group and then a team row, so a team-row-first lock
here can deadlock with them. The advisory lock cannot, since those endpoints
never take it.
"""
rows: Final = await tx.query_raw(
'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE',
'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1',
team_id,
)
if not rows:

307
tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py Normal file
View file

@ -0,0 +1,307 @@
"""
Real-Postgres coverage for the /team/member_add vs /team/delete race (LIT-5544), and for
/team/member_delete's participation in the same lock.
A member_add that validated the team before a delete began could previously still commit
its writes after the delete's reference sweeps had already run, leaving a user record and
a membership row pointing at a team id that no longer exists. Neither side of that race can
be forced by a sequential script: it needs one request to be genuinely mid-flight while the
other commits. A mocked prisma cannot arbitrate that either, since the property under test
is whether Postgres's own advisory lock actually serializes the two requests.
These tests pin the interleaving the same way test_access_group_team_sync.py does: a second
real connection holds the team's advisory lock in its own transaction, so the function under
test is provably blocked on it rather than hoping a sleep lands in the right gap.
"""
import asyncio
import json
import os
from contextlib import asynccontextmanager
from datetime import timedelta
from unittest.mock import MagicMock
import pytest
from fastapi import HTTPException
from litellm.proxy._types import (
DeleteTeamRequest,
LitellmUserRoles,
Member,
TeamMemberAddRequest,
UserAPIKeyAuth,
)
from litellm.caching.caching import DualCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
TEAM = "lit5544-race-team"
USER = "lit5544-race-user"
_DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1'
_DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1'
_DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1'
_LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
@asynccontextmanager
async def _clean_db():
"""Connects inside the running test's loop: an async fixture would be torn up on a
different loop than the test body, which prisma's engine lock refuses outright."""
from prisma import Prisma
if not os.getenv("DATABASE_URL"):
pytest.fail("DATABASE_URL is required; these tests must not silently skip")
db = Prisma()
await db.connect()
try:
await db.execute_raw(_DELETE_SEEDED, TEAM)
await db.execute_raw(_DELETE_USER, USER)
await db.execute_raw(_DELETE_TEAM, TEAM)
yield db
finally:
await db.execute_raw(_DELETE_SEEDED, TEAM)
await db.execute_raw(_DELETE_USER, USER)
await db.execute_raw(_DELETE_TEAM, TEAM)
await db.disconnect()
@asynccontextmanager
async def _real_prisma_client():
"""The full app-level PrismaClient, not the raw generated client: add_new_member reads
and writes through PrismaClient.get_data/insert_data, which the raw client doesn't have."""
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj)
await client.connect()
try:
yield client
finally:
await client.db.disconnect()
def _admin_auth():
return UserAPIKeyAuth(user_id="lit5544-admin", api_key="sk-lit5544", user_role=LitellmUserRoles.PROXY_ADMIN.value)
@pytest.mark.asyncio
async def test_member_add_blocked_by_delete_writes_no_dangling_reference():
"""
member_add re-reads the team under the advisory lock before writing anything. When a
delete already holds that lock and then removes the row, member_add's re-read must see
the row gone and raise, without ever calling the write that appends the user/membership
references, which is the only way this leaves zero trace after the delete wins.
"""
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.management_endpoints.team_endpoints import (
_add_team_members_to_team,
)
async with _clean_db() as db:
await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"})
async with _real_prisma_client() as prisma_client:
from prisma import Prisma
blocker = Prisma()
await blocker.connect()
lock_acquired = asyncio.Event()
async def add_member():
lock_acquired.set()
await _add_team_members_to_team(
data=TeamMemberAddRequest(
team_id=TEAM,
member=Member(user_id=USER, role="user"),
max_budget_in_team=5.0,
),
complete_team_data=LiteLLM_TeamTable(team_id=TEAM, members_with_roles=[]),
prisma_client=prisma_client,
user_api_key_dict=_admin_auth(),
litellm_proxy_admin_name="lit5544-admin",
)
try:
async with blocker.tx(timeout=timedelta(seconds=30)) as held:
await held.query_raw(_LOCK_SQL, TEAM)
task = asyncio.create_task(add_member())
await lock_acquired.wait()
await asyncio.sleep(0.2)
assert not task.done(), "member_add did not wait on the team's advisory lock"
# the delete wins the race: strip the team row while the lock is held
await held.execute_raw(_DELETE_TEAM, TEAM)
with pytest.raises(HTTPException) as exc_info:
await asyncio.wait_for(task, timeout=30)
assert exc_info.value.status_code == 404
finally:
await blocker.disconnect()
user_row = await db.litellm_usertable.find_unique(where={"user_id": USER})
assert user_row is None, "member_add must not have written a user row for a team that was gone under its lock"
membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER})
assert membership_row is None
@pytest.mark.asyncio
async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster():
"""
team_member_delete takes the same advisory lock and re-reads the roster under it, so a
member_add that committed while member_delete was waiting on the lock is not silently
undone. Without the re-read, member_delete would compute its new roster from the stale
snapshot it validated against before the lock, and its write would overwrite the
member_add's addition right back out even though member_add's request already succeeded.
"""
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import TeamMemberDeleteRequest
from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
other_user = f"{USER}-other"
seeded_roster = '[{"user_id": "%s", "user_email": null, "role": "user"}]' % USER
winning_add_roster = (
'[{"user_id": "%s", "user_email": null, "role": "user"}, '
'{"user_id": "%s", "user_email": null, "role": "user"}]' % (USER, other_user)
)
async with _clean_db() as db:
await db.litellm_teamtable.create(
data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": seeded_roster}
)
async with _real_prisma_client() as prisma_client:
original_prisma_client = proxy_server_module.prisma_client
proxy_server_module.prisma_client = prisma_client
try:
from prisma import Prisma
blocker = Prisma()
await blocker.connect()
lock_acquired = asyncio.Event()
async def run_delete():
lock_acquired.set()
return await team_member_delete(
data=TeamMemberDeleteRequest(team_id=TEAM, user_id=USER),
user_api_key_dict=_admin_auth(),
)
try:
async with blocker.tx(timeout=timedelta(seconds=30)) as held:
await held.query_raw(_LOCK_SQL, TEAM)
task = asyncio.create_task(run_delete())
await lock_acquired.wait()
await asyncio.sleep(0.2)
assert not task.done(), "member_delete did not wait on the team's advisory lock"
# member_add wins the race: it adds `other_user` while holding the lock
await held.litellm_teamtable.update(
where={"team_id": TEAM},
data={"members_with_roles": winning_add_roster},
)
await asyncio.wait_for(task, timeout=30)
finally:
await blocker.disconnect()
finally:
proxy_server_module.prisma_client = original_prisma_client
team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM})
raw_roster = team_row.members_with_roles
parsed_roster = json.loads(raw_roster) if isinstance(raw_roster, str) else raw_roster
remaining_ids = {m["user_id"] for m in parsed_roster}
assert remaining_ids == {other_user}, (
"member_delete must remove only the user it targeted from the roster it actually "
"committed to, not silently drop the member the winning add just committed"
)
@pytest.mark.asyncio
async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference():
"""
A member_add that wins the lock race writes its reference and releases the lock; the
delete that was waiting on it must then run its locked sweep against the row as it
actually is, not a stale snapshot, and reap that reference rather than leaving it
stranded on a team id the delete is about to remove.
"""
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.management_endpoints.team_endpoints import delete_team
async with _clean_db() as db:
await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"})
async with _real_prisma_client() as prisma_client:
proxy_logging_obj = prisma_client.proxy_logging_obj
original_prisma_client = proxy_server_module.prisma_client
original_admin_name = proxy_server_module.litellm_proxy_admin_name
original_proxy_logging_obj = proxy_server_module.proxy_logging_obj
original_cache = proxy_server_module.user_api_key_cache
original_router = proxy_server_module.llm_router
proxy_server_module.prisma_client = prisma_client
proxy_server_module.litellm_proxy_admin_name = "lit5544-admin"
proxy_server_module.proxy_logging_obj = proxy_logging_obj
proxy_server_module.user_api_key_cache = original_cache or proxy_logging_obj.internal_usage_cache
proxy_server_module.llm_router = None
async def restore():
proxy_server_module.prisma_client = original_prisma_client
proxy_server_module.litellm_proxy_admin_name = original_admin_name
proxy_server_module.proxy_logging_obj = original_proxy_logging_obj
proxy_server_module.user_api_key_cache = original_cache
proxy_server_module.llm_router = original_router
try:
from prisma import Prisma
blocker = Prisma()
await blocker.connect()
lock_acquired = asyncio.Event()
async def run_delete():
lock_acquired.set()
return await delete_team(
data=DeleteTeamRequest(team_ids=[TEAM]),
http_request=MagicMock(),
user_api_key_dict=_admin_auth(),
litellm_changed_by="lit5544-admin",
)
try:
async with blocker.tx(timeout=timedelta(seconds=30)) as held:
await held.query_raw(_LOCK_SQL, TEAM)
task = asyncio.create_task(run_delete())
await lock_acquired.wait()
await asyncio.sleep(0.3)
assert not task.done(), "delete_team did not wait on the team's advisory lock"
# member_add wins the race: write the reference while holding the lock
await held.litellm_usertable.upsert(
where={"user_id": USER},
data={
"create": {"user_id": USER, "teams": [TEAM]},
"update": {"teams": {"push": [TEAM]}},
},
)
await held.litellm_teammembership.create(data={"team_id": TEAM, "user_id": USER})
await held.litellm_teamtable.update(
where={"team_id": TEAM},
data={"members_with_roles": '[{"user_id": "%s", "role": "user"}]' % USER},
)
await asyncio.wait_for(task, timeout=30)
finally:
await blocker.disconnect()
finally:
await restore()
team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM})
assert team_row is None
user_row = await db.litellm_usertable.find_unique(where={"user_id": USER})
assert user_row is not None and TEAM not in user_row.teams, (
"delete_team's locked sweep must reap the reference member_add wrote just before losing the lock"
)
membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER})
assert membership_row is None

34
tests/proxy_unit_tests/test_proxy_server.py
View file

@ -1169,6 +1169,22 @@ async def test_create_user_default_budget(prisma_client, user_role): # noqa: F8
assert mock_client.call_args.kwargs["data"]["budget_duration"] is None
def _member_add_tx_cm(team_table):
"""Transaction whose member writes land on whatever tables are mocked on `prisma_client.db`"""
class _Tx:
query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
litellm_teamtable = team_table
def __getattr__(self, table_name):
return getattr(litellm.proxy.proxy_server.prisma_client.db, table_name)
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=_Tx())
tx_cm.__aexit__ = AsyncMock(return_value=None)
return tx_cm
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
@ -1230,7 +1246,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa
)
)
mock_litellm_usertable.upsert = mock_client
mock_litellm_usertable.find_many = AsyncMock(return_value=None)
mock_litellm_usertable.find_many = AsyncMock(return_value=[])
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
@ -1245,12 +1261,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
tx_mock = AsyncMock()
tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
tx_mock.litellm_teamtable = team_mock_client
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
tx_cm.__aexit__ = AsyncMock(return_value=None)
tx_cm = _member_add_tx_cm(team_mock_client)
original_tx = litellm.proxy.proxy_server.prisma_client.tx
litellm.proxy.proxy_server.prisma_client.tx = MagicMock(
return_value=tx_cm
@ -1432,7 +1443,7 @@ async def test_create_team_member_add_team_admin(
)
)
mock_litellm_usertable.upsert = mock_client
mock_litellm_usertable.find_many = AsyncMock(return_value=None)
mock_litellm_usertable.find_many = AsyncMock(return_value=[])
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
@ -1443,12 +1454,7 @@ async def test_create_team_member_add_team_admin(
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
tx_mock = AsyncMock()
tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
tx_mock.litellm_teamtable = team_mock_client
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
tx_cm.__aexit__ = AsyncMock(return_value=None)
tx_cm = _member_add_tx_cm(team_mock_client)
with (
patch.object(

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

@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Optional, cast
from unittest.mock import AsyncMock, MagicMock, call, patch
from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
import pytest
from fastapi import HTTPException
@ -58,6 +58,9 @@ from litellm.proxy.management_endpoints.team_endpoints import (
update_team,
validate_team_org_change,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
TeamMemberPermissionChecks,
)
@ -75,7 +78,11 @@ client = TestClient(app)
def _wire_team_create_tx(prisma_client):
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
so a mocked client has to hand its team table back out of `db.tx()`."""
so a mocked client has to hand its team table back out of `db.tx()`.
A `/team/new` carrying members then adds them under the team's advisory lock, and those
writes run on that lock's transaction, so `tx()` has to hand back the mocked tables too
for the per-table assertions on `prisma_client.db.*` to keep seeing them."""
@asynccontextmanager
async def _tx():
@ -85,18 +92,67 @@ def _wire_team_create_tx(prisma_client):
)
prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
_wire_member_add_tx(prisma_client)
def _wire_member_add_tx(prisma_client):
"""/team/member_add takes the team's advisory lock, re-reads the roster under it, and runs
the user, budget, and membership writes on that same transaction, so a mocked client has
to hand its own table mocks back out of `tx()`.
Tables resolve on access, not here, since tests routinely replace `db.<table>` after
wiring the transaction."""
class _Tx:
query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
def __getattr__(self, table_name):
return getattr(prisma_client.db, table_name)
tx = _Tx()
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)
def _wire_member_delete_tx(prisma_client):
"""/team/member_delete's four cleanups run inside one transaction, so a mocked
client has to hand back its own table mocks out of `tx()` for the existing
per-table assertions to keep seeing the calls."""
"""/team/member_delete's four cleanups, plus the advisory-lock re-read that now guards
them, run inside one transaction, so a mocked client has to hand back its own table
mocks (and a `query_raw` that answers the locked re-read from the same team row the
test already configured on `find_unique`) out of `tx()` for the existing per-table
assertions to keep seeing the calls."""
async def _query_raw(sql, team_id):
if sql != TEAM_ADVISORY_LOCK_SQL:
team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
if team_row is not None:
return [{"members_with_roles": team_row.model_dump()["members_with_roles"]}]
return []
class _Tx:
query_raw = staticmethod(_query_raw)
def __getattr__(self, table_name):
return getattr(prisma_client.db, table_name)
tx = _Tx()
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)
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."""
tx = SimpleNamespace(
litellm_teamtable=prisma_client.db.litellm_teamtable,
litellm_usertable=prisma_client.db.litellm_usertable,
litellm_teammembership=prisma_client.db.litellm_teammembership,
litellm_verificationtoken=prisma_client.db.litellm_verificationtoken,
litellm_deletedverificationtoken=prisma_client.db.litellm_deletedverificationtoken,
query_raw=AsyncMock(return_value=[]),
execute_raw=prisma_client.db.execute_raw,
)
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx)
@ -1669,6 +1725,7 @@ async def test_process_team_members_single_member():
default_team_budget_id="budget-123",
allowed_models=None,
budget_duration=None,
tx=None,
)
@ -1809,8 +1866,8 @@ async def test_update_team_members_list_duplicate_prevention():
async def test_add_team_members_reconciles_against_freshly_locked_row():
"""
Regression: _add_team_members_to_team must build the new members_with_roles
from the row it re-reads under a lock inside the write transaction, not from
the stale complete_team_data snapshot captured at the start of the request.
from the row it re-reads under the team's advisory lock, not from the stale
complete_team_data snapshot captured at the start of the request.
Two concurrent /team/member_add calls for the same team read the same
snapshot; without the locked re-read the losing write rewrites the whole
@ -1871,24 +1928,89 @@ async def test_add_team_members_reconciles_against_freshly_locked_row():
written_ids = sorted(m["user_id"] for m in json.loads(captured["data"]["members_with_roles"]))
assert written_ids == ["alice", "bob", "zed"]
lock_reads = [call for call in tx.query_raw.call_args_list if "FOR UPDATE" in str(call.args[0])]
assert lock_reads, "expected a SELECT ... FOR UPDATE row-lock read before the write"
assert tx.query_raw.call_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "test-team-lock"), (
"expected the team's advisory lock to be acquired before the members_with_roles read"
)
assert not any("FOR UPDATE" in str(call.args[0]) for call in tx.query_raw.call_args_list), (
"a row lock here can deadlock with the access-group endpoints; only the advisory lock is safe"
)
assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"]
@pytest.mark.asyncio
async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request():
async def test_add_team_members_runs_member_writes_on_the_lock_holding_transaction():
"""
Regression pin against exhausting the connection pool with advisory-lock waiters.
Every concurrent /team/member_add for one team holds a pooled connection while it waits
on the team's advisory lock. If the holder's member writes went to the regular client,
it would need a second connection to finish, so enough concurrent adds fill the pool
with waiters and the holder can never commit or release the lock. The member writes
therefore have to run on the transaction that already owns the connection.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
_add_team_members_to_team,
)
added_user = MagicMock()
added_user.user_id = "bob"
added_user.model_dump.return_value = {"user_id": "bob", "teams": ["team-pool"]}
created_budget = MagicMock()
created_budget.budget_id = "budget-pool"
membership = MagicMock()
membership.model_dump.return_value = {
"team_id": "team-pool",
"user_id": "bob",
"budget_id": "budget-pool",
"litellm_budget_table": None,
}
tx = MagicMock()
tx.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
tx.litellm_teamtable.update = AsyncMock(
return_value=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[])
)
tx.litellm_usertable.upsert = AsyncMock(return_value=added_user)
tx.litellm_usertable.update_many = AsyncMock()
tx.litellm_budgettable.create = AsyncMock(return_value=created_budget)
tx.litellm_teammembership.create = AsyncMock(return_value=membership)
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx)
tx_cm.__aexit__ = AsyncMock(return_value=None)
prisma_client = MagicMock()
prisma_client.tx = MagicMock(return_value=tx_cm)
type(prisma_client).db = PropertyMock(
side_effect=AssertionError("member writes must not reach for a second pooled connection")
)
_, updated_users, updated_team_memberships = await _add_team_members_to_team(
data=TeamMemberAddRequest(
team_id="team-pool",
member=Member(user_id="bob", role="user"),
max_budget_in_team=50.0,
),
complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]),
prisma_client=cast(object, prisma_client),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
litellm_proxy_admin_name="admin",
)
assert [user.user_id for user in updated_users] == ["bob"]
assert [tm.budget_id for tm in updated_team_memberships] == ["budget-pool"]
@pytest.mark.asyncio
async def test_add_team_members_writes_nothing_when_the_team_is_deleted_mid_request():
"""
Regression pin for the /team/member_add vs /team/delete race.
The user row and membership writes land before the reconcile takes the team
row lock, so a /team/delete that commits in between has already run its own
reference sweep and cannot see them. The empty locked SELECT is the only
signal that happened, and leaving it at that would strand the member on a
deleted team id, which authorization paths that trust `user.teams` would
treat as membership if the id were ever recreated. So the request must sweep
the references it just wrote and fail, not report success.
The advisory lock is acquired, and the team is gone, before any write is attempted:
the empty locked SELECT is proof a /team/delete already committed under the same
lock, so this request must fail without writing the user or membership rows in the
first place, rather than writing them and then trying to sweep them back out.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
_add_team_members_to_team,
@ -1907,9 +2029,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request()
prisma_client.db.execute_raw = AsyncMock()
prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
process_team_members = AsyncMock(return_value=([], []))
with patch(
"litellm.proxy.management_endpoints.team_endpoints._process_team_members",
new=AsyncMock(return_value=([], [])),
new=process_team_members,
):
with pytest.raises(HTTPException) as exc_info:
await _add_team_members_to_team(
@ -1924,14 +2047,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request()
)
assert exc_info.value.status_code == 404
process_team_members.assert_not_awaited()
tx.litellm_teamtable.update.assert_not_awaited()
assert prisma_client.db.execute_raw.await_args_list == [
call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add")
]
prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with(
where={"team_id": {"in": ("team-deleted-mid-add",)}}
)
prisma_client.db.execute_raw.assert_not_awaited()
prisma_client.db.litellm_teammembership.delete_many.assert_not_awaited()
def test_add_new_models_to_team_with_existing_models():
@ -4246,6 +4365,86 @@ async def test_team_member_delete_cleans_verification_tokens(
)
@pytest.mark.asyncio
async def test_team_member_delete_reads_on_the_lock_holding_transaction(
mock_db_client, mock_admin_auth
):
"""
Regression pin against exhausting the connection pool with advisory-lock waiters.
Every concurrent removal for one team holds a pooled connection while it waits on the
team's advisory lock, and /team/delete fans its per-member removals out concurrently.
A holder whose reads went to the regular client would need a second connection to
finish, so enough waiters fill the pool and the holder can never release the lock.
Both reads therefore have to run on the transaction that already owns the connection.
"""
from litellm.proxy._types import TeamMemberDeleteRequest
from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
test_team_id = "team-del-pool-123"
test_user_id = "user-del-pool-123"
roster_entry = {"user_id": test_user_id, "user_email": None, "role": "user"}
mock_team_row = MagicMock()
mock_team_row.model_dump.return_value = {
"team_id": test_team_id,
"members_with_roles": [roster_entry],
"team_member_permissions": [],
"metadata": {},
"models": [],
"spend": 0.0,
}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team_row
)
user_row = MagicMock()
user_row.user_id = test_user_id
user_row.teams = [test_team_id]
# Both are wired to answer, so the endpoint completes either way and the awaits below
# are what tells which connection it read on.
pooled_user_read = AsyncMock(return_value=[user_row])
pooled_token_read = AsyncMock(return_value=[])
mock_db_client.db.litellm_usertable.find_many = pooled_user_read
mock_db_client.db.litellm_verificationtoken.find_many = pooled_token_read
tx = MagicMock()
tx.query_raw = AsyncMock(return_value=[{"members_with_roles": [roster_entry]}])
tx.litellm_teamtable.update = AsyncMock(return_value=mock_team_row)
tx.litellm_usertable.find_many = AsyncMock(return_value=[user_row])
tx.litellm_usertable.update = AsyncMock()
tx.litellm_teammembership.delete_many = AsyncMock()
tx.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
tx.litellm_verificationtoken.delete_many = AsyncMock()
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx)
tx_cm.__aexit__ = AsyncMock(return_value=None)
mock_db_client.tx = MagicMock(return_value=tx_cm)
await team_member_delete(
data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id),
user_api_key_dict=mock_admin_auth,
)
tx.litellm_usertable.find_many.assert_awaited_once_with(
where={"user_id": {"in": [test_user_id]}}
)
tx.litellm_verificationtoken.find_many.assert_awaited_once_with(
where={"user_id": {"in": [test_user_id]}, "team_id": test_team_id}
)
pooled_user_read.assert_not_awaited()
pooled_token_read.assert_not_awaited()
tx.litellm_usertable.update.assert_awaited_once_with(
where={"user_id": test_user_id}, data={"teams": {"set": []}}
)
tx.litellm_teammembership.delete_many.assert_awaited_once_with(
where={"team_id": test_team_id, "user_id": test_user_id}
)
@pytest.mark.parametrize(
"roster_email",
["Alice@Example.com", "alice-invited-as@example.com"],
@ -7411,6 +7610,7 @@ async def test_delete_team_persists_deleted_teams(
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)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
@ -7481,15 +7681,14 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
cache_state_when_rows_deleted = {}
async def record_cache_state_then_delete(*args, **kwargs):
if kwargs.get("table_name") == "team":
cache_state_when_rows_deleted["doomed_still_cached"] = (
fresh_cache.get_cache(key="team_id:team-doomed") is not None
)
return {"deleted_teams": ["team-doomed"]}
cache_state_when_rows_deleted["doomed_still_cached"] = (
fresh_cache.get_cache(key="team_id:team-doomed") is not None
)
return 1
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team)
mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete)
mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0})
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
@ -7498,6 +7697,7 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
mock_prisma_client.db.execute_raw = mock_execute_raw
mock_membership_delete_many = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many
mock_prisma_client.db.litellm_teamtable.delete_many = AsyncMock(side_effect=record_cache_state_then_delete)
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
@ -7506,6 +7706,11 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
# 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)
mock_lock_tx = mock_prisma_client.tx.return_value.__aenter__.return_value
fresh_cache = UserApiKeyCache()
for cached_team_id, cached_alias in (
("team-doomed", "doomed-team"),
@ -7539,14 +7744,22 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
assert mock_execute_raw.await_args_list == [
call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"),
call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"),
], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind"
], (
"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: the second one reaps a membership row inserted while the delete was running
# 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"
)
assert fresh_cache.get_cache(key="team_id:team-doomed") is None
assert fresh_cache.get_cache(key="team_alias:doomed-team") is None
assert fresh_cache.get_cache(key="team_id:team-kept") is not None
@ -7596,6 +7809,7 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma_client)
fresh_cache = UserApiKeyCache()
fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed"))
@ -7623,14 +7837,17 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
@pytest.mark.asyncio
async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache(
async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves_the_cache_alone(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
The reconcile sweep runs after the team row is committed deleted. If it ran before cache
eviction, a sweep failure would return an error with the team gone from the db but still
served from cache, which is the exact bug this PR exists to fix.
The team delete and its post-delete reconcile sweep run inside one transaction, under the
team's advisory lock, so a sweep failure rolls the delete back with it rather than leaving
the row gone with the sweep half done. Cache eviction only runs after that transaction
commits, so a failure here must leave the team exactly as it was: still in the db, and
still cached. Evicting a cache entry for a delete that never actually committed would be
the same class of bug this PR exists to fix, just on the other side of the transaction.
"""
from litellm.proxy._types import DeleteTeamRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -7650,7 +7867,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
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 first sweep succeeds, the post-delete reconcile sweep blows up
# 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_tx = AsyncMock()
@ -7659,6 +7876,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma_client)
fresh_cache = UserApiKeyCache()
cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team")
@ -7682,9 +7900,10 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
litellm_changed_by="admin-user",
)
# the delete committed, so the cache must not still be serving the team
assert fresh_cache.get_cache(key="team_id:team-doomed") is None
assert fresh_cache.get_cache(key="team_alias:doomed-team") is None
# the transaction that deletes the row and runs the locked sweep never committed, so
# cache eviction (which only runs after that commit) must never have been reached
assert fresh_cache.get_cache(key="team_id:team-doomed") is not None
assert fresh_cache.get_cache(key="team_alias:doomed-team") is not None
@pytest.mark.asyncio
@ -7726,6 +7945,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)
published = []
@ -7796,6 +8016,7 @@ async def test_delete_team_survives_a_failing_cache_backend(
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)
exploding_logging_obj = MagicMock()
exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(
@ -7820,7 +8041,7 @@ async def test_delete_team_survives_a_failing_cache_backend(
)
assert result == {"deleted_teams": ["team-doomed"]}
mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team")
mock_prisma_client.db.litellm_teamtable.delete_many.assert_any_await(where={"team_id": {"in": ["team-doomed"]}})
assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0

78
tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
View file

@ -1120,3 +1120,81 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert():
assert upsert_data["create"]["teams"] == ["team-1"]
assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT"
assert "teams" not in upsert_data["update"]
def _member_write_tx() -> MagicMock:
tx = MagicMock()
created_user = MagicMock()
created_user.user_id = "pool-user"
created_user.model_dump.return_value = {
"user_id": "pool-user",
"user_email": "pool@example.com",
"teams": ["team-pool"],
"user_role": "internal_user",
}
created_budget = MagicMock()
created_budget.budget_id = "budget-pool"
membership = MagicMock()
membership.model_dump.return_value = {
"team_id": "team-pool",
"user_id": "pool-user",
"budget_id": "budget-pool",
"litellm_budget_table": None,
}
tx.litellm_usertable.upsert = AsyncMock(return_value=created_user)
tx.litellm_usertable.create = AsyncMock(return_value=created_user)
tx.litellm_usertable.update_many = AsyncMock()
tx.litellm_usertable.find_many = AsyncMock(return_value=[])
tx.litellm_budgettable.find_unique = AsyncMock(return_value=None)
tx.litellm_budgettable.create = AsyncMock(return_value=created_budget)
tx.litellm_teammembership.create = AsyncMock(return_value=membership)
return tx
@pytest.mark.parametrize(
"new_member",
[
Member(user_id="pool-user", role="user"),
Member(user_email="pool@example.com", role="user"),
],
ids=["by_user_id", "by_user_email"],
)
@pytest.mark.asyncio
async def test_add_new_member_runs_every_write_on_the_caller_transaction(new_member):
"""
Regression pin against exhausting the connection pool with advisory-lock waiters.
/team/member_add calls this while holding the team's advisory lock inside a transaction,
so it already owns a pooled connection. Any query issued on the regular client here needs
a second one, and enough concurrent adds for one team leave every connection parked on the
lock while the holder waits for a free one, so nothing ever commits or releases the lock.
Given a transaction, every read and write has to go through it.
"""
from litellm.proxy._types import LitellmUserRoles
tx = _member_write_tx()
prisma_client = AsyncMock()
result_user, result_membership = await add_new_member(
new_member=new_member,
max_budget_in_team=50.0,
prisma_client=prisma_client,
team_id="team-pool",
user_api_key_dict=UserAPIKeyAuth(
user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
),
litellm_proxy_admin_name="admin",
tx=tx,
)
assert result_user.user_id == "pool-user"
assert result_membership is not None
assert result_membership.budget_id == "budget-pool"
assert tx.litellm_budgettable.create.await_count == 1
assert tx.litellm_teammembership.create.await_count == 1
assert tx.litellm_usertable.upsert.await_count + tx.litellm_usertable.create.await_count == 1
prisma_client.db.assert_not_called()
prisma_client.get_data.assert_not_awaited()
prisma_client.insert_data.assert_not_awaited()

12
tests/test_litellm/repositories/test_repositories.py
View file

@ -541,17 +541,19 @@ class TestTeamRepository:
assert [m.user_id for m in members] == expected_ids
sql = tx.query_raw.call_args.args[0]
assert "FOR UPDATE" in sql
assert "FOR UPDATE" not in sql, (
"a row lock here can deadlock with the access-group endpoints; the caller must "
"already hold the team's advisory lock, so a plain read is all this needs"
)
assert tx.query_raw.call_args.args[1] == "team-1"
@pytest.mark.asyncio
async def test_get_members_with_roles_locked_missing_row(self, repo):
"""None, not [], so a caller can tell a deleted team from an empty one.
/team/member_add reconciles membership under this lock and has to fail,
and clean up the references it already wrote, when a /team/delete
committed underneath it. An empty list would look like a live team with
no members and it would carry on writing.
/team/member_add reconciles membership under the team's advisory lock and has to
fail, without writing anything, when a /team/delete committed underneath it. An
empty list would look like a live team with no members and it would carry on writing.
"""
tx = MagicMock()
tx.query_raw = AsyncMock(return_value=[])