From ab4392cd367e23d7bca29828c9b9ca0e29f6235c Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 17 Jul 2026 09:08:00 -0700 Subject: [PATCH] test(e2e): user budget across keys and team member budget isolation Fill the two live-coverage gaps in the budgets e2e suite. The internal-user enforcement test drove a single key, so it never proved a user budget is shared across that user's keys; the team-member enforcement test used one member, so it never proved isolation between members. Adds test_user_budget_across_keys_e2e.py: one user with a tiny max_budget owns two keys, neither with a budget of its own; spending the first to a budget_exceeded block makes the untouched second key reject the same way, and the recorded user spend crosses the cap. Adds test_team_member_budget_isolation_e2e.py: two members on one team, the tight member driven past its per-team cap while the roomy member keeps serving and the tight member stays blocked. Extends budget_client with a /user/info spend read, adds the two registry cells, and widens the quota grammar with the isolates_per_member and enforced_across_keys assertions. Resolves LIT-4548 --- tests/e2e/CLAUDE.md | 3 +- .../coverage_registry/quota_management.yaml | 2 + .../quota_management/budgets/budget_client.py | 26 ++++ .../test_team_member_budget_isolation_e2e.py | 119 ++++++++++++++++++ .../test_user_budget_across_keys_e2e.py | 61 +++++++++ 5 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py create mode 100644 tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e1eafb5196..3635e471faa 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -139,7 +139,8 @@ quota_management... | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking - | isolates_per_model | routes_to_fallback | reseed_matches_db | logs_cost | zero_cost + | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback + | reseed_matches_db | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index fac266149f4..1e91a2c1c47 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -8,9 +8,11 @@ - {id: quota_management.ratelimit.priority_strict.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_strict, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:53-71", rationale: "Strict mode (>=80% sat) enforces priority fairness"} - {id: quota_management.budget.key.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: key, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A key's max_budget blocks further paid calls once spend crosses it"} - {id: quota_management.budget.internal_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An internal user's max_budget governs personal keys"} +- {id: quota_management.budget.internal_user.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An internal user's max_budget governs every personal key it owns; a second untouched key is blocked once the shared user budget is exhausted"} - {id: quota_management.budget.end_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A customer (end-user) max_budget blocks calls attributed via user="} - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} +- {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} - {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"} - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} - {id: quota_management.budget.soft.alerts_without_blocking, module: quota_management, tier: P1, behavior: budget, variant: soft, assertions: [alerts_without_blocking], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "soft_budget alerts but never blocks traffic"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 7b37c3af98e..01e4d63c1c3 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -39,6 +39,19 @@ class UserNewResponse(BaseModel): user_id: str +class UserInfoParams(BaseModel): + user_id: str + + +class UserInfoRow(BaseModel): + spend: float | None = None + max_budget: float | None = None + + +class UserInfoResponse(BaseModel): + user_info: UserInfoRow | None = None + + class UserDeleteBody(BaseModel): user_ids: list[str] @@ -262,6 +275,19 @@ class BudgetClient: response_type=NoBody, ) + def user_info(self, user_id: str) -> UserInfoRow | None: + result = self.gateway.transport.get( + "/user/info", + headers=self.gateway.transport.master, + params=UserInfoParams(user_id=user_id), + response_type=UserInfoResponse, + ) + match result: + case Success(data=data): + return data.user_info + case _: + return None + # ---- customer / end-user ------------------------------------------- def create_customer(self, customer_id: str, *, max_budget: float) -> str: diff --git a/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py new file mode 100644 index 00000000000..87855a9a1c1 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py @@ -0,0 +1,119 @@ +"""Live e2e: per-team-member budgets are enforced independently between members. + +Two members share one team that has a large team budget. The tight member is capped +at a tiny per-team budget and spends past it; the roomy member has plenty of room. +Once the tight member is blocked with budget_exceeded, the roomy member still serves +on the same team, its calls land in the spend logs under its own user id, and the +tight member stays blocked. A shared or leaky member counter would either block the +roomy member too or let the tight member back through once its peer spent. +""" + +import time +from collections.abc import Iterator +from dataclasses import dataclass + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import Success, require_successful_call +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage + +pytestmark = pytest.mark.e2e + +MODEL = "gpt-5.5" +TEAM_BUDGET = 100.0 +TIGHT_MEMBER_BUDGET = 3e-6 +ROOMY_MEMBER_BUDGET = 100.0 +ROOMY_BURST = 3 + + +@dataclass(frozen=True, slots=True) +class _Pair: + team_id: str + tight_user_id: str + roomy_user_id: str + tight_key: str + roomy_key: str + + +@pytest.fixture(scope="class") +def pair(client: BudgetClient) -> Iterator[_Pair]: + """One team with a large budget and two members on it: a tight member capped at + a tiny per-team budget and a roomy member with headroom, each with their own key. + Shared across the class and torn down LIFO best-effort when it finishes.""" + resources = ResourceManager(client=client.gateway) + try: + marker = unique_marker() + team_id = client.create_team(alias=f"e2e-member-iso-{marker}", max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_team(team_id)) + tight_user = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(tight_user)) + roomy_user = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(roomy_user)) + client.add_team_member(team_id, tight_user, max_budget_in_team=TIGHT_MEMBER_BUDGET) + client.add_team_member(team_id, roomy_user, max_budget_in_team=ROOMY_MEMBER_BUDGET) + tight_key = client.generate_key(team_id=team_id, user_id=tight_user) + resources.defer(lambda: client.delete_key(tight_key)) + roomy_key = client.generate_key(team_id=team_id, user_id=roomy_user) + resources.defer(lambda: client.delete_key(roomy_key)) + yield _Pair( + team_id=team_id, + tight_user_id=tight_user, + roomy_user_id=roomy_user, + tight_key=tight_key, + roomy_key=roomy_key, + ) + finally: + resources.teardown() + + +def _roomy_send(client: BudgetClient, key: str) -> str: + """One roomy-member call that must go through; returns its request id.""" + match client.gateway.chat( + key, + ChatBody( + model=MODEL, + messages=[ChatMessage(role="user", content=f"roomy {unique_marker()}")], + max_tokens=16, + ), + ): + case Success(data=response): + assert response.id is not None, "roomy member call returned no id" + return response.id + case other: + pytest.fail(f"roomy member call failed while a peer was over budget: {other}") + + +class TestTeamMemberBudgetIsolation: + @pytest.mark.covers("quota_management.budget.team_member.isolates_per_member") + def test_blocked_member_does_not_block_peer(self, client: BudgetClient, pair: _Pair) -> None: + blocked = False + for _ in range(40): + result = client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, "tight member's per-team budget never enforced" + + sent = frozenset(_roomy_send(client, pair.roomy_key) for _ in range(ROOMY_BURST)) + + assert is_budget_block( + client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16) + ), "tight member stopped being blocked once the peer spent" + + rows = client.gateway.poll_logs_for_key( + pair.roomy_key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) + ) + logged = [row for row in rows if row.request_id in sent] + assert logged, "none of the roomy member's calls reached the spend logs" + for row in logged: + assert row.user == pair.roomy_user_id, ( + f"roomy call {row.request_id} logged under user {row.user}, not {pair.roomy_user_id}" + ) + assert row.team_id == pair.team_id, ( + f"roomy call {row.request_id} logged under team {row.team_id}, not {pair.team_id}" + ) diff --git a/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py b/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py new file mode 100644 index 00000000000..e87d6384406 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py @@ -0,0 +1,61 @@ +"""Live e2e: a per-user max_budget is enforced across ALL of that user's keys. + +An internal user's budget governs every personal key it owns, not only the one +that happened to spend it down. One user with a tiny max_budget owns two keys: +driving the first key to a budget_exceeded block then makes a fresh, untouched +second key of the same user (which carries no budget of its own, so nothing but the +shared user budget can block it) reject the same way, and the user's recorded spend +has crossed the cap. A key-scoped-only budget would leave the second key serving. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "gpt-5.5" +TINY_CAP = 3e-6 +RECORDED_SPEND_DEADLINE_SECONDS = 90 + + +def _call(client: BudgetClient, key: str): + return client.chat(key, MODEL, f"across {unique_marker()}", max_tokens=16) + + +def _drive_to_block(client: BudgetClient, key: str, subject: str) -> None: + for _ in range(40): + result = _call(client, key) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail(f"user budget never enforced on {subject} within the call budget") + + +class TestUserBudgetAcrossKeys: + @pytest.mark.covers("quota_management.budget.internal_user.enforced_across_keys") + def test_user_budget_blocks_a_second_key(self, client: BudgetClient, resources: ResourceManager) -> None: + user_id = client.create_user(max_budget=TINY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + + first_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(first_key)) + second_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(second_key)) + + _drive_to_block(client, first_key, "the first key") + _drive_to_block(client, second_key, "the second key") + + deadline = time.monotonic() + RECORDED_SPEND_DEADLINE_SECONDS + while time.monotonic() < deadline: + info = client.user_info(user_id) + if info is not None and (info.spend or 0.0) >= TINY_CAP: + return + time.sleep(5) + pytest.fail(f"user spend never reached the {TINY_CAP} cap in the recorded state")