fix(proxy): backport team member spend jsonb flush to rc/1.103.0 (#43029)

litellm/proxy/db/db_spend_update_writer.py

@@ -123,6 +123,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
@@ -203,17 +209,22 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
 
 # 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
@@ -223,15 +234,12 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
 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

@@ -982,11 +982,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
@@ -995,7 +995,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:
@@ -1022,13 +1022,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),