scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix(proxy): skip recreating membership rows for members removed before a spend flush

Browse source 
Take the team advisory lock in the spend flush transaction and read the roster
through it, so a delayed flush after /team/member_delete cannot recreate the
deleted LiteLLM_TeamMembership row. TEAM_ADVISORY_LOCK_SQL moves to
team_repository so the spend writer can import it without a circular import

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-16 02:10:29 +00:00 • committed by ryan-crabbe-berri
parent d9b48ac941
commit 6dce85c728
5 changed files with 160 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/proxy/db/db_spend_update_writer.py
View file

@ -75,7 +75,8 @@ from litellm.proxy.spend_tracking.savings import (
marks_gateway_injection,
)
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.repositories.prisma_protocols import BatchTable
from litellm.repositories.prisma_protocols import BatchTable, RawQueryTransaction
from litellm.repositories.team_repository import TEAM_ADVISORY_LOCK_SQL, TeamRepository
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
@ -133,7 +134,7 @@ class _SpendBatchManager(Protocol):
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
class _SpendTransaction(Protocol):
class _SpendTransaction(RawQueryTransaction, Protocol):
def batch_(self) -> _SpendBatchManager: ...
@ -161,6 +162,50 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
return tx
async def _lock_and_read_rosters(
prisma_client: PrismaClient, transaction: _SpendTransaction, team_ids: Sequence[str]
) -> frozenset[tuple[str, str]]:
"""Take each team's advisory lock on ``transaction`` and return its rostered (user_id, team_id) pairs.
A spend flush can land after ``/team/member_delete`` removed the member. Holding the same
lock that endpoint takes, until this transaction commits, means a member read here is on
the team for the whole flush, so only they may have a missing membership row created.
"""
repository: Final = TeamRepository(prisma_client)
async def locked_roster(team_id: str) -> tuple[tuple[str, str], ...]:
await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
roster: Final = await repository.get_members_with_roles_locked(transaction, team_id)
return tuple((member.user_id, team_id) for member in roster or () if member.user_id is not None)
rosters: Final = tuple([await locked_roster(team_id) for team_id in sorted(frozenset(team_ids))])
return frozenset(pair for roster in rosters for pair in roster)
def _queue_team_member_spend(
memberships: BatchTable, user_id: str, team_id: str, response_cost: float, rostered: bool
) -> None:
increments: Final = {
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
}
if not rostered:
memberships.update_many(where={"team_id": team_id, "user_id": user_id}, data=increments)
return
memberships.upsert(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={
"create": {
"team_id": team_id,
"user_id": user_id,
"spend": response_cost,
"total_spend": response_cost,
},
"update": increments,
},
)
def get_llm_router():
"""The proxy's router, or None outside a running proxy.
@ -1685,6 +1730,9 @@ class DBSpendUpdateWriter:
start_time = time.time()
try:
async with _spend_update_tx(prisma_client) as transaction:
rostered_members = await _lock_and_read_rosters(
prisma_client, transaction, tuple(team_id for _, team_id in team_memberships_to_invalidate)
)
async with transaction.batch_() as batcher:
# Sort by composite key for consistent lock ordering across pods to prevent deadlocks.
# Key format "team_id::<v>::user_id::<v>" makes the string sort equivalent to sorting by (team_id, user_id).
@ -1693,20 +1741,12 @@ class DBSpendUpdateWriter:
team_id = key.split("::")[1]
user_id = key.split("::")[3]
batcher.litellm_teammembership.upsert(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={
"create": {
"team_id": team_id,
"user_id": user_id,
"spend": response_cost,
"total_spend": response_cost,
},
"update": {
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
},
},
_queue_team_member_spend(
batcher.litellm_teammembership,
user_id,
team_id,
response_cost,
(user_id, team_id) in rostered_members,
)
# Transaction succeeded, break out of retry loop
break

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

@ -21,13 +21,7 @@ from typing import Final, Protocol
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 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"
from litellm.repositories.team_repository import TEAM_ADVISORY_LOCK_SQL
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'

8
litellm/repositories/prisma_protocols.py
View file

@ -7,7 +7,7 @@ private ones per file.
"""
from collections.abc import Mapping, Sequence
from typing import Protocol, TypeVar
from typing import LiteralString, Protocol, TypeVar
RowT_co = TypeVar("RowT_co", covariant=True)
@ -108,6 +108,12 @@ class PrismaRecord(Protocol):
def dict(self) -> Mapping[str, object]: ...
class RawQueryTransaction(Protocol):
"""A prisma transaction handle that can run raw SQL, e.g. an advisory lock or a locked read."""
async def query_raw(self, query: LiteralString, *args: str) -> Sequence[Mapping[str, object]]: ...
class ReadOnlyTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[PrismaRecord]: ...

12
litellm/repositories/team_repository.py
View file

