From 52f61c6c74a15efc653a264888c3e9d0de8d144b Mon Sep 17 00:00:00 2001 From: RoyVivat Date: Wed, 1 Apr 2026 13:08:19 -0700 Subject: [PATCH] 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 --- .../migration.sql | 4 + .../proxy/common_utils/reset_budget_job.py | 10 ++- .../management_endpoints/team_endpoints.py | 39 ++++----- .../common_utils/test_reset_budget_job.py | 32 ++++++++ .../test_team_endpoints.py | 80 ++++++++++++++++--- 5 files changed, 132 insertions(+), 33 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_total_spend_to_team_membership/migration.sql new file mode 100644 index 00000000000..f06c138284a --- /dev/null +++ b/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; diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 88e3cdd2134..26a68854faa 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b879fd031ac..ed75541d357 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 2625aae0e0d..396df6b177c 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b96f972eb2c..7b1a990722b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/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