From 74d5d1c3a7e4feddae5afa17ee0749310aec4262 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sat, 20 Jun 2026 12:23:13 -0700 Subject: [PATCH] fix: add new tests, split gateway --- tests/e2e/budgets/budget_client.py | 75 +++++++++-- .../budgets/test_spend_counter_reseed_e2e.py | 118 ++++++++++++++++++ .../budgets/test_team_member_budget_e2e.py | 99 +++++++++++++++ .../test_team_member_budget_reset_e2e.py | 47 +++++++ .../test_team_multi_window_budget_e2e.py | 71 +++++++++++ tests/e2e/conftest.py | 28 +++-- tests/e2e/e2e_config.py | 11 ++ tests/e2e/e2e_gateway.py | 24 ++-- tests/e2e/models.py | 14 +-- tests/e2e/transport.py | 100 +++++++++++++++ 10 files changed, 558 insertions(+), 29 deletions(-) create mode 100644 tests/e2e/budgets/test_spend_counter_reseed_e2e.py create mode 100644 tests/e2e/budgets/test_team_member_budget_e2e.py create mode 100644 tests/e2e/budgets/test_team_member_budget_reset_e2e.py create mode 100644 tests/e2e/budgets/test_team_multi_window_budget_e2e.py diff --git a/tests/e2e/budgets/budget_client.py b/tests/e2e/budgets/budget_client.py index cfe70a605f9..af8021f9b93 100644 --- a/tests/e2e/budgets/budget_client.py +++ b/tests/e2e/budgets/budget_client.py @@ -12,7 +12,7 @@ from __future__ import annotations from dataclasses import dataclass -from pydantic import BaseModel, RootModel +from pydantic import AliasPath, BaseModel, Field, RootModel from e2e_gateway import Gateway, build_gateway from e2e_http import NoBody, StreamingResponse, Success, unwrap @@ -65,6 +65,7 @@ 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): @@ -81,6 +82,29 @@ class TeamMemberAddBody(BaseModel): 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 @@ -125,9 +149,7 @@ def is_budget_block(result: StreamingResponse) -> bool: return not result.ok and "budget_exceeded" in result.body -def model_budget( - model: str, limit: float, period: str = "30d" -) -> dict[str, ModelBudgetEntry]: +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)} @@ -254,6 +276,7 @@ class BudgetClient: 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( @@ -263,6 +286,7 @@ class BudgetClient: team_alias=alias, max_budget=max_budget, organization_id=organization_id, + budget_limits=budget_limits, ), response_type=TeamNewResponse, ) @@ -276,9 +300,7 @@ class BudgetClient: response_type=NoBody, ) - def add_team_member( - self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None - ) -> None: + 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, @@ -290,6 +312,45 @@ class BudgetClient: ) 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: 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..97bf0bae422 --- /dev/null +++ b/tests/e2e/budgets/test_spend_counter_reseed_e2e.py @@ -0,0 +1,118 @@ +"""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 + +import pytest + +from budget_client import BudgetClient +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager + +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(): + import redis + + return redis.Redis( + host=os.getenv("E2E_REDIS_HOST", "localhost"), + port=int(os.getenv("E2E_REDIS_PORT", "6380")), + password=os.getenv("REDIS_PASSWORD") or None, + decode_responses=True, + socket_connect_timeout=2, + ) + + +def _spend_counter(rds, key: str) -> float | None: + """The shared spend counter for `key`, or None if it is cold. The counter key is + ``{cache namespace}:spend:key:{sha256(key)}``; matched by suffix so the configured + namespace need not be hard-coded.""" + digest = hashlib.sha256(key.encode()).hexdigest() + matches = list(rds.scan_iter(match=f"*spend:key:{digest}")) + 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: + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(lambda _: _chat(client, key), 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 E2E_REDIS_HOST/E2E_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_team_member_budget_e2e.py b/tests/e2e/budgets/test_team_member_budget_e2e.py new file mode 100644 index 00000000000..32e8a0d2513 --- /dev/null +++ b/tests/e2e/budgets/test_team_member_budget_e2e.py @@ -0,0 +1,99 @@ +"""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 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.""" + marker = unique_marker() + team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET) + user_id = client.create_user(max_budget=TEAM_BUDGET) + 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) + yield _Member(team_id=team_id, user_id=user_id, key=key) + client.delete_key(key) + client.delete_user(user_id) + client.delete_team(team_id) + + +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..8f42550be38 --- /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() + 90 + 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 90s") 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..2d4fb860045 --- /dev/null +++ b/tests/e2e/budgets/test_team_multi_window_budget_e2e.py @@ -0,0 +1,71 @@ +"""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. + +Currently fails at team creation: /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. Left failing rather than weakened - it passes once that write is +fixed. +""" + +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) + + +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 + the reset job. When + # a call flows again the window has reset; the elapsed clock must be short enough + # that this is the 30s window resetting, not the roomy 1m one. + deadline = time.monotonic() + 90 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 45, 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 90s") diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 7a3c5a7f80d..cc95c7538dd 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -21,7 +21,7 @@ from typing import Iterator import pytest import requests -from e2e_config import PROXY_BASE_URL +from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL from lifecycle import GatewayProvider, ResourceManager @@ -35,15 +35,27 @@ def pytest_configure(config: pytest.Config) -> None: ) +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.""" - try: - resp = requests.get(f"{PROXY_BASE_URL}/health/liveliness", timeout=5) - except requests.RequestException as exc: - return f"No live proxy at {PROXY_BASE_URL}: {exc}" - if resp.status_code >= 500: - return f"Proxy at {PROXY_BASE_URL} returned {resp.status_code}" + """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 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 6f4ea6a0812..3865804b08f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -10,6 +10,17 @@ 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")) diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/e2e_gateway.py index 3f8fb6fd9c6..b700145d434 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/e2e_gateway.py @@ -39,13 +39,14 @@ from models import ( SpendLogsParams, ) from e2e_config import ( + CONTROL_PLANE_BASE_URL, MASTER_KEY, POLL_INTERVAL, POLL_TIMEOUT, PROXY_BASE_URL, REQUEST_TIMEOUT, ) -from transport import HttpTransport, Transport +from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -188,13 +189,22 @@ class Gateway: def build_gateway() -> Gateway: - """The Gateway every suite's client is built from: an HttpTransport pointed at - the configured proxy, with the shared poll budget.""" + """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=HttpTransport( - base_url=PROXY_BASE_URL, - master_key=MASTER_KEY, - request_timeout=REQUEST_TIMEOUT, + 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/models.py b/tests/e2e/models.py index 0095dbde347..54095d85262 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -32,6 +32,8 @@ class KeyGenerateBody(BaseModel): 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): @@ -134,6 +136,8 @@ class SpendLogRow(BaseModel): 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 @@ -228,14 +232,10 @@ class CustomPricing(BaseModel): 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 + 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): diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 65e093993c4..37412fc0cf5 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -142,3 +142,103 @@ class HttpTransport: 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)