Skip to content

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

Merged
mubashir1osmani merged 2 commits into
litellm_internal_stagingfrom
litellm_e2e_user_team_budget_isolation
Jul 17, 2026
Merged

test(e2e): user budget across keys and team member budget isolation#33745
mubashir1osmani merged 2 commits into
litellm_internal_stagingfrom
litellm_e2e_user_team_budget_isolation

Conversation

@yassin-berriai

@yassin-berriai yassin-berriai commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

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

Both new tests pass on a live docker-compose proxy stack against the real OpenAI API with real spend, captured at commit c27515c5b7

$ docker compose -p lit4548 ps
lit4548-db-1        Up (healthy)
lit4548-jaeger-1    Up (healthy)
lit4548-litellm-1   Up
lit4548-redis-1     Up (healthy)

$ curl -s http://localhost:4548/health/readiness
{"status":"healthy","db":"connected"}

$ LITELLM_PROXY_URL=http://localhost:4548 uv run pytest \
    quota_management/budgets/test_user_budget_across_keys_e2e.py \
    quota_management/budgets/test_team_member_budget_isolation_e2e.py -v -s

quota_management/budgets/test_user_budget_across_keys_e2e.py::TestUserBudgetAcrossKeys::test_user_budget_blocks_a_second_key
  E2E_RESULT outcome=passed duration_ms=9297 covers=quota_management.budget.internal_user.enforced_across_keys
PASSED
quota_management/budgets/test_team_member_budget_isolation_e2e.py::TestTeamMemberBudgetIsolation::test_blocked_member_does_not_block_peer
  E2E_RESULT outcome=passed duration_ms=13613 covers=quota_management.budget.team_member.isolates_per_member
PASSED

============================== 2 passed in 20.41s ==============================

Gate checks, same commit

$ make lint-e2e-basedpyright
0 errors, 0 warnings, 0 notes

$ python -m coverage_registry.collector --strict    # from tests/e2e
STRICT EXIT: 0    # both @covers ids match their registry rows

Independent e2e verification (Devin)

Devin independently checked out c27515c5b7, stood up the tests/e2e docker-compose stack, sanity-checked the OpenAI gpt-5.5 key (HTTP 200), and ran both new tests to 2 passed in 20.73s against real spend. Terminal recording of the run:

e2e budget isolation tests passing on a live compose stack

Full-resolution recording: https://raw.githubusercontent.com/yassin-berriai/litellm-pr-media/main/lit4548/e2e-run.mp4

Type

✅ Test

Changes

Gap-fill for the budgets e2e suite (tests/e2e/quota_management/budgets/) against two of the Quota Management behavior contracts whose prior coverage recorded the right state but never proved the enforced behavior the contract promises

A per-user (internal user) max_budget is meant to be enforced across all of that user's keys, not only the key that happened to spend it down. The existing internal-user test drove a single key, so a key-scoped-only budget would have passed it. test_user_budget_across_keys_e2e.py gives one internal user a tiny max_budget and two keys, neither carrying a budget of its own, drives the first key to a budget_exceeded block, then asserts the untouched second key is rejected the same way (nothing but the shared user budget can block it) and that the user's recorded spend has crossed the cap

Individual team-member budgets are meant to be enforced independently; one member exhausting their per-team budget must not block a different member on the same team. The existing team-member test used one member, so it never proved isolation. test_team_member_budget_isolation_e2e.py puts two members on one team with a large team budget, caps the tight member at a tiny per-team budget and drives it past a budget_exceeded block, then asserts the roomy member still serves on the same team with its calls attributed to its own user id and the team id in the spend logs, while the tight member stays blocked after the peer spends

Supporting changes: a /user/info read-back helper on the budgets client so the R4 test can assert recorded user spend; two new rows in tests/e2e/coverage_registry/quota_management.yaml, which is what the collector --strict matches the @covers markers against; and the two new assertion tokens (isolates_per_member, enforced_across_keys) added to the naming grammar in tests/e2e/CLAUDE.md to keep the documented vocabulary in sync (the grammar itself is doc hygiene, not a --strict gate)

The R4 test uses a tiny max_budget so a single billed call exhausts it; the load-bearing assertion is the second untouched key being blocked (a key with no budget of its own can only be blocked by the shared user budget), driven to block promptly rather than by accumulating its own spend, with the recorded user spend crossing the cap as the corroborating recorded-state check

Each test asserts both sides of the contract per tests/e2e/CONTRIBUTING.md: the recorded state (user spend crossed the cap; roomy rows attributed to the right user and team) and the enforced gateway behavior (budget_exceeded on the blocked principal, a served 200 on the isolated one)

QA runbook

