Commit graph

8 commits

Author SHA1 Message Date
yuneng-jiang
7d1bd9d9f4
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.
2026-05-30 17:48:16 -07:00
Michael-RZ-Berri
02edaef50c
fix: reset org and tag budgets (#27326)
* reset org budgets

* reset tag budgets

---------

Co-authored-by: Michael Riad Zaky <michaelr@Mac.localdomain>
2026-05-09 19:15:32 -04:00
Ishaan Jaffer
e8461b5b97
style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
shivam
9f7f190670 resolved greptile comment 2026-02-16 10:24:06 -08:00
shivam
0117b35a6b added more tests, fixed tests 2026-02-16 09:59:19 -08:00
Krish Dholakia
d202ce229b
Prevent writing default user setting updates to yaml (error in non-root env) + Use central team member budget when max_budget_in_team set on UI (#12533)
* fix(proxy_setting_endpoints.py): require store model in db is enabled for setting user default settings

* test(test_proxy_server.py): update test

* fix(reset_budget_job.py): initial commit adding reset budget logic for team members

* test: update unit testing

* test(test_proxy_budget_reset.py): validate team member budget was reset

* test(test_reset_budget_job.py): update unit tests

* test: update tests
2025-07-12 10:13:07 -07:00
Laurien
0c50f8bcc9
Update enduser spend and budget reset date based on budget duration (#8460) 2025-06-08 08:39:14 -07:00
Ishaan Jaff
2753de1458
(Bug Fix + Better Observability) - BudgetResetJob: (#8562)
* use class ResetBudgetJob

* refactor reset budget job

* update reset_budget job

* refactor reset budget job

* fix LiteLLM_UserTable

* refactor reset budget job

* add telemetry for reset budget job

* dd - log service success/failure on DD

* add detailed reset budget reset info on DD

* initialize_scheduled_background_jobs

* refactor reset budget job

* trigger service failure hook when fails to reset a budget for team, key, user

* fix resetBudgetJob

* unit testing for ResetBudgetJob

* test_duration_in_seconds_basic

* testing for triggering service logging

* fix logs on test teams fail

* remove unused imports

* fix import duration in s

* duration_in_seconds
2025-02-15 16:13:08 -08:00