From 5e4b1b9df0ac766fc0640c186d994977dcf023d0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 24 Sep 2026 13:44:56 -0700 Subject: [PATCH] fix(proxy): pass team member spend rows as jsonb so a $0 flush cannot poison the pool connection (#43029) Prisma types a raw array parameter from the first batch a connection sees. After a flush in which every member cost was a whole number (a free model), the connection's cached statement expected int8[] and every later fractional batch on it failed with "improper binary format in array element", so member spend silently stopped landing while team spend kept rising. The rows now travel as one JSON document unpacked by jsonb_to_recordset with the column types declared in SQL, so Postgres types the numbers and the batch shape no longer matters. --- litellm/proxy/db/db_spend_update_writer.py | 34 +++--- .../spend/test_team_member_spend_flush.py | 106 ++++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 12 +- 3 files changed, 133 insertions(+), 19 deletions(-) create mode 100644 tests/integration/spend/test_team_member_spend_flush.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d3ac37a3e0f..17e6152bef6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -124,6 +124,12 @@ class _SpendIncrement(TypedDict): increment: ReadOnly[float] +class _MemberSpendRow(TypedDict): + user_id: ReadOnly[str] + team_id: ReadOnly[str] + cost: ReadOnly[float] + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -351,17 +357,22 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS # One statement adds every member's cost to their membership row. A missing row is created only # while the user is still on the team's roster, so a spend flush landing after a removal never -# recreates the member. +# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma +# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch +# (integers) every later fractional batch on that connection failed with "improper binary format". _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) -SELECT p.user_id, p.team_id, p.cost, p.cost -FROM unnest($1::text[], $2::text[], $3::float8[]) AS p(user_id, team_id, cost) +SELECT member.user_id, member.team_id, member.cost, member.cost +FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8) WHERE EXISTS ( SELECT 1 FROM "LiteLLM_TeamTable" t - WHERE t.team_id = p.team_id - AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id)) + WHERE t.team_id = member.team_id + AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id)) +) + OR EXISTS ( + SELECT 1 FROM "LiteLLM_TeamMembership" m + WHERE m.user_id = member.user_id AND m.team_id = member.team_id ) - OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id) ON CONFLICT (user_id, team_id) DO UPDATE SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend @@ -371,15 +382,12 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) - team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows) - for team_id in dict.fromkeys(team_ids): + for team_id in dict.fromkeys(team_id for team_id, _user_id, _cost in rows): _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) - _ = await transaction.execute_raw( - _TEAM_MEMBER_SPEND_SQL, - tuple(user_id for _team_id, user_id, _cost in rows), - team_ids, - tuple(cost for _team_id, _user_id, cost in rows), + members: Final = tuple( + _MemberSpendRow(user_id=user_id, team_id=team_id, cost=cost) for team_id, user_id, cost in rows ) + _ = await transaction.execute_raw(_TEAM_MEMBER_SPEND_SQL, json.dumps(members)) def get_llm_router(): diff --git a/tests/integration/spend/test_team_member_spend_flush.py b/tests/integration/spend/test_team_member_spend_flush.py new file mode 100644 index 00000000000..2433731f733 --- /dev/null +++ b/tests/integration/spend/test_team_member_spend_flush.py @@ -0,0 +1,106 @@ +"""Team member spend keeps landing after a flush in which every cost was a whole number. + +The proxy runs on a one-connection pool so every spend flush reuses the same database +connection. A batch of $0 requests (a free model here) is the whole-number batch, and the +fractional batches that follow it must still land on that connection. + +The $0 batch has to be flushed on its own before the paid request is sent. The spend log +row cannot prove that, since a separate monitor writes spend logs whenever they queue up, +but the daily user spend row is written by the flush cycle right after the member spend +statement, so its arrival means the whole-number batch has already been sent. +""" + +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from pydantic import JsonValue + +SINGLE_CONNECTION_CONFIG: Final = """ +model_list: [] +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 1 + proxy_batch_polling_interval: 1 + database_connection_pool_limit: 1 +router_settings: + disable_cooldowns: true +""" + + +def _member_row(team_id: str, user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', + (team_id, user_id), + ) + + +def _member_spend_is(rows: list[dict[str, JsonValue]], amount: float) -> bool: + return len(rows) == 1 and all( + float(str(rows[0][column])) == pytest.approx(amount) for column in ("spend", "total_spend") + ) + + +def _daily_user_spend_rows(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT spend FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s', (user_id,)) + + +def _logged_spend(request_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), + lambda found: len(found) == 1, + seconds=30, + ) + return float(str(rows[0]["spend"])) + + +def _chat(gateway: Gateway, key: str, model: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "member spend control"}]}, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(response.json()["id"]) + + +def test_fractional_member_spend_lands_after_a_whole_number_flush_on_the_same_connection( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = tmp_path / "single_connection_proxy.yaml" + config.write_text(SINGLE_CONNECTION_CONFIG) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + free: Final = scenario.model(input_cost_per_token=0, output_cost_per_token=0, num_retries=0) + paid: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002, num_retries=0) + team: Final = scenario.team(models=[free, paid]) + first: Final = scenario.user() + second: Final = scenario.user() + candidate.post( + "/team/member_add", + {"team_id": team, "member": [{"role": "user", "user_id": first}, {"role": "user", "user_id": second}]}, + ) + first_key: Final = scenario.key(team_id=team, user_id=first) + second_key: Final = scenario.key(team_id=team, user_id=second) + + _chat(candidate, first_key, free) + flushed: Final = eventually(lambda: _daily_user_spend_rows(first), lambda rows: len(rows) == 1, seconds=30) + assert float(str(flushed[0]["spend"])) == 0 + assert _member_spend_is(_member_row(team, first), 0) + + paid_spend: Final = _logged_spend(_chat(candidate, second_key, paid)) + assert paid_spend > 0 + eventually(lambda: _member_row(team, second), lambda rows: _member_spend_is(rows, paid_spend), seconds=30) + + repeat_spend: Final = _logged_spend(_chat(candidate, first_key, paid)) + eventually(lambda: _member_row(team, first), lambda rows: _member_spend_is(rows, repeat_spend), seconds=30) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,)), + lambda rows: float(str(rows[0]["spend"])) == pytest.approx(paid_spend + repeat_spend), + seconds=30, + ) 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 bf3b9aed234..7abb6e1ef92 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 @@ -994,11 +994,11 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster assert lock_statement is _TEAM_ADVISORY_LOCK_SQL assert locked_team_id == team_id assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement - statement, user_ids, team_ids, costs = spend_call.args + statement, members = spend_call.args assert statement is _TEAM_MEMBER_SPEND_SQL - assert (list(user_ids), list(team_ids), list(costs)) == ([user_id], [team_id], [response_cost]) + assert json.loads(members) == [{"user_id": user_id, "team_id": team_id, "cost": response_cost}] assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement - assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement + assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement @@ -1007,7 +1007,7 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster @pytest.mark.asyncio async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_user(): """ - The member spend statement touches rows in the order of its input arrays, so the batch + The member spend statement touches rows in the order of its input rows, so the batch is handed over sorted by (team_id, user_id), with each cost kept next to its member, and each distinct team is locked once, in `sorted(team_ids)` order, the order /team/delete locks in, so a concurrent flush and delete cannot deadlock. `eng` and `eng2` pin that: @@ -1034,13 +1034,13 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u ) *lock_calls, spend_call = mock_transaction.execute_raw.await_args_list - _statement, user_ids, team_ids, costs = spend_call.args + _statement, members = spend_call.args assert [lock_call.args for lock_call in lock_calls] == [ (_TEAM_ADVISORY_LOCK_SQL, "eng"), (_TEAM_ADVISORY_LOCK_SQL, "eng-b"), (_TEAM_ADVISORY_LOCK_SQL, "eng2"), ] - assert list(zip(team_ids, user_ids, costs)) == [ + assert [(row["team_id"], row["user_id"], row["cost"]) for row in json.loads(members)] == [ ("eng", "user_x", 0.3), ("eng", "user_y", 0.2), ("eng-b", "user_x", 0.4),