fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter (#29358)

* fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter

ResetBudgetJob's batched update_data path shipped the full key/user/team
model on each reset. Prisma rejects object_permission_id and budget_limits
on the update input type, so any row carrying those fields detonated the
entire batch -- spend never reset, budget_reset_at never advanced. After
v1.84.0 started populating object_permission_id on UI-created keys, this
fires routinely.

_reset_budget_common also zeroed the cross-pod spend counter before the
DB write, so failed resets left enforcement reading 0 from the counter
while the DB still held the over-budget spend, admitting requests past
the cap until the counter naturally re-saturated from new reservations.

Switch the write to per-row narrow updates ({spend, budget_reset_at})
via db.batch_, and move the counter invalidation out of
_reset_budget_common so it only fires after the DB write commits. On
DB-write failure the counter is left untouched, enforcement continues
to block, and the next scheduler tick can retry without leaving a
bypass window.

Fixes #27730.

* fix(reset_budget): address Greptile review on #29358

- Strengthen the bypass-half regression test: replace the for-loop over
  call_args_list (vacuously true when empty) with assert_not_called(),
  so the test would actually flag a re-introduction of counter-zeroing
  via any code path.
- Add the same explanatory docstring on _write_user_reset_updates and
  _write_team_reset_updates that _write_key_reset_updates already has,
  so all three helpers point future maintainers at #27730.

* test(reset_budget): update test_proxy_budget_reset for new batch-write path

Same shape as the previous test_reset_budget_job.py update: keys/users/teams
now write through prisma.db.batch_().<table>.update, not update_data, so the
tests need a batcher mock and updated assertions. Adds:

- _wire_batcher_for_test helper that returns a list which accumulates per-row
  batch updates captured from prisma_client.db.batch_().
- _attrify helper that wraps dict fixtures so getattr(item, "token") works
  alongside the dict item-access the fake_reset_* mocks rely on. The new
  narrow-write helpers use getattr to pull out the row's id, and would
  silently skip plain dicts otherwise.
- Updates 3 partial_failure tests to assert against the batch-call list
  (rows by id, payload contains only {spend, budget_reset_at}) instead of
  update_data.assert_awaited_once + data_list inspection.
- Updates test_reset_budget_continues_other_categories_on_failure: only
  budget + enduser still flow through update_data; key/user/team go through
  the batch path now.
- Wires the batcher mock into 3 service_logger_*_success tests so commit()
  is actually awaitable and the success hook fires.

These tests were silently passing locally only because the editable install
in .venv pointed at the main repo, not the worktree — running pytest with
PYTHONPATH overridden to the worktree (matching CI) reproduces the failures.

litellm/proxy/common_utils/reset_budget_job.py

@@ -414,6 +414,72 @@ class ResetBudgetJob:
         )
         return [LiteLLM_EndUserTable(**row.dict()) for row in rows]
 
