Fix bugs that bypasses per-team member budget limit - #26204
Conversation
|
Michael Riad Zaky seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Greptile SummaryThis PR fixes three related bugs that allowed users to exceed per-team-member budget limits: Redis counters were reseeding from a stale cache instead of the DB after TTL expiry, the auth check skipped enforcement when a member had no individual budget row (instead of falling back to the team-level default), and the backfill function wasn't healing pre-existing membership rows with a null Confidence Score: 5/5Safe to merge; all three bugs are correctly fixed with targeted changes and passing mock tests, and no new P0/P1 issues were found. The three bug fixes are logically sound and well-isolated. The remaining open item (direct DB query in No files require special attention beyond the previously flagged direct DB query in
|
| Filename | Overview |
|---|---|
| litellm/proxy/auth/auth_checks.py | Adds get_team_member_default_budget (direct DB query, flagged in prior review) and correctly rewires _check_team_member_budget to fall back to the team-level default when a member has no linked budget table. |
| litellm/proxy/management_endpoints/team_endpoints.py | Extends backfill_team_member_budget_entries with an unconditional update_many to heal pre-existing membership rows whose budget_id is NULL; new rows from create_many already carry the correct budget_id so no double-update occurs. |
| litellm/proxy/proxy_server.py | Introduces _reseed_spend_from_db to read authoritative spend from the DB on counter-cache miss (handling key/team/team_member/user/org prefixes and short-circuiting on window-suffix keys) and wires it into _init_and_increment_spend_counter; outstanding P2 concerns from prior threads acknowledged as in-scope or out-of-scope. |
| tests/test_litellm/proxy/auth/test_auth_checks.py | Adds two new async mock-only tests covering the team-default budget fallback path and the per-member override priority; no real network calls, good assertion coverage including cache-hit bypass of DB. |
| tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py | Adds update_many mocks to existing tests (necessary for new code path) and a new dedicated test asserting the null-budget_id healing behavior; does not weaken prior assertions. |
| tests/test_litellm/proxy/test_proxy_server.py | Three new mock-only tests cover DB-reseed on counter-cache miss, user/org prefix handling, and window-key short-circuit; correctly isolates module-level globals and restores them in finally blocks. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Request arrives] --> B[_check_team_member_budget]
B --> C[get_team_membership]
C --> D{membership.litellm_budget_table not None?}
D -- Yes --> E[team_member_budget = membership.budget_table.max_budget]
D -- No --> F{team.metadata has team_member_budget_id?}
F -- No --> G[No budget limit enforced]
F -- Yes --> H[get_team_member_default_budget cache-first DB lookup]
H --> I[team_member_budget = default_budget.max_budget]
E --> J[get_current_spend Redis counter first]
I --> J
J --> K{spend >= budget?}
K -- Yes --> L[raise BudgetExceededError]
K -- No --> M[Request proceeds]
Reviews (5): Last reviewed commit: "fix linting" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async def _reseed_spend_from_db(counter_key: str) -> float: | ||
| """ | ||
| Read the authoritative spend for a missing counter from the DB. The | ||
| counter_key prefix encodes the table to query: | ||
|
|
||
| spend:key:{token} -> LiteLLM_VerificationToken.spend | ||
| spend:team:{team_id} -> LiteLLM_TeamTable.spend | ||
| spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend | ||
|
|
||
| Returns 0.0 if prisma is unavailable, the row is missing, or the | ||
| key format is unrecognized. On failure, logs and returns 0.0 rather | ||
| than raising so the caller can still record the current increment. | ||
| """ | ||
| if prisma_client is None: | ||
| return 0.0 | ||
| try: | ||
| if counter_key.startswith("spend:key:"): | ||
| token = counter_key[len("spend:key:") :] | ||
| row = await prisma_client.db.litellm_verificationtoken.find_unique( | ||
| where={"token": token} | ||
| ) | ||
| elif counter_key.startswith("spend:team_member:"): | ||
| suffix = counter_key[len("spend:team_member:") :] | ||
| if ":" not in suffix: | ||
| return 0.0 | ||
| user_id, team_id = suffix.rsplit(":", 1) | ||
| row = await prisma_client.db.litellm_teammembership.find_unique( | ||
| where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} | ||
| ) | ||
| elif counter_key.startswith("spend:team:"): | ||
| team_id = counter_key[len("spend:team:") :] | ||
| row = await prisma_client.db.litellm_teamtable.find_unique( | ||
| where={"team_id": team_id} | ||
| ) | ||
| else: | ||
| return 0.0 | ||
| except Exception: | ||
| verbose_proxy_logger.exception( | ||
| "Failed to reseed spend counter %s from DB", counter_key | ||
| ) | ||
| return 0.0 | ||
| if row is None: | ||
| return 0.0 | ||
| return float(getattr(row, "spend", 0.0) or 0.0) |
There was a problem hiding this comment.
_reseed_spend_from_db handles the three base spend keys but the per-window counter variants (incremented directly in increment_spend_counters for budget_limits windows) have no DB reseed path. After the 5-minute TTL expires those counters reset to 0. A user inactive for more than 5 minutes during a budget window will have their window counter silently wiped, letting them spend up to their window cap again within the same period.
A follow-up could either extend _reseed_spend_from_db to handle window-suffix keys, or set each window counter's TTL to match its actual window duration rather than the shared SPEND_COUNTER_REDIS_TTL_SECONDS.
There was a problem hiding this comment.
Yes the window budget can get exceeded once on each expiry (if the primary budget is not exceeded), but this was pre-existing and is currently out of scope / will go in another PR.
cb90b8b to
ed6eb85
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
ed6eb85 to
2e4304b
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
2e4304b to
0bd49ec
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
Low: Budget enforcement fix with no new security concernsThis PR fixes a bug where per-team member budget limits were silently skipped when the membership row lacked a direct budget link. It adds a fallback path to read the team-level default budget from All new DB queries use Prisma ORM with parameterized access. Counter keys are constructed from server-side auth variables, not user input. The backfill Status: 0 open Posted by Veria AI · 2026-04-23T19:10:55.100Z |
|
|
||
| cache_key = f"team_member_default_budget:{budget_id}" | ||
|
|
||
| cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) |
There was a problem hiding this comment.
Can this be stale? I see you are caching this, but I do not see any cache invalidation on update of the budget
There was a problem hiding this comment.
Yeah, but it's at most stale for a minute and there are six other stale caches in this file that are similar
Fix bugs that bypasses per-team member budget limit
3 participants
Relevant issues
Should address LIT-2466 and LIT-2467, both having to do with users exceeding budget limits. Three bugs in particular:
This is increased to five minutes(now stays the same), and reseeding the value after expiry now reads straight from the DB (it was previously getting it from a stale cache that didn't get updated on increment).Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewScreenshots / Proof of Fix
The tests added (enforcing budget limits) fail before the fix and succeed after the fix. Also, here are some screenshots before and after:
Type
🐛 Bug Fix
✅ Test
Changes
Fixes to proxy auth_checks and proxy_server, with respective tests added.