From ba1b466480a000b559fde36c246c9d31392af1ce Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Feb 2026 18:16:01 -0800 Subject: [PATCH 1/7] fix to ensure budget duration is being inherited from budget tier for keys --- .../proxy/common_utils/reset_budget_job.py | 36 +++++ .../key_management_endpoints.py | 7 + .../common_utils/test_reset_budget_job.py | 116 +++++++++++++++ .../test_key_management_endpoints.py | 140 ++++++++++++++++++ 4 files changed, 299 insertions(+) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index fb600cee26b..6f038d127f6 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -69,6 +69,38 @@ class ResetBudgetJob: }, ) + async def reset_budget_for_keys_linked_to_budgets( + self, budgets_to_reset: List[LiteLLM_BudgetTableFull] + ): + """ + Resets the spend for keys linked to budget tiers that are being reset. + + This handles keys that have budget_id but no budget_duration set on the key + itself (e.g. keys created before the fix to inherit budget_duration from + the linked budget tier). + + Keys that have their own budget_duration are already handled by + reset_budget_for_litellm_keys() and are excluded here to avoid + double-resetting. + """ + budget_ids = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + if not budget_ids: + return + + return await self.prisma_client.db.litellm_verificationtoken.update_many( + where={ + "budget_id": {"in": budget_ids}, + "budget_duration": None, # only keys without their own reset schedule + }, + data={ + "spend": 0, + }, + ) + async def reset_budget_for_litellm_budget_table(self): """ Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired @@ -112,6 +144,10 @@ class ResetBudgetJob: budgets_to_reset=budgets_to_reset ) + await self.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + if endusers_to_reset is not None and len(endusers_to_reset) > 0: for enduser in endusers_to_reset: try: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2eb6cf65281..b62ce329548 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -594,6 +594,13 @@ async def _common_key_generation_helper( # noqa: PLR0915 if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) + elif _budget_id is not None and prisma_client is not None: + # Inherit budget_duration from linked budget tier if not explicitly set on the key + budget_row = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": _budget_id} + ) + if budget_row is not None and budget_row.budget_duration is not None: + data_json["key_budget_duration"] = budget_row.budget_duration if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id 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 a059a3adcb1..f63c77c1fc8 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 @@ -25,9 +25,21 @@ class MockLiteLLMTeamMembership: return {"count": 1} +class MockLiteLLMVerificationToken: + def __init__(self): + self.update_many_calls: List[Dict[str, Any]] = [] + + async def update_many( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Dict[str, Any]: + self.update_many_calls.append({"where": where, "data": data}) + return {"count": 1} + + class MockDB: def __init__(self): self.litellm_teammembership = MockLiteLLMTeamMembership() + self.litellm_verificationtoken = MockLiteLLMVerificationToken() class MockPrismaClient: @@ -320,3 +332,107 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): assert mock_prisma_client.updated_data["user"][0].spend == 0.0 assert mock_prisma_client.updated_data["team"][0].spend == 0.0 assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0 + + +def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_client): + """ + Test that when a budget tier is reset, keys linked to that budget + (via budget_id) that don't have their own budget_duration also get + their spend reset. + + This covers the case where keys were created with budget_id but + budget_duration was not inherited to the key (pre-fix keys). + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + + # Create a budget tier that is due for reset + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + budgets_to_reset = [test_budget] + + # Run the method + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + ) + + # Verify that update_many was called on litellm_verificationtoken + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1, f"Expected 1 update_many call, got {len(calls)}" + + # Verify the where clause filters by budget_id and null budget_duration + call = calls[0] + assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + assert call["where"]["budget_duration"] is None + + # Verify spend is reset to 0 + assert call["data"]["spend"] == 0 + + +def test_reset_budget_for_keys_linked_to_budgets_empty( + reset_budget_job, mock_prisma_client +): + """ + Test that when there are no budgets to reset, no update is performed + on the verification token table. + """ + # Run with empty list + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=[] + ) + ) + + # Verify no update_many calls were made + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 0 + + +def test_budget_table_reset_also_resets_linked_keys( + reset_budget_job, mock_prisma_client +): + """ + Integration-style test: when reset_budget_for_litellm_budget_table runs, + it should also reset spend for keys linked to the expiring budget tiers + (in addition to end-users and team members). + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + + # Run the full budget table reset + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Verify that keys linked to the budget were also reset + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1, ( + "Expected reset_budget_for_litellm_budget_table to also reset keys " + f"linked to expiring budgets, but got {len(calls)} update_many calls" + ) + assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + assert calls[0]["data"]["spend"] == 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 39f8d1cccb0..472504871ed 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5559,3 +5559,143 @@ async def test_validate_key_list_check_key_hash_not_found(): assert exc_info.value.code == "403" or exc_info.value.code == 403 assert "Key Hash not found" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_key_inherits_budget_duration_from_budget_tier(): + """ + Test that when a key is created with budget_id pointing to a budget tier + that has budget_duration, the key inherits budget_duration from the tier + even when budget_duration is not explicitly set on the key request. + + This verifies the fix for the bug where keys created with budget_id + would have null budget_duration and budget_reset_at, causing the + budget reset job to never reset their spend. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + # Mock the budget tier lookup to return a budget with budget_duration="7d" + mock_budget_row = MagicMock() + mock_budget_row.budget_duration = "7d" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_budget_row + ) + + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), patch( + "litellm.proxy.proxy_server.llm_router", None + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + budget_id="7d-budget-tier", + max_budget=10.0, + # NOTE: budget_duration is intentionally NOT set here + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ), + litellm_changed_by=None, + team_table=None, + ) + + # Verify generate_key_helper_fn was called + mock_generate_key.assert_awaited_once() + call_kwargs = mock_generate_key.call_args.kwargs + + # The key should have inherited key_budget_duration from the budget tier + assert call_kwargs.get("key_budget_duration") == "7d", ( + "key_budget_duration should be inherited from the linked budget tier " + f"but got: {call_kwargs.get('key_budget_duration')}" + ) + + # Verify the budget tier was looked up with the correct budget_id + mock_prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( + where={"budget_id": "7d-budget-tier"} + ) + + +@pytest.mark.asyncio +async def test_key_does_not_override_explicit_budget_duration(): + """ + Test that when a key is created with both budget_id and an explicit + budget_duration, the explicit budget_duration takes precedence over + the budget tier's budget_duration. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma = MagicMock() + # The budget tier has budget_duration="7d" + mock_budget_row = MagicMock() + mock_budget_row.budget_duration = "7d" + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_budget_row + ) + + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), patch( + "litellm.proxy.proxy_server.llm_router", None + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + budget_id="7d-budget-tier", + max_budget=10.0, + budget_duration="30d", # explicit budget_duration should take precedence + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ), + litellm_changed_by=None, + team_table=None, + ) + + mock_generate_key.assert_awaited_once() + call_kwargs = mock_generate_key.call_args.kwargs + + # The explicit budget_duration should take precedence + assert call_kwargs.get("key_budget_duration") == "30d", ( + "Explicit budget_duration should take precedence over the budget tier's value " + f"but got: {call_kwargs.get('key_budget_duration')}" + ) + + # The budget tier should NOT have been looked up since budget_duration was explicit + mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited() From cf14f0c8214951b63593e2e19e50652f253ba5c9 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 14 Feb 2026 10:53:58 -0800 Subject: [PATCH 2/7] change logic to match Kriish's input --- .../key_management_endpoints.py | 11 ++---- .../test_key_management_endpoints.py | 38 +++++++------------ 2 files changed, 18 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b62ce329548..ab610a77b4d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -592,15 +592,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 if _budget_id is not None: data_json["budget_id"] = _budget_id + # Only set budget_duration on key when explicitly provided. Keys with budget_id + # but no explicit budget_duration follow their linked budget tier's schedule; + # reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + # This avoids duplicating budget_duration on keys so tier updates apply automatically. if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - elif _budget_id is not None and prisma_client is not None: - # Inherit budget_duration from linked budget tier if not explicitly set on the key - budget_row = await prisma_client.db.litellm_budgettable.find_unique( - where={"budget_id": _budget_id} - ) - if budget_row is not None and budget_row.budget_duration is not None: - data_json["key_budget_duration"] = budget_row.budget_duration if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 472504871ed..21be2ad69e8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5562,26 +5562,19 @@ async def test_validate_key_list_check_key_hash_not_found(): @pytest.mark.asyncio -async def test_key_inherits_budget_duration_from_budget_tier(): +async def test_key_with_budget_id_does_not_store_budget_duration(): """ - Test that when a key is created with budget_id pointing to a budget tier - that has budget_duration, the key inherits budget_duration from the tier - even when budget_duration is not explicitly set on the key request. + Test that when a key is created with budget_id but without explicit + budget_duration, the key does NOT get budget_duration stored on it. - This verifies the fix for the bug where keys created with budget_id - would have null budget_duration and budget_reset_at, causing the - budget reset job to never reset their spend. + Keys with budget_id follow their linked budget tier's reset schedule; + reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + This avoids duplicating budget_duration on keys so tier updates apply + automatically to all linked keys. """ from unittest.mock import AsyncMock, MagicMock, patch - # Mock the budget tier lookup to return a budget with budget_duration="7d" - mock_budget_row = MagicMock() - mock_budget_row.budget_duration = "7d" - mock_prisma = MagicMock() - mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( - return_value=mock_budget_row - ) mock_generate_key = AsyncMock( return_value={ @@ -5619,20 +5612,17 @@ async def test_key_inherits_budget_duration_from_budget_tier(): team_table=None, ) - # Verify generate_key_helper_fn was called mock_generate_key.assert_awaited_once() call_kwargs = mock_generate_key.call_args.kwargs - # The key should have inherited key_budget_duration from the budget tier - assert call_kwargs.get("key_budget_duration") == "7d", ( - "key_budget_duration should be inherited from the linked budget tier " - f"but got: {call_kwargs.get('key_budget_duration')}" + # Key should NOT have key_budget_duration - it follows the budget tier's schedule + assert call_kwargs.get("key_budget_duration") is None, ( + "key_budget_duration should be None for budget-linked keys without explicit " + f"budget_duration; got: {call_kwargs.get('key_budget_duration')}" ) - # Verify the budget tier was looked up with the correct budget_id - mock_prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( - where={"budget_id": "7d-budget-tier"} - ) + # No budget tier lookup - we don't copy budget_duration onto the key + mock_prisma.db.litellm_budgettable.find_unique.assert_not_called() @pytest.mark.asyncio @@ -5698,4 +5688,4 @@ async def test_key_does_not_override_explicit_budget_duration(): ) # The budget tier should NOT have been looked up since budget_duration was explicit - mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited() + mock_prisma.db.litellm_budgettable.find_unique.assert_not_called() From f79a8f7809ceaa1b2650c98ebdb95bc126a5beac Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 14:42:17 -0800 Subject: [PATCH 3/7] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 6f038d127f6..c63309da1dc 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -95,11 +95,13 @@ class ResetBudgetJob: where={ "budget_id": {"in": budget_ids}, "budget_duration": None, # only keys without their own reset schedule + "spend": {"gt": 0}, # only reset keys that have accumulated spend }, data={ "spend": 0, }, ) + ) async def reset_budget_for_litellm_budget_table(self): """ From ea87dd216281155bc60da8a2751ae69122a97cce Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 19:01:35 -0800 Subject: [PATCH 4/7] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index c63309da1dc..b926fe28bed 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -101,7 +101,6 @@ class ResetBudgetJob: "spend": 0, }, ) - ) async def reset_budget_for_litellm_budget_table(self): """ From 2d9508ec96253557c97e6d6cc229d0a1d9e702d8 Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 19:09:26 -0800 Subject: [PATCH 5/7] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index b926fe28bed..4933e679d71 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -76,13 +76,14 @@ class ResetBudgetJob: Resets the spend for keys linked to budget tiers that are being reset. This handles keys that have budget_id but no budget_duration set on the key - itself (e.g. keys created before the fix to inherit budget_duration from - the linked budget tier). + itself. Keys with budget_id rely on their linked budget tier's reset schedule + rather than having their own budget_duration. Keys that have their own budget_duration are already handled by reset_budget_for_litellm_keys() and are excluded here to avoid double-resetting. """ + """ budget_ids = [ budget.budget_id for budget in budgets_to_reset From 0117b35a6bc1c2877f074794664e9ece1294114d Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 16 Feb 2026 09:59:19 -0800 Subject: [PATCH 6/7] added more tests, fixed tests --- .../proxy/common_utils/reset_budget_job.py | 3 +- .../test_proxy_budget_reset.py | 16 +++++++ .../common_utils/test_reset_budget_job.py | 43 +++++++++++++++++++ 3 files changed, 60 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 4933e679d71..8ce73d29c84 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -83,7 +83,6 @@ class ResetBudgetJob: reset_budget_for_litellm_keys() and are excluded here to avoid double-resetting. """ - """ budget_ids = [ budget.budget_id for budget in budgets_to_reset @@ -617,4 +616,4 @@ class ResetBudgetJob: await ResetBudgetJob._reset_budget_common( item=key, current_time=current_time, item_type="key" ) - return key + return key \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 7cddde30421..34423a88da4 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -229,6 +229,10 @@ async def test_reset_budget_endusers_partial_failure(): prisma_client.get_data.side_effect = get_data_mock prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -389,6 +393,10 @@ async def test_reset_budget_continues_other_categories_on_failure(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -863,6 +871,10 @@ async def test_service_logger_endusers_success(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -938,6 +950,10 @@ async def test_service_logger_endusers_failure(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() 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 f63c77c1fc8..f975460836a 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 @@ -382,6 +382,49 @@ def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 +def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_duration( + reset_budget_job, mock_prisma_client +): + """ + Test that keys with BOTH budget_id AND budget_duration are excluded from + reset_budget_for_keys_linked_to_budgets. Such keys have their own reset + schedule and are handled only by reset_budget_for_litellm_keys(). The + budget_duration=None filter ensures they are NOT double-reset when the + linked budget tier expires. + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + budgets_to_reset = [test_budget] + + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + ) + + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1 + call = calls[0] + + # Critical: budget_duration must be None so keys with their own budget_duration + # (e.g. key has budget_id="X" AND budget_duration=60) are excluded. + # Those keys are reset only by reset_budget_for_litellm_keys() - no double-reset. + assert call["where"]["budget_duration"] is None + assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + + def test_reset_budget_for_keys_linked_to_budgets_empty( reset_budget_job, mock_prisma_client ): From 9f7f19067079977ae6d8fd3ecbd82a6f39362f02 Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 16 Feb 2026 10:24:06 -0800 Subject: [PATCH 7/7] resolved greptile comment --- tests/litellm_utils_tests/test_proxy_budget_reset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 34423a88da4..34b2043261c 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -1042,6 +1042,9 @@ async def test_reset_budget_for_litellm_team_members_called(): prisma_client.db.litellm_teammembership.update_many = AsyncMock( return_value={"count": 2} ) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock()