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

@@ -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::<value>::user_id::<value>"; 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():

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,
+        )

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),