From d566d2e01967d2b2ee6c940c461e3e2b86e849f6 Mon Sep 17 00:00:00 2001 From: jibanez-staticduo Date: Mon, 20 Apr 2026 09:43:03 +0200 Subject: [PATCH] fix(proxy): avoid invalid Prisma JsonFilter in reset_budget_windows (cherry picked from commit e3af3d168fbd5e65b8f989e19fb8e5b165fce689) (cherry picked from commit d6809667df49442f060bdad02f4b4e0976071a97) --- .../proxy/common_utils/reset_budget_job.py | 8 +- .../common_utils/test_reset_budget_job.py | 100 ++++++++++++++++++ 2 files changed, 102 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index b11af04e2b1..6c769b0ff54 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -634,9 +634,7 @@ class ResetBudgetJob: # --- Keys --- try: - all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( - where={"budget_limits": {"not": None}} # type: ignore[arg-type] - ) + all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many() for key in all_keys: raw = key.budget_limits # type: ignore[attr-defined] if not raw: @@ -663,9 +661,7 @@ class ResetBudgetJob: # --- Teams --- try: - all_teams = await self.prisma_client.db.litellm_teamtable.find_many( - where={"budget_limits": {"not": None}} # type: ignore[arg-type] - ) + all_teams = await self.prisma_client.db.litellm_teamtable.find_many() for team in all_teams: raw = team.budget_limits # type: ignore[attr-defined] if not raw: 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 8206079cb8d..17eb882ea21 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 @@ -4,6 +4,7 @@ import sys import time from datetime import datetime, timedelta, timezone from typing import Any, Dict, List +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -27,7 +28,21 @@ class MockLiteLLMTeamMembership: class MockLiteLLMVerificationToken: def __init__(self): + self.find_many_calls: List[Dict[str, Any]] = [] + self.update_calls: List[Dict[str, Any]] = [] self.update_many_calls: List[Dict[str, Any]] = [] + self._find_many_results: List[Any] = [] + + def set_find_many_results(self, results: List[Any]): + self._find_many_results = results + + async def find_many(self, where: Dict[str, Any] | None = None) -> List[Any]: + self.find_many_calls.append({"where": where}) + return self._find_many_results + + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: + self.update_calls.append({"where": where, "data": data}) + return {"count": 1} async def update_many( self, where: Dict[str, Any], data: Dict[str, Any] @@ -36,6 +51,24 @@ class MockLiteLLMVerificationToken: return {"count": 1} +class MockLiteLLMTeamTable: + def __init__(self): + self.find_many_calls: List[Dict[str, Any]] = [] + self.update_calls: List[Dict[str, Any]] = [] + self._find_many_results: List[Any] = [] + + def set_find_many_results(self, results: List[Any]): + self._find_many_results = results + + async def find_many(self, where: Dict[str, Any] | None = None) -> List[Any]: + self.find_many_calls.append({"where": where}) + return self._find_many_results + + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: + self.update_calls.append({"where": where, "data": data}) + return {"count": 1} + + class MockLiteLLMEndUserTable: def __init__(self): self.find_many_calls: List[Dict[str, Any]] = [] @@ -53,6 +86,7 @@ class MockDB: def __init__(self): self.litellm_teammembership = MockLiteLLMTeamMembership() self.litellm_verificationtoken = MockLiteLLMVerificationToken() + self.litellm_teamtable = MockLiteLLMTeamTable() self.litellm_endusertable = MockLiteLLMEndUserTable() @@ -615,6 +649,72 @@ def test_budget_table_reset_also_resets_linked_keys( assert calls[0]["data"]["spend"] == 0 +def test_reset_budget_windows_filters_null_budget_limits_in_python( + reset_budget_job, mock_prisma_client +): + """ + reset_budget_windows() should fetch keys and teams without Prisma JSON null + filters and skip empty budget_limits records in Python. + """ + key_with_windows = type( + "LiteLLM_VerificationToken", + (), + { + "token": "sk-key-with-windows", + "budget_limits": '[{"budget_duration": "24h", "reset_at": null}]', + }, + ) + key_without_windows = type( + "LiteLLM_VerificationToken", + (), + {"token": "sk-key-without-windows", "budget_limits": None}, + ) + team_with_windows = type( + "LiteLLM_TeamTable", + (), + { + "team_id": "team-with-windows", + "budget_limits": [{"budget_duration": "7d", "reset_at": None}], + }, + ) + team_without_windows = type( + "LiteLLM_TeamTable", + (), + {"team_id": "team-without-windows", "budget_limits": None}, + ) + + mock_prisma_client.db.litellm_verificationtoken.set_find_many_results( + [key_with_windows, key_without_windows] + ) + mock_prisma_client.db.litellm_teamtable.set_find_many_results( + [team_with_windows, team_without_windows] + ) + + with patch( + "litellm.proxy.proxy_server.spend_counter_cache", new=MagicMock() + ), patch.object( + ResetBudgetJob, + "_reset_expired_window", + new=AsyncMock(return_value=True), + ): + asyncio.run(reset_budget_job.reset_budget_windows()) + + assert mock_prisma_client.db.litellm_verificationtoken.find_many_calls == [ + {"where": None} + ] + assert mock_prisma_client.db.litellm_teamtable.find_many_calls == [ + {"where": None} + ] + + key_updates = mock_prisma_client.db.litellm_verificationtoken.update_calls + assert len(key_updates) == 1 + assert key_updates[0]["where"] == {"token": "sk-key-with-windows"} + + team_updates = mock_prisma_client.db.litellm_teamtable.update_calls + assert len(team_updates) == 1 + assert team_updates[0]["where"] == {"team_id": "team-with-windows"} + + def test_reset_budget_resets_endusers_with_null_budget_id( reset_budget_job, mock_prisma_client ):