test(e2e): user budget across keys and team member budget isolation

[yassin-berriai]

Relevant issues

Linear ticket

Resolves LIT-4548

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

All proof captured at commit ab4392cd36 against a live docker-compose proxy stack (ghcr.io/berriai/litellm:main-latest, real OpenAI calls, real spend). The chat model is gpt-5.5 because that is the model with a working provider key on this stack; the existing budget suite uses the same literal per its coverage matrix note

Both new tests passing:

$ LITELLM_PROXY_URL=http://localhost:4620 python -m pytest \
    tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py \
    tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py -v

tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py::TestUserBudgetAcrossKeys::test_user_budget_blocks_a_second_key PASSED [ 50%]
tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py::TestTeamMemberBudgetIsolation::test_blocked_member_does_not_block_peer PASSED [100%]

============================== 2 passed in 19.37s ==============================

R4 money assertion, raw curl (a per-user budget blocks a second, untouched key of the same user):

### R4: per-user budget enforced across ALL of a user's keys ###
user_id=b99e3b68-c67a-4460-be20-f28406f320e8 (max_budget=3e-6)
key1=sk-kkex8lEeNK2...  key2=sk-l4s3onxCbS_...  (neither key carries its own budget)
--- drive key1 until the shared user budget blocks it ---
key1 call #1 -> HTTP 200 (ok, accruing spend)
key1 call #2 -> HTTP 429 budget_exceeded
--- now the UNTOUCHED key2 of the same user ---
key2 first call -> HTTP 429
  body: ExceededBudget: User=b99e3b68-c67a-4460-be20-f28406f320e8 over budget. Spend=0.000585, Budget=3e-06
--- recorded user spend vs cap (poll for the batch-write flush) ---
  user spend=0.000585  max_budget=3e-06  crossed=True
(cleaned up)

R2 money assertion, raw curl (a blocked member does not block a peer on the same team):

### R2: team-member budgets enforced independently between members ###
team=b3e00345-... (max_budget=100)  memberA=a820bee5-... (in-team cap 3e-6)  memberB=b7e621ff-... (in-team cap 100)
--- drive member A over its tiny per-team cap ---
memberA call #2 -> HTTP 429 budget_exceeded
  A block body: Budget has been exceeded! Current cost: 0.000585, Max budget: 3e-06