+    async def _write_key_reset_updates(
+        self, updated_keys: List[LiteLLM_VerificationToken]
+    ) -> None:
+        """
+        Write per-row {spend, budget_reset_at} updates for keys.
+
+        Avoids the batched full-model update path, which trips
+        prisma.errors.DataError on any row carrying object_permission_id or
+        budget_limits (see #27730). Both fields are rejected by Prisma's
+        update input type for LiteLLM_VerificationToken, and the failure
+        aborts the entire batch — silently leaving spend over the cap and
+        budget_reset_at unchanged forever.
+        """
+        batcher = self.prisma_client.db.batch_()
+        for k in updated_keys:
+            token = getattr(k, "token", None)
+            if token is None:
+                continue
+            batcher.litellm_verificationtoken.update(
+                where={"token": token},
+                data={"spend": 0, "budget_reset_at": k.budget_reset_at},
+            )
+        await batcher.commit()
+
+    async def _write_user_reset_updates(
+        self, updated_users: List[LiteLLM_UserTable]
+    ) -> None:
+        """
+        Write per-row {spend, budget_reset_at} updates for users.
+
+        Mirrors _write_key_reset_updates — avoids the full-model update path
+        that trips Prisma's DataError on rows carrying unrecognised fields
+        (see #27730).
+        """
+        batcher = self.prisma_client.db.batch_()
+        for u in updated_users:
+            user_id = getattr(u, "user_id", None)
+            if user_id is None:
+                continue
+            batcher.litellm_usertable.update(
+                where={"user_id": user_id},
+                data={"spend": 0, "budget_reset_at": u.budget_reset_at},
+            )
+        await batcher.commit()
+
+    async def _write_team_reset_updates(
+        self, updated_teams: List[LiteLLM_TeamTable]
+    ) -> None:
+        """
+        Write per-row {spend, budget_reset_at} updates for teams.
+
+        Mirrors _write_key_reset_updates — avoids the full-model update path
+        that trips Prisma's DataError on rows carrying unrecognised fields
+        (see #27730).
+        """
+        batcher = self.prisma_client.db.batch_()
+        for t in updated_teams:
+            team_id = getattr(t, "team_id", None)
+            if team_id is None:
+                continue
+            batcher.litellm_teamtable.update(
+                where={"team_id": team_id},
+                data={"spend": 0, "budget_reset_at": t.budget_reset_at},
+            )
+        await batcher.commit()
+
     async def reset_budget_for_litellm_keys(self):
         """
         Resets the budget for all the litellm keys
@@ -455,11 +521,7 @@ class ResetBudgetJob:
                 )
 
                 if updated_keys:
-                    await self.prisma_client.update_data(
-                        query_type="update_many",
-                        data_list=updated_keys,
-                        table_name="key",
-                    )
+                    await self._write_key_reset_updates(updated_keys=updated_keys)
                     for k in updated_keys:
                         token = getattr(k, "token", None)
                         if token:
@@ -544,11 +606,7 @@ class ResetBudgetJob:
                     "Updated users %s", json.dumps(updated_users, indent=4, default=str)
                 )
                 if updated_users:
-                    await self.prisma_client.update_data(
-                        query_type="update_many",
-                        data_list=updated_users,
-                        table_name="user",
-                    )
+                    await self._write_user_reset_updates(updated_users=updated_users)
                     for u in updated_users:
                         user_id = getattr(u, "user_id", None)
                         if user_id:
@@ -641,11 +699,7 @@ class ResetBudgetJob:
                     "Updated teams %s", json.dumps(updated_teams, indent=4, default=str)
                 )
                 if updated_teams:
-                    await self.prisma_client.update_data(
-                        query_type="update_many",
-                        data_list=updated_teams,
-                        table_name="team",
-                    )
+                    await self._write_team_reset_updates(updated_teams=updated_teams)
                     for t in updated_teams:
                         team_id = getattr(t, "team_id", None)
                         if team_id:
@@ -816,49 +870,16 @@ class ResetBudgetJob:
         """
         In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration
 
