From 6dce85c7285e7e7a6ea1b0b450ff4041815470ba Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:10:29 +0000 Subject: [PATCH] fix(proxy): skip recreating membership rows for members removed before a spend flush 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> --- litellm/proxy/db/db_spend_update_writer.py | 72 +++++++--- .../access_group_team_sync.py | 8 +- litellm/repositories/prisma_protocols.py | 8 +- litellm/repositories/team_repository.py | 12 +- .../proxy/db/test_db_spend_update_writer.py | 134 ++++++++++++------ 5 files changed, 160 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 913675705e3..5c86c52c539 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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::::user_id::" 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 diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py index 664e36c9f10..cfe207e66ea 100644 --- a/litellm/proxy/management_helpers/access_group_team_sync.py +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -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' diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 0919ae9f808..c742f3f5f3f 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -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]: ... diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 5ff07d76b5d..1b9cba48e41 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -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. diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index bed903d7a1f..b12dc68b7b1 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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),