From 3db26a5faa455f5bfa6ae0d25cada06f0bbd9730 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaan@berri.ai>
Date: Fri, 24 Apr 2026 18:56:09 -0700
Subject: [PATCH] fix: use atomic upsert+increment on
 LiteLLM_TeamMemberModelSpend instead of non-atomic JSON read-modify-write

---
 litellm/proxy/db/db_spend_update_writer.py | 38 ++++++++++------------
 1 file changed, 17 insertions(+), 21 deletions(-)

diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 9ee08e3fe2d..dadfc29bcce 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1816,13 +1816,11 @@ class DBSpendUpdateWriter:
         team_member_model_list_transactions: Dict[str, float],
     ) -> None:
         """
-        Persist per-model spend increments to LiteLLM_TeamMembership.model_spend so that
-        the spend counter cache can be seeded from the DB after a restart.
+        Persist per-model spend increments to LiteLLM_TeamMemberModelSpend using
+        atomic upsert+increment — no read-modify-write race condition.
 
         Key format: "team_id::<tid>::user_id::<uid>::model::<model>"
         """
-        import json as _json
-
         from litellm.proxy.utils import _raise_failed_update_spend_exception
 
         for key, cost in team_member_model_list_transactions.items():
@@ -1843,25 +1841,23 @@ class DBSpendUpdateWriter:
             for attempt in range(n_retry_times + 1):
                 start_time = time.time()
                 try:
-                    row = await prisma_client.db.litellm_teammembership.find_unique(
+                    await prisma_client.db.litellm_teammembermodeispend.upsert(
                         where={
-                            "user_id_team_id": {"user_id": user_id, "team_id": team_id}
-                        }
-                    )
-                    if row is None:
-                        break
-                    existing: dict = {}
-                    raw = getattr(row, "model_spend", None)
-                    if isinstance(raw, str):
-                        existing = _json.loads(raw)
-                    elif isinstance(raw, dict):
-                        existing = raw
-                    existing[model_name] = float(existing.get(model_name) or 0.0) + cost
-                    await prisma_client.db.litellm_teammembership.update(
-                        where={
-                            "user_id_team_id": {"user_id": user_id, "team_id": team_id}
+                            "user_id_team_id_model": {
+                                "user_id": user_id,
+                                "team_id": team_id,
+                                "model": model_name,
+                            }
+                        },
+                        data={
+                            "create": {
+                                "user_id": user_id,
+                                "team_id": team_id,
+                                "model": model_name,
+                                "spend": cost,
+                            },
+                            "update": {"spend": {"increment": cost}},
                         },
-                        data={"model_spend": _json.dumps(existing)},
                     )
                     break
                 except DB_CONNECTION_ERROR_TYPES as e: