mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
4d33964898
commit
ab4392cd36
5 changed files with 210 additions and 1 deletions
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
@ -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")
|
||||
Loading…
Add table
Reference in a new issue