From 5c02d627103f9010bdeb0016fd104e5b05bdd4d1 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaan@berri.ai>
Date: Fri, 24 Apr 2026 18:19:40 -0700
Subject: [PATCH] feat: persist per-model team member spend to DB via
 _commit_team_member_model_spend_to_db

---
 litellm/proxy/db/db_spend_update_writer.py | 99 ++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index c06e1850d9f..9ee08e3fe2d 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -361,6 +361,7 @@ class DBSpendUpdateWriter:
                 response_cost=response_cost,
                 team_id=team_id,
                 user_id=user_id,
+                model=payload_copy.get("model_group") or payload_copy.get("model"),
                 prisma_client=prisma_client,
             )
         except Exception:
@@ -556,6 +557,7 @@ class DBSpendUpdateWriter:
         team_id: Optional[str],
         user_id: Optional[str],
         prisma_client: Optional[PrismaClient],
+        model: Optional[str] = None,
     ):
         try:
             if team_id is None or prisma_client is None:
@@ -584,6 +586,19 @@ class DBSpendUpdateWriter:
                             response_cost=response_cost,
                         )
                     )
+                    # Track per-model spend for per-model budget enforcement
+                    if model is not None:
+                        # key is "team_id::<value>::user_id::<value>::model::<value>"
+                        team_member_model_key = (
+                            f"team_id::{team_id}::user_id::{user_id}::model::{model}"
+                        )
+                        await self.spend_update_queue.add_update(
+                            update=SpendUpdateQueueItem(
+                                entity_type=Litellm_EntityType.TEAM_MEMBER_MODEL,
+                                entity_id=team_member_model_key,
+                                response_cost=response_cost,
+                            )
+                        )
             except Exception as e:
                 verbose_proxy_logger.error(
                     "Spend tracking - failed to enqueue team member spend update. "
@@ -1337,6 +1352,21 @@ class DBSpendUpdateWriter:
                             f"Invalidated team membership cache for user_id={user_id}, team_id={team_id}"
                         )
 
+        ### UPDATE TEAM MEMBER MODEL SPEND ###
+        team_member_model_list_transactions = db_spend_update_transactions.get(
+            "team_member_model_list_transactions"
+        )
+        if (
+            team_member_model_list_transactions is not None
+            and len(team_member_model_list_transactions) > 0
+        ):
+            await self._commit_team_member_model_spend_to_db(
+                prisma_client=prisma_client,
+                n_retry_times=n_retry_times,
+                proxy_logging_obj=proxy_logging_obj,
+                team_member_model_list_transactions=team_member_model_list_transactions,
+            )
+
         ### UPDATE ORG TABLE ###
         org_list_transactions = db_spend_update_transactions["org_list_transactions"]
         verbose_proxy_logger.debug(
@@ -1778,6 +1808,75 @@ class DBSpendUpdateWriter:
                 e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
             )
 
+    async def _commit_team_member_model_spend_to_db(
+        self,
+        prisma_client: PrismaClient,
+        n_retry_times: int,
+        proxy_logging_obj: ProxyLogging,
+        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.
+
+        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():
+            parts = key.split("::")
+            # Expected: ["team_id", "<tid>", "user_id", "<uid>", "model", "<model>"]
+            if (
+                len(parts) != 6
+                or parts[0] != "team_id"
+                or parts[2] != "user_id"
+                or parts[4] != "model"
+            ):
+                verbose_proxy_logger.debug(
+                    "Skipping malformed team_member_model key: %s", key
+                )
+                continue
+            team_id, user_id, model_name = parts[1], parts[3], parts[5]
+
+            for attempt in range(n_retry_times + 1):
+                start_time = time.time()
+                try:
+                    row = await prisma_client.db.litellm_teammembership.find_unique(
+                        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}
+                        },
+                        data={"model_spend": _json.dumps(existing)},
+                    )
+                    break
+                except DB_CONNECTION_ERROR_TYPES as e:
+                    if attempt >= n_retry_times:
+                        _raise_failed_update_spend_exception(
+                            e=e,
+                            start_time=start_time,
+                            proxy_logging_obj=proxy_logging_obj,
+                        )
+                    await asyncio.sleep(2**attempt)
+                except Exception as e:
+                    _raise_failed_update_spend_exception(
+                        e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
+                    )
+
     @staticmethod
     async def update_daily_user_spend(
         n_retry_times: int,