Address PR review: fix test patch target, migration, multi-period reset, wire-up scope

- Fix test_upsert_team_member_budget_table: patch target was
  team_endpoints.prisma_client (module attr) but production code does a
  lazy import from proxy_server; correct target is
  litellm.proxy.proxy_server.prisma_client
- Split test into two: create path wires up null-budget members (expected),
  update path does NOT (new assertion for backwards-compat fix)
- Move null-budget wire-up to create-only path in upsert_team_member_budget_table
  to avoid silently constraining members who were intentionally unlimited
  on every settings save
- Add migration 20260401000000_add_total_spend_to_team_membership with
  ALTER TABLE ... ADD COLUMN IF NOT EXISTS "total_spend" for existing deployments
- Fix _reset_budget_reset_at_date else branch: use while loop (matching the
  None path) so proxies down for multiple periods advance the timestamp all
  the way to the next future date in one scheduler tick
- Add test_reset_budget_reset_at_date_multi_period_catchup to cover this

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql

@@ -0,0 +1,4 @@
+-- Add total_spend column to LiteLLM_TeamMembership
+-- Tracks lifetime (never-zeroed) spend for a user within a team,
+-- independent of the current-period spend that resets periodically.
+ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

litellm/proxy/common_utils/reset_budget_job.py

@@ -602,9 +602,13 @@ class ResetBudgetJob:
                     budget.budget_reset_at = anchor
                 else:
                     # Normal roll-forward: advance from the previous reset time.
-                    budget.budget_reset_at = budget.budget_reset_at + timedelta(
-                        seconds=duration_s
-                    )
+                    # Use a while loop so proxies that were down for multiple
+                    # periods catch up in one pass rather than firing the
+                    # spend-zero operation repeatedly over many scheduler ticks.
+                    while budget.budget_reset_at <= current_time:
+                        budget.budget_reset_at = budget.budget_reset_at + timedelta(
+                            seconds=duration_s
+                        )
         except Exception as e:
             verbose_proxy_logger.exception(
                 "Error resetting budget_reset_at for budget: %s. Item: %s", e, budget

litellm/proxy/management_endpoints/team_endpoints.py

@@ -229,7 +229,7 @@ class TeamMemberBudgetHandler:
                 updated_kv["metadata"] = {}
             updated_kv["metadata"]["team_member_budget_id"] = budget_row.budget_id
 
-        else:  # budget does not exist
+        else:  # budget does not exist — newly creating the template
             updated_kv = await TeamMemberBudgetHandler.create_team_member_budget_table(
                 data=team_table,
                 new_team_data_json=updated_kv,
@@ -240,25 +240,26 @@ class TeamMemberBudgetHandler:
                 team_member_budget_duration=team_member_budget_duration,
             )
 
-        # Wire up any existing members whose budget_id is null so they
-        # immediately see the template budget (and its budget_reset_at).
-        # This handles members who were added before the team's budget
-        # duration was ever configured.
-        if team_table.team_id is not None:
-            final_budget_id: Optional[str] = (
-                updated_kv.get("metadata", {}).get("team_member_budget_id")
-            )
-            if final_budget_id is not None:
-                from litellm.proxy.proxy_server import prisma_client as _prisma_client
+            # Wire up existing members whose budget_id is null to the
+            # newly-created template budget so they immediately inherit the
+            # reset schedule. Only do this on creation (not on every update)
+            # to avoid silently constraining members who were intentionally
+            # left unlimited.
+            if team_table.team_id is not None:
+                new_budget_id: Optional[str] = (
+                    updated_kv.get("metadata", {}).get("team_member_budget_id")
+                )
+                if new_budget_id is not None:
+                    from litellm.proxy.proxy_server import prisma_client as _prisma_client
 
-                if _prisma_client is not None:
-                    await _prisma_client.db.litellm_teammembership.update_many(
-                        where={
-                            "team_id": team_table.team_id,
-                            "budget_id": None,
-                        },
-                        data={"budget_id": final_budget_id},
-                    )
+                    if _prisma_client is not None:
+                        await _prisma_client.db.litellm_teammembership.update_many(
+                            where={
+                                "team_id": team_table.team_id,
+                                "budget_id": None,
+                            },
+                            data={"budget_id": new_budget_id},
+                        )
 
         # Remove team member fields from updated_kv
         TeamMemberBudgetHandler._clean_team_member_fields(updated_kv)

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -688,6 +688,38 @@ async def test_reset_budget_reset_at_date_none_reset_at_already_past():
     )
 
 