--- member B, on the SAME team, must still serve after A is blocked ---
memberB call #1 -> HTTP 200 OK
memberB call #2 -> HTTP 200 OK
memberB call #3 -> HTTP 200 OK
--- member A is still blocked (its cap not lifted by B's traffic) ---
memberA recheck -> HTTP 429 budget_exceeded
(cleaned up)

Stack brought up namespaced so it does not collide with concurrent suites:

docker compose -p e2ebud -f docker-compose.yml -f override.yml up -d   # override remaps litellm 4620:4000, jaeger 16620:16686
curl -fs http://localhost:4620/health/liveliness   # -> "I'm alive!"

Registry collector accepts the two new cells:

$ cd tests/e2e && PYTHONPATH=. python -m coverage_registry.collector --strict   # exit 0

Type

✅ Test

Changes

The budgets e2e suite had two contracts asserted only halfway. The internal-user enforcement test (test_budget_enforcement_e2e.py) drove a single key, so it proved a user budget blocks that key but never that the budget is shared across all of the user's keys. The team-member enforcement test (test_team_member_budget_e2e.py) used one member, so it proved a member cap blocks that member but never that members are isolated from each other

test_user_budget_across_keys_e2e.py closes the first gap. One internal user gets a tiny max_budget and owns two keys, neither carrying a budget of its own. The test spends the first key to a budget_exceeded block, then the untouched second key is rejected the same way; since the second key has no budget of its own, only the shared user budget can block it. It then polls /user/info and asserts the recorded user spend has crossed the cap, so both the recorded state and the enforced gateway behavior are checked

test_team_member_budget_isolation_e2e.py closes the second. Two members share one team with a large team budget; the tight member is capped at a tiny per-team budget and driven past it. Once the tight member is blocked, the roomy member still serves on the same team, its calls land in the spend logs attributed to its own user id, and the tight member stays blocked, proving one member's exhausted budget does not leak onto another's

Supporting changes: budget_client.py gains a typed /user/info spend read; coverage_registry/quota_management.yaml gains the quota_management.budget.internal_user.enforced_across_keys and quota_management.budget.team_member.isolates_per_member cells; and the Quota Management naming grammar in tests/e2e/CLAUDE.md is widened with the enforced_across_keys and isolates_per_member assertions so the collector grammar stays exhaustive

The other three requirements were already covered and are left alone: virtual-key over-budget rejection by test_budget_enforcement_e2e.py::test_budget_enforcement[KeyBudgetCase] and test_budget_reset_advances_e2e.py::test_key_spend_blocks_at_cap; team-member reset by test_team_member_budget_reset_e2e.py and test_budget_reset_advances_e2e.py::test_team_member_budget_reset_at_advances; and budget_duration reset by test_budget_reset_e2e.py::test_key_budget_resets_after_duration and test_budget_reset_advances_e2e.py::test_key_budget_reset_at_advances_after_window

QA runbook

• tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py::TestUserBudgetAcrossKeys::test_user_budget_blocks_a_second_key - a per-user budget, once exhausted through one key, also rejects a second untouched key of the same user, and the recorded user spend has crossed the cap
  • Create a user with a tiny budget: curl -X POST http://localhost:4000/user/new -H "Authorization: Bearer sk-1234" -d '{"max_budget": 0.000003}'
  • Generate two keys under that user id (no per-key budget): curl -X POST http://localhost:4000/key/generate -H "Authorization: Bearer sk-1234" -d '{"user_id": "<user_id>"}' twice
  • Send /chat/completions on the first key with a unique prompt each call until it returns 429 budget_exceeded (response caching means identical prompts do not accrue spend)
  • Send one /chat/completions on the second key and expect 429 naming the user over budget
  • GET /user/info?user_id=<user_id> and expect spend >= the max_budget
  • Sanity check: this test makes sense to add and is not hand-wavey (the second key carries no budget of its own, so only the shared user budget can block it) or potentially flaky
• tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py::TestTeamMemberBudgetIsolation::test_blocked_member_does_not_block_peer - one team member exhausting their per-team budget does not block a different member on the same team
  • Create a team with a large budget: curl -X POST http://localhost:4000/team/new -H "Authorization: Bearer sk-1234" -d '{"team_alias": "iso", "max_budget": 100.0}'
  • Add two members: /team/member_add the first with max_budget_in_team 0.000003 and the second with max_budget_in_team 100.0
  • Generate one key per member (team_id + user_id), then send /chat/completions on the tight member's key with unique prompts until it returns 429 budget_exceeded (member spend enforces off the batch write, so poll)
  • Send /chat/completions on the roomy member's key and expect 200; recheck the tight member's key and expect it still 429
  • Sanity check: this test makes sense to add and is not hand-wavey (the roomy member still serving after the tight member is blocked is the isolation proof) or potentially flaky

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[yassin-berriai]

@greptileai

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

This PR closes two gaps in the budget e2e suite by adding tests for user budget sharing across keys and per-team-member budget isolation, along with the supporting registry and grammar updates.

• test_user_budget_across_keys_e2e.py creates one user with a tiny max_budget and two budget-free keys, drives the first key to budget_exceeded, then asserts the untouched second key is also rejected — proving the shared user budget (not per-key tracking) is the enforcement layer. It also polls /user/info to confirm the recorded spend crossed the cap.
• test_team_member_budget_isolation_e2e.py places two members on a team: one with a tiny per-team cap (driven to exhaustion) and one with ample headroom. After the tight member is blocked it verifies the roomy member still serves and that spend rows are attributed to the roomy member's user/team IDs.
• budget_client.py gains a typed user_info() helper; quota_management.yaml and CLAUDE.md register the two new assertion tokens to keep the collector grammar exhaustive.

Confidence Score: 5/5

Safe to merge — purely additive test and registry changes with no production code touched.

All five changed files are under tests/e2e/ and the coverage registry. No production logic is modified. The new test files follow established harness conventions (typed Pydantic bodies, Result-tagged-union handling, ResourceManager cleanup, unique_marker isolation), the spend-poll deadline is generous, and the registry cells match the naming grammar. The PR description includes live proof-of-pass against a real proxy stack with raw curl output confirming the gateway behavior being tested.

No files require special attention.

Important Files Changed

Filename	Overview
tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py	New e2e test proving per-user max_budget is enforced across all keys; test logic, fixture setup, and spend-poll deadline are sound
tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py	New e2e test for per-team-member budget isolation; fixture, log-polling assertion, and tight/roomy member split are correctly structured
tests/e2e/quota_management/budgets/budget_client.py	Adds UserInfoParams/UserInfoRow/UserInfoResponse models and user_info() method; follows existing patterns for typed GET helpers
tests/e2e/coverage_registry/quota_management.yaml	Two new registry cells added (enforced_across_keys, isolates_per_member); IDs, tiers, and rationale are consistent with the existing entries
tests/e2e/CLAUDE.md	Naming grammar extended with isolates_per_member and enforced_across_keys assertion tokens; consistent with how existing tokens are listed

_{Reviews (5): Last reviewed commit: "test(e2e): annotate _call return type in..." | Re-trigger Greptile}

[greptile-apps]

PR description does not match the diff

The PR description (and attached screenshots/runbook) describes a new tests/e2e/mcp/ suite testing MCP authorization via test_mcp_key_access_e2e.py, with Docker compose bring-up steps, MCP server registration, and references to KeyGenerateBody.object_permission. None of those files or changes are present in this diff. The actual changes are entirely budget-related (test_user_budget_across_keys_e2e.py and test_team_member_budget_isolation_e2e.py). The mismatch makes it difficult to verify the pre-submission checklist and understand the intent of the PR during review.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR adds two new live e2e tests under tests/e2e/quota_management/budgets/ and the supporting plumbing to exercise two budget isolation properties: (1) a user's max_budget is enforced across every key it owns, not just the key that spent the budget, and (2) one team member's exhausted per-team budget does not block a different member of the same team.

• test_user_budget_across_keys_e2e.py drives a first key to a budget_exceeded block and then confirms a fresh second key belonging to the same user is also blocked, polling the recorded user spend to verify the cap was genuinely crossed.
• test_team_member_budget_isolation_e2e.py creates two team members with different per-team caps, exhausts the tight member's cap, and then proves the roomy member's calls still succeed and land in spend logs under the correct user and team IDs.
• budget_client.py gains a user_info helper that calls GET /user/info and returns the spend/budget row, used by the cross-key test to confirm the server-side recorded spend.

Confidence Score: 4/5

Safe to merge — the change is entirely additive new e2e tests with no production code touched.

The test logic for both isolation scenarios is sound: fixture teardown is guarded with try/finally, lambda captures are not in loops so closure semantics are correct, and the spend-log assertions use the right predicate. The only nit is a missing return type annotation on the private _call helper in the cross-keys test, which is inconsistent with the strict basedpyright config and the explicit annotations on every other helper in the suite.

tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py — the _call helper needs an explicit return type to stay consistent with the zero-error basedpyright gate.

Important Files Changed

Filename	Overview
tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py	New e2e test proving user max_budget is enforced across all keys; logic is correct but the private helper _call is missing an explicit return type annotation, inconsistent with strict basedpyright mode and all other helpers in the suite.
tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py	New e2e test for per-team-member budget isolation; fixture teardown, lambda captures, and spend-log assertions are all correct.
tests/e2e/quota_management/budgets/budget_client.py	Adds UserInfoParams, UserInfoRow, UserInfoResponse models and the user_info helper method; correctly uses the transport layer and follows existing patterns.
tests/e2e/coverage_registry/quota_management.yaml	Adds two new coverage registry entries for the new tests; IDs match the @pytest.mark.covers markers in the test files.
tests/e2e/CLAUDE.md	Expands the assertion vocabulary in the quota_management naming grammar with isolates_per_member and enforced_across_keys; change is accurate and self-consistent.

_{Reviews (2): Last reviewed commit: "test(e2e): user budget across keys and t..." | Re-trigger Greptile}

[yassin-berriai]

@greptileai addressed both findings: corrected the PR description (it now describes this budgets diff, the earlier body was mistakenly the wrong content) and added the explicit return type annotation on _call. Please re-review.

[yassin-berriai]

@greptileai the return-type fix is on HEAD 9c878f2 now (the prior review raced the push). Please re-review the latest commit.

[yassin-berriai]

@greptileai please review the current head 9c878f2

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_e2e_budget_gap_fill (9c878f2) with litellm_internal_staging (4d33964)}

[yassin-berriai]

Closing in favor of #33745, which supersedes this with the same two tests plus a tightened second-key assertion (block the untouched key promptly rather than driving it to accumulate its own spend) and corrected PR framing. Continuing the work there