diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 180639b53e3..abfd198b7a7 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -150,14 +150,15 @@ quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict key | internal_user | end_user | organization | team | team_member | tag - | model_max | soft | key_multi_window | team_multi_window + | model_max | team_model_max | soft | key_multi_window | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user | per_model | failure | 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 | isolates_per_member | enforced_across_keys | routes_to_fallback + | isolates_per_model | isolates_per_member | enforced_across_keys | key_override_wins + | 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 diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index eb620395c46..db6ac7939c8 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -19,6 +19,8 @@ - {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.team_model_max.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: team_model_max, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "A team's model_max_budget entry is one shared cap: exhausting it on one key blocks every other team key for that model while other models stay open"} +- {id: quota_management.budget.team_model_max.key_override_wins, module: quota_management, tier: P1, behavior: budget, variant: team_model_max, assertions: [key_override_wins], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "A key's own model_max_budget entry takes precedence over the team's cap for that model, so the key keeps serving after the team cap is exhausted"} - {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"} - {id: quota_management.budget.key.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: key, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes key spend after the window; a blocked key serves again"} - {id: quota_management.budget.team.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes a team's spend after the window; every key on the team serves again"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 83e8f27b597..ecb94342eeb 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -97,6 +97,7 @@ class TeamNewBody(BaseModel): budget_duration: str | None = None organization_id: str | None = None budget_limits: list[BudgetWindow] | None = None + model_max_budget: dict[str, ModelBudgetEntry] | None = None class TeamNewResponse(BaseModel): @@ -385,6 +386,7 @@ class BudgetClient: budget_duration: str | None = None, organization_id: str | None = None, budget_limits: list[BudgetWindow] | None = None, + model_max_budget: dict[str, ModelBudgetEntry] | None = None, ) -> str: team_id = unwrap( self.proxy.transport.post( @@ -396,6 +398,7 @@ class BudgetClient: budget_duration=budget_duration, organization_id=organization_id, budget_limits=budget_limits, + model_max_budget=model_max_budget, ), response_type=TeamNewResponse, ) diff --git a/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py new file mode 100644 index 00000000000..21514fee753 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py @@ -0,0 +1,89 @@ +"""Live e2e: a team's `model_max_budget` is one shared per-model cap for every +team key, and a key carrying its own entry for that model is exempt from it.""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block, model_budget +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +CAPPED_MODEL = "claude-haiku-4-5" +FREE_MODEL = "gemini-2.5-flash" +BLOCK_DEADLINE_SECONDS = 60 + + +def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse: + result = client.chat(key, model, f"hi {unique_marker()}", max_tokens=16) + if not result.ok and not is_budget_block(result): + require_successful_call(result) + return result + + +def _drive_to_team_block(client: BudgetClient, key: str, team_id: str) -> None: + deadline = time.monotonic() + BLOCK_DEADLINE_SECONDS + while time.monotonic() < deadline: + result = _call(client, key, CAPPED_MODEL) + if is_budget_block(result): + assert team_id in result.body, result.body + return + time.sleep(1) + pytest.fail(f"team model_max_budget on {CAPPED_MODEL} never enforced") + + +# User flow (TLDR^2) +# 1. Admin caps claude-haiku tiny on the team, leaves gemini roomy +# 2. Alice's team key burns through the claude cap +# 3. Bob's untouched team key is refused claude on its first try +# 4. Bob's same key still gets gemini answers +@pytest.mark.covers("quota_management.budget.team_model_max.enforced_across_keys") +def test_team_model_cap_blocks_sibling_key(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-model-max-{unique_marker()}", + model_max_budget={ + **model_budget(CAPPED_MODEL, 1e-6), + **model_budget(FREE_MODEL, 1000.0), + }, + ) + resources.defer(lambda: client.delete_team(team_id)) + spender = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(spender)) + bystander = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(bystander)) + + _drive_to_team_block(client, spender, team_id) + + first = _call(client, bystander, CAPPED_MODEL) + assert is_budget_block(first), ( + f"sibling team key served {CAPPED_MODEL} after the team cap was exhausted: " + f"HTTP {first.status_code} {first.body}" + ) + assert team_id in first.body, first.body + + require_successful_call(_call(client, bystander, FREE_MODEL)) + + +# User flow (TLDR^2) +# 1. Admin caps claude-haiku tiny on the team +# 2. Admin issues Carol a team key with its own big claude budget +# 3. Dave's plain team key exhausts the team cap and is refused +# 4. Carol keeps getting claude answers on the same team +@pytest.mark.covers("quota_management.budget.team_model_max.key_override_wins") +def test_key_override_exempts_from_team_cap(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-model-max-override-{unique_marker()}", + model_max_budget=model_budget(CAPPED_MODEL, 1e-6), + ) + resources.defer(lambda: client.delete_team(team_id)) + inherit = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(inherit)) + override = client.generate_key(team_id=team_id, model_max_budget=model_budget(CAPPED_MODEL, 1000.0)) + resources.defer(lambda: client.delete_key(override)) + + _drive_to_team_block(client, inherit, team_id) + + require_successful_call(_call(client, override, CAPPED_MODEL))