@ -15,10 +15,9 @@ from litellm.repositories.base_repository import (
DbRecord,
record_to_dict,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.prisma_protocols import RawQueryTransaction, TableActions
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
@ -40,6 +39,13 @@ def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays:
return team
# hashtext collisions only cost two unrelated teams a little serialization, and the
# 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,
# the access-group mirror and the team member spend flush all reuse this exact statement to
# serialize against each other, 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"
_MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member])
_JSON_ENCODED_TEAM_FIELDS: Final = (
"metadata",
@ -78,7 +84,7 @@ 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:
async def get_members_with_roles_locked(self, tx: RawQueryTransaction, team_id: str) -> list[Member] | None:
"""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.

134
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -20,6 +20,7 @@ from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
from litellm.repositories.team_repository import TEAM_ADVISORY_LOCK_SQL
@pytest.mark.asyncio
@ -913,38 +914,22 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total_spend():
"""
Verify that _commit_spend_updates_to_db increments BOTH spend (cycle-scoped)
and total_spend (non-resetting) on LiteLLM_TeamMembership in a single
upsert call, using the same response_cost.
Regression (LIT-5502): members added without a budget had no membership row, and
the previous update_many matched zero rows, so their spend was silently dropped.
The upsert has to create the row seeded with this call's cost in that case.
"""
db_writer = DBSpendUpdateWriter()
def _team_member_flush_fixtures(team_id: str, rostered_user_ids: list[str]) -> tuple[MagicMock, AsyncMock, MagicMock]:
"""A batcher, transaction and prisma client whose locked roster read for `team_id` lists `rostered_user_ids`."""
mock_batcher = MagicMock()
mock_batcher.litellm_verificationtoken = MagicMock()
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
mock_batcher.litellm_usertable = MagicMock()
mock_batcher.litellm_usertable.update_many = MagicMock()
mock_batcher.litellm_teamtable = MagicMock()
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_teammembership = MagicMock()
mock_batcher.litellm_teammembership.upsert = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_tagtable = MagicMock()
mock_batcher.litellm_tagtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
mock_batcher.litellm_teammembership.update_many = MagicMock()
roster_row = {"members_with_roles": json.dumps([{"user_id": uid, "role": "user"} for uid in rostered_user_ids])}
async def query_raw(query: str, *args: str) -> list[dict[str, object]]:
return [] if query == TEAM_ADVISORY_LOCK_SQL else [roster_row]
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.query_raw = AsyncMock(side_effect=query_raw)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
@ -955,16 +940,11 @@ async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
return mock_batcher, mock_transaction, mock_prisma_client
mock_proxy_logging = MagicMock()
# Skip team-membership cache invalidation — out of scope for this test.
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
team_id = "team-abc"
user_id = "user-xyz"
response_cost = 0.75
entity_id = f"team_id::{team_id}::user_id::{user_id}"
db_spend_update_transactions = {
def _team_member_only_transactions(entity_id: str, response_cost: float) -> dict[str, dict[str, float]]:
return {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
@ -975,14 +955,38 @@ async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total
"agent_list_transactions": {},
}
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=db_spend_update_transactions,
)
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total_spend():
"""
Verify that _commit_spend_updates_to_db increments BOTH spend (cycle-scoped)
and total_spend (non-resetting) on LiteLLM_TeamMembership in a single
upsert call, using the same response_cost.
Regression (LIT-5502): members added without a budget had no membership row, and
the previous update_many matched zero rows, so their spend was silently dropped.
For a member still on the team roster the upsert has to create the row seeded
with this call's cost in that case.
"""
db_writer = DBSpendUpdateWriter()
team_id = "team-abc"
user_id = "user-xyz"
response_cost = 0.75
mock_batcher, mock_transaction, mock_prisma_client = _team_member_flush_fixtures(team_id, [user_id])
mock_proxy_logging = MagicMock()
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=_team_member_only_transactions(
f"team_id::{team_id}::user_id::{user_id}", response_cost
),
)
assert mock_transaction.query_raw.await_args_list[0] == call(TEAM_ADVISORY_LOCK_SQL, team_id)
mock_batcher.litellm_teammembership.upsert.assert_called_once()
mock_batcher.litellm_teammembership.update_many.assert_not_called()
call_kwargs = mock_batcher.litellm_teammembership.upsert.call_args.kwargs
@ -1001,6 +1005,43 @@ async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total
}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_does_not_recreate_membership_of_removed_team_member():
"""
A spend flush that lands after /team/member_delete must not resurrect the deleted
membership row: a user missing from the team roster, read under the team's advisory
lock inside the flush transaction, only gets an increment on whatever row still
exists, never a create.
"""
db_writer = DBSpendUpdateWriter()
team_id = "team-abc"
removed_user_id = "user-removed"
response_cost = 0.75
mock_batcher, mock_transaction, mock_prisma_client = _team_member_flush_fixtures(team_id, ["user-still-here"])
mock_proxy_logging = MagicMock()
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=_team_member_only_transactions(
f"team_id::{team_id}::user_id::{removed_user_id}", response_cost
),
)
assert mock_transaction.query_raw.await_args_list[0] == call(TEAM_ADVISORY_LOCK_SQL, team_id)
mock_batcher.litellm_teammembership.upsert.assert_not_called()
mock_batcher.litellm_teammembership.update_many.assert_called_once_with(
where={"team_id": team_id, "user_id": removed_user_id},
data={
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
},
)
@pytest.mark.asyncio
async def test_org_spend_increments_organization_membership_row_for_the_calling_user():
"""A request made with a user_id inside an org must increment that user's
@ -2232,13 +2273,9 @@ async def test_commit_daily_tag_spend_no_requeue_on_success():
"team_id::team_b::user_id::user_x": 0.3,
},
"litellm_teammembership",
"upsert",
"user_id_team_id",
[
{"user_id": "user_x", "team_id": "team_a"},
{"user_id": "user_x", "team_id": "team_b"},
{"user_id": "user_x", "team_id": "team_c"},
],
"update_many",
"team_id",
["team_a", "team_b", "team_c"],
id="team_member",
),
pytest.param(
@ -2312,6 +2349,8 @@ async def test_commit_spend_updates_iterates_in_sorted_order(
)
)
mock_transaction.query_raw = AsyncMock(return_value=[])
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
@ -3071,6 +3110,7 @@ def _good_tx(mock_batcher):
tx = AsyncMock()
tx.__aenter__ = AsyncMock(return_value=tx)
tx.__aexit__ = AsyncMock(return_value=False)
tx.query_raw = AsyncMock(return_value=[])
tx.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),