-        Common logic for resetting budget for a team, user, or key
+        Common logic for resetting budget for a team, user, or key.
+
+        Spend-counter invalidation happens in the caller, AFTER the DB write
+        commits. Zeroing the counter here would open a bypass window when the
+        DB write fails: get_current_spend reads 0 from Redis while the DB
+        still holds the pre-reset value, admitting requests past the cap.
         """
         try:
             item.spend = 0.0
-
-            # Reset the cross-pod spend counter.
-            # Reset Redis directly (not via DualCache) so a Redis failure
-            # doesn't silently leave a stale counter that get_current_spend
-            # would read as authoritative, permanently blocking the user.
-            from litellm.proxy.proxy_server import spend_counter_cache
-
-            counter_key = None
-            if item_type == "key" and hasattr(item, "token") and item.token is not None:  # type: ignore[union-attr]
-                counter_key = f"spend:key:{item.token}"  # type: ignore[union-attr]
-            elif (
-                item_type == "team"
-                and hasattr(item, "team_id")
-                and item.team_id is not None  # type: ignore[union-attr]
-            ):
-                counter_key = f"spend:team:{item.team_id}"  # type: ignore[union-attr]
-
-            if counter_key is not None:
-                # Always reset in-memory (local fallback)
-                spend_counter_cache.in_memory_cache.set_cache(
-                    key=counter_key, value=0.0
-                )
-                # Explicitly reset Redis with warning on failure
-                if spend_counter_cache.redis_cache is not None:
-                    try:
-                        await spend_counter_cache.redis_cache.async_set_cache(
-                            key=counter_key, value=0.0
-                        )
-                    except Exception as redis_err:
-                        verbose_proxy_logger.warning(
-                            "Failed to reset spend counter in Redis for %s key=%s: %s. "
-                            "Budget may be over-enforced until counter expires.",
-                            item_type,
-                            counter_key,
-                            redis_err,
-                        )
-
             if hasattr(item, "budget_duration") and item.budget_duration is not None:
-                # Get standardized reset time based on budget duration
                 from litellm.proxy.common_utils.timezone_utils import (
                     get_budget_reset_time,
                 )

tests/litellm_utils_tests/test_proxy_budget_reset.py

@@ -22,6 +22,60 @@ from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
 # In a real-world scenario, these would be instances of LiteLLM_VerificationToken, LiteLLM_UserTable, etc.
 
 
+def _attrify(d: dict):
+    """
+    Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`,
+    etc.) works alongside the existing item-access the fake_reset_* helpers
+    rely on. The reset job's narrow-write helpers use `getattr(item, "token",
+    None)` (et al), which returns None for plain dicts — that would silently
+    skip the row.
+    """
+    class _AttrDict(dict):
+        def __getattr__(self, k):
+            try:
+                return self[k]
+            except KeyError:
+                raise AttributeError(k)
+
+        def __setattr__(self, k, v):
+            self[k] = v
+
+    return _AttrDict(d)
+
+
+def _wire_batcher_for_test(prisma_client):
+    """
+    Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is
+    awaitable and whose per-table .update() calls get captured. The reset job
+    writes key/user/team resets via prisma.db.batch_().<table>.update — not via
+    prisma_client.update_data — so tests must let that batch path complete.
+
+    Returns the list that will accumulate {table, where, data} dicts from
+    each captured update call.
+    """
+    batch_calls = []
+
+    def make_batcher():
+        class _Table:
+            def __init__(self, table_name):
+                self._table_name = table_name
+
+            def update(self, where=None, data=None):
+                batch_calls.append(
+                    {"table": self._table_name, "where": where, "data": data}
+                )
+
+        batcher = MagicMock()
+        batcher.litellm_verificationtoken = _Table("key")
+        batcher.litellm_usertable = _Table("user")
+        batcher.litellm_teamtable = _Table("team")
+        batcher.commit = AsyncMock(return_value=None)
+        return batcher
+
+    prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
+    return batch_calls
+
+
 @pytest.mark.asyncio
 async def test_reset_budget_keys_partial_failure():
     """
@@ -45,6 +99,9 @@ async def test_reset_budget_keys_partial_failure():
         return_value=[key1, key2, key3, key4, key5, key6]
     )
     prisma_client.update_data = AsyncMock()
+    # Reset job writes key resets via prisma.db.batch_().<table>.update — not
+    # via update_data — so wire that path.
+    batch_calls = _wire_batcher_for_test(prisma_client)
 
     # Using a dummy logging object with async hooks mocked out.
     proxy_logging_obj = MagicMock()
@@ -56,6 +113,15 @@ async def test_reset_budget_keys_partial_failure():
 
     now = datetime.utcnow()
 
+    # token is needed because the new write path uses where={"token": ...}
+    # and _AttrDict makes getattr work alongside item access used by fake_reset_key.
+    for k in [key1, key2, key3, key4, key5, key6]:
+        k.setdefault("token", k["id"])
+    key1, key2, key3, key4, key5, key6 = (
+        _attrify(k) for k in [key1, key2, key3, key4, key5, key6]
+    )
+    prisma_client.get_data = AsyncMock(return_value=[key1, key2, key3, key4, key5, key6])
+
     async def fake_reset_key(key, current_time):
         if key["id"] == "key1":
             # Simulate a failure on key1 (for example, this might be due to an invariant check)
@@ -80,17 +146,17 @@ async def test_reset_budget_keys_partial_failure():
     # Assert that the helper was called for 6 keys
     assert mock_reset_key.call_count == 6
 
-    # Assert that update_data was called once with a list containing all 6 keys
-    prisma_client.update_data.assert_awaited_once()
-    update_call = prisma_client.update_data.call_args
-    assert update_call.kwargs.get("table_name") == "key"
-    updated_keys = update_call.kwargs.get("data_list", [])
-    assert len(updated_keys) == 5
-    assert updated_keys[0]["id"] == "key2"
-    assert updated_keys[1]["id"] == "key3"
-    assert updated_keys[2]["id"] == "key4"
-    assert updated_keys[3]["id"] == "key5"
-    assert updated_keys[4]["id"] == "key6"
+    # Assert that the new narrow write path got 5 batched updates (key1 failed).
+    # update_data must NOT have been called for keys.
+    prisma_client.update_data.assert_not_awaited()
+    key_writes = [c for c in batch_calls if c["table"] == "key"]
+    assert len(key_writes) == 5
+    written_ids = [c["where"]["token"] for c in key_writes]
+    assert written_ids == ["key2", "key3", "key4", "key5", "key6"]
+    # And every write must carry only {spend, budget_reset_at} — never the full row.
+    for c in key_writes:
+        assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+        assert c["data"]["spend"] == 0
 
     # Verify that the failure logging hook was scheduled (due to the failure for key1)
     failure_hook_calls = (
@@ -125,6 +191,7 @@ async def test_reset_budget_users_partial_failure():
         return_value=[user1, user2, user3, user4, user5, user6]
     )
     prisma_client.update_data = AsyncMock()
+    batch_calls = _wire_batcher_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -133,6 +200,15 @@ async def test_reset_budget_users_partial_failure():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
+    # user_id required for the new write path's where clause; _AttrDict so
+    # getattr(u, 'user_id') works alongside the dict access fake_reset_user uses.
+    for u in [user1, user2, user3, user4, user5, user6]:
+        u.setdefault("user_id", u["id"])
+    user1, user2, user3, user4, user5, user6 = (
+        _attrify(u) for u in [user1, user2, user3, user4, user5, user6]
+    )
+    prisma_client.get_data = AsyncMock(return_value=[user1, user2, user3, user4, user5, user6])
+
     async def fake_reset_user(user, current_time):
         if user["id"] == "user1":
             raise Exception("Simulated failure for user1")
@@ -150,16 +226,14 @@ async def test_reset_budget_users_partial_failure():
         await asyncio.sleep(0.1)
 
     assert mock_reset_user.call_count == 6
-    prisma_client.update_data.assert_awaited_once()
-    update_call = prisma_client.update_data.call_args
-    assert update_call.kwargs.get("table_name") == "user"
-    updated_users = update_call.kwargs.get("data_list", [])
-    assert len(updated_users) == 5
-    assert updated_users[0]["id"] == "user2"
-    assert updated_users[1]["id"] == "user3"
-    assert updated_users[2]["id"] == "user4"
-    assert updated_users[3]["id"] == "user5"
-    assert updated_users[4]["id"] == "user6"
+    prisma_client.update_data.assert_not_awaited()
+    user_writes = [c for c in batch_calls if c["table"] == "user"]
+    assert len(user_writes) == 5
+    written_ids = [c["where"]["user_id"] for c in user_writes]
+    assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
+    for c in user_writes:
+        assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+        assert c["data"]["spend"] == 0
 
     failure_hook_calls = (
         proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -308,6 +382,7 @@ async def test_reset_budget_teams_partial_failure():
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(return_value=[team1, team2])
     prisma_client.update_data = AsyncMock()
+    batch_calls = _wire_batcher_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -316,6 +391,12 @@ async def test_reset_budget_teams_partial_failure():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
+    # team_id required for the new write path's where clause; _AttrDict for getattr.
+    for t in [team1, team2]:
+        t.setdefault("team_id", t["id"])
+    team1, team2 = _attrify(team1), _attrify(team2)
+    prisma_client.get_data = AsyncMock(return_value=[team1, team2])
+
     async def fake_reset_team(team, current_time):
         if team["id"] == "team1":
             raise Exception("Simulated failure for team1")
@@ -333,12 +414,12 @@ async def test_reset_budget_teams_partial_failure():
         await asyncio.sleep(0.1)
 
     assert mock_reset_team.call_count == 2
-    prisma_client.update_data.assert_awaited_once()
-    update_call = prisma_client.update_data.call_args
-    assert update_call.kwargs.get("table_name") == "team"
-    updated_teams = update_call.kwargs.get("data_list", [])
-    assert len(updated_teams) == 1
-    assert updated_teams[0]["id"] == "team2"
+    prisma_client.update_data.assert_not_awaited()
+    team_writes = [c for c in batch_calls if c["table"] == "team"]
+    assert len(team_writes) == 1
+    assert team_writes[0]["where"] == {"team_id": "team2"}
+    assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+    assert team_writes[0]["data"]["spend"] == 0
 
     failure_hook_calls = (
         proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -402,6 +483,18 @@ 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()
+    batch_calls = _wire_batcher_for_test(prisma_client)
+    # ID fields required by the new write path's where clauses; _AttrDict
+    # lets getattr() see them alongside the item-access fake_reset_* helpers use.
+    for k in [key1, key2]:
+        k.setdefault("token", k["id"])
+    for u in [user1, user2]:
+        u.setdefault("user_id", u["id"])
+    for t in [team1, team2]:
+        t.setdefault("team_id", t["id"])
+    key1, key2 = _attrify(key1), _attrify(key2)
+    user1, user2 = _attrify(user1), _attrify(user2)
+    team1, team2 = _attrify(team1), _attrify(team2)
     # 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}
@@ -488,32 +581,29 @@ async def test_reset_budget_continues_other_categories_on_failure():
         "team_membership",
     }
 
-    # Verify that update_data was called three times (one per category, enduser update includes two)
-    assert prisma_client.update_data.await_count == 5
+    # After the fix, keys/users/teams write via prisma.db.batch_().<table>.update,
+    # so only budget + enduser still go through update_data.
     calls = prisma_client.update_data.await_args_list
-
-    # Check keys update: both keys succeed.
-    keys_call = calls[0]
-    assert keys_call.kwargs.get("table_name") == "key"
-    assert len(keys_call.kwargs.get("data_list", [])) == 2
-
-    # Check users update: only user2 succeeded.
-    users_call = calls[1]
-    assert users_call.kwargs.get("table_name") == "user"
-    users_updated = users_call.kwargs.get("data_list", [])
-    assert len(users_updated) == 1
-    assert users_updated[0]["id"] == "user2"
-
-    # Check teams update: both teams succeed.
-    teams_call = calls[2]
-    assert teams_call.kwargs.get("table_name") == "team"
-    assert len(teams_call.kwargs.get("data_list", [])) == 2
+    update_data_tables = [c.kwargs.get("table_name") for c in calls]
+    assert sorted(update_data_tables) == ["budget", "enduser"]
 
     # Check enduser update: enduser succeed.
-    enduser_call = calls[4]
-    assert enduser_call.kwargs.get("table_name") == "enduser"
+    enduser_call = next(c for c in calls if c.kwargs.get("table_name") == "enduser")
     assert len(enduser_call.kwargs.get("data_list", [])) == 1
 
+    # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams.
+    key_writes = [c for c in batch_calls if c["table"] == "key"]
+    user_writes = [c for c in batch_calls if c["table"] == "user"]
+    team_writes = [c for c in batch_calls if c["table"] == "team"]
+    assert len(key_writes) == 2
+    assert len(user_writes) == 1
+    assert user_writes[0]["where"] == {"user_id": "user2"}
+    assert len(team_writes) == 2
+    # Every batched write must carry only the two reset fields, never the full row.
+    for c in key_writes + user_writes + team_writes:
+        assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
+        assert c["data"]["spend"] == 0
+
 
 # ---------------------------------------------------------------------------
 # Additional tests for service logger behavior (keys, users, teams, endusers)
@@ -527,12 +617,13 @@ async def test_service_logger_keys_success():
     logger success hook is called with the correct event metadata and no exception is logged.
     """
     keys = [
-        {"id": "key1", "spend": 10.0, "budget_duration": 60},
-        {"id": "key2", "spend": 15.0, "budget_duration": 60},
+        {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"},
+        {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"},
     ]
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(return_value=keys)
     prisma_client.update_data = AsyncMock()
+    _wire_batcher_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -644,12 +735,13 @@ async def test_service_logger_users_success():
     the correct metadata and no exception is logged.
     """
     users = [
-        {"id": "user1", "spend": 20.0, "budget_duration": 120},
-        {"id": "user2", "spend": 25.0, "budget_duration": 120},
+        {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"},
+        {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"},
     ]
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(return_value=users)
     prisma_client.update_data = AsyncMock()
+    _wire_batcher_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -756,12 +848,13 @@ async def test_service_logger_teams_success():
     the proper metadata and nothing is logged as an exception.
     """
     teams = [
-        {"id": "team1", "spend": 30.0, "budget_duration": 180},
-        {"id": "team2", "spend": 35.0, "budget_duration": 180},
+        {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"},
+        {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"},
     ]
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(return_value=teams)
     prisma_client.update_data = AsyncMock()
+    _wire_batcher_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -92,6 +92,37 @@ class MockLiteLLMEndUserTable:
         return self._find_many_results
 
 
+class MockBatcher:
+    """Captures per-row update calls and exposes them after commit().
+
+    Mirrors prisma's `db.batch_()` ergonomics enough that the reset job's
+    narrow-write helpers (`_write_key_reset_updates` et al) can run against
+    the mock and the test can assert on what would have been written.
+    """
+
+    def __init__(self):
+        self.calls: List[Dict[str, Any]] = []
+        self.committed: bool = False
+
+        class _Table:
+            def __init__(_self, table_name: str, outer: "MockBatcher"):
+                _self._table_name = table_name
+                _self._outer = outer
+
+            def update(_self, where, data):
+                _self._outer.calls.append(
+                    {"table": _self._table_name, "where": where, "data": data}
+                )
+
+        self.litellm_verificationtoken = _Table("key", self)
+        self.litellm_usertable = _Table("user", self)
+        self.litellm_teamtable = _Table("team", self)
+
+    async def commit(self):
+        self.committed = True
+        return self.calls
+
+
 class MockDB:
     def __init__(self):
         self.litellm_teammembership = MockLiteLLMTeamMembership()
@@ -99,6 +130,19 @@ class MockDB:
         self.litellm_endusertable = MockLiteLLMEndUserTable()
         self.litellm_organizationtable = MockLiteLLMOrganizationTable()
         self.litellm_tagtable = MockLiteLLMTagTable()
+        self.batch_calls: List[Dict[str, Any]] = []
+
+    def batch_(self):
+        batcher = MockBatcher()
+        # Aggregate calls across all batches so tests can assert on cumulative writes.
+        original_commit = batcher.commit
+
+        async def _record_and_commit():
+            self.batch_calls.extend(batcher.calls)
+            return await original_commit()
+
+        batcher.commit = _record_and_commit  # type: ignore[assignment]
+        return batcher
 
 
 class MockPrismaClient:
@@ -205,6 +249,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
             "budget_duration": "30d",
             "budget_reset_at": now,
             "id": "test-key-1",
+            "token": "tok-key-1",
         },
     )
 
@@ -213,11 +258,16 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
     # Run the test
     asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
 
-    # Verify results
-    assert len(mock_prisma_client.updated_data["key"]) == 1
-    updated_key = mock_prisma_client.updated_data["key"][0]
-    assert updated_key.spend == 0.0
-    assert updated_key.budget_reset_at > now
+    # The reset writes only {spend, budget_reset_at} per row via batch_().
+    # Full-row writes would re-detonate the Prisma DataError on rows carrying
+    # object_permission_id / budget_limits (see #27730).
+    key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
+    assert len(key_writes) == 1
+    write = key_writes[0]
+    assert write["where"] == {"token": "tok-key-1"}
+    assert write["data"]["spend"] == 0
+    assert write["data"]["budget_reset_at"] > now
+    assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
 
 
 def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
@@ -231,6 +281,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
             "budget_duration": "7d",
             "budget_reset_at": now,
             "id": "test-user-1",
+            "user_id": "uid-1",
         },
     )
 
@@ -239,11 +290,13 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
     # Run the test
     asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
 
-    # Verify results
-    assert len(mock_prisma_client.updated_data["user"]) == 1
-    updated_user = mock_prisma_client.updated_data["user"][0]
-    assert updated_user.spend == 0.0
-    assert updated_user.budget_reset_at > now
+    user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"]
+    assert len(user_writes) == 1
+    write = user_writes[0]
+    assert write["where"] == {"user_id": "uid-1"}
+    assert write["data"]["spend"] == 0
+    assert write["data"]["budget_reset_at"] > now
+    assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
 
 
 def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
@@ -257,6 +310,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
             "budget_duration": "1mo",
             "budget_reset_at": now,
             "id": "test-team-1",
+            "team_id": "tid-1",
         },
     )
 
@@ -265,11 +319,13 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
     # Run the test
     asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
 
-    # Verify results
-    assert len(mock_prisma_client.updated_data["team"]) == 1
-    updated_team = mock_prisma_client.updated_data["team"][0]
-    assert updated_team.spend == 0.0
-    assert updated_team.budget_reset_at > now
+    team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"]
+    assert len(team_writes) == 1
+    write = team_writes[0]
+    assert write["where"] == {"team_id": "tid-1"}
+    assert write["data"]["spend"] == 0
+    assert write["data"]["budget_reset_at"] > now
+    assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
 
 
 def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
@@ -324,6 +380,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
             "budget_duration": "30d",
             "budget_reset_at": now,
             "id": "test-key-1",
+            "token": "tok-all-1",
         },
     )
 
@@ -335,6 +392,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
             "budget_duration": "7d",
             "budget_reset_at": now,
             "id": "test-user-1",
+            "user_id": "uid-all-1",
         },
     )
 
@@ -346,6 +404,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
             "budget_duration": "1mo",
             "budget_reset_at": now,
             "id": "test-team-1",
+            "team_id": "tid-all-1",
         },
     )
 
@@ -379,17 +438,22 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
     # Run the test
     asyncio.run(reset_budget_job.reset_budget())
 
-    # Verify results
-    assert len(mock_prisma_client.updated_data["key"]) == 1
-    assert len(mock_prisma_client.updated_data["user"]) == 1
-    assert len(mock_prisma_client.updated_data["team"]) == 1
+    # key/user/team rows are written via batch_().<table>.update — verify each
+    # one fired exactly once with the narrow {spend, budget_reset_at} payload.
+    for table_name, where in [
+        ("key", {"token": "tok-all-1"}),
+        ("user", {"user_id": "uid-all-1"}),
+        ("team", {"team_id": "tid-all-1"}),
+    ]:
+        writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == table_name]
+        assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
+        assert writes[0]["where"] == where
+        assert writes[0]["data"]["spend"] == 0
+        assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
+
+    # Enduser + budget rows still go through update_data (not narrowed; different path).
     assert len(mock_prisma_client.updated_data["enduser"]) == 1
     assert len(mock_prisma_client.updated_data["budget"]) == 1
-
-    # Check that all spends were reset to 0
-    assert mock_prisma_client.updated_data["key"][0].spend == 0.0
-    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
 
 
@@ -1399,6 +1463,105 @@ def test_reset_budget_for_teams_invalidates_redis_counter(
     )
 
 
+def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch):
+    """
+    Regression for #27730 (the bypass-half).
+
+    If the DB write inside the reset job raises (e.g. Prisma DataError on a
+    row carrying object_permission_id or budget_limits), the Redis spend
+    counter MUST NOT be zeroed — that would let get_current_spend admit
+    requests past the cap while the DB row still holds the over-budget
+    spend.
+
+    Pre-fix: _reset_budget_common pre-zeroed the counter before the DB
+    write attempt, opening the bypass window.
+    Post-fix: counter invalidation lives in the caller, AFTER the DB write
+    commits. If the write raises, the post-write invalidation never runs.
+    """
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    now = datetime.now(timezone.utc)
+    prisma_client = MagicMock()
+
+    matching_key = type(
+        "Key",
+        (),
+        {
+            "spend": 100.0,
+            "budget_duration": "30d",
+            "budget_reset_at": now - timedelta(seconds=1),
+            "token": "sk-failing",
+        },
+    )
+
+    # get_data returns one key needing reset; the batched DB write then explodes.
+    async def fake_get_data(table_name, query_type, **kwargs):
+        if table_name == "key":
+            return [matching_key]
+        return []
+
+    prisma_client.get_data = fake_get_data
+
+    batcher = MagicMock()
+    batcher.litellm_verificationtoken.update = MagicMock()
+
+    async def failing_commit():
+        raise RuntimeError("simulated Prisma DataError on update")
+
+    batcher.commit = failing_commit
+    prisma_client.db.batch_ = MagicMock(return_value=batcher)
+
+    job = ResetBudgetJob(
+        proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client
+    )
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    # CRITICAL: counter invalidation must NOT have been called at all —
+    # the DB write raised before the post-write invalidation loop. Using
+    # assert_not_called() instead of iterating call_args_list, because the
+    # latter is vacuously true when the list is empty (would pass even if
+    # the bypass were re-introduced via a different code path).
+    counter_cache.in_memory_cache.set_cache.assert_not_called()
+
+
+def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, mock_prisma_client):
+    """
+    Regression for #27730 (the trigger-half).
+
+    The reset job must write only {spend, budget_reset_at} per row — never
+    the full key object. Sending the full object via the old update_data
+    batcher path made Prisma reject any row carrying object_permission_id
+    or budget_limits (both became non-NULL on UI-created keys after v1.84.0).
+    """
+    now = datetime.now(timezone.utc)
+    key_with_problematic_fields = type(
+        "LiteLLM_VerificationToken",
+        (),
+        {
+            "spend": 50.0,
+            "budget_duration": "30d",
+            "budget_reset_at": now,
+            "token": "sk-problematic",
+            "object_permission_id": "perm-abc",   # would be rejected on update
+            "budget_limits": [{"max_budget": 5}],   # would be rejected on update
+            "metadata": {"some": "thing"},
+        },
+    )
+    mock_prisma_client.data["key"] = [key_with_problematic_fields]
+
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
+
+    key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
+    assert len(key_writes) == 1
+    payload_keys = set(key_writes[0]["data"].keys())
+    assert payload_keys == {"spend", "budget_reset_at"}, (
+        f"reset payload must not include any field besides spend / budget_reset_at, "
+        f"got: {payload_keys}. Any extra field (object_permission_id, budget_limits, etc.) "
+        f"trips Prisma DataError and detonates the whole batch."
+    )
+
+
 def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
     """Resetting keys via budget tier must clear each linked key's counter."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)