Bring up the stack from tests/e2e/ with a .env holding a working OPENAI_API_KEY: docker compose up -d then curl -fs http://localhost:4000/health/liveliness. Both tests drive gpt-5.5 (real OpenAI, real spend)

  • tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py::TestUserBudgetAcrossKeys::test_user_budget_blocks_a_second_key - one internal user's max_budget blocks a second key it owns that carries no budget of its own
    • Create an internal user with a tiny cap: curl -X POST http://localhost:4000/user/new -H "Authorization: Bearer sk-1234" -d '{"max_budget": 0.000003}'
    • Mint two keys for that user id with no per-key budget: curl -X POST http://localhost:4000/key/generate -H "Authorization: Bearer sk-1234" -d '{"user_id": "<uid>"}' (twice)
    • Send /v1/chat/completions calls on the first key until one returns a budget_exceeded block
    • Send one /v1/chat/completions call on the untouched second key and expect it rejected with budget_exceeded too
    • Read GET /user/info?user_id=<uid> and expect user_info.spend >= 0.000003
    • Sanity check: this test makes sense to add and is not hand-wavey (the second key has 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 its 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 '{"max_budget": 100}'
    • Create two internal users, add both to the team with /team/member_add, capping the tight member at max_budget_in_team: 0.000003 and the roomy member at 100
    • Mint a key per member (team_id + user_id), drive the tight member's key until it returns budget_exceeded
    • Send a few /v1/chat/completions calls on the roomy member's key and expect 200s
    • Re-send on the tight member's key and expect it still budget_exceeded; poll /spend/logs and expect the roomy calls attributed to the roomy user id and the team id
    • Sanity check: this test makes sense to add and is not hand-wavey (asserts the roomy member both serves and is attributed correctly, and the tight member stays blocked after the peer spends) 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

…solation

Gap-fill for the budgets e2e suite against two Quota Management behavior
contracts that had incomplete coverage.

A per-user max_budget must be enforced across all of that user's keys, not
only the key that spent it down. The prior internal-user test drove a single
key, so it never proved the budget is shared. The new test gives one internal
user a tiny max_budget and two keys carrying no budget of their own, drives the
first key to a budget_exceeded block, then shows the untouched second key is
rejected the same way and the user's recorded spend has crossed the cap.

Per-team-member budgets must be enforced independently between members. The
prior team-member test used one member, so it never proved isolation. The new
test puts two members on one team with a large team budget, caps the tight
member at a tiny per-team budget and drives it past the block, then shows the
roomy member still serves on the same team with its calls attributed to its own
user id in the spend logs while the tight member stays blocked.

Adds a /user/info read-back helper, two coverage-registry rows, and the two new
assertion tokens to the naming grammar so the collector --strict passes.

Resolves LIT-4548
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
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.

@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@yassin-berriai

Copy link
Copy Markdown
Contributor Author

osv-scan is red here but it is base drift, not from this change: it fails identically on every recent open PR (a CVE published after the base branch's last lockfile bump). This PR touches only tests/e2e/** and does not modify uv.lock, so there is nothing to fix here; the dependency bump belongs in its own PR on litellm_internal_staging

@codecov

codecov Bot commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds two new live e2e tests that close gaps in the quota management budget suite: one proves that a per-user max_budget is enforced across all of a user's keys (not just the key that spent it down), and the other proves that one team member exhausting their per-team budget does not block a peer on the same team.

  • test_user_budget_across_keys_e2e.py: Drives first_key to a budget_exceeded block, then uses the new _expect_prompt_block helper (bounded to 6 attempts) to assert the untouched second_key — which carries no per-key budget — is also rejected. Because second_key has no per-key budget, a block within 6 calls can only originate from the shared user budget, cleanly isolating what the test proves. A 90-second spend-propagation poll then confirms recorded user spend crossed the cap.
  • test_team_member_budget_isolation_e2e.py: Creates two team members (tight cap at 3e-6, roomy cap at 100), drives the tight member to block, then asserts the roomy member still serves and that its calls are attributed to the correct user and team in the spend logs.
  • budget_client.py: Adds user_info helper and supporting Pydantic models (UserInfoRow, UserInfoResponse) for the spend read-back assertion in the cross-keys test.

Confidence Score: 5/5

Test-only PR; no production code is touched. Both tests are well-scoped and clean up after themselves.

All five changed files are under tests/e2e/. The previous concern about the second key looping up to 40 times is directly addressed by _expect_prompt_block, which bounds the retry to 6 calls — small enough that the second key could never block itself through key-scoped spend accumulation since it carries no per-key budget. Both new tests include the required dual-sided assertions. No logic gaps or structural issues found.

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 for shared user budget enforcement across keys. _expect_prompt_block (6 bounded attempts) replaces the previous 40-call loop on the second key, cleanly separating shared-budget signals from key-scoped accumulation.
tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py New e2e test proving per-team-member budget isolation. Two-member fixture with fixture-scoped teardown; tight member blocked first, then roomy member served 3 calls, spend attribution asserted via poll_logs_for_key. Logic is sound.
tests/e2e/quota_management/budgets/budget_client.py Adds UserInfoParams, UserInfoRow, UserInfoResponse models and a user_info helper that reads back recorded spend. Consistent with existing GET-param helpers; silently returns None on transport failure.
tests/e2e/coverage_registry/quota_management.yaml Two new registry rows added for enforced_across_keys and isolates_per_member; both match the @covers markers in the new test files.
tests/e2e/CLAUDE.md Doc-only: two new assertion tokens (isolates_per_member, enforced_across_keys) appended to the naming grammar to keep the vocabulary list in sync with the new tests.

Reviews (2): Last reviewed commit: "test(e2e): block the second key promptly..." | Re-trigger Greptile

…udget

Drive only the first key to exhaust the shared user budget, then require the
second key (which has no budget of its own) to be rejected within a small
bounded number of calls rather than driving it to accumulate its own spend. The
block can then only come from the shared user budget, not the second key's own
usage.
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai please review the current head c27515c. It addresses the 4/5 concern on the second-key blocking strategy: the second key is no longer driven through a 40-call loop. Since the shared user budget is already exhausted when the first key blocks, the untouched second key (which carries no budget of its own) is now required to block within a small bounded number of calls via _expect_prompt_block, so a block there can only come from the shared user budget, not the key's own accumulated spend

@codspeed-hq

codspeed-hq Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_e2e_user_team_budget_isolation (c27515c) with litellm_internal_staging (a7d01cb)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (e5a9f3f) during the generation of this report, so a7d01cb was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@mubashir1osmani
mubashir1osmani merged commit 62207ac into litellm_internal_staging Jul 17, 2026
79 of 80 checks passed
@mubashir1osmani
mubashir1osmani deleted the litellm_e2e_user_team_budget_isolation branch July 17, 2026 21:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants