Track per-member total spend on team memberships

litellm-proxy-extras/litellm_proxy_extras/migrations/20260421135425_add_team_membership_total_spend/migration.sql

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN     "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
+

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
   user_id    String
   team_id    String
   spend      Float    @default(0.0)
+  total_spend Float @default(0.0)
   budget_id String?
   litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
   @@id([user_id, team_id])

litellm/proxy/_types.py

@@ -1997,7 +1997,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
 
 
 class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
-    """Represents user-controllable params for a LiteLLM_BudgetTable record"""
+    """Represents user-controllable params for a LiteLLM_BudgetTable record.
+
+    Budget-write paths use `model_fields.keys()` on this class as an allowlist
+    for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
+    `LiteLLM_BudgetTableFull` so they aren't user-settable.
+    """
 
     budget_id: Optional[str] = None
     soft_budget: Optional[float] = None
@@ -2015,7 +2020,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
 
 
 class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
-    """Represents all params for a LiteLLM_BudgetTable record"""
+    """LiteLLM_BudgetTable + server-managed fields returned on API responses."""
 
     budget_reset_at: Optional[datetime] = None
     created_at: datetime
@@ -3695,7 +3700,11 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
     team_id: str
     budget_id: Optional[str] = None
     spend: Optional[float] = 0.0
-    litellm_budget_table: Optional[LiteLLM_BudgetTable]
+    total_spend: Optional[float] = 0.0
+    # Union so Pydantic picks Full when data has server-managed fields
+    # (/team/info) and Base when callers/tests construct with only
+    # user-settable fields.
+    litellm_budget_table: Optional[Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]]
 
     def safe_get_team_member_rpm_limit(self) -> Optional[int]:
         if self.litellm_budget_table is not None:

litellm/proxy/db/db_spend_update_writer.py

@@ -1300,7 +1300,10 @@ async def _commit_spend_updates_to_db(  # noqa: PLR0915
 
                                 batcher.litellm_teammembership.update_many(  # 'update_many' prevents error from being raised if no row exists
                                     where={"team_id": team_id, "user_id": user_id},
-                                    data={"spend": {"increment": response_cost}},
+                                    data={
+                                        "spend": {"increment": response_cost},
+                                        "total_spend": {"increment": response_cost},
+                                    },
                                 )
                     # Transaction succeeded, break out of retry loop
                     break

litellm/proxy/schema.prisma

@@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
   user_id    String
   team_id    String
   spend      Float    @default(0.0)
+  total_spend Float @default(0.0)
   budget_id String?
   litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
   @@id([user_id, team_id])

schema.prisma

@@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
   user_id    String
   team_id    String
   spend      Float    @default(0.0)
+  total_spend Float @default(0.0)
   budget_id String?
   litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
   @@id([user_id, team_id])

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -4,6 +4,7 @@
 import time
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -784,3 +785,37 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li
     assert len(find_many_calls) == 0
 
     litellm.max_end_user_budget_id = None
+
+
+def test_reset_budget_for_team_members_preserves_total_spend():
+    """Regression guard: reset_budget_for_litellm_team_members must zero `spend`
+    but leave `total_spend` untouched.
+
+    The reset writes `data={"spend": 0}` explicitly. If a future refactor adds
+    `"total_spend": 0` to that dict, this test fails immediately.
+    """
+    expired_budget = type(
+        "LiteLLM_BudgetTableFull",
+        (),
+        {"budget_id": "budget-1"},
+    )
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[])
+    mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(
+        return_value={"count": 1}
+    )
+
+    job = ResetBudgetJob(
+        proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client
+    )
+
+    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
+
+    mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once()
+    call_kwargs = (
+        mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs
+    )
+    assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"]
+    assert call_kwargs["data"] == {"spend": 0}
+    assert "total_spend" not in call_kwargs["data"]

tests/test_litellm/proxy/db/test_db_spend_update_writer.py

@@ -642,6 +642,81 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
     assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
 
 
+@pytest.mark.asyncio
+async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total_spend():
+    """
+    Verify that _commit_spend_updates_to_db increments BOTH spend (cycle-scoped)
+    and total_spend (non-resetting) on LiteLLM_TeamMembership in a single
+    update_many call, using the same response_cost.
+    """
+    db_writer = DBSpendUpdateWriter()
+
+    mock_batcher = MagicMock()
+    mock_batcher.litellm_verificationtoken = MagicMock()
+    mock_batcher.litellm_verificationtoken.update_many = MagicMock()
+    mock_batcher.litellm_usertable = MagicMock()
+    mock_batcher.litellm_usertable.update_many = MagicMock()
+    mock_batcher.litellm_teamtable = MagicMock()
+    mock_batcher.litellm_teamtable.update_many = MagicMock()
+    mock_batcher.litellm_teammembership = MagicMock()
+    mock_batcher.litellm_teammembership.update_many = MagicMock()
+    mock_batcher.litellm_organizationtable = MagicMock()
+    mock_batcher.litellm_organizationtable.update_many = MagicMock()
+    mock_batcher.litellm_tagtable = MagicMock()
+    mock_batcher.litellm_tagtable.update_many = MagicMock()
+    mock_batcher.litellm_agentstable = MagicMock()
+    mock_batcher.litellm_agentstable.update_many = MagicMock()
+
+    mock_transaction = AsyncMock()
+    mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
+    mock_transaction.__aexit__ = AsyncMock(return_value=False)
+    mock_transaction.batch_ = MagicMock(
+        return_value=AsyncMock(
+            __aenter__=AsyncMock(return_value=mock_batcher),
+            __aexit__=AsyncMock(return_value=False),
+        )
+    )
+
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db = MagicMock()
+    mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
+
+    mock_proxy_logging = MagicMock()
+    # Skip team-membership cache invalidation — out of scope for this test.
+    mock_proxy_logging.call_details.get = MagicMock(return_value=None)
+
+    team_id = "team-abc"
+    user_id = "user-xyz"
+    response_cost = 0.75
+    entity_id = f"team_id::{team_id}::user_id::{user_id}"
+    db_spend_update_transactions = {
+        "user_list_transactions": {},
+        "end_user_list_transactions": {},
+        "key_list_transactions": {},
+        "team_list_transactions": {},
+        "team_member_list_transactions": {entity_id: response_cost},
+        "org_list_transactions": {},
+        "tag_list_transactions": {},
+        "agent_list_transactions": {},
+    }
+
+    with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
+        await db_writer._commit_spend_updates_to_db(
+            prisma_client=mock_prisma_client,
+            n_retry_times=0,
+            proxy_logging_obj=mock_proxy_logging,
+            db_spend_update_transactions=db_spend_update_transactions,
+        )
+
+    mock_batcher.litellm_teammembership.update_many.assert_called_once()
+    call_kwargs = mock_batcher.litellm_teammembership.update_many.call_args[1]
+    assert call_kwargs["where"] == {"team_id": team_id, "user_id": user_id}
+    assert call_kwargs["data"] == {
+        "spend": {"increment": response_cost},
+        "total_spend": {"increment": response_cost},
+    }
+
+
 @pytest.mark.asyncio
 async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
     """