diff --git a/tests/e2e/budgets/BUDGET_CODE_MATRIX.md b/tests/e2e/budgets/BUDGET_CODE_MATRIX.md new file mode 100644 index 00000000000..63cb41228d2 --- /dev/null +++ b/tests/e2e/budgets/BUDGET_CODE_MATRIX.md @@ -0,0 +1,93 @@ +# Budget Code Matrix + +What LiteLLM actually implements for budgets: every entity that can carry a dollar +budget, how the limit is enforced, and where in the code it happens. This is the +"what we support" reference; the companion `BUDGET_TEST_COVERAGE_MATRIX.md` maps +each row to its tests and the e2e gaps. + +Over-budget surfaces as a `budget_exceeded` error (the live suite +`tests/otel_tests/test_e2e_budgeting.py` asserts `type == "budget_exceeded"`, +`code == "429"`); the underlying `BudgetExceededError` is defined in +`litellm/exceptions.py` (`status_code=400`). Enforcement runs in `common_checks()` +/ `auth_checks.py` at auth time, plus pre-call reservation in +`budget_reservation.py`. + +Legend for "Enforced": **block** = request rejected; **filter** = router skips the +deployment; **alert** = notify only, request proceeds. + +--- + +## 1. Per-entity dollar budgets + +| Entity | Budget stored | Hard `max_budget` | Soft budget | Per-window | Model budget | Reset by `budget_duration` | +|--------|---------------|-------------------|-------------|------------|--------------|----------------------------| +| API key | `LiteLLM_VerificationToken` (direct cols + `budget_id` FK) | block (`_virtual_key_max_budget_check`) | alert (`_virtual_key_soft_budget_check`) + 80% alert | block (`_virtual_key_multi_budget_check`) | block (`model_max_budget_limiter.is_key_within_model_budget`) | keys reset job | +| Internal user | `LiteLLM_UserTable` (direct cols) | block (`common_checks`, only when not on a team) | - | - | via `model_max_budget` json | users reset job | +| Team | `LiteLLM_TeamTable` (direct cols) | block (`_team_max_budget_check`) | alert (`_team_soft_budget_check`) | block (`_team_multi_budget_check`) | via `model_max_budget` | teams reset job | +| Team member | `LiteLLM_TeamMembership` -> `LiteLLM_BudgetTable` | block (`_check_team_member_budget`) | - | - | - | budget-table reset job | +| End-user / customer | `LiteLLM_EndUserTable` -> `LiteLLM_BudgetTable` | block (`_check_end_user_budget`) | - | - | block (`is_end_user_within_model_budget`) | budget-table reset job | +| Organization | `LiteLLM_OrganizationTable` -> `LiteLLM_BudgetTable` | block (`_organization_max_budget_check`) | - | - | via budget-table | budget-table reset job | +| Tag | `LiteLLM_TagTable` -> `LiteLLM_BudgetTable` | block (`_tag_max_budget_check`) | - | - | via budget-table | budget-table reset job | +| Project | `LiteLLM_ProjectTable` -> `LiteLLM_BudgetTable` | block (`_project_max_budget_check`) | alert (`_project_soft_budget_check`) | - | - | budget-table reset job | +| Provider (router) | config `provider_budget_config` (in-memory) | filter (`router_strategy/budget_limiter`) | - | yes (time window) | - | window TTL | +| Global proxy | `litellm.max_budget` (config) | block (`_global_proxy_budget_check`) | - | - | - | - | + +Notes / flags from the code: +- **User budget only enforced off-team**: `common_checks` skips the personal-user + budget when the key belongs to a team (team budget governs instead). +- **Comparison operators are inconsistent**: key/user use `>=`, team/end-user main + budget use `>`. Spend exactly at `max_budget` blocks a key but not a team. +- **Provider budgets are filter-only**: an over-budget provider is removed from + routing; if all are over budget the router raises + `no_deployments_with_provider_budget_routing` (not a per-entity block). +- **Enforcement timing differs by entity**: key / user / org / team-member / tag / + model enforce off real-time reservation counters (block within ~2 calls); + **end-user** enforcement reads `EndUserTable.spend`, which only updates on the + `proxy_batch_write_at` flush, so it lags by that interval (verified live). + +## 2. Budget mechanisms + +| Mechanism | What it does | Code | +|-----------|--------------|------| +| Pre-call reservation | Estimates max request cost, atomically reserves against redis spend counters for key/team/user/end_user/tag/team_member/org before the call; blocks if a counter would exceed | `spend_tracking/budget_reservation.py` | +| Post-call reconciliation | Adjusts the reservation to the actual cost once known | `reconcile_budget_reservation` | +| Read-time enforcement | Auth-time check of current spend vs `max_budget` | `auth_checks.common_checks` + per-entity `_*_max_budget_check` | +| Soft budget / alerts | At `soft_budget` (or 80% of max) fire Slack/email alert, do not block | `_virtual_key_soft_budget_check`, `_team_soft_budget_check`, `budget_alerts` | +| Multi-window budgets | `budget_limits` list of `{budget_duration, max_budget}`; each window enforced + reset independently | `_virtual_key_multi_budget_check`, `reset_budget_windows` | +| Model-level budgets | `model_max_budget` dict (per model: `budget_limit` + `time_period`) on key/user/end_user | `hooks/model_max_budget_limiter.py` | +| Reset by duration | Job zeros `spend`, recomputes `budget_reset_at = now + duration_in_seconds(budget_duration)`, invalidates redis counters | `common_utils/reset_budget_job.py`, `duration_parser.duration_in_seconds` | +| Zero-cost bypass | Models with no configured price bypass budget reservation | `budget_reservation` zero-cost path | + +## 3. Budget management surface (endpoints) + +| Action | Endpoint | Handler | +|--------|----------|---------| +| Create budget | `POST /budget/new` | `new_budget` | +| Update budget | `POST /budget/update` | `update_budget` | +| Budget info | `POST /budget/info` (`{"budgets": [id]}`) | `info_budget` | +| Budget settings | `GET /budget/settings` | `budget_settings` | +| List budgets | `GET /budget/list` | `list_budget` | +| Delete budget | `POST /budget/delete` (`{"id": id}`) | `delete_budget` | +| Set on key | `POST /key/generate`, `/key/update` (`max_budget`, `soft_budget`, `budget_duration`, `model_max_budget`, `budget_id`) | key mgmt | +| Set on user | `POST /user/new` (`max_budget`, `budget_duration`) | internal user | +| Set on team | `POST /team/new` (`max_budget`, `soft_budget`, `team_member_budget`) | team | +| Set on team member | `POST /team/member_add` (`max_budget_in_team`) | team | +| Set on org | `POST /organization/new` (`max_budget`, `soft_budget`, `model_max_budget`) | org | +| Set on customer | `POST /customer/new`, `/customer/update` (`max_budget`, `budget_id`) | customer | +| Set on tag | `POST /tag/new`, `/tag/update` (`max_budget`) | tag mgmt | +| Read budget+spend | `/key/info`, `/user/info`, `/team/info`, `/organization/info`, `/customer/info`, `/budget/info` | per-entity info | + +Endpoint method/shape gotchas verified live: `/organization/delete` is **DELETE** +with `{"organization_ids": [id]}`; `/budget/info` takes `{"budgets": [id]}`; +`model_max_budget` entries use `{"budget_limit", "time_period"}`. + +## 4. Config knobs + +| Setting | Effect | +|---------|--------| +| `litellm.max_budget` | proxy-wide hard cap (global proxy budget) | +| `max_internal_user_budget` / `default_max_internal_user_budget` | default `max_budget` for internal users | +| `internal_user_budget_duration` | default reset duration for internal users | +| `max_end_user_budget` / `max_end_user_budget_id` | default budget for end-users | +| `default_team_params` | default `max_budget` / `budget_duration` / limits for teams | +| `provider_budget_config` (router) | per-provider spend caps + windows | diff --git a/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md b/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..62bfc1fdd41 --- /dev/null +++ b/tests/e2e/budgets/BUDGET_TEST_COVERAGE_MATRIX.md @@ -0,0 +1,78 @@ +# Budget Test Coverage Matrix + +Maps every row of `BUDGET_CODE_MATRIX.md` (what LiteLLM implements) to its tests +and level, then marks the live e2e coverage this suite adds. + +Levels: `unit` mocked (`AsyncMock` on `get_current_spend`/prisma); `router` live +router with fake deployments; `live-e2e` real proxy, real key/team, real requests +until blocked. Status: `covered` / `partial` / `gap`. + +Pre-existing live coverage outside this suite: +- `tests/otel_tests/test_e2e_budgeting.py` - key + team enforcement, budget update. +- `tests/local_testing/test_router_budget_limiter.py` - provider / tag / deployment + budgets at the router. + +This suite (`tests/e2e/budgets/`) adds the missing live coverage and runs +on the shared lifecycle (every entity it creates is deleted on teardown). + +--- + +## Per-entity enforcement + +| Entity | Unit | Pre-existing live | This suite (live) | Status | +|--------|------|-------------------|-------------------|--------| +| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** | +| Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** | +| Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** | +| Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** | +| End-user / customer | `test_custom_auth_end_user_budget.py` | - | `test_end_user_budget_blocks` | **covered (new)** | +| Organization | `test_organization_budget_enforcement.py` (flagged weak) | - | `test_organization_budget_blocks` | **covered (new)** | +| Tag (proxy-level) | - | router only | `test_tag_budget_e2e::test_tag_budget_blocks_tagged_requests` | **covered (new)** | +| Model-level (`model_max_budget`) | `test_unit_test_max_model_budget_limiter.py` | - | `test_model_max_budget_e2e::test_model_max_budget_isolates_per_model` | **covered (new)** | +| Provider (router) | `test_budget_limiter_hotpath.py` | `test_router_budget_limiter.py` | - | **covered** (router) | +| Global proxy (`litellm.max_budget`) | unit | - | - | **gap** (needs a config-level cap; not key-settable) | + +## Budget mechanisms + +| Mechanism | Unit | This suite (live) | Status | +|-----------|------|-------------------|--------| +| Pre-call reservation | `test_budget_reservation.py` | exercised by every enforcement test | **partial** | +| Soft budget / alerts | `SlackAlerting/test_budget_alert_types.py` | `test_soft_budget_e2e::test_soft_budget_does_not_block` | **covered (new)** (block-vs-alert; the alert side-effect itself stays unit) | +| Budget CRUD | `test_budget_endpoints.py` | `test_budget_crud_e2e` (roundtrip + delete) | **covered (new)** | +| Reset scheduling | `test_proxy_budget_reset.py` | `test_budget_crud_e2e::test_budget_duration_schedules_reset_on_key` | **covered (new)** (scheduling; actual zeroing is time-dependent -> unit) | +| Multi-window budgets | `test_multi_budget_windows.py` | - | **gap** (window setup is fiddly; left to unit for now) | +| Read budget+spend | `test_spend_management_endpoints.py` | `/key/info` asserted in CRUD + enforcement | **partial** | + +## Remaining gaps (intentionally not live-tested) + +- **Global proxy budget** (`litellm.max_budget`): set via proxy config, not a + per-key API, so it needs a dedicated proxy boot with that config rather than a + runtime-created entity. Out of scope for the per-entity suite. +- **Multi-window budgets**: the `budget_limits` list shape and per-window reset are + covered by `test_multi_budget_windows.py` (unit); a live version would need to + wait out a short window to see the reset, which is time-dependent. +- **Soft-budget alert delivery**: whether the Slack/email actually fires is not + observable from the proxy API; unit tests own that. The live test pins the + load-bearing behavior (soft does not block). +- **Reset zeroing after the window elapses**: time-dependent; unit tests own the + reset-job logic. The live test pins that `budget_reset_at` is scheduled. + +## This suite's files + +| File | Covers | +|------|--------| +| `test_budget_enforcement_e2e.py` | key / internal-user / end-user / organization / team-member hard enforcement | +| `test_model_max_budget_e2e.py` | per-model caps isolate by model | +| `test_soft_budget_e2e.py` | soft budget alerts but does not block | +| `test_tag_budget_e2e.py` | proxy-level tag budget blocks tagged requests, spares others | +| `test_budget_crud_e2e.py` | `/budget/*` CRUD roundtrip + delete + `budget_reset_at` scheduling | + +## Pattern + timing + +Create the entity with a tiny `max_budget`, drive spend until a `budget_exceeded` +block. The enforcement helper is two-phase: a fast warmup (key/user/org/member/tag/ +model block within ~2 calls off real-time counters), then a poll across the ~60s +batch-write window (end-user enforcement reads table spend that lags). Skip on a +non-budget error (provider down / key missing); fail if the budget is never +enforced. Chat tests use `gpt-5.5` (the model with a working key on the reference +proxy); swap the literal if your proxy differs. diff --git a/tests/e2e/budgets/budget_client.py b/tests/e2e/budgets/budget_client.py new file mode 100644 index 00000000000..af8021f9b93 --- /dev/null +++ b/tests/e2e/budgets/budget_client.py @@ -0,0 +1,418 @@ +"""Client for budget e2e tests: the shared Gateway plus budget-bearing entity +management (user / team / team-member / org / customer / tag / budget-table) and +info reads. + +Over-budget surfaces as a ``budget_exceeded`` error; ``is_budget_block`` detects it +on a chat outcome. Create methods return the new id and raise on failure; tests +register the matching delete with ``resources.defer(...)`` for cleanup. The request +and response models are co-located here because only this suite uses them. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic import AliasPath, BaseModel, Field, RootModel + +from e2e_gateway import Gateway, build_gateway +from e2e_http import NoBody, StreamingResponse, Success, unwrap +from models import ( + BudgetWindow, + ChatBody, + ChatMessage, + ChatMetadata, + KeyGenerateBody, + ModelBudgetEntry, +) + + +class UserNewBody(BaseModel): + max_budget: float + + +class UserNewResponse(BaseModel): + user_id: str + + +class UserDeleteBody(BaseModel): + user_ids: list[str] + + +class CustomerNewBody(BaseModel): + user_id: str + max_budget: float + + +class OrgNewBody(BaseModel): + organization_alias: str + max_budget: float + + +class OrgNewResponse(BaseModel): + organization_id: str + + +class OrgDeleteBody(BaseModel): + organization_ids: list[str] + + +class TeamMember(BaseModel): + role: str + user_id: str + + +class TeamNewBody(BaseModel): + team_alias: str + max_budget: float | None = None + organization_id: str | None = None + budget_limits: list[BudgetWindow] | None = None + + +class TeamNewResponse(BaseModel): + team_id: str + + +class TeamDeleteBody(BaseModel): + team_ids: list[str] + + +class TeamMemberAddBody(BaseModel): + team_id: str + member: TeamMember + max_budget_in_team: float | None = None + + +class TeamMemberUpdateBody(BaseModel): + team_id: str + user_id: str + max_budget_in_team: float | None = None + budget_duration: str | None = None + + +class TeamMembershipRow(BaseModel): + user_id: str | None = None + budget_reset_at: str | None = Field( + default=None, + validation_alias=AliasPath("litellm_budget_table", "budget_reset_at"), + ) + + +class TeamInfoParams(BaseModel): + team_id: str + + +class TeamInfoResponse(BaseModel): + team_memberships: list[TeamMembershipRow] = [] + + +class TagNewBody(BaseModel): + name: str + max_budget: float + + +class TagDeleteBody(BaseModel): + name: str + + +class BudgetNewBody(BaseModel): + max_budget: float + soft_budget: float | None = None + budget_duration: str | None = None + + +class BudgetNewResponse(BaseModel): + budget_id: str + + +class BudgetDeleteBody(BaseModel): + id: str + + +class BudgetInfoBody(BaseModel): + budgets: list[str] + + +class BudgetRow(BaseModel): + budget_id: str | None = None + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: str | None = None + + +class BudgetInfoResponse(RootModel[list[BudgetRow]]): + pass + + +def is_budget_block(result: StreamingResponse) -> bool: + """True if the call was rejected for being over budget (vs a provider error).""" + return not result.ok and "budget_exceeded" in result.body + + +def model_budget(model: str, limit: float, period: str = "30d") -> dict[str, ModelBudgetEntry]: + """A model_max_budget entry: per-model cap with a reset window.""" + return {model: ModelBudgetEntry(budget_limit=limit, time_period=period)} + + +@dataclass(frozen=True, slots=True) +class BudgetClient: + gateway: Gateway + + # ---- generic key ops (delegate to the shared Gateway) --------------- + + def generate_key( + self, + *, + models: list[str] | None = None, + max_budget: float | None = None, + soft_budget: float | None = None, + budget_duration: str | None = None, + budget_id: str | None = None, + user_id: str | None = None, + team_id: str | None = None, + model_max_budget: dict[str, ModelBudgetEntry] | None = None, + budget_limits: list[BudgetWindow] | None = None, + ) -> str: + return self.gateway.generate_key( + KeyGenerateBody( + models=models or [], + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + budget_id=budget_id, + user_id=user_id, + team_id=team_id, + model_max_budget=model_max_budget, + budget_limits=budget_limits, + ) + ) + + def delete_key(self, key: str) -> None: + self.gateway.delete_key(key) + + def delete_customers(self, user_ids: list[str]) -> None: + self.gateway.delete_customers(user_ids) + + # ---- chat (raw HTTP outcome: a budget block surfaces as a non-2xx) -- + + def chat( + self, + key: str, + model: str, + content: str, + *, + max_tokens: int | None = None, + user: str | None = None, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + "/chat/completions", + headers=self.gateway.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + user=user, + metadata=ChatMetadata(tags=tags) if tags else None, + ), + ) + + # ---- internal user -------------------------------------------------- + + def create_user(self, *, max_budget: float) -> str: + return unwrap( + self.gateway.transport.post( + "/user/new", + headers=self.gateway.transport.master, + json=UserNewBody(max_budget=max_budget), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = self.gateway.transport.post( + "/user/delete", + headers=self.gateway.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + + # ---- customer / end-user ------------------------------------------- + + def create_customer(self, customer_id: str, *, max_budget: float) -> str: + resp = self.gateway.transport.send( + "/customer/new", + headers=self.gateway.transport.master, + json=CustomerNewBody(user_id=customer_id, max_budget=max_budget), + ) + assert resp.ok, resp.body + return customer_id + + # ---- organization --------------------------------------------------- + + def create_org(self, *, max_budget: float, alias: str) -> str: + return unwrap( + self.gateway.transport.post( + "/organization/new", + headers=self.gateway.transport.master, + json=OrgNewBody(organization_alias=alias, max_budget=max_budget), + response_type=OrgNewResponse, + ) + ).organization_id + + def delete_org(self, org_id: str) -> None: + _ = self.gateway.transport.delete( + "/organization/delete", + headers=self.gateway.transport.master, + json=OrgDeleteBody(organization_ids=[org_id]), + response_type=NoBody, + ) + + # ---- team ----------------------------------------------------------- + + def create_team( + self, + *, + alias: str, + max_budget: float | None = None, + organization_id: str | None = None, + budget_limits: list[BudgetWindow] | None = None, + ) -> str: + return unwrap( + self.gateway.transport.post( + "/team/new", + headers=self.gateway.transport.master, + json=TeamNewBody( + team_alias=alias, + max_budget=max_budget, + organization_id=organization_id, + budget_limits=budget_limits, + ), + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + _ = self.gateway.transport.post( + "/team/delete", + headers=self.gateway.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None: + resp = self.gateway.transport.send( + "/team/member_add", + headers=self.gateway.transport.master, + json=TeamMemberAddBody( + team_id=team_id, + member=TeamMember(role="user", user_id=user_id), + max_budget_in_team=max_budget_in_team, + ), + ) + assert resp.ok, resp.body + + def update_team_member( + self, + team_id: str, + user_id: str, + *, + max_budget_in_team: float | None = None, + budget_duration: str | None = None, + ) -> None: + resp = self.gateway.transport.send( + "/team/member_update", + headers=self.gateway.transport.master, + json=TeamMemberUpdateBody( + team_id=team_id, + user_id=user_id, + max_budget_in_team=max_budget_in_team, + budget_duration=budget_duration, + ), + ) + assert resp.ok, resp.body + + def member_budget_reset_at(self, team_id: str, user_id: str) -> str | None: + """The member's per-team budget_reset_at as /team/info reports it, or None if + no reset is scheduled. The reset job advances this each time the window + elapses; a job that skips the row leaves it pinned forever.""" + result = self.gateway.transport.get( + "/team/info", + headers=self.gateway.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + match result: + case Success(data=data): + return next( + (row.budget_reset_at for row in data.team_memberships if row.user_id == user_id), + None, + ) + case _: + return None + + # ---- tag ------------------------------------------------------------ + + def create_tag(self, name: str, *, max_budget: float) -> str: + resp = self.gateway.transport.send( + "/tag/new", + headers=self.gateway.transport.master, + json=TagNewBody(name=name, max_budget=max_budget), + ) + assert resp.ok, resp.body + return name + + def delete_tag(self, name: str) -> None: + _ = self.gateway.transport.post( + "/tag/delete", + headers=self.gateway.transport.master, + json=TagDeleteBody(name=name), + response_type=NoBody, + ) + + # ---- budget table --------------------------------------------------- + + def create_budget( + self, + *, + max_budget: float, + soft_budget: float | None = None, + budget_duration: str | None = None, + ) -> str: + return unwrap( + self.gateway.transport.post( + "/budget/new", + headers=self.gateway.transport.master, + json=BudgetNewBody( + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + ), + response_type=BudgetNewResponse, + ) + ).budget_id + + def delete_budget(self, budget_id: str) -> None: + _ = self.gateway.transport.post( + "/budget/delete", + headers=self.gateway.transport.master, + json=BudgetDeleteBody(id=budget_id), + response_type=NoBody, + ) + + def budget_info(self, budget_id: str) -> tuple[BudgetRow, ...]: + result = self.gateway.transport.post( + "/budget/info", + headers=self.gateway.transport.master, + json=BudgetInfoBody(budgets=[budget_id]), + response_type=BudgetInfoResponse, + ) + match result: + case Success(data=data): + return tuple(data.root) + case _: + return () + + +def build_client() -> BudgetClient: + return BudgetClient(gateway=build_gateway()) diff --git a/tests/e2e/budgets/conftest.py b/tests/e2e/budgets/conftest.py new file mode 100644 index 00000000000..236822f4309 --- /dev/null +++ b/tests/e2e/budgets/conftest.py @@ -0,0 +1,16 @@ +"""Budgets suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. BudgetClient holds the shared Gateway, +so the `resources` fixture cleans up keys through it; tests register entity deletes +via `resources.defer(...)`. +""" + +import pytest + +from budget_client import BudgetClient, build_client + + +@pytest.fixture(scope="session") +def client() -> BudgetClient: + return build_client() diff --git a/tests/e2e/budgets/test_budget_crud_e2e.py b/tests/e2e/budgets/test_budget_crud_e2e.py new file mode 100644 index 00000000000..e697eca0051 --- /dev/null +++ b/tests/e2e/budgets/test_budget_crud_e2e.py @@ -0,0 +1,60 @@ +"""Live e2e for the budget management surface (no LLM calls, fast). + +Covers the budget-table CRUD round-trip and that `budget_duration` schedules a +`budget_reset_at`. The actual zeroing after the window is time-dependent, so we +assert the reset is *scheduled* (now + duration), not waited out. +""" + +from datetime import datetime, timezone + +import pytest + +from budget_client import BudgetClient +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + + +def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager) -> None: + budget_id = client.create_budget(max_budget=12.5, soft_budget=10.0, budget_duration="30d") + resources.defer(lambda: client.delete_budget(budget_id)) + + rows = client.budget_info(budget_id) + assert rows, f"/budget/info returned nothing for {budget_id}" + row = rows[0] + assert row.max_budget == 12.5 + assert row.soft_budget == 10.0 + assert row.budget_reset_at, "budget_duration did not schedule a reset" + + # Attach the budget to a key and confirm the key reflects it. + key = client.generate_key(budget_id=budget_id) + resources.defer(lambda: client.delete_key(key)) + info = client.gateway.key_info(key) + linked = info.litellm_budget_table + assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), ( + f"key does not reflect attached budget: {info.budget_id}, {linked}" + ) + + +def test_budget_delete_removes_it(client: BudgetClient, resources: ResourceManager) -> None: + budget_id = client.create_budget(max_budget=1.0) + resources.defer(lambda: client.delete_budget(budget_id)) + client.delete_budget(budget_id) + assert not client.budget_info(budget_id), "budget still present after delete" + + +def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=10.0, budget_duration="30d") + resources.defer(lambda: client.delete_key(key)) + + reset_at = client.gateway.key_info(key).budget_reset_at + assert reset_at, "budget_duration did not set budget_reset_at on the key" + + # budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the + # proxy may align the reset to a calendar boundary (e.g. start of next month), + # so "30d" can land ~12 days out mid-month. Assert it's scheduled ahead. + + # get current time -> assert budget from days_left - budget_duration == days_left + reset_dt = datetime.fromisoformat(str(reset_at).replace("Z", "+00:00")) + days_out = (reset_dt - datetime.now(timezone.utc)).total_seconds() / 86400 + assert 0 < days_out < 40, f"reset should be scheduled ahead, got {days_out:.1f}d out" diff --git a/tests/e2e/budgets/test_budget_enforcement_e2e.py b/tests/e2e/budgets/test_budget_enforcement_e2e.py new file mode 100644 index 00000000000..d03288b637a --- /dev/null +++ b/tests/e2e/budgets/test_budget_enforcement_e2e.py @@ -0,0 +1,148 @@ +"""Live e2e: a tiny max_budget on an entity actually blocks requests. + +Each entity is an E2ECase (lifecycle.E2ECase) driven by run_case: init() creates +the budgeted entity + a key, run() drives spend until a `budget_exceeded` block, +teardown() deletes everything init() created (always runs, even on failure/skip). +Covers the entities with no prior live coverage - internal user, end-user, +organization, team member. See BUDGET_TEST_COVERAGE_MATRIX.md. + +A non-budget error fails hard (never a skip); if calls never get blocked, budget +enforcement is broken -> fail. +""" + +import time +from dataclasses import dataclass, field +from typing import Callable, List, Type + +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 run_case + +pytestmark = pytest.mark.e2e + +def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> None: + """Send paid calls until the entity's budget blocks one. Key/user/org/member + block within a couple calls off real-time reservation counters; the end-user + budget enforces off table spend that lands on the batch write, so it takes a + few more. A non-budget error fails hard (never a skip).""" + for _ in range(40): + result = client.chat( + key, + "claude-haiku-4-5", + f"spend {unique_marker()}", + max_tokens=16, + user=user or None, + ) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail("budget never enforced within the call budget") + + +@dataclass +class _BudgetCase: + """Base E2ECase: a key under some budgeted entity must get blocked. + + Subclasses set up the budgeted entity in init() and register every created id + in `_undo` (run LIFO in teardown so a key is deleted before its team/org). + """ + + client: BudgetClient + key: str = "" + _undo: List[Callable[[], None]] = field( + default_factory=list + ) # mutable-ok: per-case teardown registry + + def init(self) -> None: + raise NotImplementedError + + def run(self) -> None: + _assert_budget_blocks(self.client, self.key) + + def teardown(self) -> None: + for undo in reversed(self._undo): + undo() + + +class KeyBudgetCase(_BudgetCase): + def init(self) -> None: + self.key = self.client.generate_key(max_budget=3e-6) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class InternalUserBudgetCase(_BudgetCase): + def init(self) -> None: + user_id = self.client.create_user(max_budget=3e-6) + self._undo.append(lambda: self.client.delete_user(user_id)) + # personal key (no team) -> the user budget governs + self.key = self.client.generate_key(user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class EndUserBudgetCase(_BudgetCase): + def init(self) -> None: + customer = f"e2e-budget-cust-{unique_marker()}" + self.client.create_customer(customer, max_budget=3e-6) + self._undo.append(lambda: self.client.delete_customers([customer])) + self.key = self.client.generate_key(models=["claude-haiku-4-5"]) + self._undo.append(lambda: self.client.delete_key(self.key)) + self._customer = customer + + def run(self) -> None: + _assert_budget_blocks(self.client, self.key, user=self._customer) + + +class OrganizationBudgetCase(_BudgetCase): + def init(self) -> None: + # Org carries the tiny budget; the team under it has none, so a block here + # is org-level enforcement (the historically weak link). + org_id = self.client.create_org( + max_budget=3e-6, alias=f"e2e-budget-org-{unique_marker()}" + ) + self._undo.append(lambda: self.client.delete_org(org_id)) + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", organization_id=org_id + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + self.key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +class TeamMemberBudgetCase(_BudgetCase): + def init(self) -> None: + # Member's per-team budget is tiny while the team has a large budget, so a + # block proves member-level (not team-level) enforcement. + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", max_budget=100.0 + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + user_id = self.client.create_user(max_budget=100.0) + self._undo.append(lambda: self.client.delete_user(user_id)) + self.client.add_team_member(team_id, user_id, max_budget_in_team=3e-6) + self.key = self.client.generate_key(team_id=team_id, user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + + +def _case_id(case_cls: Type[_BudgetCase]) -> str: + return case_cls.__name__ + + +@pytest.mark.parametrize( + "case_cls", + [ + KeyBudgetCase, + InternalUserBudgetCase, + EndUserBudgetCase, + OrganizationBudgetCase, + TeamMemberBudgetCase, + ], + ids=_case_id, +) +def test_budget_enforcement( + client: BudgetClient, case_cls: Type[_BudgetCase] +) -> None: + run_case(case_cls(client)) diff --git a/tests/e2e/budgets/test_budget_reset_e2e.py b/tests/e2e/budgets/test_budget_reset_e2e.py new file mode 100644 index 00000000000..dcf776db9a2 --- /dev/null +++ b/tests/e2e/budgets/test_budget_reset_e2e.py @@ -0,0 +1,59 @@ +"""Live e2e: a key budget resets (zeroes spend) after its budget_duration. + +Short budget_duration (30s) + the fast-rescheduled reset job: a key blocked for +exceeding its max_budget starts succeeding again once the duration elapses and the +reset job zeroes key.spend. Closes the reset-zeroing gap in +BUDGET_TEST_COVERAGE_MATRIX.md (reset_budget_for_litellm_keys), which the unit +suite covers but no live test did - distinct from the per-window reset in +test_multi_window_budget_e2e.py. +""" + +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 + + +def _call(client: BudgetClient, key: str): + return client.chat( + key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16 + ) + + +def test_key_budget_resets_after_duration( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key(max_budget=3e-6, budget_duration="30s") + resources.defer(lambda: client.delete_key(key)) + + # 1. exceed the budget -> litellm returns budget_exceeded + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, "key budget never enforced" + + # 2. once the 30s duration elapses + the reset job runs, key.spend zeroes and + # calls flow again. The window is wall-clock-aligned, so the reset lands up to + # a window later, then the rescheduler (~15-20s) zeroes the spend; allow + # generous headroom over that. A stuck rescheduler is caught by the wait-loop + # timeout, not this elapsed bound. + start = time.monotonic() + while time.monotonic() < start + 150: + time.sleep(5) + result = _call(client, key) + if result.ok: + assert time.monotonic() - start < 120, "reset too slow for a 30s budget" + return + assert is_budget_block(result), f"non-budget error: {result.body[:200]}" + pytest.fail("key budget never reset within 150s") diff --git a/tests/e2e/budgets/test_model_max_budget_e2e.py b/tests/e2e/budgets/test_model_max_budget_e2e.py new file mode 100644 index 00000000000..44e6a333ef0 --- /dev/null +++ b/tests/e2e/budgets/test_model_max_budget_e2e.py @@ -0,0 +1,57 @@ +"""Live e2e: per-model budgets (`model_max_budget`) isolate by model. + +A key caps one model tiny and leaves another generous. Exhausting the capped +model must block *that* model while the other still works - proving the per-model +cap is enforced independently, not as a key-wide budget. Closes the +model_max_budget gap in BUDGET_TEST_COVERAGE_MATRIX.md. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block, model_budget +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +CAPPED_MODEL = "claude-haiku-4-5" +FREE_MODEL = "gemini-2.5-flash" + + +def _call(client: BudgetClient, key: str, model: str): + 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 test_model_max_budget_isolates_per_model( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key( + model_max_budget={ + **model_budget(CAPPED_MODEL, 1e-6), + **model_budget(FREE_MODEL, 1000.0), + } + ) + resources.defer(lambda: client.delete_key(key)) + + # Exhaust the capped model. + blocked = False + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if is_budget_block(_call(client, key, CAPPED_MODEL)): + blocked = True + break + time.sleep(1) + assert blocked, f"{CAPPED_MODEL} per-model budget never enforced" + + # The other model shares the key but has its own (large) cap -> still works. + other = _call(client, key, FREE_MODEL) + assert not is_budget_block(other), ( + f"{FREE_MODEL} was blocked by {CAPPED_MODEL}'s budget; per-model caps not isolated" + ) + require_successful_call(other) diff --git a/tests/e2e/budgets/test_multi_window_budget_e2e.py b/tests/e2e/budgets/test_multi_window_budget_e2e.py new file mode 100644 index 00000000000..553ad1ce701 --- /dev/null +++ b/tests/e2e/budgets/test_multi_window_budget_e2e.py @@ -0,0 +1,70 @@ +"""Live e2e: multi-window budgets (budget_limits) enforce AND reset per window. + +Short windows make the time limit reachable inside a test: a tight 30s window and +a roomy 1m window. The 30s window blocks once its tiny cap is exceeded, then - once +its 30s elapses and the reset job runs (rescheduled fast via +PROXY_BUDGET_RESCHEDULER_* in docker-compose) - the window resets and calls flow +again. Closes the multi-window gap (enforcement + per-window reset) in +BUDGET_TEST_COVERAGE_MATRIX.md, which the unit suite covered but no live test did. +""" + +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 +from models import BudgetWindow + +pytestmark = pytest.mark.e2e + +WINDOW_SECONDS = 30 # the tight window; calls succeed again only after it elapses + + +def _call(client: BudgetClient, key: str): + return client.chat( + key, "claude-haiku-4-5", f"window {unique_marker()}", max_tokens=16 + ) + + +def test_short_window_blocks_then_resets( + client: BudgetClient, resources: ResourceManager +) -> None: + key = client.generate_key( + budget_limits=[ + BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ] + ) + resources.defer(lambda: client.delete_key(key)) + + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, f"{WINDOW_SECONDS}s window never enforced" + + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, ( + f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + ) + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s") diff --git a/tests/e2e/budgets/test_soft_budget_e2e.py b/tests/e2e/budgets/test_soft_budget_e2e.py new file mode 100644 index 00000000000..407de7ae467 --- /dev/null +++ b/tests/e2e/budgets/test_soft_budget_e2e.py @@ -0,0 +1,35 @@ +"""Live e2e: soft_budget alerts but does NOT block. + +A key with a tiny `soft_budget` well under a large `max_budget`: spend crosses the +soft threshold within a couple calls, but requests keep succeeding (soft budget is +advisory). Closes the soft_budget gap in BUDGET_TEST_COVERAGE_MATRIX.md. The alert +side-effect (Slack/email) is not observable from the proxy API, so we assert the +load-bearing behavior: soft != block. +""" + +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 + + +def test_soft_budget_does_not_block( + client: BudgetClient, resources: ResourceManager +) -> None: + # soft far below max: spend crosses soft immediately, stays under max. + key = client.generate_key(max_budget=1000.0, soft_budget=1e-9) + resources.defer(lambda: client.delete_key(key)) + + for _ in range(3): + result = client.chat( + key, "claude-haiku-4-5", f"hi {unique_marker()}", max_tokens=16 + ) + assert not is_budget_block(result), ( + "soft_budget blocked a request; it must alert only, not block " + f"(body={result.body[:200]})" + ) + require_successful_call(result) # any other non-2xx (e.g. provider down) is a hard fail diff --git a/tests/e2e/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py new file mode 100644 index 00000000000..a6860aeef43 --- /dev/null +++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py @@ -0,0 +1,144 @@ +"""Live e2e: concurrent cold-counter reseeds keep the spend counter equal to DB spend (#26829). + +Regression for the cross-pod spend-counter multiplication. Real requests build a key's +DB spend through the spend writer; the Redis spend counter then expires (the e2e proxy +sets a short default_redis_ttl) and goes cold. The proxy runs several workers sharing one +Redis, so a concurrent burst makes more than one worker reseed the same cold counter at +once. The fix seeds with SET NX - one worker initializes the counter at the DB spend and +the rest read it back - so the counter still equals the DB spend (plus the burst's own +small cost). The pre-#26829 additive reseed stacked the DB spend once per worker, leaving +the counter at ~N x the real spend. + +The test reads the shared counter straight from Redis and asserts it equals the DB spend, +not a multiple. It also asserts the counter actually went cold before the burst, so a proxy +that never expires the counter (no short TTL) fails loudly instead of passing vacuously. +Skipped when the e2e Redis is not reachable. +""" + +import hashlib +import os +import time +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from typing import TYPE_CHECKING + +import pytest + +from budget_client import BudgetClient +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager + +if TYPE_CHECKING: + import redis + from redis.cluster import RedisCluster + +pytestmark = pytest.mark.e2e + +MODEL = "claude-haiku-4-5" +ACCUMULATE_CALLS = 24 +BURST = 6 +# proxy_batch_write_at (60s) flushes the spend to the DB and default_redis_ttl (20s) +# expires the counter; this waits out both. +COLD_WAIT_SECONDS = 80 + + +def _redis() -> "redis.Redis[str] | RedisCluster[str]": + """The proxy's Redis. The deployed runner sets REDIS_HOST to the serverless + ElastiCache, which is always TLS + cluster-mode; without it, fall back to a + local standalone redis for docker-compose runs.""" + import redis + + host = os.getenv("REDIS_HOST") + if not host: + return redis.Redis(host="localhost", port=6380, decode_responses=True, socket_connect_timeout=2) + + from redis.cluster import RedisCluster + + return RedisCluster( + host=host, + port=int(os.getenv("REDIS_PORT", "6379")), + ssl=True, + decode_responses=True, + socket_connect_timeout=2, + ) + + +def _spend_counter(rds: "redis.Redis[str] | RedisCluster[str]", key: str) -> float | None: + """The shared spend counter for `key`, or None if it is cold. A cluster client + can't run a keyspace SCAN that spans shards, so read the key directly - the stage + gateway sets no cache namespace, so the key is the bare ``spend:key:{sha256(key)}``. + A standalone client matches by suffix, so the local cache namespace (litellm.caching) + need not be hard-coded here.""" + from redis.cluster import RedisCluster + + digest = hashlib.sha256(key.encode()).hexdigest() + suffix = f"spend:key:{digest}" + if isinstance(rds, RedisCluster): + raw = rds.get(suffix) + return float(raw) if raw is not None else None + + matches = list(rds.scan_iter(match=f"*{suffix}")) + if not matches: + return None + raw = rds.get(matches[0]) + return float(raw) if raw is not None else None + + +def _chat(client: BudgetClient, key: str) -> StreamingResponse: + return client.chat(key, MODEL, f"reseed {unique_marker()}", max_tokens=16) + + +def _accumulate(client: BudgetClient, key: str, count: int) -> None: + def one(_: int) -> StreamingResponse: + return _chat(client, key) + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(one, range(count))) + + +def _burst(client: BudgetClient, key: str, count: int) -> None: + """Fire `count` requests that start together, so multiple workers reseed the cold + counter concurrently rather than one warming it before the others arrive.""" + barrier = Barrier(count) + + def one(_: int) -> StreamingResponse: + barrier.wait() + return _chat(client, key) + + with ThreadPoolExecutor(max_workers=count) as pool: + list(pool.map(one, range(count))) + + +def test_cold_counter_reseed_keeps_counter_equal_to_db_spend( + client: BudgetClient, resources: ResourceManager +) -> None: + try: + rds = _redis() + rds.ping() + except Exception as exc: # noqa: BLE001 - any connect failure means skip + pytest.skip(f"e2e redis not reachable (set REDIS_HOST/REDIS_PORT): {exc}") + + key = client.generate_key(max_budget=1.0, models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + _accumulate(client, key, ACCUMULATE_CALLS) + time.sleep(COLD_WAIT_SECONDS) + + assert _spend_counter(rds, key) is None, ( + "the spend counter never went cold; default_redis_ttl must be short enough for it " + "to expire, otherwise the burst reads a warm counter and the reseed is never exercised" + ) + db_spend = client.gateway.key_info(key).spend or 0.0 + assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}" + + _burst(client, key, BURST) + time.sleep(3) + + counter = _spend_counter(rds, key) + assert counter is not None, "the burst did not reseed the cold counter" + assert db_spend * 0.95 <= counter < db_spend * 1.7, ( + f"redis spend counter {counter} does not equal DB spend {db_spend} (expected ~equal " + f"plus the burst's small cost); a near-multiple means the cold-counter reseed stacked " + f"the DB spend once per worker instead of seeding it once (#26829)" + ) diff --git a/tests/e2e/budgets/test_tag_budget_e2e.py b/tests/e2e/budgets/test_tag_budget_e2e.py new file mode 100644 index 00000000000..7cec5bc96c1 --- /dev/null +++ b/tests/e2e/budgets/test_tag_budget_e2e.py @@ -0,0 +1,59 @@ +"""Live e2e: proxy-level tag budgets block tagged requests. + +A tag with a tiny budget: requests carrying that tag get blocked once the tag's +spend is exceeded, while a request with a different tag (no budget) still works. +Closes the proxy-level tag-budget gap in BUDGET_TEST_COVERAGE_MATRIX.md (today +only router-level tag budgets are tested). +""" + +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 + +TINY_BUDGET = 1e-6 + + +def _tagged_call(client: BudgetClient, key: str, tag: str): + result = client.chat( + key, + "claude-haiku-4-5", + f"hi {unique_marker()}", + tags=[tag], + max_tokens=16, + ) + if not result.ok and not is_budget_block(result): + require_successful_call(result) + return result + + +def test_tag_budget_blocks_tagged_requests( + client: BudgetClient, scoped_key: str, resources: ResourceManager +) -> None: + budgeted_tag = f"e2e-budget-tag-{unique_marker()}" + client.create_tag(budgeted_tag, max_budget=TINY_BUDGET) + resources.defer(lambda: client.delete_tag(budgeted_tag)) + + # Requests under the budgeted tag get blocked once its spend is exceeded. + blocked = False + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if is_budget_block(_tagged_call(client, scoped_key, budgeted_tag)): + blocked = True + break + time.sleep(1) + assert blocked, f"tag budget for {budgeted_tag!r} never enforced" + + # A request with an unbudgeted tag on the same key is unaffected. + free_tag = f"e2e-free-tag-{unique_marker()}" + other = _tagged_call(client, scoped_key, free_tag) + assert not is_budget_block(other), ( + f"unbudgeted tag {free_tag!r} was blocked by {budgeted_tag!r}'s budget" + ) + require_successful_call(other) diff --git a/tests/e2e/budgets/test_team_member_budget_e2e.py b/tests/e2e/budgets/test_team_member_budget_e2e.py new file mode 100644 index 00000000000..301617bfdca --- /dev/null +++ b/tests/e2e/budgets/test_team_member_budget_e2e.py @@ -0,0 +1,107 @@ +"""Live e2e: a team member's per-team budget attributes spend and enforces a cap. + +The team carries a large budget while the one enrolled member is capped at a tiny +per-team budget, so any block is member-level, not team-level. Two scenarios share +that single member: +- attribution: the member's calls land in the spend logs tagged with both the team_id + and the member's user_id, so per-member spend can be billed back +- enforcement: once the member's spend passes the per-team budget, calls are blocked + with budget_exceeded while the team's own budget is nowhere near exhausted + +Per-member budgets enforce off batch-written spend (~60s), so a quick burst all goes +through; the block only lands once that spend flushes. +""" + +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 = "claude-haiku-4-5" +TEAM_BUDGET = 100.0 +MEMBER_BUDGET = 3e-6 +BURST = 6 + + +@dataclass(frozen=True, slots=True) +class _Member: + team_id: str + user_id: str + key: str + + +@pytest.fixture(scope="class") +def member(client: BudgetClient) -> Iterator[_Member]: + """A team with a large budget plus one member capped at a tiny per-team budget, + and that member's key. Shared across the class; torn down when it finishes. + Cleanups register progressively and run LIFO best-effort through ResourceManager, + so a partial-setup failure still releases what came before and one failed delete + never strands the rest on the shared proxy.""" + resources = ResourceManager(client=client.gateway) + try: + marker = unique_marker() + team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_team(team_id)) + user_id = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(user_id)) + client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET) + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + yield _Member(team_id=team_id, user_id=user_id, key=key) + finally: + resources.teardown() + + +def _send(client: BudgetClient, key: str) -> str | None: + """One member call; its response id (== the spend-log request_id) if it went + through, else None.""" + match client.gateway.chat( + key, + ChatBody( + model=MODEL, + messages=[ChatMessage(role="user", content=f"hi {unique_marker()}")], + max_tokens=16, + ), + ): + case Success(data=response): + return response.id + case _: + return None + + +class TestTeamMemberBudget: + def test_member_spend_attributed_to_team_and_user(self, client: BudgetClient, member: _Member) -> None: + sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid) + assert sent, "no member call went through; cannot check attribution" + + rows = client.gateway.poll_logs_for_key( + member.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, f"none of the member's {len(sent)} calls reached the spend logs" + + for row in logged: + assert row.team_id == member.team_id, ( + f"call {row.request_id} logged under team {row.team_id}, not the member's team {member.team_id}" + ) + assert row.user == member.user_id, ( + f"call {row.request_id} logged under user {row.user}, not member {member.user_id}" + ) + + def test_member_spend_over_budget_is_blocked(self, client: BudgetClient, member: _Member) -> None: + for _ in range(40): + result = client.chat(member.key, MODEL, f"spend {unique_marker()}", max_tokens=16) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail("per-member budget never enforced within the call budget") diff --git a/tests/e2e/budgets/test_team_member_budget_reset_e2e.py b/tests/e2e/budgets/test_team_member_budget_reset_e2e.py new file mode 100644 index 00000000000..2749f16a26e --- /dev/null +++ b/tests/e2e/budgets/test_team_member_budget_reset_e2e.py @@ -0,0 +1,47 @@ +import time +from datetime import datetime + +import pytest + +from budget_client import BudgetClient +from e2e_config import unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MEMBER_BUDGET = 1.0 # default member budget is $50, we're testing with a smaller value + +def _as_datetime(value: str) -> datetime: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def test_team_member_budget_reset_keeps_advancing(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-member-reset-{unique_marker()}", max_budget=100.0) + resources.defer(lambda: client.delete_team(team_id)) + user_id = client.create_user(max_budget=100.0) + resources.defer(lambda: client.delete_user(user_id)) + + # add the member, then update them onto a short per-team budget window + client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET) + client.update_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET, budget_duration="30s") + + scheduled = client.member_budget_reset_at(team_id, user_id) + assert scheduled, "updating the member with a budget_duration set no budget_reset_at" + first_reset = _as_datetime(scheduled) + + # the member can spend within the team while the window is live + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + require_successful_call(client.chat(key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16)) + + # once the window elapses the reset job must move budget_reset_at forward; a job + # that skips the member's budget row (the #25109 regression) leaves it pinned at + # first_reset forever + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + current = client.member_budget_reset_at(team_id, user_id) + if current and _as_datetime(current) > first_reset: + return + pytest.fail(f"member budget_reset_at never advanced past {first_reset.isoformat()} in 150s") diff --git a/tests/e2e/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py new file mode 100644 index 00000000000..946ab8ee1f2 --- /dev/null +++ b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py @@ -0,0 +1,79 @@ +"""Live e2e: a team's multi-window budgets (budget_limits) enforce AND reset per window. + +The team analog of test_multi_window_budget_e2e.py (which covers keys). A team is +created with a tight 30s window and a roomy 1m window; a key on that team blocks once +the tight window's cap is exceeded, then - once the 30s elapses and the reset job runs +(rescheduled fast via PROXY_BUDGET_RESCHEDULER_* in docker-compose) - the window resets +and calls flow again. This exercises the reset_budget_windows TEAM branch (raw SQL over +LiteLLM_TeamTable.budget_limits, the literal #25109 path), which had no live coverage. + +Fails at team creation today: /team/new writes the raw window list straight to the +Json? column, where Prisma rejects it (500), unlike the key path and /team/update which +json.dumps it first. Marked xfail(strict=True) so the suite stays green while the bug +persists and flips to a failure the moment the write is fixed and the marker should be +removed. +""" + +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 +from models import BudgetWindow + +pytestmark = pytest.mark.e2e + +WINDOW_SECONDS = 30 + + +def _call(client: BudgetClient, key: str): + return client.chat(key, "claude-haiku-4-5", f"team-window {unique_marker()}", max_tokens=16) + + +@pytest.mark.xfail( + strict=True, + reason="known proxy bug: /team/new writes budget_limits straight to the Json? " + "column and Prisma rejects it (500), unlike the key path and /team/update which " + "json.dumps first; remove this marker once that write is fixed", +) +def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + blocked = False + for _ in range(20): + result = _call(client, key) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, f"team {WINDOW_SECONDS}s window never enforced" + + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py new file mode 100644 index 00000000000..cc95c7538dd --- /dev/null +++ b/tests/e2e/conftest.py @@ -0,0 +1,120 @@ +"""Shared fixtures for all live e2e suites under tests/e2e/. + +Design rule: skip on environment, fail on behavior. Live tests (marked `e2e`) +skip when no proxy answers; once a request reaches the proxy, behavior is +asserted. Pure unit coverage of the harness itself carries no `e2e` marker and +runs regardless of whether a proxy is up. + +Lifecycle: the `resources` fixture maps the init -> run -> teardown contract +(lifecycle.E2ECase) onto pytest - setup is init(), the test body is run(), and +teardown deletes every resource the test created on the long-lived proxy. + +Each suite provides its own `client` fixture (a lifecycle.ResourceClient); these +shared fixtures build on it. +""" + +import functools +import sys +from pathlib import Path +from typing import Iterator + +import pytest +import requests + +from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL +from lifecycle import GatewayProvider, ResourceManager + + +_E2E_TEST_RAN = pytest.StashKey[bool]() + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "markers", + "e2e: live test that requires a running proxy and real provider keys", + ) + + +def _liveness_reason(label: str, base_url: str) -> str | None: + """None if `base_url` answers its liveness probe, else a skip reason.""" + try: + resp = requests.get(f"{base_url}/health/liveliness", timeout=5) + except requests.RequestException as exc: + return f"No live {label} at {base_url}: {exc}" + if resp.status_code >= 500: + return f"{label} at {base_url} returned {resp.status_code}" + return None + + +@functools.lru_cache(maxsize=1) +def _proxy_skip_reason() -> str | None: + """Probe the proxy once per session. None if it answers, else a skip reason. In + a split deployment the management/admin control plane is a separate service, so + require it too (when it differs) - else its tests would fail rather than skip.""" + reason = _liveness_reason("proxy", PROXY_BASE_URL) + if reason is not None: + return reason + if CONTROL_PLANE_BASE_URL != PROXY_BASE_URL: + return _liveness_reason("control plane", CONTROL_PLANE_BASE_URL) + return None + + +def pytest_runtest_setup(item: pytest.Item) -> None: + """Skip `e2e`-marked tests unless a proxy answers its liveness probe. Unmarked + tests (unit coverage of the harness) don't touch the proxy, so they run even + when none is up.""" + if item.get_closest_marker("e2e") is None: + return + reason = _proxy_skip_reason() + if reason is not None: + pytest.skip(reason) + + +def pytest_runtest_call(item: pytest.Item) -> None: + """Mark that an e2e test body actually ran (not skipped at setup). Skipped + sessions never reach this hook, so the session-finish cleanup can use it as a + guard before truncating the spend-log DB. Tests under `tests/e2e/` without the + `e2e` marker (pure unit coverage for the harness itself) never hit the proxy, + so they must not arm the destructive DB truncate.""" + if item.get_closest_marker("e2e") is None: + return + item.session.stash[_E2E_TEST_RAN] = True + + +def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: + """Once the whole e2e session is done (all suites), truncate the spend logs so + the DB doesn't accumulate test rows. Skipped sessions (no live proxy, no test + actually executed) leave the DB alone so a `DATABASE_URL` pointing at a shared + instance is never wiped without an e2e run. Best-effort: a cleanup failure (no + DB reachable) must not fail the run. The spend_tracking dir goes on sys.path + only for this import and is removed after, so a broader `pytest tests/` run is + not left with a mutated path.""" + if not session.stash.get(_E2E_TEST_RAN, False): + return + spend_dir = str(Path(__file__).parent / "spend_tracking") + sys.path.insert(0, spend_dir) + try: + from spend_e2e_client import reset_spend_logs # pyright: ignore + + reset_spend_logs() + except Exception as exc: # noqa: BLE001 - cleanup is best-effort + print(f"spend-log cleanup skipped: {exc}") + finally: + if spend_dir in sys.path: + sys.path.remove(spend_dir) + + +@pytest.fixture +def resources(client: GatewayProvider) -> Iterator[ResourceManager]: + """init -> run -> teardown: create a manager, run the test, release resources. + Cleanup goes through the shared Gateway, whatever the suite's client adds.""" + manager = ResourceManager(client=client.gateway) + manager.init() + yield manager + manager.teardown() + + +@pytest.fixture +def scoped_key(resources: ResourceManager) -> str: + """A fresh all-models key per test, auto-deleted by the resources teardown.""" + return resources.key() diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py new file mode 100644 index 00000000000..3865804b08f --- /dev/null +++ b/tests/e2e/e2e_config.py @@ -0,0 +1,34 @@ +"""Generic configuration for live e2e tests against a running LiteLLM proxy. + +Shared by every e2e suite under tests/e2e/. Values come from the +environment so the same tests run against localhost or a deployed proxy. +""" + +import os +import uuid + +PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000").rstrip("/") +MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY", "sk-1234") + +# Control-plane (management/admin) base URL. In a split control-plane/data-plane +# deployment the LLM data plane (PROXY_BASE_URL: /chat, /embeddings, native +# passthrough) and the management API (keys, users, teams, orgs, budgets, spend, +# model info, /openapi.json) are served by *different* services. The suite drives +# both through one Transport that routes by path (see transport.SplitTransport). +# Defaults to PROXY_BASE_URL so a monolithic proxy serving everything on one URL +# behaves exactly as before. +CONTROL_PLANE_BASE_URL = os.environ.get( + "LITELLM_CONTROL_PLANE_URL", PROXY_BASE_URL +).rstrip("/") + +# Writes on the proxy are eventually consistent (e.g. spend rows flush on +# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. +POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) +POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) +REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) + + +def unique_marker() -> str: + """A short unique token per call/run, so concurrent runs and the shared + response cache never collide on prompts, tags, or customer ids.""" + return uuid.uuid4().hex[:12] diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py new file mode 100644 index 00000000000..b700145d434 --- /dev/null +++ b/tests/e2e/e2e_gateway.py @@ -0,0 +1,211 @@ +"""Gateway: the shared proxy operations, DI'd into every client (composition). + +A frozen-slots dataclass holding a Transport plus poll config. Clients hold a +Gateway and add their own route methods; the lifecycle ResourceManager uses the +Gateway's key/customer methods for cleanup. Read-backs are eventually consistent +(proxy_batch_write_at ~60s) so they poll to a deadline. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass + +from e2e_http import ( + NoBody, + ProbeResult, + Result, + StreamingResponse, + Success, + unwrap, +) +from models import ( + ChatBody, + ChatResponse, + CustomerDeleteBody, + EmbedBody, + EmbedResponse, + KeyDeleteBody, + KeyGenerateBody, + KeyGenerateResponse, + KeyInfo, + KeyInfoParams, + KeyInfoResponse, + ModelInfoEntry, + ModelInfoResponse, + SpendLogRow, + SpendLogs, + SpendLogsParams, +) +from e2e_config import ( + CONTROL_PLANE_BASE_URL, + MASTER_KEY, + POLL_INTERVAL, + POLL_TIMEOUT, + PROXY_BASE_URL, + REQUEST_TIMEOUT, +) +from transport import HttpTransport, SplitTransport, Transport + +RowsPredicate = Callable[[list[SpendLogRow]], bool] + + +@dataclass(frozen=True, slots=True) +class Gateway: + transport: Transport + poll_timeout: float = 120.0 + poll_interval: float = 5.0 + + # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- + + def generate_key(self, body: KeyGenerateBody) -> str: + return unwrap( + self.transport.post( + "/key/generate", + headers=self.transport.master, + json=body, + response_type=KeyGenerateResponse, + ) + ).key + + def delete_key(self, key: str) -> None: + _ = self.transport.post( + "/key/delete", + headers=self.transport.master, + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, + ) + + def delete_customers(self, user_ids: list[str]) -> None: + if not user_ids: + return + _ = self.transport.post( + "/customer/delete", + headers=self.transport.master, + json=CustomerDeleteBody(user_ids=user_ids), + response_type=NoBody, + ) + + def key_info(self, key: str) -> KeyInfo: + return unwrap( + self.transport.get( + "/key/info", + headers=self.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ).info + + def model_info(self) -> list[ModelInfoEntry]: + """Every configured deployment with the price the proxy resolved for it + (config override merged over cost-map defaults).""" + return unwrap( + self.transport.get( + "/model/info", + headers=self.transport.master, + params=NoBody(), + response_type=ModelInfoResponse, + ) + ).data + + # ---- LLM calls ------------------------------------------------------ + + def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]: + return self.transport.post( + "/chat/completions", + headers=self.transport.bearer(key), + json=body, + response_type=ChatResponse, + ) + + def chat_stream(self, key: str, body: ChatBody) -> StreamingResponse: + return self.transport.stream( + "/chat/completions", headers=self.transport.bearer(key), json=body + ) + + def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]: + return self.transport.post( + "/embeddings", + headers=self.transport.bearer(key), + json=body, + response_type=EmbedResponse, + ) + + # ---- spend read-back ------------------------------------------------ + + def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]: + result = self.transport.get( + "/spend/logs", + headers=self.transport.master, + params=params, + response_type=SpendLogs, + ) + match result: + case Success(data=logs): + return logs.root + case _: + return [] + + def poll_logs_for_key( + self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None + ) -> list[SpendLogRow]: + return self._poll( + lambda: self.spend_logs(SpendLogsParams(api_key=key)), min_rows, predicate + ) + + def poll_logs_for_request_id( + self, + request_id: str, + *, + min_rows: int = 1, + predicate: RowsPredicate | None = None, + ) -> list[SpendLogRow]: + return self._poll( + lambda: self.spend_logs(SpendLogsParams(request_id=request_id)), + min_rows, + predicate, + ) + + def _poll( + self, + fetch: Callable[[], list[SpendLogRow]], + min_rows: int, + predicate: RowsPredicate | None, + ) -> list[SpendLogRow]: + deadline = time.monotonic() + self.poll_timeout + rows: list[SpendLogRow] = [] + while time.monotonic() < deadline: + rows = fetch() + if len(rows) >= min_rows and (predicate is None or predicate(rows)): + return rows + time.sleep(self.poll_interval) + return rows + + # ---- route probe ---------------------------------------------------- + + def probe(self, path: str, *, params: NoBody) -> ProbeResult: + return self.transport.probe(path, params=params) + + +def build_gateway() -> Gateway: + """The Gateway every suite's client is built from: a SplitTransport that routes + LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the + control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two + base URLs are the same for a monolithic proxy, so routing is then a no-op.""" + return Gateway( + transport=SplitTransport( + data=HttpTransport( + base_url=PROXY_BASE_URL, + master_key=MASTER_KEY, + request_timeout=REQUEST_TIMEOUT, + ), + control=HttpTransport( + base_url=CONTROL_PLANE_BASE_URL, + master_key=MASTER_KEY, + request_timeout=REQUEST_TIMEOUT, + ), + ), + poll_timeout=POLL_TIMEOUT, + poll_interval=POLL_INTERVAL, + ) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py new file mode 100644 index 00000000000..7458f316852 --- /dev/null +++ b/tests/e2e/e2e_http.py @@ -0,0 +1,306 @@ +"""The ONLY module permitted to call ``requests.*``. + +Enforced by tests/code_coverage_tests/check_e2e_no_raw_requests.py. Every request +body / query / header / response is a pydantic model; outcomes are a tagged union +(``Result[R]``) so callers ``match`` on them instead of catching exceptions. + +Named e2e_http (not http) so it does not shadow the stdlib ``http`` package that +requests itself imports. +""" + +from __future__ import annotations + +from typing import Generic, Iterator, Literal, NewType, TypeVar, cast + +import pytest +import requests +from pydantic import BaseModel, ConfigDict, Field + +URL = NewType("URL", str) + + +class Headers(BaseModel): + """Base for header models. Subclasses may alias to hyphenated header names + (e.g. ``x-litellm-api-key``); serialization uses by_alias.""" + + model_config = ConfigDict(populate_by_name=True) + + +class AuthHeaders(Headers): + # litellm accepts either; set whichever the call needs, leave the other None. + authorization: str | None = None + x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key") + + +class NoBody(BaseModel): + """Empty body/query for routes that take none.""" + + +# ---------- Result types ---------- + +R = TypeVar("R", bound=BaseModel) + + +class Success(BaseModel, Generic[R]): + kind: Literal["success"] = "success" + data: R + + +class NetworkError(BaseModel): + kind: Literal["network"] = "network" + message: str + + +class UnauthorizedError(BaseModel): + kind: Literal["unauthorized"] = "unauthorized" + + +class RateLimitedError(BaseModel): + kind: Literal["rate_limited"] = "rate_limited" + retry_after_seconds: int | None = None + # litellm overloads 429 for budget_exceeded too, so keep the body to tell them apart. + body: str = "" + + +class ValidationError(BaseModel): + kind: Literal["validation"] = "validation" + message: str + + +class UnknownApiError(BaseModel): + kind: Literal["unknown"] = "unknown" + status_code: int + body: str + + +type Result[R: BaseModel] = ( + Success[R] + | NetworkError + | UnauthorizedError + | RateLimitedError + | ValidationError + | UnknownApiError +) + + +class ProbeResult(BaseModel): + """A route's reachability: status + body, no schema validation. Healthy == + route exists (not 404) and the handler did not crash (not 5xx).""" + + status_code: int + body: str + + @property + def healthy(self) -> bool: + return 200 <= self.status_code < 500 and self.status_code != 404 + + +class StreamingResponse(BaseModel): + """Raw outcome for calls whose body is provider-native or streamed: status, the + x-litellm-call-id header (== SpendLogs.request_id), the content-type (which + tells streaming `text/event-stream` from non-streaming `application/json`), and + the body. Used by passthrough and streaming, where one validated JSON model + does not fit.""" + + status_code: int + call_id: str | None = None # x-litellm-call-id header + content_type: str | None = None + body: str + chunks: int = 0 # streamed events (0 for non-streaming) + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + @property + def is_streaming(self) -> bool: + return "text/event-stream" in (self.content_type or "") + + +def _hdr(resp: requests.Response, name: str) -> str | None: + value = resp.headers.get(name) + return value if isinstance(value, str) else None + + +def unwrap[R: BaseModel](result: Result[R]) -> R: + match result: + case Success(data=data): + return data + case _: + raise AssertionError(result) + + +def is_ok[R: BaseModel](result: Result[R]) -> bool: + match result: + case Success(): + return True + case _: + return False + + +def require_successful_call(result: StreamingResponse) -> None: + """A call that should have succeeded but didn't is a hard failure, never a skip: + if the proxy can't make a call it's expected to, the test must fail.""" + if result.ok: + return + pytest.fail( + f"upstream call failed (status {result.status_code}); body={result.body[:300]}" + ) + + +def _headers(headers: BaseModel) -> dict[str, str]: + dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) + return {key: str(value) for key, value in dumped.items()} + + +def _params(params: BaseModel | None) -> dict[str, str]: + if params is None: + return {} + dumped: dict[str, object] = params.model_dump(by_alias=True, exclude_none=True) + return {key: str(value) for key, value in dumped.items()} + + +def _classify[R: BaseModel]( + resp: requests.Response, response_type: type[R] +) -> Result[R]: + if resp.status_code == 401: + return UnauthorizedError() + if resp.status_code == 429: + return RateLimitedError(body=resp.text) + if not resp.ok: + return UnknownApiError(status_code=resp.status_code, body=resp.text) + try: + return Success(data=response_type.model_validate(resp.json())) + except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value + return ValidationError(message=str(exc)) + + +def post[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def get[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def delete[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.delete( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + +def probe( + url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0 +) -> ProbeResult: + try: + resp = requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return ProbeResult(status_code=-1, body=str(exc)) + return ProbeResult(status_code=resp.status_code, body=resp.text) + + +def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingResponse: + call_id = _hdr(resp, "x-litellm-call-id") + content_type = _hdr(resp, "content-type") + if not stream or not (200 <= resp.status_code < 300): + return StreamingResponse( + status_code=resp.status_code, + call_id=call_id, + content_type=content_type, + body=resp.text, + ) + lines = cast("Iterator[bytes]", resp.iter_lines()) + chunks = sum(1 for line in lines if line) + return StreamingResponse( + status_code=resp.status_code, + call_id=call_id, + content_type=content_type, + body="", + chunks=chunks, + ) + + +def send( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + timeout: float = 60.0, +) -> StreamingResponse: + """Raw POST returning the unparsed HTTP outcome: status, full body, and the + x-litellm-call-id header. For native/passthrough bodies and for calls judged by + status rather than a typed JSON model (e.g. a budget block is a non-2xx). With + ``stream=True`` the SSE body is consumed and its events counted instead.""" + try: + resp = requests.post( + str(url), + headers=_headers(headers), + params=_params(params), + json=json.model_dump(by_alias=True, exclude_none=True), + stream=stream, + timeout=timeout, + ) + except requests.RequestException as exc: + return StreamingResponse(status_code=-1, body=str(exc)) + return _streaming_outcome(resp, stream) + + +def stream( + url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0 +) -> StreamingResponse: + """Streaming (SSE) call: consumes the stream counting events, and captures the + x-litellm-call-id + content-type headers. Body is elided.""" + return send(url, headers=headers, json=json, stream=True, timeout=timeout) diff --git a/tests/e2e/gateway/litellm-config.yml b/tests/e2e/gateway/litellm-config.yml new file mode 100644 index 00000000000..f4ca48cfee0 --- /dev/null +++ b/tests/e2e/gateway/litellm-config.yml @@ -0,0 +1,170 @@ +# This default config file aims to support most popular model providers out of the box + +#In general, the model name used by the client will be the same as the ones from the provider (For example, you will use "anthropic.claude-3-5-sonnet-20240620-v1:0" when you're calling LiteLLM just like you would when calling Amazon Bedrock directly) +#In the case where there are model name conflicts, a prefix will be used (For example, the Azure and the openAI model names conflict, so when you are using Azure, you will use "azure/gpt-4o-realtime-preview-2024-10-01") + +#Some model providers require additional user-specific configuration (such as Azure which requires you to specify your own api_base with your resource name, and your api_version). +#In this case, the provider is commented out, and you should uncomment it and provide your specific info + +#For more detailed information about each provider, refer to the docs: https://docs.litellm.ai/docs/providers + +#If you are not interested in a particular provider, just remove it from your config.yaml, and redeploy, and it will no longer show up in your LiteLLM deployment + +#If a particular provider is not working, double check your .env file, and make sure you have provided a valid api key for that provider, and then redeploy + +#Full details on guardrails here: https://docs.litellm.ai/docs/proxy/guardrails/bedrock +general_settings: + store_prompts_in_spend_logs: true + master_key: os.environ/LITELLM_MASTER_KEY + proxy_batch_write_at: 60 + database_connection_pool_limit: 10 + # disable_error_logs: True + forward_client_headers_to_llm_api: false + maximum_spend_logs_retention_period: "60d" # GSE-13389: Cleanup logs older than 60 days + maximum_spend_logs_cleanup_cron: "0 1 * * *" # 01:00 UTC daily = 18:00 PDT + database_url: os.environ/DATABASE_URL + control_plane_url: os.environ/CONTROL_PLANE_URL + alerts: ["email"] + proxy_budget_rescheduler_min_time: 15 + proxy_budget_rescheduler_max_time: 20 + +# fallbacks: [{"gpt-4": ["anthropic.claude-3-5-sonnet-20240620-v1:0"]}] #Configure fallbacks for context window exeeded errors (In this example, we will fall back to Claude Sonnet if over 8000 tokens, which is gpt-4's limit) + # default_fallbacks: ["anthropic.claude-3-haiku-20240307-v1:0"] #Configure fallbacks for any error for every model (the above fallback configurations override this one) +# environment_variables: +# STORE_MODEL_IN_DB: 'True' +# LITELLM_LOG: "DEBUG" +litellm_settings: + drop_params: True + # Spend counters inherit this as their Redis TTL, so an idle counter goes cold and + # the next request reseeds it from the DB; kept short to exercise the cross-pod + # reseed path in test_spend_counter_reseed_e2e. Response-cache writes pass their own + # ttl and are unaffected. + default_redis_ttl: 20 + request_timeout: 600 + num_retries: 3 + json_logs: true + store_audit_logs: True + cache: true + cache_params: + type: redis + host: redis + port: 6379 + password: os.environ/REDIS_PASSWORD + namespace: litellm.caching + ttl: 16600 + # max_budget: 1000000000.0 # (float) sets max budget in dollars across the entire proxy across all API keys. Note, the budget does not apply to the master key. That is the only exception. + # budget_duration: 1mo # (str) frequency of budget reset - You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + # max_internal_user_budget: 1000000000.0 # (float) sets default budget in dollars for each internal user. (Doesn't apply to Admins. Doesn't apply to Teams. Doesn't apply to master key) + # internal_user_budget_duration: "1mo" # (str) frequency of budget reset - You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + # success_callback: ["s3_v2"] + # failure_callback: ["s3_v2"] + # service_callback: ["datadog"] + callbacks: ["arize_phoenix", "datadog", "smtp_email", "prometheus", "otel"] + require_auth_for_metrics_endpoint: false + #type: redis-semantic + #similarity_threshold: 0.8 # similarity threshold for semantic cache + #redis_semantic_cache_embedding_model: text-embedding-ada-002 # only works with text-embedding-ada-002 for now... https://github.com/BerriAI/litellm/issues/4001 + +router_settings: + routing_strategy: simple-shuffle + num_retries: 3 + allowed_fails: 5 + cooldown_time: 30 + # When gemini deployments are exhausted (provider 429 / auth), cross over to + # working models. Exercised by tests/e2e/router/test_rate_limiter.py. + fallbacks: + - gemini-2.5-flash: ["gpt-5.5", "claude-haiku-4-5"] + +#ttl: Optional[float] +#default_in_memory_ttl: Optional[float] +#default_in_redis_ttl: Optional[float] + +model_list: + - model_name: gpt-5.5 + litellm_params: + model: openai/gpt-5.5 + api_key: os.environ/OPENAI_API_KEY + + - model_name: claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + + # Same underlying model via Vertex AI — distinct routing/auth path + # # (service-account JSON), so it gets its own model_name. + - model_name: gemini-2.5-flash-vertex + litellm_params: + model: vertex_ai/gemini-2.5-flash + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + vertex_credentials: os.environ/VERTEXAI_CREDENTIALS + + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + + # load balancing to a different deployment, if gemini gets rate limited. + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + + # Custom per-token pricing exercised by llm_translation/test_custom_pricing_e2e.py. + # Rates deliberately exceed canonical gemini-2.5-flash (input 3e-7 / output 2.5e-6) + # so an override that is ignored or under-applied reports spend at the base rate + # and fails that test. The test reads these same rates back from this file. + - model_name: custom-priced-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + input_cost_per_token: 0.00005 + output_cost_per_token: 0.0001 + + # embedding models + - model_name: openai-text-embedding-3-small + litellm_params: + model: openai/text-embedding-3-small + api_key: os.environ/OPENAI_API_KEY + + - model_name: gemini-2-embedding + litellm_params: + model: gemini/gemini-2-embedding + api_key: os.environ/GEMINI_API_KEY + + # realtime models + - model_name: openai-realtime + litellm_params: + model: openai/realtime-2 + api_key: os.environ/OPENAI_API_KEY + model_info: + mode: realtime + + +mcp_servers: + deepwiki_mcp: + url: "https://mcp.deepwiki.com/mcp" + auth_type: none + description: "just a test" + + atlassian: + url: "https://mcp.atlassian.com/v1/mcp" + auth_type: oauth2 + authorization_url: https://auth.atlassian.com/authorize + + +guardrails: + - guardrail_name: "presidio-pii" + litellm_params: + guardrail: presidio + mode: pre_call + presidio_analyzer_api_base: os.environ/PRESIDIO_ANALYZER_API_BASE + presidio_anonymizer_api_base: os.environ/PRESIDIO_ANONYMIZER_API_BASE + default_on: false + pii_entities_config: + EMAIL_ADDRESS: BLOCK + CREDIT_CARD: BLOCK + US_SSN: BLOCK + PHONE_NUMBER: BLOCK + + diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py new file mode 100644 index 00000000000..fdf2137584e --- /dev/null +++ b/tests/e2e/lifecycle.py @@ -0,0 +1,117 @@ +"""Lifecycle contract and resource cleanup for stateful e2e tests. + +Shared by every e2e suite under tests/e2e/. The proxy under test is +long-lived and never reset between tests, so anything a test creates (keys, +customers, teams, orgs, users, guardrails, budgets, ...) persists unless +explicitly deleted. Every check follows an init -> run -> teardown lifecycle; +teardown releases each resource init() created, even when run() raises. + +In pytest terms (see conftest.py): the `resources` fixture's setup is init(), +the test body is run(), and the fixture's teardown is teardown(). +""" + +from dataclasses import dataclass, field +from typing import Callable, List, Protocol, runtime_checkable + +from e2e_gateway import Gateway +from models import KeyGenerateBody + + +@runtime_checkable +class E2ECase(Protocol): + """A stateful e2e check run against a long-lived proxy. + + init() acquires resources, run() exercises behaviour and asserts, teardown() + releases everything init() created. teardown() must run even if init() fails + partway or run() raises. + """ + + def init(self) -> None: ... + + def run(self) -> None: ... + + def teardown(self) -> None: ... + + +def run_case(case: E2ECase) -> None: + """Drive a case through its lifecycle: init -> run -> teardown. + + teardown always runs - even when init() fails partway or run() raises (or + skips) - so resources the case already registered on the long-lived proxy are + released. init() is inside the try because cases register cleanups + progressively (e.g. create team, then user, then key), and a failure after + the first creation must still release what came before. + """ + try: + case.init() + case.run() + finally: + case.teardown() + + +@runtime_checkable +class ResourceClient(Protocol): + """Proxy operations the convenience creators use. Resource types without a + creator here are handled generically via ResourceManager.defer(). The Gateway + satisfies this.""" + + def generate_key(self, body: KeyGenerateBody) -> str: ... + + def delete_key(self, key: str) -> None: ... + + def delete_customers(self, user_ids: List[str]) -> None: ... + + +@runtime_checkable +class GatewayProvider(Protocol): + """Every suite's client exposes the shared Gateway, which the resources fixture + uses for cleanup. The client adds its own route methods on top.""" + + @property + def gateway(self) -> Gateway: ... + + +@dataclass +class ResourceManager: + """Registry of teardown actions for resources a test creates on the stateful + proxy. + + Not limited to any resource type: register a cleanup with ``defer()`` for a + key, customer, team, org, user, guardrail, budget, MCP server - anything with + a delete. The two most common resources have sugar (``key``, ``customer``); + everything else is ``resources.defer(lambda: client.delete_team(team_id))``. + + Cleanups run LIFO (so a resource is removed before whatever it depends on) and + best-effort (one failing cleanup never blocks the rest). + """ + + client: ResourceClient + _cleanups: List[Callable[[], None]] = field( + default_factory=list + ) # mutable-ok: append-only teardown registry + + def init(self) -> None: + """No global setup needed today; present for lifecycle symmetry.""" + return None + + def defer(self, cleanup: Callable[[], None]) -> None: + """Register a teardown action for any resource the test just created.""" + self._cleanups.append(cleanup) + + def key(self) -> str: + """Create an all-models virtual key; delete it on teardown.""" + key = self.client.generate_key(KeyGenerateBody(models=[])) + self.defer(lambda: self.client.delete_key(key)) + return key + + def customer(self, customer_id: str) -> str: + """Track an end-user id (from the `user` param); delete it on teardown.""" + self.defer(lambda: self.client.delete_customers([customer_id])) + return customer_id + + def teardown(self) -> None: + for cleanup in reversed(self._cleanups): + try: + cleanup() + except Exception: + pass # best-effort: a failed cleanup must not block the rest diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..5e4a448857f --- /dev/null +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -0,0 +1,84 @@ +# LLM Translation Test Coverage Matrix + +Scope: the proxy's two translation surfaces, end to end against a live proxy. + +1. **Passthrough** - the client speaks the provider's NATIVE API (Gemini + `generateContent`, Anthropic `/v1/messages`); the proxy forwards it and still + logs a costed `SpendLogs` row (`call_type="pass_through_endpoint"`). Routes: + `/gemini`, `/anthropic`, `/vertex_ai`, `/openai`, `/bedrock`, `/cohere`, + `/mistral`, `/vllm`. +2. **Non-passthrough** - the client speaks OpenAI format + (`/chat/completions`, `/embeddings`); litellm translates to/from the provider. + +The two axes that must work in production for each: **passthrough vs +non-passthrough** and **streaming vs non-streaming**, with **cost logged** and +**tool calls** working in every cell. + +Companion: live suite `test_passthrough_e2e.py` (this directory). The +non-passthrough chat/embedding cells are exercised by `../spend_tracking/`. + +Levels: `live` real provider + proxy + SpendLogs row; `unit` mocked. +Status: `covered` / `partial` / `gap`. + +--- + +## Passthrough endpoints (native provider format) + +| Provider | Non-streaming | Streaming | Tool calls | Cost logged | Status | +|----------|---------------|-----------|------------|-------------|--------| +| Gemini (`/gemini/v1beta/models/{m}:generateContent` / `:streamGenerateContent`) | live | live | live | live | **covered** | +| Anthropic (`/anthropic/v1/messages`) | live | live | live | live | **covered** | +| Vertex AI (`/vertex_ai/...`) | - | - | - | - | gap (gcloud auth) | +| OpenAI / Bedrock / Cohere / Mistral / VLLM | - | - | - | - | gap | + +Each covered cell asserts: `call_type == "pass_through_endpoint"`, `spend > 0`, +`status == "success"`, correct `custom_llm_provider`/`model`, row correlated by the +`x-litellm-call-id` header. Gemini non-streaming also pins `request_tags` +propagation; streaming pins `chunks > 0` then a costed row; tool tests assert the +provider emitted a tool call (`functionCall` / `tool_use`) and it was costed. + +Cost on passthrough is computed in the success handler by transforming the native +response to a `ModelResponse` and calling `litellm.completion_cost()`; for +streaming, chunks are buffered and costed after the stream ends. This is the path +most likely to silently break and the one a mock can't prove works. + +## Non-passthrough endpoints (OpenAI-compatible translation) + +| Modality | Non-streaming | Streaming | Tool calls | Cost logged | Status | +|----------|---------------|-----------|------------|-------------|--------| +| Chat | live (spend suite) | live (spend suite) | gap | live | partial | +| Embeddings | live (spend suite) | n/a | n/a | live | covered | +| Responses / image / audio / rerank / realtime | - | - | - | - | gap | + +## This suite's files + +| Test | Cell | +|------|------| +| `test_gemini_passthrough_nonstreaming_logs_cost` | gemini native, non-stream, cost + tags | +| `test_gemini_passthrough_streaming_logs_cost` | gemini native, stream, cost | +| `test_gemini_passthrough_tool_call_logs_cost` | gemini native, tool call, cost | +| `test_anthropic_passthrough_nonstreaming_logs_cost` | anthropic native, non-stream, cost | +| `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | +| `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | + +## Gaps + +- Vertex / OpenAI / Bedrock / Cohere passthrough (same shape; add once the + provider credential is configured; Vertex is closest - route exists, auth stale). +- Non-passthrough tool calls over `/chat/completions` end to end with cost. +- Image / audio / rerank / responses / realtime translation + cost. +- Streaming cost-injection (`include_cost_in_streaming_usage`); passthrough on + client disconnect (partial-usage logging). + +## Adding a provider/modality + +Extend `PassthroughClient` with the native call (it inherits keys, cleanup, and +SpendLogs polling from `ProxyClient`), then add a test that calls it, +`require_successful_call(result)`, and `_costed_row(...)`. + +## Timing + +Passthrough spend is logged asynchronously after the response and lands on the +`proxy_batch_write_at` (~60s) cycle, so cost assertions poll +`/spend/logs?request_id=` to a deadline. Streaming cost is only +known after the stream is fully consumed. diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py new file mode 100644 index 00000000000..fbf008cf085 --- /dev/null +++ b/tests/e2e/llm_translation/conftest.py @@ -0,0 +1,15 @@ +"""LLM-translation suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared +Gateway, so the `resources` fixture cleans up keys this suite creates. +""" + +import pytest + +from passthrough_client import PassthroughClient, build_client + + +@pytest.fixture(scope="session") +def client() -> PassthroughClient: + return build_client() diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py new file mode 100644 index 00000000000..fff4064a328 --- /dev/null +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -0,0 +1,163 @@ +"""Client for LLM-translation e2e tests over the proxy's passthrough endpoints. + +A passthrough request is sent in the PROVIDER's native format (Gemini +generateContent, Anthropic /v1/messages) to the proxy, which forwards it to the +provider and still logs a SpendLogs row (call_type="pass_through_endpoint"). The +litellm virtual key is passed as the provider key; the proxy swaps in the real env +credential. SpendLogs.request_id == the x-litellm-call-id response header. The +native request models are co-located here because only this suite uses them. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic import BaseModel, Field + +from e2e_gateway import Gateway, build_gateway +from e2e_http import Headers, StreamingResponse +from models import ChatMessage + + +class JsonSchemaProperty(BaseModel): + type: str + + +class JsonSchema(BaseModel): + type: str + properties: dict[str, JsonSchemaProperty] + required: list[str] + + +class GeminiHeaders(Headers): + x_goog_api_key: str = Field(serialization_alias="x-goog-api-key") + content_type: str = Field( + default="application/json", serialization_alias="Content-Type" + ) + tags: str | None = None + + +class AnthropicHeaders(Headers): + x_api_key: str = Field(serialization_alias="x-api-key") + anthropic_version: str = Field( + default="2023-06-01", serialization_alias="anthropic-version" + ) + content_type: str = Field( + default="application/json", serialization_alias="Content-Type" + ) + tags: str | None = None + + +class AltSseParams(BaseModel): + alt: str = "sse" + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + role: str = "user" + parts: list[GeminiPart] + + +class GeminiFunctionDeclaration(BaseModel): + name: str + description: str + parameters: JsonSchema + + +class GeminiTool(BaseModel): + function_declarations: list[GeminiFunctionDeclaration] = Field( + serialization_alias="functionDeclarations" + ) + + +class GeminiGenerateBody(BaseModel): + contents: list[GeminiContent] + tools: list[GeminiTool] | None = None + + +class AnthropicTool(BaseModel): + name: str + description: str + input_schema: JsonSchema + + +class AnthropicMessageBody(BaseModel): + model: str + max_tokens: int + messages: list[ChatMessage] + tools: list[AnthropicTool] | None = None + stream: bool = False + + +def _tags_header(tags: list[str] | None) -> str | None: + return ",".join(tags) if tags else None + + +@dataclass(frozen=True, slots=True) +class PassthroughClient: + gateway: Gateway + + # ---- Gemini native passthrough (/gemini/v1beta/...) ----------------- + + def gemini_generate( + self, + key: str, + model: str, + text: str, + *, + tools: list[GeminiTool] | None = None, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + f"/gemini/v1beta/models/{model}:generateContent", + headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=text)])], tools=tools + ), + ) + + def gemini_stream( + self, key: str, model: str, text: str, *, tags: list[str] | None = None + ) -> StreamingResponse: + return self.gateway.transport.send( + f"/gemini/v1beta/models/{model}:streamGenerateContent", + headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), + json=GeminiGenerateBody( + contents=[GeminiContent(parts=[GeminiPart(text=text)])] + ), + params=AltSseParams(), + stream=True, + ) + + # ---- Anthropic native passthrough (/anthropic/v1/messages) ---------- + + def anthropic_message( + self, + key: str, + model: str, + text: str, + *, + max_tokens: int = 64, + tools: list[AnthropicTool] | None = None, + stream: bool = False, + tags: list[str] | None = None, + ) -> StreamingResponse: + return self.gateway.transport.send( + "/anthropic/v1/messages", + headers=AnthropicHeaders(x_api_key=key, tags=_tags_header(tags)), + json=AnthropicMessageBody( + model=model, + max_tokens=max_tokens, + messages=[ChatMessage(role="user", content=text)], + tools=tools, + stream=stream, + ), + stream=stream, + ) + + +def build_client() -> PassthroughClient: + return PassthroughClient(gateway=build_gateway()) diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py new file mode 100644 index 00000000000..4b3e87b78e1 --- /dev/null +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -0,0 +1,219 @@ +"""Live e2e: a model's custom per-token pricing is loaded, billed, and isolated. + +The gateway config declares ``custom-priced-flash`` (gemini-2.5-flash underneath) +with input/output rates deliberately far above the canonical gemini price, read +back here from the same config file. Three behaviors are checked independently: + +- billing: a real call's logged cost breakdown charges input and output tokens at + the custom rates, each component checked separately (a base-rate bill lands + ~100x lower; a swapped input/output rate passes a total-only check but not this) +- reporting: /model/info surfaces those rates for the model +- isolation: gemini-2.5-flash shares the same underlying gemini/gemini-2.5-flash + but sets no override, so it must keep its own price; an override that leaks into + the shared cost map misprices it. This fails on a real proxy gap today, so it is + marked xfail(strict=True): the suite stays green while the leak persists and + flips to a failure the moment isolation is fixed and the marker should be removed. +""" + +import time +from dataclasses import dataclass +from pathlib import Path + +import pytest +import yaml +from pydantic import BaseModel, RootModel + +from e2e_config import unique_marker +from e2e_http import Success, unwrap +from models import ChatBody, ChatMessage, CustomPricing, ModelInfoEntry, SpendLogsParams +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +CUSTOM_MODEL = "custom-priced-flash" +BASE_MODEL = "gemini-2.5-flash" +CONFIG_PATH = Path(__file__).resolve().parents[1] / "gateway" / "litellm-config.yml" + + +@dataclass(frozen=True, slots=True) +class _Rates: + input_per_token: float + output_per_token: float + + +class _ConfiguredModel(BaseModel): + model_name: str + litellm_params: CustomPricing + + +class _GatewayConfig(BaseModel): + model_list: list[_ConfiguredModel] + + +class _CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + +class _RowMetadata(BaseModel): + cost_breakdown: _CostBreakdown | None = None + + +class _SpendRow(BaseModel): + request_id: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + metadata: _RowMetadata | None = None + + +class _SpendRows(RootModel[list[_SpendRow]]): + pass + + +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +def _configured_pricing(model_name: str) -> _Rates: + """The custom rates declared for `model_name` in the gateway config the proxy + runs with - the source of truth the billed and reported prices are checked + against.""" + config = _GatewayConfig.model_validate(yaml.safe_load(CONFIG_PATH.read_text())) + for entry in config.model_list: + if entry.model_name == model_name: + pricing = entry.litellm_params + assert pricing.input_cost_per_token and pricing.output_cost_per_token, ( + f"{model_name} declares no custom per-token rates in {CONFIG_PATH.name}" + ) + return _Rates(pricing.input_cost_per_token, pricing.output_cost_per_token) + pytest.fail(f"{model_name} not found in {CONFIG_PATH.name}") + + +def _model_info_entry( + entries: list[ModelInfoEntry], model_name: str +) -> ModelInfoEntry: + for entry in entries: + if entry.model_name == model_name: + return entry + pytest.fail(f"{model_name} absent from /model/info; the override did not load") + + +def _poll_breakdown_row( + client: PassthroughClient, key: str, response_id: str | None +) -> _SpendRow: + """Poll /spend/logs until the call's row lands with a cost breakdown (rows + flush ~60s behind the call via proxy_batch_write_at).""" + deadline = time.monotonic() + client.gateway.poll_timeout + while time.monotonic() < deadline: + result = client.gateway.transport.get( + "/spend/logs", + headers=client.gateway.transport.master, + params=SpendLogsParams(api_key=key), + response_type=_SpendRows, + ) + match result: + case Success(data=data): + rows = data.root + case _: + rows = [] + priced = [ + row + for row in rows + if row.metadata + and row.metadata.cost_breakdown + and row.metadata.cost_breakdown.input_cost is not None + ] + for row in priced: + if response_id and row.request_id == response_id: + return row + if priced and response_id is None: + return priced[0] + time.sleep(client.gateway.poll_interval) + pytest.fail("no spend row with a cost breakdown landed before the deadline") + + +def test_custom_pricing_is_billed_at_configured_rate( + client: PassthroughClient, scoped_key: str +) -> None: + rates = _configured_pricing(CUSTOM_MODEL) + + chat = unwrap( + client.gateway.chat( + scoped_key, + ChatBody( + model=CUSTOM_MODEL, + messages=[ + ChatMessage( + role="user", content=f"reply with one word {unique_marker()}" + ) + ], + max_tokens=16, + ), + ) + ) + + row = _poll_breakdown_row(client, scoped_key, chat.id) + assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll + breakdown = row.metadata.cost_breakdown + + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + assert prompt > 0 and completion > 0, f"call tokens not logged on the row: {row}" + + input_cost = breakdown.input_cost + output_cost = breakdown.output_cost + assert input_cost is not None and output_cost is not None, ( + f"row cost breakdown missing input/output cost: {breakdown}" + ) + assert _approx_equal(input_cost, prompt * rates.input_per_token), ( + f"input_cost {input_cost} != {prompt} tokens * {rates.input_per_token} " + f"= {prompt * rates.input_per_token}" + ) + assert _approx_equal(output_cost, completion * rates.output_per_token), ( + f"output_cost {output_cost} != {completion} tokens * {rates.output_per_token} " + f"= {completion * rates.output_per_token}" + ) + + +def test_model_info_reports_custom_pricing(client: PassthroughClient) -> None: + rates = _configured_pricing(CUSTOM_MODEL) + entry = _model_info_entry(client.gateway.model_info(), CUSTOM_MODEL) + + assert entry.litellm_params.input_cost_per_token == rates.input_per_token, ( + f"/model/info litellm_params input rate " + f"{entry.litellm_params.input_cost_per_token} != configured " + f"{rates.input_per_token}" + ) + assert entry.litellm_params.output_cost_per_token == rates.output_per_token, ( + f"/model/info litellm_params output rate " + f"{entry.litellm_params.output_cost_per_token} != configured " + f"{rates.output_per_token}" + ) + + +@pytest.mark.xfail( + strict=True, + reason="known proxy bug: a deployment's custom per-token pricing leaks into the " + "shared cost map for sibling deployments of the same underlying model; remove " + "this marker once isolation is fixed", +) +def test_custom_pricing_is_isolated_from_sibling_deployment( + client: PassthroughClient, +) -> None: + entries = {entry.model_name: entry for entry in client.gateway.model_info()} + custom = entries.get(CUSTOM_MODEL) + base = entries.get(BASE_MODEL) + assert custom is not None, f"{CUSTOM_MODEL} absent from /model/info" + assert base is not None, f"{BASE_MODEL} absent from /model/info" + + # custom-priced-flash overrides pricing; gemini-2.5-flash shares the same + # underlying gemini/gemini-2.5-flash but sets no override, so it must keep its + # own price. Equal rates mean the override leaked into the shared cost map. + assert ( + base.model_info.input_cost_per_token != custom.model_info.input_cost_per_token + ), ( + f"{BASE_MODEL} input rate {base.model_info.input_cost_per_token} matches " + f"{CUSTOM_MODEL}'s override {custom.model_info.input_cost_per_token}; " + f"per-deployment custom pricing is not isolated" + ) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py new file mode 100644 index 00000000000..37d55c665b3 --- /dev/null +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -0,0 +1,159 @@ +"""Live e2e for LLM-translation passthrough endpoints. + +Each test sends a NATIVE provider request through the proxy's passthrough route +and verifies the proxy still logged a costed SpendLogs row +(call_type="pass_through_endpoint"), correlated by the x-litellm-call-id header. + +Covered: gemini ("gemini-2.5-flash") + anthropic ("claude-haiku-4-5"), streaming + +non-streaming, plus native tool calls. See LLM_TRANSLATION_COVERAGE_MATRIX.md. + +A passthrough call returning non-2xx fails hard (never a skip); once it returns +2xx, a missing or zero-cost SpendLogs row fails too. +""" + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from models import SpendLogRow +from passthrough_client import ( + AnthropicTool, + GeminiFunctionDeclaration, + GeminiTool, + JsonSchema, + JsonSchemaProperty, + PassthroughClient, +) + +pytestmark = pytest.mark.e2e + + +def _fetch_cost_breakdown(client: PassthroughClient, result: StreamingResponse) -> SpendLogRow: + """The passthrough call's logged row, polled until it carries a cost. + + Asserts (not skips) that a 2xx passthrough call produced a costed row - the + whole point of passthrough spend tracking. + """ + assert result.call_id, "passthrough response had no x-litellm-call-id header" + rows = client.gateway.poll_logs_for_request_id( + result.call_id, + predicate=lambda rs: (rs[0].spend or 0) > 0, + ) + assert rows, f"no SpendLogs row for passthrough call_id {result.call_id}" + row = rows[0] + assert row.call_type == "pass_through_endpoint" + assert (row.spend or 0) > 0, f"passthrough call was not costed: {row}" + assert row.status == "success" + return row + + +# ---- Gemini passthrough ------------------------------------------------ + + +def test_gemini_passthrough_nonstreaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + tag = f"e2e-passthrough-{unique_marker()}" + result = client.gemini_generate( + scoped_key, "gemini-2.5-flash", "Say hello in one word", tags=[tag, "gemini"] + ) + require_successful_call(result) + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + assert "gemini" in (row.model or "") + assert tag in (row.request_tags or []), f"tags not logged: {row.request_tags}" + + +def test_gemini_passthrough_streaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.gemini_stream(scoped_key, "gemini-2.5-flash", "Count to five") + require_successful_call(result) + assert result.chunks > 0, "streaming passthrough produced no events" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + + +def test_gemini_passthrough_tool_call_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.gemini_generate( + scoped_key, + "gemini-2.5-flash", + "What is the weather in Paris? Use the get_weather tool.", + tools=[ + GeminiTool( + function_declarations=[ + GeminiFunctionDeclaration( + name="get_weather", + description="Get the weather for a city", + parameters=JsonSchema( + type="object", + properties={"city": JsonSchemaProperty(type="string")}, + required=["city"], + ), + ) + ] + ) + ], + ) + require_successful_call(result) + assert "functionCall" in result.body, "gemini did not emit a tool call" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "gemini" + + +# ---- Anthropic passthrough --------------------------------------------- + + +def test_anthropic_passthrough_nonstreaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message(scoped_key, "claude-haiku-4-5", "Say hello") + require_successful_call(result) + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" + assert "claude" in (row.model or "") + + +def test_anthropic_passthrough_streaming_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message( + scoped_key, "claude-haiku-4-5", "Count to five", stream=True + ) + require_successful_call(result) + assert result.chunks > 0, "streaming passthrough produced no events" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" + + +def test_anthropic_passthrough_tool_call_logs_cost( + client: PassthroughClient, scoped_key: str +) -> None: + result = client.anthropic_message( + scoped_key, + "claude-haiku-4-5", + "What is the weather in Paris? Use the get_weather tool.", + tools=[ + AnthropicTool( + name="get_weather", + description="Get the weather for a city", + input_schema=JsonSchema( + type="object", + properties={"city": JsonSchemaProperty(type="string")}, + required=["city"], + ), + ) + ], + ) + require_successful_call(result) + assert "tool_use" in result.body, "anthropic did not emit a tool call" + + row = _fetch_cost_breakdown(client, result) + assert row.custom_llm_provider == "anthropic" diff --git a/tests/e2e/models.py b/tests/e2e/models.py new file mode 100644 index 00000000000..fbeb3d44fa5 --- /dev/null +++ b/tests/e2e/models.py @@ -0,0 +1,240 @@ +"""Shared pydantic request/response models for the e2e gateway. + +Only the fields the tests read are modelled; pydantic ignores the rest, so a +response validates without mirroring every proxy field. No untyped dicts. +""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, RootModel + +# ---------- keys ---------- + + +class ModelBudgetEntry(BaseModel): + budget_limit: float + time_period: str + + +class BudgetWindow(BaseModel): + budget_duration: str + max_budget: float + + +class KeyGenerateBody(BaseModel): + models: list[str] = [] + duration: str | None = None + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + user_id: str | None = None + team_id: str | None = None + budget_id: str | None = None + model_max_budget: dict[str, ModelBudgetEntry] | None = None + budget_limits: list[BudgetWindow] | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + + +class KeyGenerateResponse(BaseModel): + key: str + + +class KeyDeleteBody(BaseModel): + keys: list[str] + + +class KeyInfoParams(BaseModel): + key: str + + +class LiteLLMBudgetTable(BaseModel): + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: str | None = None + + +class KeyInfo(BaseModel): + spend: float | None = None + max_budget: float | None = None + budget_reset_at: str | None = None + budget_id: str | None = None + litellm_budget_table: LiteLLMBudgetTable | None = None + + +class KeyInfoResponse(BaseModel): + info: KeyInfo + + +# ---------- customers ---------- + + +class CustomerDeleteBody(BaseModel): + user_ids: list[str] + + +# ---------- chat / embeddings ---------- + + +class ChatMetadata(BaseModel): + tags: list[str] | None = None + + +class ChatMessage(BaseModel): + role: str + content: str + + +class ChatBody(BaseModel): + model: str + messages: list[ChatMessage] + stream: bool = False + max_tokens: int | None = None + user: str | None = None + metadata: ChatMetadata | None = None + + +class OutMessage(BaseModel): + content: str | None = None + + +class ChatChoice(BaseModel): + message: OutMessage | None = None + + +class Usage(BaseModel): + prompt_tokens: int | None = None + completion_tokens: int | None = None + total_tokens: int | None = None + + +class ChatResponse(BaseModel): + id: str | None = None + model: str | None = None + choices: list[ChatChoice] = [] + usage: Usage | None = None + + +class EmbedBody(BaseModel): + model: str + input: str + + +class EmbedResponse(BaseModel): + model: str | None = None + + +# ---------- spend logs ---------- + + +class SpendLogRow(BaseModel): + request_id: str | None = None + model: str | None = None + spend: float | None = None + status: str | None = None + cache_hit: str | None = None + call_type: str | None = None + custom_llm_provider: str | None = None + team_id: str | None = None + user: str | None = None + end_user: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + total_tokens: int | None = None + request_tags: list[str] | None = None + + +class SpendLogs(RootModel[list[SpendLogRow]]): + pass + + +class SpendLogsParams(BaseModel): + request_id: str | None = None + api_key: str | None = None + + +# ---------- spend calculate ---------- + + +class SpendCalculateBody(BaseModel): + model: str + messages: list[ChatMessage] + + +class SpendCalculateResponse(BaseModel): + cost: float + + +# ---------- route probing ---------- + + +class DateRangeParams(BaseModel): + start_date: str + end_date: str + + +class RouteSpec(RootModel[dict[str, object]]): + """One /openapi.json path entry: a map of HTTP method -> operation. Only the + method names are read, so the operation specs stay opaque.""" + + @property + def methods(self) -> frozenset[str]: + return frozenset(method.lower() for method in self.root) + + +class OpenAPISchema(BaseModel): + paths: dict[str, RouteSpec] = {} + + +# ---------- model info / custom pricing ---------- + + +class CustomPricing(BaseModel): + """The per-token custom-pricing fields a deployment can override in + litellm_params - the token-cost subset of litellm's CustomPricingLiteLLMParams + the proxy applies to chat spend. All optional: a config sets only what it + overrides, and /model/info echoes the rates the proxy resolved.""" + + model_config = ConfigDict(extra="ignore") + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + + def overrides(self) -> dict[str, float]: + """The rates actually declared (non-null) - e.g. those a config.yml sets.""" + declared = { + "input_cost_per_token": self.input_cost_per_token, + "output_cost_per_token": self.output_cost_per_token, + "cache_read_input_token_cost": self.cache_read_input_token_cost, + "cache_creation_input_token_cost": self.cache_creation_input_token_cost, + } + return {field: rate for field, rate in declared.items() if rate is not None} + + def token_cost(self, prompt_tokens: int, completion_tokens: int) -> float: + """Spend for a fresh (uncached) call under these rates: the proxy's + custom-pricing formula (prompt * input + completion * output).""" + assert ( + self.input_cost_per_token is not None + and self.output_cost_per_token is not None + ), "custom pricing has no per-token rates" + return ( + prompt_tokens * self.input_cost_per_token + + completion_tokens * self.output_cost_per_token + ) + + +class ModelInfoEntry(BaseModel): + """One /model/info row. `litellm_params` is the configured deployment (carries + any custom-pricing override); `model_info` is the price the proxy resolved for + it - the override merged over the cost-map defaults.""" + + model_config = ConfigDict(protected_namespaces=()) + model_name: str + litellm_params: CustomPricing = CustomPricing() + model_info: CustomPricing = CustomPricing() + + +class ModelInfoResponse(BaseModel): + data: list[ModelInfoEntry] = [] diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini new file mode 100644 index 00000000000..7799f6b16a2 --- /dev/null +++ b/tests/e2e/pytest.ini @@ -0,0 +1,7 @@ +[pytest] +# Config when any e2e suite under tests/e2e/ is run directly, e.g. +# uv run pytest tests/e2e/spend_tracking/ -v +# The e2e marker is also registered in conftest.py for runs rooted elsewhere. +addopts = --strict-markers --strict-config +markers = + e2e: live test that requires a running proxy and real provider keys diff --git a/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md new file mode 100644 index 00000000000..53c4d4ace83 --- /dev/null +++ b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md @@ -0,0 +1,78 @@ +# Spend Tracking Test Coverage Matrix + +Scope: every distinct spend-tracking code path, mapped to the test that exercises +it and the level it runs at. Highlights where a live e2e check is the only thing +that would catch a regression. + +Companion: live suite `test_spend_tracking_e2e.py` + route breadth +`test_spend_routes.py` (this directory). Offline regression suite: +`tests/test_litellm/proxy/spend_tracking/`. Reference PR: BerriAI/litellm#29956. + +Levels: `unit` mocked; `integration` real DB/cost-map; `live` real provider + +proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. + +--- + +## SpendLogs row construction (`spend_tracking_utils.get_logging_payload`) + +| Path | Existing | Level | Status | Live e2e | +|------|----------|-------|--------|----------| +| `_get_status_for_spend_log` | `test_spend_tracking_utils.py` | unit | covered | yes (status read off the row) | +| cache-hit `request_id` suffix | `test_spend_tracking_utils.py` | unit | covered | yes (`test_cache_hit_is_zero_cost_and_suffixed`) | +| failure status + zero spend | `test_spend_tracking_utils.py` | unit | covered | no (live failure logging is non-deterministic across providers) | +| per-model / per-provider attribution | `test_spend_tracking_utils.py` | unit | covered | yes (`test_each_model_on_a_shared_key_gets_its_own_row`) | +| field population (model/tokens/api_key/team/org) | `test_spend_tracking_utils.py` | unit | partial | yes (asserts real values) | +| `request_tags` propagation | `test_db_spend_update_writer.py` | unit | partial | yes (`test_request_tags_round_trip`) | +| `end_user` attribution | unit | unit | partial | yes (`test_end_user_spend_attributed_on_row`) | + +## Cost calculation by modality + +| Modality | Existing | Status | Live e2e | +|----------|----------|--------|----------| +| Chat (non-stream) | `test_cost_calculator.py`, `local_testing/test_completion_cost.py` | covered | yes (`test_chat_completion_writes_nonzero_spend_row`) | +| Chat (streaming) | `test_streaming_interrupt_spend_tracking.py` | partial | yes (`test_streaming_chat_completion_tracks_spend`) | +| Embedding | `test_cost_calculator.py` (#29956) | partial | yes (`test_embedding_writes_nonzero_spend_row`) | +| Pass-through (gemini/anthropic) | `pass_through_tests/*.test.js` + `llm_translation/` suite | covered | yes (llm_translation suite) | +| Image / audio / rerank / responses / realtime | per-provider unit cost tests | partial/gap | gap | + +## Entity spend aggregation + +| Entity | Existing | Status | Live e2e | +|--------|----------|--------|----------| +| API key | `test_db_spend_update_writer.py`, `test_spend_counters.py` | covered | yes (`test_key_spend_equals_sum_of_logs`) | +| Tag | `test_update_daily_tag_spend.py` | partial | yes (`test_request_tags_round_trip`, propagation only) | +| End-user | `test_proxy_update_spend.py` | covered | yes | +| Spend == sum(logs) consistency | none | gap | yes (key aggregate == sum of rows) | + +## Spend read endpoints (verification surface) + +| Endpoint | Existing | Status | Live e2e | +|----------|----------|--------|----------| +| `/spend/logs` (request_id / api_key) | `test_spend_management_endpoints.py` | covered | yes (primary read path; `test_spend_logs_endpoint_returns_spend` asserts 200 + spend, never 5xx) | +| `/spend/calculate` | `local_testing/test_spend_calculate_endpoint.py` | covered | yes (`test_spend_calculate_returns_nonzero_cost`) | +| `/spend/tags` | `test_spend_management_endpoints.py` | partial | yes (`test_spend_routes.py` route probe) | +| whole spend GET surface (22 routes) | unit per-handler | partial | yes (`test_spend_routes.py` probes each for 404/5xx) | + +## What this suite pins + +| Test | Invariant | +|------|-----------| +| `test_chat_completion_writes_nonzero_spend_row` | nonzero cost, token arithmetic, status, row findable by `response.id` | +| `test_streaming_chat_completion_tracks_spend` | streamed responses still costed | +| `test_embedding_writes_nonzero_spend_row` | embedding cost, `completion_tokens == 0` | +| `test_cache_hit_is_zero_cost_and_suffixed` | cache hits not double-charged; `_cache_hit` suffix | +| `test_key_spend_equals_sum_of_logs` | key aggregate == sum of rows | +| `test_request_tags_round_trip` | tags persist onto the row | +| `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed | +| `test_each_model_on_a_shared_key_gets_its_own_row` | per-model/provider rows, correct model + cost, distinct request_ids matching response id | +| `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) | +| `test_spend_logs_endpoint_returns_spend` | `/spend/logs` returns 200 + the key's spend, never a 5xx (intermittent-500 regression) | +| `test_spend_routes.py` (23) | no spend route 404s or 5xxs | + +## Design + timing + +`proxy_batch_write_at` (~60s) means rows land late; every read polls to a deadline. +Fresh scoped key per test (isolation, xdist-safe, cleaned up). Assert invariants +(`spend > 0`, `total == prompt + completion`, aggregate == sum), not literal +$/token values, so pricing drift is not a failure. Skip on environment (no proxy / +no provider key), fail on behavior (a real 2xx call with a wrong/missing row). diff --git a/tests/e2e/spend_tracking/conftest.py b/tests/e2e/spend_tracking/conftest.py new file mode 100644 index 00000000000..1d01ab3d17a --- /dev/null +++ b/tests/e2e/spend_tracking/conftest.py @@ -0,0 +1,16 @@ +"""Spend-tracking suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway +(GatewayProvider), so the `resources` fixture cleans up keys and customers this +suite creates. +""" + +import pytest + +from spend_e2e_client import SpendClient, build_client + + +@pytest.fixture(scope="session") +def client() -> SpendClient: + return build_client() diff --git a/tests/e2e/spend_tracking/spend_e2e_client.py b/tests/e2e/spend_tracking/spend_e2e_client.py new file mode 100644 index 00000000000..d749d69f1a4 --- /dev/null +++ b/tests/e2e/spend_tracking/spend_e2e_client.py @@ -0,0 +1,167 @@ +"""Spend-tracking e2e client: a Gateway plus the spend-specific read endpoints. + +Generic proxy operations (keys, customers, chat/embed, route probing, SpendLogs +polling) come from the shared Gateway, DI'd in (composition, not inheritance). +This client adds only the spend surface: /spend/calculate, key-spend +polling, and the route probes the breadth test uses. + +Re-exports unwrap / is_ok / unique_marker / SpendLogRow so the tests import their +helpers from one place. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import dataclass + +from e2e_config import unique_marker +from e2e_http import ( + NoBody, + ProbeResult, + Result, + StreamingResponse, + is_ok, + unwrap, +) +from e2e_gateway import Gateway, build_gateway +from models import ( + ChatBody, + ChatMessage, + ChatMetadata, + ChatResponse, + DateRangeParams, + EmbedBody, + EmbedResponse, + OpenAPISchema, + SpendCalculateBody, + SpendCalculateResponse, + SpendLogRow, +) + +__all__ = [ + "SpendClient", + "build_client", + "reset_spend_logs", + "unique_marker", + "unwrap", + "is_ok", + "SpendLogRow", + "ProbeResult", +] + + +def reset_spend_logs() -> None: + """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes + spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses + DATABASE_URL (default: the local docker postgres on its mapped host port; note + the in-container `@db` host isn't resolvable from the host, so default to + localhost). + """ + import psycopg + + url = os.environ.get( + "DATABASE_URL", + "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", + ) + with psycopg.connect(url) as conn: + _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') + + +def _chat_body( + model: str, + content: str, + *, + max_tokens: int | None = None, + tags: list[str] | None = None, + user: str | None = None, + stream: bool = False, +) -> ChatBody: + return ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=max_tokens, + stream=stream, + user=user, + metadata=ChatMetadata(tags=tags) if tags else None, + ) + + +@dataclass(frozen=True, slots=True) +class SpendClient: + gateway: Gateway + + def chat( + self, + key: str, + model: str, + content: str, + *, + max_tokens: int | None = None, + tags: list[str] | None = None, + user: str | None = None, + ) -> Result[ChatResponse]: + return self.gateway.chat( + key, _chat_body(model, content, max_tokens=max_tokens, tags=tags, user=user) + ) + + def chat_stream( + self, key: str, model: str, content: str, *, max_tokens: int | None = None + ) -> StreamingResponse: + return self.gateway.chat_stream( + key, _chat_body(model, content, max_tokens=max_tokens, stream=True) + ) + + def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]: + return self.gateway.embed(key, EmbedBody(model=model, input=content)) + + def poll_logs_for_key( + self, + key: str, + *, + min_rows: int = 1, + predicate: Callable[[list[SpendLogRow]], bool] | None = None, + ) -> list[SpendLogRow]: + return self.gateway.poll_logs_for_key( + key, min_rows=min_rows, predicate=predicate + ) + + def calculate_spend(self, model: str, content: str) -> float: + return unwrap( + self.gateway.transport.post( + "/spend/calculate", + headers=self.gateway.transport.master, + json=SpendCalculateBody( + model=model, messages=[ChatMessage(role="user", content=content)] + ), + response_type=SpendCalculateResponse, + ) + ).cost + + def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float: + deadline = time.monotonic() + self.gateway.poll_timeout + spend = 0.0 + while time.monotonic() < deadline: + spend = self.gateway.key_info(key).spend or 0.0 + if spend > minimum: + return spend + time.sleep(self.gateway.poll_interval) + return spend + + def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: + return self.gateway.transport.probe(path, params=params) + + def openapi(self) -> OpenAPISchema: + return unwrap( + self.gateway.transport.get( + "/openapi.json", + headers=self.gateway.transport.master, + params=NoBody(), + response_type=OpenAPISchema, + ) + ) + + +def build_client() -> SpendClient: + return SpendClient(gateway=build_gateway()) diff --git a/tests/e2e/spend_tracking/test_spend_routes.py b/tests/e2e/spend_tracking/test_spend_routes.py new file mode 100644 index 00000000000..e3c96a4d578 --- /dev/null +++ b/tests/e2e/spend_tracking/test_spend_routes.py @@ -0,0 +1,96 @@ +"""Breadth check: query every route on the spend read surface and show what it +returns. + +Spend tracking sprawls across many routes (model-cost / key / user / team / org / +customer aggregation, tags, and activity reports). Most are served with +`include_in_schema=False`, so they do NOT appear in `/openapi.json` - discovery +from the schema alone misses ~70% of the surface. So we probe a curated, verified +list directly, plus any spend route the schema does list (to auto-catch new ones). + +Each probe captures status AND body, so a failure shows the proxy's actual error +(a 500 traceback, a 404 meaning the route was removed) rather than a bare code. +Run with `-rA` (or `-s`) to print every route's response, not just failures. + +Healthy == route exists (not 404) and handler did not crash (not 5xx). A 4xx +(missing params / auth nuance) still means the route is wired and ran. Cheap and +fast: no batch-write wait, no provider calls. +""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from models import DateRangeParams +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +# Verified present and responsive on a live proxy. One per row of the spend +# surface: key / user / team / org / customer aggregation, model-cost, tags, +# activity. +SPEND_ROUTES = ( + "/spend/keys", + "/spend/users", + "/spend/tags", + "/spend/logs", + "/spend/logs/ui", + "/global/spend", + "/global/spend/keys", + "/global/spend/teams", + "/global/spend/models", + "/global/spend/provider", + "/global/spend/report", + "/global/spend/tags", + "/global/spend/logs", + "/global/spend/all_tag_names", + "/global/activity", + "/global/activity/model", + "/global/activity/exceptions", + "/key/list", + "/user/list", + "/team/list", + "/organization/list", + "/customer/list", +) + +_SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") + + +def _date_range() -> DateRangeParams: + # Satisfies date-required endpoints (report/activity/provider); ignored elsewhere. + end = datetime.now(timezone.utc).date() + start = end - timedelta(days=1) + return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) + + +@pytest.mark.parametrize("route", SPEND_ROUTES) +def test_spend_route_responsive(client: SpendClient, route: str) -> None: + result = client.probe(route, params=_date_range()) + print(f"{route} -> {result.status_code}\n{result.body[:600]}") + assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}" + + +def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None: + """Probe any spend GET route the schema lists that isn't in SPEND_ROUTES.""" + schema = client.openapi() + assert schema.paths, "/openapi.json had no paths" + + discovered = [ + path + for path, spec in schema.paths.items() + if "get" in spec.methods + and "{" not in path + and any(path.startswith(prefix) for prefix in _SPEND_PREFIXES) + ] + extras = [path for path in discovered if path not in SPEND_ROUTES] + + params = _date_range() + results = [(path, client.probe(path, params=params)) for path in extras] + for path, result in results: + print(f"{path} -> {result.status_code}") + offenders = [ + f"{path} -> {result.status_code}\n{result.body[:600]}" + for path, result in results + if not result.healthy + ] + assert not offenders, "non-responsive schema spend routes:\n" + "\n".join(offenders) diff --git a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py new file mode 100644 index 00000000000..8c9e913b10f --- /dev/null +++ b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py @@ -0,0 +1,328 @@ +"""Live end-to-end spend-tracking tests against a running proxy. + +Run against a proxy started with the gateway config. Coverage rationale: +SPEND_TRACKING_COVERAGE_MATRIX.md. + +Model names are literals from that config: chat tests hit "gemini-2.5-flash", +embedding tests hit "openai-text-embedding-3-small". + +Every test: fresh scoped key (isolation) -> real provider call -> unwrap (hard +fail if the proxy couldn't make a call it should) -> poll /spend/logs to a +deadline (rows land ~60s later via proxy_batch_write_at) -> assert invariants on +the real row (spend, token arithmetic, status, cache). + +Assertions target invariants, not literals: a regression in the spend pipeline +fails the test; a pricing or token-count drift does not. +""" + +import time +from collections.abc import Callable + +import pytest + +from e2e_http import Success +from lifecycle import ResourceManager +from models import SpendLogs, SpendLogsParams +from spend_e2e_client import SpendClient, SpendLogRow, unique_marker, unwrap + +pytestmark = pytest.mark.e2e + + +def _approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +def _summarize(rows: list[SpendLogRow]) -> list[dict[str, object]]: + fields = { + "request_id", + "model", + "spend", + "status", + "cache_hit", + "prompt_tokens", + "completion_tokens", + "total_tokens", + } + return [row.model_dump(include=fields) for row in rows] + + +def _require_row( + rows: list[SpendLogRow], predicate: Callable[[SpendLogRow], bool], what: str +) -> SpendLogRow: + matches = [r for r in rows if predicate(r)] + assert matches, ( + f"no SpendLogs row {what} after polling; saw {len(rows)} row(s): " + f"{_summarize(rows)}" + ) + return matches[0] + + +def test_chat_completion_writes_nonzero_spend_row( + client: SpendClient, scoped_key: str +) -> None: + chat = unwrap( + client.chat( + scoped_key, + "gemini-2.5-flash", + f"reply with one word {unique_marker()}", + max_tokens=16, + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.status == "success" for r in rs) + ) + row = _require_row(rows, lambda r: r.status == "success", "for the chat call") + + assert (row.spend or 0) > 0, f"chat row should cost > 0: {_summarize(rows)}" + assert row.status == "success" + assert row.cache_hit != "True", "fresh call must not be a cache hit" + assert "gemini-2.5-flash" in (row.model or "") + + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + total = row.total_tokens or 0 + assert prompt > 0 and completion > 0 + assert total == prompt + completion, f"token arithmetic broken: {_summarize(rows)}" + + if chat.id: + assert any( + r.request_id == chat.id for r in rows + ), f"row request_id != client response.id ({chat.id})" + + +def test_streaming_chat_completion_tracks_spend( + client: SpendClient, scoped_key: str +) -> None: + result = client.chat_stream( + scoped_key, + "gemini-2.5-flash", + f"count to three {unique_marker()}", + max_tokens=64, + ) + assert ( + result.ok + ), f"stream failed (status {result.status_code}): {result.body[:300]}" + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + row = _require_row( + rows, lambda r: (r.spend or 0) > 0, "with nonzero spend for the stream" + ) + prompt = row.prompt_tokens or 0 + completion = row.completion_tokens or 0 + assert ( + prompt > 0 and completion > 0 + ), f"streaming tokens not tracked: {_summarize(rows)}" + assert (row.total_tokens or 0) == prompt + completion + + +def test_embedding_writes_nonzero_spend_row( + client: SpendClient, scoped_key: str +) -> None: + _ = unwrap( + client.embed( + scoped_key, + "openai-text-embedding-3-small", + f"vectorize this sentence {unique_marker()}", + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + row = _require_row( + rows, lambda r: (r.spend or 0) > 0, "with nonzero spend for the embedding" + ) + assert (row.prompt_tokens or 0) > 0 + assert (row.completion_tokens or 0) == 0, "embeddings have no completion tokens" + assert "text-embedding-3-small" in (row.model or "") + + +def test_cache_hit_is_zero_cost_and_suffixed( + client: SpendClient, scoped_key: str +) -> None: + # Unique marker shared by both calls: call 1 is a guaranteed cache MISS (fresh + # content, paid), call 2 repeats the identical request and HITS the cache just + # populated. The marker keeps each run isolated - a fixed prompt would persist + # in the shared response cache across runs and make both calls hit (flaky). + prompt = f"What is the capital of France? Answer in one word. {unique_marker()}" + _ = unwrap(client.chat(scoped_key, "gemini-2.5-flash", prompt, max_tokens=16)) + _ = unwrap(client.chat(scoped_key, "gemini-2.5-flash", prompt, max_tokens=16)) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.cache_hit == "True" for r in rs) + ) + cache_rows = [r for r in rows if r.cache_hit == "True"] + if not cache_rows: + pytest.skip( + "no cache-hit row observed; caching may be disabled on this proxy. " + f"rows seen: {_summarize(rows)}" + ) + + cache_row = cache_rows[0] + assert ( + cache_row.spend or 0 + ) == 0.0, f"cache hit was charged (double-charge regression): {_summarize(rows)}" + assert "_cache_hit" in (cache_row.request_id or ""), ( + "cache-hit row missing the _cache_hit request_id suffix; " + "duplicate-key collisions will silently drop rows" + ) + paid_rows = [r for r in rows if r.cache_hit != "True"] + assert any( + (r.spend or 0) > 0 for r in paid_rows + ), f"the non-cached call should still be charged: {_summarize(rows)}" + + +def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> None: + for _ in range(2): + _ = unwrap( + client.chat( + scoped_key, + "gemini-2.5-flash", + f"say hi {unique_marker()}", + max_tokens=16, + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, + min_rows=2, + predicate=lambda rs: sum((r.spend or 0) for r in rs) > 0, + ) + assert len(rows) >= 2, f"expected >=2 rows for the key, saw {_summarize(rows)}" + logs_total = sum((r.spend or 0) for r in rows) + assert logs_total > 0 + + key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999) + assert _approx_equal( + key_spend, logs_total + ), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" + + +def test_request_tags_round_trip(client: SpendClient, scoped_key: str) -> None: + tag = f"e2e-spend-{unique_marker()}" + _ = unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", "tagged request", tags=[tag], max_tokens=16 + ) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(tag in (r.request_tags or []) for r in rs) + ) + _require_row( + rows, lambda r: tag in (r.request_tags or []), f"carrying request tag {tag!r}" + ) + + +def test_end_user_spend_attributed_on_row( + client: SpendClient, scoped_key: str, resources: ResourceManager +) -> None: + customer = resources.customer(f"e2e-cust-{unique_marker()}") + _ = unwrap( + client.chat(scoped_key, "gemini-2.5-flash", "hi", user=customer, max_tokens=16) + ) + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any(r.end_user == customer for r in rs) + ) + row = _require_row( + rows, lambda r: r.end_user == customer, f"attributed to end_user {customer!r}" + ) + assert (row.spend or 0) > 0, f"end-user row should cost > 0: {_summarize(rows)}" + + +def test_each_model_on_a_shared_key_gets_its_own_row( + client: SpendClient, scoped_key: str +) -> None: + """One key calling two different models, on two providers, gets one spend row per + call - each carrying its own model and a nonzero cost, under distinct request_ids + that match the call's response id. Pins per-model/per-provider attribution: a + regression that stamps the wrong model on the row, bills a call's cost to the + sibling deployment, or collapses both calls onto one request_id fails here.""" + gemini = unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", f"one word {unique_marker()}", max_tokens=16 + ) + ) + claude = unwrap( + client.chat( + scoped_key, "claude-haiku-4-5", f"one word {unique_marker()}", max_tokens=16 + ) + ) + + def both_models_costed(rows: list[SpendLogRow]) -> bool: + costed = [r.model or "" for r in rows if (r.spend or 0) > 0] + return any("gemini-2.5-flash" in m for m in costed) and any( + "claude-haiku-4-5" in m for m in costed + ) + + rows = client.poll_logs_for_key(scoped_key, min_rows=2, predicate=both_models_costed) + gemini_row = _require_row( + rows, lambda r: "gemini-2.5-flash" in (r.model or ""), "for the gemini call" + ) + claude_row = _require_row( + rows, lambda r: "claude-haiku-4-5" in (r.model or ""), "for the claude call" + ) + + assert (gemini_row.spend or 0) > 0, f"gemini row should cost > 0: {_summarize(rows)}" + assert (claude_row.spend or 0) > 0, f"claude row should cost > 0: {_summarize(rows)}" + assert ( + gemini_row.request_id != claude_row.request_id + ), f"two distinct calls collapsed onto one request_id: {_summarize(rows)}" + if gemini.id: + assert ( + gemini_row.request_id == gemini.id + ), f"gemini row request_id {gemini_row.request_id} != response id {gemini.id}" + if claude.id: + assert ( + claude_row.request_id == claude.id + ), f"claude row request_id {claude_row.request_id} != response id {claude.id}" + + +def test_spend_calculate_returns_nonzero_cost(client: SpendClient) -> None: + cost = client.calculate_spend( + "gemini-2.5-flash", "estimate the cost of this request" + ) + assert cost > 0, ( + "/spend/calculate returned 0 for gemini-2.5-flash; " + "cost map may be missing this model" + ) + + +def test_spend_logs_endpoint_returns_spend( + client: SpendClient, scoped_key: str +) -> None: + """The /spend/logs read endpoint returns a 200 carrying the key's spend, never a + 5xx. Regression for intermittent 500s (DB query / serialization errors under load) + on this endpoint: every poll asserts a success response, not just a truthy row + list, so a 500 fails loudly instead of being swallowed as 'no rows yet'; the + call's nonzero spend must surface before the deadline.""" + unwrap( + client.chat( + scoped_key, "gemini-2.5-flash", f"spend logs {unique_marker()}", max_tokens=16 + ) + ) + + gateway = client.gateway + deadline = time.monotonic() + gateway.poll_timeout + while True: + result = gateway.transport.get( + "/spend/logs", + headers=gateway.transport.master, + params=SpendLogsParams(api_key=scoped_key), + response_type=SpendLogs, + ) + assert isinstance(result, Success), f"/spend/logs did not return 200 OK: {result}" + rows = result.data.root + if sum((r.spend or 0) for r in rows) > 0: + return + if time.monotonic() >= deadline: + pytest.fail( + f"/spend/logs never surfaced the key's spend before the deadline; " + f"saw {_summarize(rows)}" + ) + time.sleep(gateway.poll_interval) diff --git a/tests/e2e/test_lifecycle.py b/tests/e2e/test_lifecycle.py new file mode 100644 index 00000000000..d3c559dd2ed --- /dev/null +++ b/tests/e2e/test_lifecycle.py @@ -0,0 +1,46 @@ +"""Unit coverage for the lifecycle harness (lifecycle.run_case). + +Cases register cleanups progressively during init() (create team, then user, then +key), so a failure partway through init() must still release whatever was already +created on the long-lived shared proxy. This guards that contract. +""" + +from dataclasses import dataclass, field +from typing import Callable, List + +import pytest + +from lifecycle import run_case + + +@dataclass +class _PartialInitCase: + """init() registers a cleanup, then raises before finishing - mirroring a real + case that creates a resource, registers its delete, then fails on the next + step.""" + + released: List[str] = field(default_factory=list) + _undo: List[Callable[[], None]] = field(default_factory=list) + + def init(self) -> None: + self._undo.append(lambda: self.released.append("first")) + raise RuntimeError("init failed after registering the first resource") + + def run(self) -> None: + raise AssertionError("run() must not execute when init() failed") + + def teardown(self) -> None: + for undo in reversed(self._undo): + undo() + + +def test_run_case_releases_resources_when_init_fails_partway() -> None: + case = _PartialInitCase() + + with pytest.raises(RuntimeError, match="init failed"): + run_case(case) + + assert case.released == ["first"], ( + "a resource registered before init() failed must still be released, or it " + "leaks on the long-lived shared proxy" + ) diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py new file mode 100644 index 00000000000..37412fc0cf5 --- /dev/null +++ b/tests/e2e/transport.py @@ -0,0 +1,244 @@ +"""Transport: the typed request primitives clients use, behind a Protocol. + +`Transport` is what each client depends on (composition + DI); `HttpTransport` is +the concrete frozen-slots dataclass that fulfils it via the e2e_http wrapper. No +client touches requests.* or builds raw dicts; they pass pydantic models here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from pydantic import BaseModel + +import e2e_http +from e2e_http import URL, AuthHeaders, ProbeResult, Result, StreamingResponse + + +class Transport(Protocol): + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: ... + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: ... + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: ... + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: ... + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: ... + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... + + def bearer(self, key: str) -> AuthHeaders: ... + + @property + def master(self) -> AuthHeaders: ... + + +@dataclass(frozen=True, slots=True) +class HttpTransport: + base_url: str + master_key: str + request_timeout: float = 60.0 + + def _url(self, path: str) -> URL: + return URL(f"{self.base_url.rstrip('/')}{path}") + + def bearer(self, key: str) -> AuthHeaders: + return AuthHeaders(authorization=f"Bearer {key}") + + @property + def master(self) -> AuthHeaders: + return self.bearer(self.master_key) + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return e2e_http.post( + self._url(path), + headers=headers, + json=json, + response_type=response_type, + timeout=self.request_timeout, + ) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: + return e2e_http.get( + self._url(path), + headers=headers, + params=params, + response_type=response_type, + timeout=self.request_timeout, + ) + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return e2e_http.delete( + self._url(path), + headers=headers, + json=json, + response_type=response_type, + timeout=self.request_timeout, + ) + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: + return e2e_http.stream( + self._url(path), headers=headers, json=json, timeout=self.request_timeout + ) + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + return e2e_http.send( + self._url(path), + headers=headers, + json=json, + params=params, + stream=stream, + timeout=self.request_timeout, + ) + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + return e2e_http.probe( + self._url(path), + headers=self.master, + params=params, + timeout=self.request_timeout, + ) + + +# Top-level management/admin route groups. In a split deployment these are served +# by the control plane (a different service from the LLM data plane). LLM routes +# (/chat, /embeddings, and native passthrough like /gemini, /anthropic) are NOT +# here and fall through to the data plane. Matched as path prefixes. +CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( + "/key", + "/user", + "/team", + "/organization", + "/customer", + "/tag", + "/budget", + "/model/info", + "/spend", + "/global", + "/openapi.json", +) + + +def is_control_plane_path(path: str) -> bool: + """True if `path` is a management/admin route (served by the control plane in a + split deployment), false for LLM data-plane routes.""" + return path.startswith(CONTROL_PLANE_PREFIXES) + + +@dataclass(frozen=True, slots=True) +class SplitTransport: + """A Transport that dispatches each call by path to one of two backends: the + management/admin control plane or the LLM data plane. + + Litellm can run as a split control-plane/data-plane deployment where the two + surfaces live on different services. Clients here stay plane-agnostic — they + keep calling ``transport.post("/budget/new", ...)`` or + ``transport.send("/chat/completions", ...)`` — and routing happens in one place + by path (see ``CONTROL_PLANE_PREFIXES``). When ``control`` and ``data`` share a + base URL (the monolithic default), routing is a no-op. ``bearer``/``master`` + are plane-agnostic (same master key both planes), so they come from ``data``. + """ + + data: HttpTransport + control: HttpTransport + + def _route(self, path: str) -> HttpTransport: + return self.control if is_control_plane_path(path) else self.data + + def bearer(self, key: str) -> AuthHeaders: + return self.data.bearer(key) + + @property + def master(self) -> AuthHeaders: + return self.data.master + + def post[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return self._route(path).post( + path, headers=headers, json=json, response_type=response_type + ) + + def get[R: BaseModel]( + self, + path: str, + *, + headers: BaseModel, + params: BaseModel, + response_type: type[R], + ) -> Result[R]: + return self._route(path).get( + path, headers=headers, params=params, response_type=response_type + ) + + def delete[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return self._route(path).delete( + path, headers=headers, json=json, response_type=response_type + ) + + def stream( + self, path: str, *, headers: BaseModel, json: BaseModel + ) -> StreamingResponse: + return self._route(path).stream(path, headers=headers, json=json) + + def send( + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + params: BaseModel | None = None, + stream: bool = False, + ) -> StreamingResponse: + return self._route(path).send( + path, headers=headers, json=json, params=params, stream=stream + ) + + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + return self._route(path).probe(path, params=params) diff --git a/tests/pyrightconfig.json b/tests/pyrightconfig.json new file mode 100644 index 00000000000..5757c97f812 --- /dev/null +++ b/tests/pyrightconfig.json @@ -0,0 +1,11 @@ +{ + "include": ["e2e"], + "exclude": ["**/node_modules", "**/__pycache__"], + "pythonVersion": "3.12", + "typeCheckingMode": "strict", + "enableTypeIgnoreComments": false, + "reportMissingImports": false, + "reportPrivateImportUsage": false, + "reportExplicitAny": "error", + "reportAny": "error" +} \ No newline at end of file