+@pytest.mark.asyncio
+async def test_reset_budget_reset_at_date_multi_period_catchup():
+    """
+    When budget_reset_at is non-null but multiple periods have elapsed (e.g.
+    the proxy was down for 3 weeks and the period is 7 days), the while loop
+    must advance the timestamp all the way to the next future date in a single
+    call rather than requiring the scheduler to fire 3 times.
+    """
+    from litellm.proxy._types import LiteLLM_BudgetTableFull
+
+    now = datetime.now(timezone.utc)
+    # reset_at is 22 days in the past; 7d period → 3 full periods have elapsed
+    previous_reset = now - timedelta(days=22)
+
+    budget = LiteLLM_BudgetTableFull(
+        budget_id="b-multi",
+        budget_duration="7d",
+        budget_reset_at=previous_reset,
+        created_at=now - timedelta(days=30),
+    )
+
+    result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now)
+
+    assert result.budget_reset_at > now, "reset must be in the future after catch-up"
+    # Should be previous_reset + 28d (4 periods) = now + 6d
+    expected = previous_reset + timedelta(days=28)
+    assert result.budget_reset_at == expected, (
+        f"Expected {expected}, got {result.budget_reset_at}. "
+        "Multi-period catch-up must advance through all elapsed periods in one call."
+    )
+
+
 @pytest.mark.asyncio
 async def test_reset_budget_at_startup_advances_expired_budget_reset_at(
     reset_budget_job, mock_prisma_client

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -1662,15 +1662,15 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
 
 
 @pytest.mark.asyncio
-async def test_upsert_team_member_budget_table_wires_up_null_budget_members():
+async def test_upsert_team_member_budget_table_wires_up_null_budget_members_on_create():
     """
-    Regression test: when team_member_budget_duration is set (create or update path),
-    existing team members whose budget_id is null must be assigned the template
-    budget so that their 'Next Budget Reset' column reflects the new reset time.
+    Regression test: when a team budget is being CREATED for the first time,
+    existing team members whose budget_id is null must be assigned the new
+    template budget so that their 'Next Budget Reset' column reflects the reset time.
 
     Previously, members added before the team's budget duration was configured had
-    budget_id=null. Changing the duration created/updated the template but left
-    those memberships disconnected — they never showed a reset date.
+    budget_id=null. Creating the template budget left those memberships disconnected
+    — they never showed a reset date.
     """
     from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -1681,14 +1681,17 @@ async def test_upsert_team_member_budget_table_wires_up_null_budget_members():
         user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
     )
 
+    # No existing budget_id — triggers the create path
     team_table = MagicMock(spec=LiteLLM_TeamTable)
-    team_table.metadata = {"team_member_budget_id": "template-budget-abc"}
+    team_table.metadata = {}
     team_table.team_id = "team-xyz"
+    team_table.team_alias = "Test Team"
+    team_table.budget_duration = None
 
     updated_kv: dict = {"team_id": "team-xyz"}
 
     mock_budget_response = MagicMock()
-    mock_budget_response.budget_id = "template-budget-abc"
+    mock_budget_response.budget_id = "new-template-budget-abc"
 
     mock_db = MagicMock()
     mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 2})
@@ -1698,12 +1701,12 @@ async def test_upsert_team_member_budget_table_wires_up_null_budget_members():
 
     with (
         patch(
-            "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget",
+            "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
             new_callable=AsyncMock,
             return_value=mock_budget_response,
         ),
         patch(
-            "litellm.proxy.management_endpoints.team_endpoints.prisma_client",
+            "litellm.proxy.proxy_server.prisma_client",
             mock_prisma,
         ),
     ):
@@ -1719,7 +1722,62 @@ async def test_upsert_team_member_budget_table_wires_up_null_budget_members():
     call_kwargs = mock_db.litellm_teammembership.update_many.call_args.kwargs
     assert call_kwargs["where"]["team_id"] == "team-xyz"
     assert call_kwargs["where"]["budget_id"] is None
-    assert call_kwargs["data"]["budget_id"] == "template-budget-abc"
+    assert call_kwargs["data"]["budget_id"] == "new-template-budget-abc"
+
+
+@pytest.mark.asyncio
+async def test_upsert_team_member_budget_table_update_does_not_wire_up_null_members():
+    """
+    When the team already has a template budget and the settings are UPDATED,
+    members with budget_id=null must NOT be silently assigned the budget.
+    Those members were intentionally left unlimited; assigning a budget on
+    every settings save would be a backwards-incompatible behaviour change.
+    """
+    from unittest.mock import AsyncMock, MagicMock, patch
+
+    from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth
+    from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
+
+    mock_user = UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
+    )
+
+    # Existing budget_id — triggers the update path
+    team_table = MagicMock(spec=LiteLLM_TeamTable)
+    team_table.metadata = {"team_member_budget_id": "template-budget-abc"}
+    team_table.team_id = "team-xyz"
+
+    updated_kv: dict = {"team_id": "team-xyz"}
+
+    mock_budget_response = MagicMock()
+    mock_budget_response.budget_id = "template-budget-abc"
+
+    mock_db = MagicMock()
+    mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 0})
+
+    mock_prisma = MagicMock()
+    mock_prisma.db = mock_db
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget",
+            new_callable=AsyncMock,
+            return_value=mock_budget_response,
+        ),
+        patch(
+            "litellm.proxy.proxy_server.prisma_client",
+            mock_prisma,
+        ),
+    ):
+        await TeamMemberBudgetHandler.upsert_team_member_budget_table(
+            team_table=team_table,
+            user_api_key_dict=mock_user,
+            updated_kv=updated_kv,
+            team_member_budget_duration="7d",
+        )
+
+    # update_many must NOT have been called — no silent budget assignment on update
+    mock_db.litellm_teammembership.update_many.assert_not_awaited()
 
 
 @pytest.mark.asyncio