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
This commit is contained in:
Yassin Kortam 2026-07-17 09:08:00 -07:00
parent 4d33964898
commit ab4392cd36
5 changed files with 210 additions and 1 deletions

View file

@ -139,7 +139,8 @@ quota_management.<behavior>.<variant>.<assertion>
| 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]

View file

@ -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"}

View file

@ -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:

View file

@ -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}"
)

View file

@ -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")