test(e2e): perf rig (out-of-process mocks, PerformanceTestCase runner, per-request limits) and the small_bodies_5rps case (#46713)

* test(e2e): out-of-process mock OpenAI, Anthropic, Bedrock and guardrail servers for the perf cases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): coding-agent request corpus and real provider captures for the perf mocks

Request bodies come from a deterministic coding-agent session (system prompt, Read/Edit/Write/Bash/Grep/Glob tools, tool_use and tool_result turns carrying real source files picked by the request marker) in OpenAI chat, Responses and Anthropic formats. captures/ holds one real response per provider, route, shape and stream mode (gpt-6-luna, claude-sonnet-5, us.anthropic.claude-sonnet-5), cold prompt cache, body only, raw stream chunks with time offsets

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): replay module for the perf mocks, captured provider responses picked by route, mode and request size

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf mocks replay the captured provider responses, knobs shrink to flake and speed

The OpenAI, Anthropic and Bedrock mocks answer with the captured body or stream chunks for their route, mode and nearest request size, ids swapped per request, on the captured clock scaled by the p knob. The synthetic reply builders, the per-token stream pacing and the hand-rolled Bedrock event-stream encoder are gone

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf mock request log names its provider

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf guardrail mock remembers the request marker by call id for the response hook

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf mock captures retaken cold, one marker and system-prompt prefix per cell

Run 2 shared one marker across cells, so every stream capture was served from the provider's prompt cache (cache_read_input_tokens up to 523946) and the json cells had small prefix hits. Every cell now carries its own marker, also prepended to the system prompt, and all 30 captures report zero cached tokens

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf corpus reads tests/e2e files so the lane's runner image has them

* test(e2e): perf long_context json captures for Claude retaken with a 16k output cap, no longer cut mid-thinking

* test(e2e): perf mocks keep only what the replay path and the test read

* test(e2e): perf mocks interpolate the provider's time to first byte from the captures around the request size

* test(e2e): perf captures at 100k (short) and 500k (long) tokens with tools, the fixed-size set removed

* test(e2e): PerformanceTestCase schema, deterministic request plan and wire records for the perf cases

RequestShape names a corpus session (input_tokens, tools, output_tokens) and wire.request_body builds it with setup/corpus/session.py, so a 300k request is 85 file reads of real source in the endpoint's own format. The plan no longer predicts reply text, the mocks replay captured provider responses and the runner joins what the client saw against the mock's log by marker

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf Metric gains time_to_first_token_seconds

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf User loses the unused role field, every JWT user is an internal_user

* test(e2e): perf RequestShape sizes requests from an input-token range, short and long stock shapes

* test(e2e): perf case runner (provision, Locust plan replay, observe) and the 50 rps mixed-provider case

provision.py checks the deployed config against the case and builds the callers, run.py replays the plan from a separate Locust process (locustfile.py) recording status, header, first-byte, first-token and total time per request, observe.py joins the records with what the mocks received and the workers' RSS to score every Limit, and compares each served body with the capture the mock replayed for it as a diagnostic. The case sends 15,000 requests at 50 rps over three model groups, two guardrails and the three request shapes

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf case asserts time_to_first_token_seconds at percentile 50 and 99

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf provision creates every JWT user as internal_user

* test(e2e): perf observe reads only the mock log fields that remain

* test(e2e): 50 rps case sends short and long coding-agent sessions sized per request

* test(e2e): perf run keeps the per-request records as perf-records.jsonl under E2E_ARTIFACT_DIR

* test(e2e): the mixed small-gateway case runs at 25 rps (50 saturated one 4-worker pod on main and on v1.105.1)

* test(e2e): small_bodies_high_rate perf case, the github issue 46276 shape (11 to 17 KB chat bodies, keys only, 200 rps)

* test(e2e): the mixed small-gateway case runs at 10 rps for 6,000 requests (25 rps still saturated one pod on v1.105.1)

* ci: python lint job gets 20 minutes, the e2e harness phase (625 tests plus the replay harness suite) pushes a run past 15 when tests/e2e changes

* test(e2e): small-bodies perf case runs at 5 rps so gateway_added_seconds is per-request overhead, not queue wait

* test(e2e): mixed-provider perf case runs at 5 rps x 3,000 requests from 100 users

* test(e2e): mixed-provider perf case sends happy-path traffic only (no flaky deployment, no guardrail interventions or blocks)

* test(e2e): perf gateway model_list is three happy-path deployments, one per provider

* test(e2e): Postgres and Redis work per request as perf case limits

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): perf mocks wait a case-level latency with jitter instead of replaying captured timings, cpu_seconds_per_request limit from metrics-server

* test(e2e): perf suite PR carries the small_bodies_5rps case only, the mixed-provider case moves to its own PR

* test(e2e): small_bodies_5rps limits from the ten v1.105.1 baselines (1.5x on datastore counts, 2x on latency and cpu, absolute on errors and memory)

* test(e2e): perf mocks keep only the short captures the small case reads, captured timings and the Bedrock invoke route removed

* test(e2e): perf mocks generate each reply from the request (provider schema, usage, ids, event order, chunk sizes) instead of replaying captures

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-15 13:14:40 -07:00 • committed by GitHub
parent b4a63d3b6a
commit fbd078bcbe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 3687 additions and 37 deletions

View file

@ -21,7 +21,7 @@ jobs:
pull-requests: read
actions: read
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 20
services:
redis:
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f

View file

@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection
- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak)
- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in
- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres transactions and Redis commands each request cost (`postgres_calls_per_request`, `redis_calls_per_request`, from `pg_stat_database` and Redis `INFO` read by `datastore.py` before and after the run). Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe`
- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite
- `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/<name>` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=<system>` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager_<system>_ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_<system>.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up <system>` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only)
- `gateway/` - proxy configuration only (`litellm-config.yml`). No tests

View file

@ -481,15 +481,40 @@ def get_external[R: BaseModel](
response_type: type[R],
headers: BaseModel | None = None,
timeout: float = 30.0,
ca_bundle: str | None = None,
) -> Result[R]:
"""GET an absolute URL outside the proxy (e.g. a public /.well-known document).
Unlike the transport wrappers there is no proxy base url and no proxy auth; the
response still gets the same tagged-union classification as every other call."""
response still gets the same tagged-union classification as every other call.
`ca_bundle` is a PEM path for a server the default trust store does not cover."""
try:
resp = requests.get(
url,
headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
timeout=timeout,
verify=ca_bundle if ca_bundle is not None else True,
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return classify(resp, response_type)
def post_external[R: BaseModel](
url: str,
*,
json: BaseModel,
response_type: type[R],
headers: BaseModel | None = None,
timeout: float = 30.0,
) -> Result[R]:
"""POST JSON to an absolute URL outside the proxy (e.g. a mock server's control
route), classified the same way as every other call."""
try:
resp = requests.post(
url,
headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
json=wire_body(json),
timeout=timeout,
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))

View file

@ -0,0 +1,77 @@
# Gateway config for the perf cases (tests/e2e/perf/test_perf_cases.py), the local twin of the
# proxy_config block in project-releaser's .buildkite/e2e-perf-gateway-values.yaml: the same model
# groups over the mock providers, the same two guardrails, so a laptop run and the lane read the same
# config. One deployment per mock provider pod; how long the mocks wait before answering is set by the
# case at setup (mock_latency), not here. Teams, users and keys are created by the case at run time.
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
store_model_in_db: true
proxy_config_reload_interval_seconds: 7
enable_jwt_auth: true
litellm_jwtauth:
user_id_jwt_field: sub
user_email_jwt_field: email
team_ids_jwt_field: groups
user_id_upsert: true
proxy_batch_write_at: 5
database_connection_pool_limit: 10
litellm_settings:
cache: true
cache_params:
type: redis
host: os.environ/REDIS_HOST
port: os.environ/REDIS_PORT
supported_call_types: []
callbacks:
- prometheus
- otel
require_auth_for_metrics_endpoint: false
drop_params: true
request_timeout: 120
router_settings:
num_retries: 2
retry_after: 0
allowed_fails: 3
cooldown_time: 5
routing_strategy: simple-shuffle
model_list:
- model_name: gpt-6-luna
litellm_params:
model: openai/gpt-6-luna
api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL
api_key: perf-mock-not-checked
- model_name: claude-sonnet-5
litellm_params:
model: anthropic/claude-sonnet-5
api_base: os.environ/E2E_PERF_MOCK_ANTHROPIC_URL
api_key: perf-mock-not-checked
- model_name: claude-sonnet-5-bedrock
litellm_params:
model: bedrock/anthropic.claude-sonnet-5
aws_bedrock_runtime_endpoint: os.environ/E2E_PERF_MOCK_BEDROCK_URL
aws_region_name: us-east-1
aws_access_key_id: perf-mock-not-checked
aws_secret_access_key: perf-mock-not-checked
guardrails:
- guardrail_name: perf-pre
litellm_params:
guardrail: generic_guardrail_api
mode: pre_call
api_base: os.environ/E2E_PERF_MOCK_GUARDRAIL_URL
default_on: true
streaming_end_of_stream_only: true
additional_provider_specific_params:
latency_s: 0.05
- guardrail_name: perf-post
litellm_params:
guardrail: generic_guardrail_api
mode: post_call
api_base: os.environ/E2E_PERF_MOCK_GUARDRAIL_URL
default_on: true
streaming_end_of_stream_only: true
additional_provider_specific_params:
latency_s: 0.05

View file

@ -0,0 +1,22 @@
"""Every module in this package that defines a `CASE` is one perf test."""
from __future__ import annotations
import importlib
import pkgutil
from typing import Final
from setup.performance_test_case import PerformanceTestCase
def _load() -> tuple[PerformanceTestCase, ...]:
found: list[PerformanceTestCase] = []
for module_info in pkgutil.iter_modules(__path__):
module = importlib.import_module(f"{__name__}.{module_info.name}")
case: object = getattr(module, "CASE", None)
if isinstance(case, PerformanceTestCase):
found.append(case)
return tuple(sorted(found, key=lambda case: case.name))
CASES: Final = _load()

View file

@ -0,0 +1,55 @@
"""Small bodies at a low rate on the small gateway: the shape of github issue #46276 (short chat
requests, master key and virtual keys, no JWT, no guardrails) at 5 rps on one 4-worker pod, far
under what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own
per-request overhead. A slower middleware chain or an extra Redis round trip per request shows up
directly instead of hiding behind queue wait."""
from __future__ import annotations
from typing import Final
from setup.performance_test_case import (
Limit,
MockLatency,
PerformanceTestCase,
RequestShape,
Team,
Traffic,
VirtualKey,
)
TINY: Final = RequestShape("tiny", input_tokens=(2_500, 4_000))
CASE: Final = PerformanceTestCase(
name="small_bodies_5rps",
teams=(Team("team-a", models=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock")),),
users=(),
virtual_keys=(
VirtualKey("key-a1", team="team-a"),
VirtualKey("key-a2", team="team-a"),
VirtualKey("key-plain"),
),
model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"),
mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25),
guardrails=(),
traffic=Traffic(
seed=11,
rps=5.0,
total_requests=3000,
endpoints={"chat_completions": 0.6, "responses": 0.2, "messages": 0.2},
shapes={TINY: 1.0},
stream_share=0.5,
intervene_share=0.0,
block_share=0.0,
concurrency=20,
),
limits=(
Limit("unexpected_status_count", max=3),
Limit("gateway_added_seconds", percentile=50, max=0.38),
Limit("time_to_first_token_seconds", percentile=99, max=3.0),
Limit("memory_growth_megabytes", max=100),
Limit("postgres_calls_per_request", max=13),
Limit("redis_calls_per_request", max=35),
Limit("cpu_seconds_per_request", max=0.19),
),
)

View file

@ -70,37 +70,24 @@ class AnomalyReport:
@property
def warm_cache_read_share(self) -> float:
billed = (
self.warm_uncached_input_tokens
+ self.warm_cache_read_tokens
+ self.warm_cache_creation_tokens
)
billed = self.warm_uncached_input_tokens + self.warm_cache_read_tokens + self.warm_cache_creation_tokens
return self.warm_cache_read_tokens / billed if billed else 0.0
def _system_prefix_block(marker: str) -> TextBlock:
text = " ".join(
f"Project context paragraph {index} for session {marker}." for index in range(300)
)
text = " ".join(f"Project context paragraph {index} for session {marker}." for index in range(300))
return TextBlock(text=text, cache_control=CacheControl())
def _user_turn_text(marker: str, turn_index: int) -> str:
notes = " ".join(
f"Working note {index} of turn {turn_index} in session {marker}."
for index in range(80)
)
notes = " ".join(f"Working note {index} of turn {turn_index} in session {marker}." for index in range(80))
return f"Reply with one short sentence.\n{notes}"
def _reminder_turn() -> RichMessage:
return RichMessage(
role="system",
content=[
TextBlock(
text="<system-reminder>Keep the answer to one short sentence.</system-reminder>"
)
],
content=[TextBlock(text="<system-reminder>Keep the answer to one short sentence.</system-reminder>")],
)
@ -127,9 +114,7 @@ def retried(
return retried(call, attempts - 1, backoff_seconds, sleep)
def _metric(
result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float
) -> TurnMetric:
def _metric(result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float) -> TurnMetric:
if isinstance(result, Success):
usage = result.data.usage
return TurnMetric(
@ -167,11 +152,7 @@ def _drive_turns(
return ()
user_turn = RichMessage(
role="user",
content=[
TextBlock(
text=_user_turn_text(marker, turn_index), cache_control=CacheControl()
)
],
content=[TextBlock(text=_user_turn_text(marker, turn_index), cache_control=CacheControl())],
)
started = time.monotonic()
result = retried(
@ -190,9 +171,7 @@ def _drive_turns(
turn = _metric(result, turn_index, time.monotonic() - started)
if not isinstance(result, Success):
return (turn,)
assistant_turn = RichMessage(
role="assistant", content=[TextBlock(text=result.data.text or "Understood.")]
)
assistant_turn = RichMessage(role="assistant", content=[TextBlock(text=result.data.text or "Understood.")])
return (
turn,
*_drive_turns(
@ -242,9 +221,7 @@ def run_concurrent_sessions(
) -> tuple[TurnMetric, ...]:
with ThreadPoolExecutor(max_workers=sessions) as pool:
futures = [
pool.submit(
run_session, transport, key, model, turns_per_session, attempts_per_turn
)
pool.submit(run_session, transport, key, model, turns_per_session, attempts_per_turn)
for _ in range(sessions)
]
return tuple(turn for future in futures for turn in future.result())
@ -296,7 +273,5 @@ def summarize(turns: tuple[TurnMetric, ...], planned_turns: int) -> AnomalyRepor
warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm),
warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm),
warm_cache_creation_tokens=sum(turn.cache_creation_tokens for turn in warm),
p95_turn_seconds=_p95(
tuple(turn.latency_seconds for turn in turns if turn.ok)
),
p95_turn_seconds=_p95(tuple(turn.latency_seconds for turn in turns if turn.ok)),
)

View file

View file

@ -0,0 +1,254 @@
"""A coding-agent session as the gateway sees it. The system prompt and tool set come from this
directory, the turns are the assistant reading files and the files coming back as tool results, and the
files are real Python files from tests/e2e, chosen by the request's marker. The input budget is spent
on turns, so a 300k request is a longer session over bigger files than a 100k one, in the wire format of
the endpoint it goes to."""
from __future__ import annotations
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import cache
from pathlib import Path
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, TypeAdapter
Endpoint = Literal["chat", "responses", "messages"]
CORPUS_DIR: Final = Path(__file__).resolve().parent
E2E_ROOT: Final = CORPUS_DIR.parents[2]
CHARS_PER_TOKEN: Final = 4
_FILE_BYTES: Final = (2_000, 80_000)
_STRIDE: Final = 7919
_QUESTION: Final = (
"Which of the files above raise an exception that the caller is not told about? Answer with the paths only."
)
_CACHE_CONTROL: Final = {"type": "ephemeral"}
class ToolSpec(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
description: str
input_schema: dict[str, object]
@dataclass(frozen=True, slots=True)
class Step:
"""One tool round trip: the assistant asks to read `path`, the user turn carries `result`."""
call_id: str
path: str
offset: int
limit: int
result: str
@dataclass(frozen=True, slots=True)
class Session:
system: str
tools: tuple[ToolSpec, ...]
steps: tuple[Step, ...]
question: str
def estimate_tokens(text: str) -> int:
return max(1, len(text) // CHARS_PER_TOKEN)
@cache
def system_prompt() -> str:
return (CORPUS_DIR / "system_prompt.md").read_text()
@cache
def tool_specs() -> tuple[ToolSpec, ...]:
return tuple(TypeAdapter(list[ToolSpec]).validate_json((CORPUS_DIR / "tools.json").read_bytes()))
@cache
def corpus_files() -> tuple[Path, ...]:
"""Python files under tests/e2e between 2 KB and 80 KB, in path order so every run sees the same list.
The e2e tree is what the runner image ships, so a laptop and the lane read the same files at the same sha.
"""
low, high = _FILE_BYTES
found: Final = tuple(path for path in sorted(E2E_ROOT.rglob("*.py")) if low <= path.stat().st_size <= high)
assert found, f"no corpus files under {E2E_ROOT}"
return found
def _digest(seed: str) -> int:
return int.from_bytes(hashlib.sha1(seed.encode()).digest()[:6], "big")
def _numbered(lines: tuple[str, ...], offset: int) -> str:
return "\n".join(f"{offset + index:>6}\t{line}" for index, line in enumerate(lines, start=1))
def _step(marker: str, index: int, budget_tokens: int) -> Step:
files: Final = corpus_files()
path: Final = files[(_digest(marker) + index * _STRIDE) % len(files)]
lines: Final = tuple(path.read_text(errors="replace").splitlines())
whole: Final = _numbered(lines, 1)
if estimate_tokens(whole) <= budget_tokens:
return Step(_call_id(marker, index), f"tests/e2e/{path.relative_to(E2E_ROOT)}", 1, len(lines), whole)
keep: Final = max(1, len(lines) * budget_tokens * CHARS_PER_TOKEN // max(1, len(whole)))
return Step(_call_id(marker, index), f"tests/e2e/{path.relative_to(E2E_ROOT)}", 1, keep, _numbered(lines[:keep], 1))
def _call_id(marker: str, index: int) -> str:
return hashlib.sha1(f"{marker}:call{index}".encode()).hexdigest()[:24]
def build_session(marker: str, input_tokens: int, *, tools: bool) -> Session:
"""Spend `input_tokens` on file-reading turns after the system prompt, the tool set and the question."""
specs: Final = tool_specs() if tools else ()
fixed: Final = estimate_tokens(system_prompt()) + sum(estimate_tokens(spec.model_dump_json()) for spec in specs)
question: Final = f"{marker} {_QUESTION}"
steps: list[Step] = [] # mutable-ok: each step's size decides whether another fits
spent = fixed + estimate_tokens(question) # rebind-ok: running total over the steps
while spent < input_tokens:
steps.append(step := _step(marker, len(steps), input_tokens - spent))
spent += estimate_tokens(step.result) + 40
return Session(system_prompt(), specs, tuple(steps), question)
def _read_arguments(step: Step) -> dict[str, object]:
return {"file_path": f"/workspace/litellm/{step.path}", "offset": step.offset, "limit": step.limit}
def _chat_turns(session: Session) -> Sequence[Mapping[str, object]]:
if not session.tools:
return [
turn
for step in session.steps
for turn in (
{"role": "user", "content": f"Review this file.\n\n{step.path}\n{step.result}"},
{"role": "assistant", "content": f"Read {step.path}, {step.limit} lines. Continue."},
)
]
return [
turn
for step in session.steps
for turn in (
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": f"call_{step.call_id}",
"type": "function",
"function": {"name": "Read", "arguments": json.dumps(_read_arguments(step))},
}
],
},
{"role": "tool", "tool_call_id": f"call_{step.call_id}", "content": step.result},
)
]
def _responses_input(session: Session) -> Sequence[Mapping[str, object]]:
if not session.tools:
return _chat_turns(session)
return [
turn
for step in session.steps
for turn in (
{
"type": "function_call",
"call_id": f"call_{step.call_id}",
"name": "Read",
"arguments": json.dumps(_read_arguments(step)),
},
{"type": "function_call_output", "call_id": f"call_{step.call_id}", "output": step.result},
)
]
def _messages_turns(session: Session) -> Sequence[Mapping[str, object]]:
if not session.tools:
return _chat_turns(session)
last: Final = len(session.steps) - 1
return [
turn
for index, step in enumerate(session.steps)
for turn in (
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{step.call_id}", "name": "Read", "input": _read_arguments(step)}
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": f"toolu_{step.call_id}",
"content": step.result,
**({"cache_control": _CACHE_CONTROL} if index == last else {}),
}
],
},
)
]
def request_body(
*, endpoint: Endpoint, model: str, marker: str, input_tokens: int, tools: bool, output_tokens: int, stream: bool
) -> dict[str, object]:
"""The session as one request in the endpoint's wire format, the marker in the final user message."""
session: Final = build_session(marker, input_tokens, tools=tools)
question: Final = {"role": "user", "content": session.question}
if endpoint == "responses":
return {
"model": model,
"instructions": session.system,
"input": [*_responses_input(session), question],
"max_output_tokens": output_tokens,
"stream": stream,
**(
{
"tools": [
{"type": "function", "name": s.name, "description": s.description, "parameters": s.input_schema}
for s in session.tools
]
}
if session.tools
else {}
),
}
if endpoint == "messages":
return {
"model": model,
"system": [{"type": "text", "text": session.system, "cache_control": _CACHE_CONTROL}],
"messages": [*_messages_turns(session), question],
"max_tokens": output_tokens,
"stream": stream,
**({"tools": [spec.model_dump() for spec in session.tools]} if session.tools else {}),
}
return {
"model": model,
"messages": [{"role": "system", "content": session.system}, *_chat_turns(session), question],
"max_completion_tokens": output_tokens,
**({"reasoning_effort": "none"} if session.tools else {}),
"stream": stream,
**(
{
"tools": [
{
"type": "function",
"function": {"name": s.name, "description": s.description, "parameters": s.input_schema},
}
for s in session.tools
]
}
if session.tools
else {}
),
}

View file

@ -0,0 +1,132 @@
You are an interactive CLI tool that helps users with software engineering tasks. Use the instructions below and the tools available to you to assist the user.
IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation.
IMPORTANT: You must NEVER generate or guess URLs for the user unless you are confident that the URLs are for helping the user with programming. You may use URLs provided by the user in their messages or local files.
If the user asks for help or wants to give feedback inform them of the following:
- /help: Get help with using the tool
- To give feedback, users should report the issue at the project's issue tracker
When the user directly asks about the tool (eg 'can the tool do...', 'does the tool have...') or asks in second person (eg 'are you able...', 'can you do...'), first use the WebFetch tool to gather information to answer the question from the documentation.
# Tone and style
You should be concise, direct, and to the point. When you run a non-trivial bash command, you should explain what the command does and why you are running it, to make sure the user understands what you are doing (this is especially important when you are running a command that will make changes to the user's system).
Remember that your output will be displayed on a command line interface. Your responses can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification.
Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session.
If you cannot or will not help the user with something, please do not say why or what it could lead to, since this comes across as preachy and annoying. Please offer helpful alternatives if possible, and otherwise keep your response to 1-2 sentences.
Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked.
IMPORTANT: You should minimize output tokens as much as possible while maintaining helpfulness, quality, and accuracy. Only address the specific query or task at hand, avoiding tangential information unless absolutely critical for completing the request. If you can answer in 1-3 sentences or a short paragraph, please do.
IMPORTANT: You should NOT answer with unnecessary preamble or postamble (such as explaining your code or summarizing your action), unless the user asks you to.
IMPORTANT: Keep your responses short, since they will be displayed on a command line interface. You MUST answer concisely with fewer than 4 lines (not including tool use or code generation), unless user asks for detail. Answer the user's question directly, without elaboration, explanation, or details. One word answers are best. Avoid introductions, conclusions, and explanations. You MUST avoid text before/after your response, such as "The answer is <answer>.", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...". Here are some examples to demonstrate appropriate verbosity:
<example>
user: 2 + 2
assistant: 4
</example>
<example>
user: what is 2+2?
assistant: 4
</example>
<example>
user: is 11 a prime number?
assistant: Yes
</example>
<example>
user: what command should I run to list files in the current directory?
assistant: ls
</example>
<example>
user: what command should I run to watch files in the current directory?
assistant: [use the ls tool to list the files in the current directory, then read docs/commands in the relevant file to find out how to watch files]
npm run dev
</example>
<example>
user: How many golf balls fit inside a jetta?
assistant: 150000
</example>
<example>
user: what files are in the directory src/?
assistant: [runs ls and sees foo.c, bar.c, baz.c]
user: which file contains the implementation of foo?
assistant: src/foo.c
</example>
<example>
user: write tests for new feature
assistant: [uses grep and glob search tools to find where similar tests are defined, uses concurrent read file tool use blocks in one single message to read relevant files at the same time, uses edit file tool to write new tests]
</example>
# Proactiveness
You are allowed to be proactive, but only when the user asks you to do something. You should strive to strike a balance between:
1. Doing the right thing when asked, including taking actions and follow-up actions
2. Not surprising the user with actions you take without asking
For example, if the user asks you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions.
3. Do not add additional code explanation summary unless requested by the user. After working on a file, just stop, rather than providing an explanation of what you did.
# Following conventions
When making changes to files, first understand the file's code conventions. Mimic code style, use existing libraries and utilities, and follow existing patterns.
- NEVER assume that a given library is available, even if it is well known. Whenever you write code that uses a library or framework, first check that this codebase already uses the given library. For example, you might look at neighboring files, or check the package.json (or cargo.toml, and so on depending on the language).
- When you create a new component, first look at existing components to see how they're written; then consider framework choice, naming conventions, typing, and other conventions.
- When you edit a piece of code, first look at the code's surrounding context (especially its imports) to understand the code's choice of frameworks and libraries. Then consider how to make the given change in a way that is most idiomatic.
- Always follow security best practices. Never introduce code that exposes or logs secrets and keys. Never commit secrets or keys to the repository.
# Code style
- IMPORTANT: DO NOT ADD ***ANY*** COMMENTS unless asked
# Task Management
You have access to the TodoWrite and TodoRead tools to help you manage and plan tasks. Use these tools VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress.
These tools are also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable.
It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed.
# Doing tasks
The user will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended:
- Use the TodoWrite tool to plan the task if required
- Use the available search tools to understand the codebase and the user's query. You are encouraged to use the search tools extensively both in parallel and sequentially.
- Implement the solution using all tools available to you
- Verify the solution if possible with tests. NEVER assume specific test framework or test script. Check the README or search codebase to determine the testing approach.
- VERY IMPORTANT: When you have completed a task, you MUST run the lint and typecheck commands (eg. npm run lint, npm run typecheck, ruff, etc.) with Bash if they were provided to you to ensure your code is correct. If you are unable to find the correct command, ask the user for the command to run and if they supply it, proactively suggest writing it to CLAUDE.md so that you will know to run it next time.
NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise the user will feel that you are committing without their permission.
- Tool results and user messages may include <system-reminder> tags. <system-reminder> tags contain useful information and reminders. They are NOT part of the user's provided input or the tool result.
# Tool usage policy
- When doing file search, prefer to use the Task tool in order to reduce context usage.
- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description.
- When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response.
- You have the capability to call multiple tools in a single response. When multiple independent pieces of information are requested, batch your tool calls together for optimal performance. When making multiple bash tool calls, you MUST send a single message with multiple tools calls to run the calls in parallel. For example, if you need to run "git status" and "git diff", send a single message with two tool calls to run the calls in parallel.
You MUST answer concisely with fewer than 4 lines of text (not including tool use or code generation), unless user asks for detail.
Here is useful information about the environment you are running in:
<env>
Working directory: /workspace/litellm
Is directory a git repo: Yes
Platform: linux
OS Version: Linux 6.8.0
Today's date: 2026-10-14
</env>
You are powered by the model named the configured model. The exact model ID is the one the gateway routes this request to.
IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation.
# Code References
When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow the user to easily navigate to the source code location.
<example>
user: Where are errors from the client handled?
assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712.
</example>
gitStatus: This is the git status at the start of the conversation. Note that this status is a snapshot in time, and will not update during the conversation.
Current branch: main
Main branch (you will usually use this for PRs): main
Status: (clean)
Recent commits:
abcdef0 fix(router): keep the stream retry inside the same deployment group
1234567 docs: router_settings reference for the new fallback flag
89abcde test(e2e): perf cases on an out-of-process rig

View file

@ -0,0 +1,95 @@
[
{
"name": "Read",
"description": "Reads a file from the local filesystem. You can access any file directly by using this tool. Assume this tool is able to read all files on the machine. If the user provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned. The file_path parameter must be an absolute path, not a relative path. By default, it reads up to 2000 lines starting from the beginning of the file. You can optionally specify a line offset and limit (especially handy for long files), but it is recommended to read the whole file by not providing these parameters. Any lines longer than 2000 characters will be truncated. Results are returned using cat -n format, with line numbers starting at 1. This tool can only read files, not directories. To read a directory, use an ls command via the Bash tool. You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful. If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents.",
"input_schema": {
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "The absolute path to the file to read"},
"offset": {"type": "number", "description": "The line number to start reading from. Only provide if the file is too large to read at once"},
"limit": {"type": "number", "description": "The number of lines to read. Only provide if the file is too large to read at once."}
},
"required": ["file_path"],
"additionalProperties": false
}
},
{
"name": "Edit",
"description": "Performs exact string replacements in files. You must use your Read tool at least once in the conversation before editing. This tool will error if you attempt an edit without reading the file. When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears AFTER the line number prefix. The line number prefix format is: spaces + line number + tab. Everything after that tab is the actual file content to match. Never include any part of the line number prefix in the old_string or new_string. ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked. The edit will FAIL if old_string is not unique in the file. Either provide a larger string with more surrounding context to make it unique or use replace_all to change every instance of old_string. Use replace_all for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.",
"input_schema": {
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "The absolute path to the file to modify"},
"old_string": {"type": "string", "description": "The text to replace"},
"new_string": {"type": "string", "description": "The text to replace it with (must be different from old_string)"},
"replace_all": {"type": "boolean", "default": false, "description": "Replace all occurences of old_string (default false)"}
},
"required": ["file_path", "old_string", "new_string"],
"additionalProperties": false
}
},
{
"name": "Write",
"description": "Writes a file to the local filesystem. This tool will overwrite the existing file if there is one at the provided path. If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first. ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User. Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.",
"input_schema": {
"type": "object",
"properties": {
"file_path": {"type": "string", "description": "The absolute path to the file to write (must be absolute, not relative)"},
"content": {"type": "string", "description": "The content to write to the file"}
},
"required": ["file_path", "content"],
"additionalProperties": false
}
},
{
"name": "Bash",
"description": "Executes a given bash command in a persistent shell session with optional timeout, ensuring proper handling and security measures. Before executing the command, please follow these steps: 1. Directory Verification: If the command will create new directories or files, first use the LS tool to verify the parent directory exists and is the correct location. 2. Command Execution: Always quote file paths that contain spaces with double quotes. After ensuring proper quoting, execute the command. Capture the output of the command. Usage notes: The command argument is required. You can specify an optional timeout in milliseconds (up to 600000ms / 10 minutes). If not specified, commands will timeout after 120000ms (2 minutes). It is very helpful if you write a clear, concise description of what this command does in 5-10 words. If the output exceeds 30000 characters, output will be truncated before being returned to you. VERY IMPORTANT: You MUST avoid using search commands like find and grep. Instead use Grep, Glob, or Task to search. You MUST avoid read tools like cat, head, tail, and ls, and use Read and LS to read files. If you still need to run grep, STOP. ALWAYS USE ripgrep at rg first, which all Claude Code users have pre-installed. When issuing multiple commands, use the ; or && operator to separate them. DO NOT use newlines to separate commands (newlines are ok in quoted strings). Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd. You may use cd if the User explicitly requests it.",
"input_schema": {
"type": "object",
"properties": {
"command": {"type": "string", "description": "The command to execute"},
"timeout": {"type": "number", "description": "Optional timeout in milliseconds (max 600000)"},
"description": {"type": "string", "description": "Clear, concise description of what this command does in 5-10 words."},
"run_in_background": {"type": "boolean", "description": "Set to true to run this command in the background. Use BashOutput to read the output later."}
},
"required": ["command"],
"additionalProperties": false
}
},
{
"name": "Grep",
"description": "A powerful search tool built on ripgrep. ALWAYS use Grep for search tasks. NEVER invoke grep or rg as a Bash command. The Grep tool has been optimized for correct permissions and access. Supports full regex syntax (e.g., \"log.*Error\", \"function\\s+\\w+\"). Filter files with glob parameter (e.g., \"*.js\", \"**/*.tsx\") or type parameter (e.g., \"js\", \"py\", \"rust\"). Output modes: \"content\" shows matching lines, \"files_with_matches\" shows only file paths (default), \"count\" shows match counts. Use Task tool for open-ended searches requiring multiple rounds. Pattern syntax: Uses ripgrep (not grep), so literal braces need escaping (use interface\\{\\} to find interface{} in Go code). Multiline matching: By default patterns match within single lines only. For cross-line patterns like struct \\{[\\s\\S]*?field, use multiline: true.",
"input_schema": {
"type": "object",
"properties": {
"pattern": {"type": "string", "description": "The regular expression pattern to search for in file contents"},
"path": {"type": "string", "description": "File or directory to search in (rg PATH). Defaults to current working directory."},
"glob": {"type": "string", "description": "Glob pattern to filter files (e.g. \"*.js\", \"*.{ts,tsx}\") - maps to rg --glob"},
"output_mode": {"type": "string", "enum": ["content", "files_with_matches", "count"], "description": "Output mode: \"content\" shows matching lines, \"files_with_matches\" shows file paths (supports head_limit), \"count\" shows match counts (supports head_limit). Defaults to \"files_with_matches\"."},
"-B": {"type": "number", "description": "Number of lines to show before each match (rg -B). Requires output_mode: \"content\", ignored otherwise."},
"-A": {"type": "number", "description": "Number of lines to show after each match (rg -A). Requires output_mode: \"content\", ignored otherwise."},
"-C": {"type": "number", "description": "Number of lines to show before and after each match (rg -C). Requires output_mode: \"content\", ignored otherwise."},
"-n": {"type": "boolean", "description": "Show line numbers in output (rg -n). Requires output_mode: \"content\", ignored otherwise. Defaults to true."},
"-i": {"type": "boolean", "description": "Case insensitive search (rg -i)"},
"type": {"type": "string", "description": "File type to search (rg --type). Common types: js, py, rust, go, java, etc. More efficient than include for standard file types."},
"head_limit": {"type": "number", "description": "Limit output to first N lines/entries, equivalent to \"| head -N\". Works across all output modes: content (limits output lines), files_with_matches (limits file paths), count (limits count entries). Defaults to 250 when unspecified."},
"multiline": {"type": "boolean", "description": "Enable multiline mode where . matches newlines and patterns can span lines (rg -U --multiline-dotall). Default: false."}
},
"required": ["pattern"],
"additionalProperties": false
}
},
{
"name": "Glob",
"description": "Fast file pattern matching tool that works with any codebase size. Supports glob patterns like \"**/*.js\" or \"src/**/*.ts\". Returns matching file paths sorted by modification time. Use this tool when you need to find files by name patterns. When you are doing an open ended search that may require multiple rounds of globbing and grepping, use the Task tool instead. You have the capability to call multiple tools in a single response. It is always better to speculatively perform multiple searches as a batch that are potentially useful.",
"input_schema": {
"type": "object",
"properties": {
"pattern": {"type": "string", "description": "The glob pattern to match files against"},
"path": {"type": "string", "description": "The directory to search in. If not specified, the current working directory will be used. IMPORTANT: Omit this parameter if you do not provide a value - this is the default."}
},
"required": ["pattern"],
"additionalProperties": false
}
}
]

View file

@ -0,0 +1,64 @@
"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling
on how much database and cache work the gateway does per request. Both come from the servers' own
statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them.
`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the
counters are None and any Limit on them reports as missing."""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Final
from pydantic import TypeAdapter
_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()"
_DB_ROW: Final = TypeAdapter(tuple[int])
_INFO: Final = TypeAdapter(dict[str, object])
_INT: Final = TypeAdapter(int)
@dataclass(frozen=True, slots=True)
class DbCounters:
commits: int
@dataclass(frozen=True, slots=True)
class RedisCounters:
commands: int
connected_clients: int
def read_db_counters() -> DbCounters | None:
url: Final = os.environ.get("DATABASE_URL")
if url is None:
return None
import psycopg
with psycopg.connect(url) as conn:
row: Final = conn.execute(_DB_QUERY).fetchone()
(commits,) = _DB_ROW.validate_python(row)
return DbCounters(commits=commits)
def read_redis_counters() -> RedisCounters | None:
"""Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are."""
host: Final = os.environ.get("REDIS_HOST")
if host is None:
return None
import redis
client: Final = redis.Redis(
host=host,
port=int(os.environ.get("REDIS_PORT", "6379")),
ssl=os.environ.get("REDIS_SSL", "false").lower() == "true",
password=os.environ.get("REDIS_PASSWORD") or None,
decode_responses=True,
socket_connect_timeout=2,
)
stats: Final = _INFO.validate_python(client.info("stats"))
clients: Final = _INFO.validate_python(client.info("clients"))
return RedisCounters(
commands=_INT.validate_python(stats["total_commands_processed"]),
connected_clients=_INT.validate_python(clients["connected_clients"]),
)

View file

@ -0,0 +1,266 @@
"""The Locust side of a perf case: simulated users take the next request off the pre-generated plan,
send it at the planned rate, and write one record per request with what the client saw: status,
time to headers, time to first byte, time to first token (streams), total time, the text read back
and whether it matched. Runs in its own process (`setup/run.py` starts it), never in the gateway's.
Reads `PERF_PLAN_FILE` (JSON list of wire requests), writes `PERF_RECORDS_FILE` (one JSON object per
line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last request is recorded.
"""
from __future__ import annotations
import os
import time
from collections.abc import Iterator
from pathlib import Path
from typing import Final, cast
import gevent
from locust import FastHttpUser, constant_throughput, task
from locust.contrib.fasthttp import FastResponse, ResponseContextManager
from setup.wire import Endpoint, WireRecord, WireRequest, request_body
from pydantic import TypeAdapter
_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest])
_JSON: Final = TypeAdapter[object](object)
_PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes())
_RECORDS_PATH: Final = Path(os.environ["PERF_RECORDS_FILE"])
_RPS: Final = float(os.environ["PERF_RPS"])
_USERS: Final = int(os.environ["PERF_USERS"])
_NEXT: Final[Iterator[WireRequest]] = iter(_PLAN)
_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8")
_written = 0
BLOCK_REASON: Final = "perf-test-block"
_TOKEN_MAX_AGE_S: Final = 240.0
_tokens: dict[str, tuple[str, float]] = {}
def _as_dict(value: object) -> dict[str, object]:
return cast(dict[str, object], value) if isinstance(value, dict) else {}
def _as_list(value: object) -> list[object]:
return cast(list[object], value) if isinstance(value, list) else []
def _first_choice_field(payload: dict[str, object], holder: str) -> str:
choices: Final = _as_list(payload.get("choices"))
if not choices:
return ""
content: Final = _as_dict(_as_dict(choices[0]).get(holder)).get("content")
return content if isinstance(content, str) else ""
def text_of_json(endpoint: Endpoint, payload: dict[str, object]) -> str:
"""The assistant text in a non-streaming body of each endpoint."""
if endpoint == "chat_completions":
return _first_choice_field(payload, "message")
if endpoint == "responses":
parts: list[str] = []
for item in _as_list(payload.get("output")):
for part in _as_list(_as_dict(item).get("content")):
text = _as_dict(part).get("text")
if _as_dict(part).get("type") == "output_text" and isinstance(text, str):
parts.append(text)
return "".join(parts)
blocks: list[str] = []
for block in _as_list(payload.get("content")):
text = _as_dict(block).get("text")
if isinstance(text, str):
blocks.append(text)
return "".join(blocks)
def delta_of_event(endpoint: Endpoint, payload: dict[str, object]) -> str:
"""The text delta one SSE event of each endpoint carries, empty when it carries none."""
if endpoint == "chat_completions":
return _first_choice_field(payload, "delta")
if endpoint == "responses":
delta: Final = payload.get("delta")
return delta if payload.get("type") == "response.output_text.delta" and isinstance(delta, str) else ""
if payload.get("type") == "content_block_delta":
text: Final = _as_dict(payload.get("delta")).get("text")
return text if isinstance(text, str) else ""
return ""
def _events(buffer: bytes) -> tuple[list[dict[str, object]], bytes]:
"""Complete SSE events in `buffer` as their JSON data, plus the unfinished tail."""
events: list[dict[str, object]] = []
while b"\n\n" in buffer:
raw, buffer = buffer.split(b"\n\n", 1)
for line in raw.split(b"\n"):
if line.startswith(b"data:"):
data = line[5:].strip()
if data and data != b"[DONE]":
events.append(_as_dict(_JSON.validate_json(data)))
return events, buffer
def _header(response: FastResponse, name: str) -> str | None:
headers: Final = response.headers
if headers is None:
return None
value: object = headers.get(name) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
return value if isinstance(value, str) else None
def _quit(user: FastHttpUser) -> None:
"""End the run once every record is written. Quitting from a separate greenlet: `runner.quit`
kills the user greenlets, and the user calling it would be killed mid-shutdown."""
runner = user.environment.runner # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
if runner is not None:
gevent.spawn(runner.quit) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportUnusedCallResult]
def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str | None]:
if item.blocked:
return (
status == 400 and BLOCK_REASON in body_text,
None if status == 400 else f"expected a 400 block, got {status}",
)
if status != 200:
return False, f"status {status}: {body_text[:200]}"
return True, None
def _record(record: WireRecord) -> None:
global _written
_ = _RECORDS.write(record.model_dump_json() + "\n")
_RECORDS.flush()
_written += 1
def _authorization(user: FastHttpUser, item: WireRequest) -> str:
"""The bearer for this request: the key as given, or the caller's JWT, re-minted from Keycloak
once it is older than the realm's access token lifespan allows."""
if item.token_source is None:
assert item.authorization is not None, f"request {item.seq} has neither a key nor a token source"
return item.authorization
source: Final = item.token_source
cached: Final = _tokens.get(source.username)
if cached is not None and time.monotonic() - cached[1] < _TOKEN_MAX_AGE_S:
return f"Bearer {cached[0]}"
minted = ""
with user.client.request( # pyright: ignore[reportUnknownMemberType]
"POST",
source.token_url,
data={
"grant_type": "password",
"client_id": source.client_id,
"username": source.username,
"password": source.password,
},
headers={"Content-Type": "application/x-www-form-urlencoded"},
catch_response=True,
name="keycloak token",
) as response:
payload: Final = _as_dict(_JSON.validate_json(response.text or "{}"))
token: Final = payload.get("access_token")
if not isinstance(token, str):
response.failure(f"no access_token in {response.status_code}") # pyright: ignore[reportUnknownMemberType]
raise RuntimeError(
f"could not mint a token for {source.username}: {response.status_code} {response.text!r}"
)
response.success()
minted = token
_tokens[source.username] = (minted, time.monotonic())
return f"Bearer {minted}"
class PlanUser(FastHttpUser):
"""One simulated client; all users share the plan and together hold the planned rate."""
wait_time = constant_throughput(_RPS / _USERS)
network_timeout = 120.0
connection_timeout = 30.0
@task
def send(self) -> None:
item: Final = next(_NEXT, None)
if item is None:
if _written >= len(_PLAN):
_quit(self)
return
started_at: Final = time.time()
started: Final = time.perf_counter()
first_byte_s: float | None = None
first_token_s: float | None = None
text = ""
body_text = ""
status = 0
ok = False
error: str | None = None
header_s = 0.0
response: FastResponse | None = None
try:
headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"}
context: ResponseContextManager
with self.client.request( # pyright: ignore[reportUnknownMemberType]
"POST",
item.path,
json=request_body(item),
headers=headers,
stream=item.stream,
catch_response=True,
name=f"{item.endpoint}{' stream' if item.stream else ''}",
) as context:
response = context
header_s = time.perf_counter() - started
status = int(context.status_code)
if item.stream and status == 200:
buffer = b""
raw = context.stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the socket-backed response
while True:
line = cast(bytes, raw.readline(b"\n")) # pyright: ignore[reportUnknownMemberType] # returns as soon as a line is complete
if not line:
break
now = time.perf_counter() - started
if first_byte_s is None:
first_byte_s = now
buffer += line
events, buffer = _events(buffer)
for event in events:
delta = delta_of_event(item.endpoint, event)
if delta:
if first_token_s is None:
first_token_s = now
text += delta
else:
body_text = context.text or ""
first_byte_s = time.perf_counter() - started
if status == 200:
text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text)))
ok, error = _judge(item, status, body_text)
if ok:
context.success()
else:
context.failure(error or "failed") # pyright: ignore[reportUnknownMemberType]
except Exception as exc: # noqa: BLE001 # anything the client raised is the record's error
ok = False
error = f"{type(exc).__name__}: {exc}"
total_s: Final = time.perf_counter() - started
delay_header: Final = _header(response, "llm_provider-x-mock-delay-s") if response is not None else None
_record(
WireRecord(
seq=item.seq,
marker=item.marker,
caller=item.caller,
endpoint=item.endpoint,
stream=item.stream,
blocked=item.blocked,
started_at=started_at,
status=status,
header_s=header_s,
first_byte_s=first_byte_s,
first_token_s=first_token_s,
total_s=total_s,
text=text,
ok=ok,
mock_delay_s=float(delay_header) if delay_header else None,
call_id=_header(response, "x-litellm-call-id") if response is not None else None,
error=error,
)
)
if _written >= len(_PLAN):
_quit(self)

View file

View file

@ -0,0 +1,40 @@
"""The Anthropic-shaped perf mock: /v1/messages, JSON and SSE, answering each request with a generated reply in
api.anthropic.com's wire format (thinking blocks, tool_use, cache usage, the message_start to message_stop event
order and its chunk sizes)."""
from __future__ import annotations
from typing import Final
from fastapi import FastAPI, Request
from fastapi.responses import Response
from mock_common import (
RequestLog,
bool_field,
json_reply,
mount_control,
parse_json,
str_field,
stream_reply,
)
SSE: Final = "text/event-stream"
def build_app() -> FastAPI:
app: Final = FastAPI(title="perf mock anthropic")
log: Final = RequestLog("anthropic")
mount_control(app, log)
async def messages(request: Request) -> Response:
body: Final = parse_json(await request.body())
stream: Final = bool_field(body, "stream")
decision: Final = log.decide(
route="messages", stream=stream, body=body, model=str_field(body, "model", "claude-sonnet-5")
)
return stream_reply(decision, SSE) if stream else await json_reply(decision)
app.add_api_route("/v1/messages", messages, methods=["POST"])
app.add_api_route("/messages", messages, methods=["POST"])
return app

View file

@ -0,0 +1,51 @@
"""The Bedrock-shaped perf mock: Converse, ConverseStream, InvokeModel and InvokeModelWithResponseStream under
/model/{model_id}/..., answering each request with a generated reply in bedrock-runtime's wire format. Streams are application/vnd.amazon.eventstream
frames (prelude, headers, CRCs), so the gateway's botocore-based decoder runs against the real framing."""
from __future__ import annotations
from typing import Final, Literal
from fastapi import FastAPI, Request
from fastapi.responses import Response
from mock_common import (
RequestLog,
json_reply,
mount_control,
parse_json,
stream_reply,
)
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
def build_app() -> FastAPI:
app: Final = FastAPI(title="perf mock bedrock")
log: Final = RequestLog("bedrock")
mount_control(app, log)
async def handle(model_id: str, request: Request, route: Literal["converse", "invoke"], stream: bool) -> Response:
body: Final = parse_json(await request.body())
decision: Final = log.decide(route=route, stream=stream, body=body, model=model_id)
return stream_reply(decision, EVENT_STREAM) if stream else await json_reply(decision)
async def converse(model_id: str, request: Request) -> Response:
return await handle(model_id, request, "converse", False)
async def converse_stream(model_id: str, request: Request) -> Response:
return await handle(model_id, request, "converse", True)
async def invoke(model_id: str, request: Request) -> Response:
return await handle(model_id, request, "invoke", False)
async def invoke_with_response_stream(model_id: str, request: Request) -> Response:
return await handle(model_id, request, "invoke", True)
app.add_api_route("/model/{model_id:path}/converse", converse, methods=["POST"])
app.add_api_route("/model/{model_id:path}/converse-stream", converse_stream, methods=["POST"])
app.add_api_route("/model/{model_id:path}/invoke", invoke, methods=["POST"])
app.add_api_route(
"/model/{model_id:path}/invoke-with-response-stream", invoke_with_response_stream, methods=["POST"]
)
return app

View file

@ -0,0 +1,227 @@
"""Shared behaviour of the perf mock providers.
Every reply is generated for the request just received (see `replies.py`): sized from its prompt and output
cap, a tool call when tools were offered, in the provider's own wire format with its usage fields, ids,
event order and chunk sizes, so the gateway parses, prices and streams it as it would the real thing. The
model id stays a real one (`gpt-6-luna`, `anthropic.claude-sonnet-5`) so every name-based branch in the
gateway runs, and one mock pod per provider serves any number of deployments. How long a mock waits before answering is the one setting a case
controls, `POST /mock/latency {"seconds": 1.0, "jitter_seconds": 0.25}`, set by the test's setup directly
on the pod, never through the gateway. The wait, the reply's size and its ids are pure functions of the request marker (`perfreq-000123-n`, put in the prompt by the plan
generator), so the test can join what each provider received against what each client saw.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import re
from collections import deque
from dataclasses import dataclass
from typing import Final, Literal, cast
from fastapi import FastAPI
from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel, ConfigDict, Field
from replies import Route, render_chunks, render_json, reply_for
from reply_common import Provider, Reply, chunks_after, unit
MARKER: Final = re.compile(r"perfreq-\d{6}-[nib]")
DELAY_HEADER: Final = "x-mock-delay-s"
Flag = Literal["n", "i", "b"]
def marker_in(text: str) -> str | None:
match: Final = MARKER.search(text)
return None if match is None else match.group(0)
def flag_of(marker: str) -> Flag:
last: Final = marker[-1]
assert last in ("n", "i", "b"), marker
return "n" if last == "n" else "i" if last == "i" else "b"
def collect_text(value: object) -> list[str]:
"""Every prompt string, whatever the wire format nests it in: `text` and `content` parts, the
Responses API's bare `input`, and `system` prompts."""
if isinstance(value, dict):
items: Final = cast(dict[object, object], value).items()
return [
part
for key, item in items
for part in (
[item]
if key in ("text", "content", "input", "system") and isinstance(item, str)
else collect_text(item)
)
] # comprehension-ok: flattening one nesting level of the request body
if isinstance(value, list):
return [part for item in cast(list[object], value) for part in collect_text(item)] # comprehension-ok: flatten
return []
def token_count(text: str) -> int:
return max(1, len(text) // 4)
def str_field(body: object, key: str, default: str) -> str:
if isinstance(body, dict):
value: object = cast(dict[object, object], body).get(key)
if isinstance(value, str):
return value
return default
def bool_field(body: object, key: str) -> bool:
if isinstance(body, dict):
value: object = cast(dict[object, object], body).get(key)
return value is True
return False
class MockRequest(BaseModel):
model_config = ConfigDict(frozen=True)
provider: Provider
marker: str | None
route: str
stream: bool
output_tokens: int
text_sha1: str
delay_s: float
class MockRequestsResponse(BaseModel):
model_config = ConfigDict(frozen=True)
requests: list[MockRequest]
class ResetResponse(BaseModel):
model_config = ConfigDict(frozen=True)
cleared: int
class Latency(BaseModel):
"""Seconds a mock waits before its first byte: `seconds` plus or minus up to `jitter_seconds`, drawn
from the request marker so the same request always waits the same."""
model_config = ConfigDict(frozen=True)
seconds: float = Field(default=0.0, ge=0.0)
jitter_seconds: float = Field(default=0.0, ge=0.0)
def for_seed(self, seed: str) -> float:
return max(0.0, self.seconds + self.jitter_seconds * (2 * unit(f"{seed}:latency") - 1))
@dataclass(frozen=True, slots=True)
class Decision:
"""What this mock does with one request: the reply it generates, on which route, after how long."""
marker: str | None
reply: Reply
route: Route
delay_s: float
"""Seconds before the first byte, what `gateway_added_seconds` subtracts from the client's time."""
@property
def seed(self) -> str:
return self.marker or "unmarked"
@property
def headers(self) -> dict[str, str]:
return {DELAY_HEADER: f"{self.delay_s:.4f}"}
class RequestLog:
"""What one mock pod received, in order, readable by the test at GET /mock/requests."""
def __init__(self, provider: Provider, capacity: int = 100_000) -> None:
self._provider: Final = provider
self._entries: deque[MockRequest] = deque(maxlen=capacity)
self._latency = Latency()
@property
def latency(self) -> Latency:
return self._latency
def set_latency(self, latency: Latency) -> Latency:
self._latency = latency
return self._latency
def decide(self, *, route: Route, stream: bool, body: object, model: str) -> Decision:
prompt_text: Final = " ".join(collect_text(body))
marker: Final = marker_in(prompt_text)
seed: Final = marker or "unmarked"
reply: Final = reply_for(
provider=self._provider, body=body, seed=seed, model=model, prompt_tokens=token_count(prompt_text)
)
decision: Final = Decision(marker=marker, reply=reply, route=route, delay_s=self._latency.for_seed(seed))
self._entries.append(
MockRequest(
provider=self._provider,
marker=marker,
route=route,
stream=stream,
output_tokens=reply.output_tokens,
text_sha1=hashlib.sha1(reply.text.encode()).hexdigest(),
delay_s=decision.delay_s,
)
)
return decision
def snapshot(self) -> list[MockRequest]:
return list(self._entries)
def reset(self) -> int:
count: Final = len(self._entries)
self._entries.clear()
return count
async def json_reply(decision: Decision) -> Response:
"""The generated body after the case's latency."""
if decision.delay_s > 0:
await asyncio.sleep(decision.delay_s)
return Response(
content=render_json(decision.reply, decision.route),
media_type="application/json",
headers=decision.headers,
)
def stream_reply(decision: Decision, media_type: str) -> Response:
"""The generated chunks back to back after the case's latency."""
return StreamingResponse(
chunks_after(render_chunks(decision.reply, decision.route), decision.delay_s),
media_type=media_type,
headers=decision.headers,
)
def parse_json(raw: bytes) -> object:
parsed: object = json.loads(raw) if raw else {}
return parsed
def mount_control(app: FastAPI, log: RequestLog) -> None:
async def health() -> dict[str, str]:
return {"status": "ok"}
async def requests_seen() -> MockRequestsResponse:
return MockRequestsResponse(requests=log.snapshot())
async def reset() -> ResetResponse:
return ResetResponse(cleared=log.reset())
async def latency(request: Latency) -> Latency:
return log.set_latency(request)
app.add_api_route("/health", health, methods=["GET"])
app.add_api_route("/mock/requests", requests_seen, methods=["GET"])
app.add_api_route("/mock/reset", reset, methods=["POST"])
app.add_api_route("/mock/latency", latency, methods=["POST"])

View file

@ -0,0 +1,115 @@
"""The perf mock guardrail: a server speaking litellm's Generic Guardrail API contract
(POST /beta/litellm_basic_guardrail_api), registered once as pre_call and once as post_call.
What it does with a request is decided by the marker's flag: `-n` passes it through
(`action: NONE`), `-i` rewrites the text (`[pre]` on the request, `[post]` on the response) so the
test can see the rewrite reach the provider and the client, and `-b` blocks the request with the
reason `perf-test-block`. Its own latency comes from `additional_provider_specific_params`."""
from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import Mapping
from typing import Final, Literal
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from mock_common import flag_of, marker_in
from pydantic import BaseModel, ConfigDict
BLOCK_REASON: Final = "perf-test-block"
Action = Literal["NONE", "BLOCKED", "GUARDRAIL_INTERVENED"]
class GuardrailRequest(BaseModel):
model_config = ConfigDict(extra="allow")
input_type: Literal["request", "response"]
texts: list[str] | None = None
litellm_call_id: str | None = None
additional_provider_specific_params: dict[str, object] | None = None
def latency_s(self) -> float:
value: Final = (self.additional_provider_specific_params or {}).get("latency_s")
return float(value) if isinstance(value, (int, float)) else 0.0
class GuardrailCall(BaseModel):
marker: str | None
input_type: Literal["request", "response"]
action: Action
class GuardrailCallsResponse(BaseModel):
calls: list[GuardrailCall]
class GuardrailResetResponse(BaseModel):
cleared: int
def marker_of(body: GuardrailRequest, remembered: Mapping[str, str]) -> str | None:
"""The request's marker: in the prompt on the request hook, remembered by `litellm_call_id` on the response
hook, since a generated provider reply does not contain it."""
texts: Final = body.texts or []
in_text: Final = next((found for text in texts if (found := marker_in(text)) is not None), None)
if in_text is not None or body.litellm_call_id is None:
return in_text
return remembered.get(body.litellm_call_id)
def decide(body: GuardrailRequest, marker: str | None) -> tuple[Action, dict[str, object]]:
texts: Final = body.texts or []
if marker is None:
return "NONE", {"action": "NONE"}
flag: Final = flag_of(marker)
if flag == "b" and body.input_type == "request":
return "BLOCKED", {"action": "BLOCKED", "blocked_reason": BLOCK_REASON}
if flag == "i":
suffix: Final = " [pre]" if body.input_type == "request" else " [post]"
rewritten: Final = (
[text + suffix if marker_in(text) else text for text in texts]
if body.input_type == "request"
else [text + suffix for text in texts]
)
return "GUARDRAIL_INTERVENED", {"action": "GUARDRAIL_INTERVENED", "texts": rewritten}
return "NONE", {"action": "NONE"}
def build_app() -> FastAPI:
app: Final = FastAPI(title="perf mock guardrail")
calls: deque[GuardrailCall] = deque(maxlen=200_000)
markers: dict[str, str] = {} # mutable-ok: the request hook's marker, looked up by the response hook
async def health() -> dict[str, str]:
return {"status": "ok"}
async def guard(body: GuardrailRequest) -> JSONResponse:
latency: Final = body.latency_s()
if latency > 0:
await asyncio.sleep(latency)
marker: Final = marker_of(body, markers)
if marker is not None and body.litellm_call_id is not None and body.input_type == "request":
markers[body.litellm_call_id] = marker
action, content = decide(body, marker)
calls.append(
GuardrailCall(
marker=marker,
input_type=body.input_type,
action=action,
)
)
return JSONResponse(content=content)
async def calls_seen() -> GuardrailCallsResponse:
return GuardrailCallsResponse(calls=list(calls))
async def reset() -> GuardrailResetResponse:
count: Final = len(calls)
calls.clear()
return GuardrailResetResponse(cleared=count)
app.add_api_route("/health", health, methods=["GET"])
app.add_api_route("/beta/litellm_basic_guardrail_api", guard, methods=["POST"])
app.add_api_route("/mock/calls", calls_seen, methods=["GET"])
app.add_api_route("/mock/reset", reset, methods=["POST"])
return app

View file

@ -0,0 +1,48 @@
"""The OpenAI-shaped perf mock: /v1/chat/completions and /v1/responses, JSON and SSE, answering each request
with a generated reply in api.openai.com's wire format, so the gateway's real OpenAI client and stream parsers
run against the provider's field set and chunk sizes."""
from __future__ import annotations
from typing import Final, Literal
from fastapi import FastAPI, Request
from fastapi.responses import Response
from mock_common import (
RequestLog,
bool_field,
json_reply,
mount_control,
parse_json,
str_field,
stream_reply,
)
SSE: Final = "text/event-stream"
def build_app() -> FastAPI:
app: Final = FastAPI(title="perf mock openai")
log: Final = RequestLog("openai")
mount_control(app, log)
async def handle(request: Request, route: Literal["chat", "responses"]) -> Response:
body: Final = parse_json(await request.body())
stream: Final = bool_field(body, "stream")
decision: Final = log.decide(
route=route, stream=stream, body=body, model=str_field(body, "model", "gpt-6-luna")
)
return stream_reply(decision, SSE) if stream else await json_reply(decision)
async def chat_completions(request: Request) -> Response:
return await handle(request, "chat")
async def responses(request: Request) -> Response:
return await handle(request, "responses")
app.add_api_route("/v1/chat/completions", chat_completions, methods=["POST"])
app.add_api_route("/chat/completions", chat_completions, methods=["POST"])
app.add_api_route("/v1/responses", responses, methods=["POST"])
app.add_api_route("/responses", responses, methods=["POST"])
return app

View file

@ -0,0 +1,80 @@
"""A generated provider reply for one request: `plan_reply` reads the request (model, prompt size, output
cap, tools offered) and `render_json` / `render_chunks` write it in the wire format of the provider and
route. The words are synthetic; the schema, ids, event order, usage fields and delta and chunk sizes follow
what each provider sends, so the gateway's parsers, cost tracking and stream handling do the same work."""
from __future__ import annotations
from typing import Final, Literal, cast
import reply_anthropic
import reply_bedrock
import reply_openai
from reply_common import Provider, Reply, plan_reply
Route = Literal["chat", "responses", "messages", "converse", "invoke"]
_MAX_TOKEN_KEYS: Final = ("max_output_tokens", "max_completion_tokens", "max_tokens")
DEFAULT_MAX_TOKENS: Final = 1024
def _mapping(value: object) -> dict[str, object]:
return cast(dict[str, object], value) if isinstance(value, dict) else {}
def _items(value: object) -> tuple[object, ...]:
return tuple(cast(list[object], value)) if isinstance(value, list) else ()
def max_tokens_of(body: object) -> int:
mapping: Final = _mapping(body)
found: Final = tuple(mapping[key] for key in _MAX_TOKEN_KEYS if isinstance(mapping.get(key), int))
inference: Final = _mapping(mapping.get("inferenceConfig")).get("maxTokens")
if found:
return cast(int, found[0])
return inference if isinstance(inference, int) else DEFAULT_MAX_TOKENS
def _tool_name(tool: object) -> str | None:
mapping: Final = _mapping(tool)
nested: Final = _mapping(mapping.get("function")) or _mapping(mapping.get("toolSpec"))
name: Final = nested.get("name", mapping.get("name"))
return name if isinstance(name, str) else None
def tool_names_of(body: object) -> tuple[str, ...]:
mapping: Final = _mapping(body)
tools: Final = _items(mapping.get("tools")) or _items(_mapping(mapping.get("toolConfig")).get("tools"))
return tuple(name for name in map(_tool_name, tools) if name is not None)
def reply_for(*, provider: Provider, body: object, seed: str, model: str, prompt_tokens: int) -> Reply:
return plan_reply(
provider=provider,
seed=seed,
model=model,
prompt_tokens=prompt_tokens,
max_tokens=max_tokens_of(body),
tool_names=tool_names_of(body),
)
def render_json(reply: Reply, route: Route) -> bytes:
if route == "chat":
return reply_openai.chat_json(reply)
if route == "responses":
return reply_openai.responses_json(reply)
if route in ("messages", "invoke"):
return reply_anthropic.messages_json(reply)
return reply_bedrock.converse_json(reply)
def render_chunks(reply: Reply, route: Route) -> tuple[bytes, ...]:
if route == "chat":
return reply_openai.chat_chunks(reply)
if route == "responses":
return reply_openai.responses_chunks(reply)
if route == "messages":
return reply_anthropic.messages_chunks(reply)
if route == "invoke":
return reply_bedrock.invoke_chunks(reply)
return reply_bedrock.converse_chunks(reply)

View file

@ -0,0 +1,132 @@
"""Anthropic wire format for a generated `Reply`: `/v1/messages` JSON and the SSE event sequence
(`message_start`, `ping`, content blocks of `thinking`, `text` or `tool_use` with their deltas,
`message_delta`, `message_stop`) with the usage fields, delta and chunk sizes api.anthropic.com uses."""
from __future__ import annotations
import json
from typing import Final
from reply_common import Reply, grouped, ident, pieces, sse
DELTA_CHARS: Final = 40
CHUNK_BYTES: Final = 1000
def _usage(reply: Reply) -> dict[str, object]:
return {
"input_tokens": reply.uncached_tokens,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": reply.cached_tokens,
"cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
"output_tokens": reply.output_tokens,
"service_tier": "standard",
}
def _signature(reply: Reply) -> str:
return ident(reply.seed, "", 64, hex_digits=True)
def _blocks(reply: Reply) -> tuple[dict[str, object], ...]:
thinking: Final[tuple[dict[str, object], ...]] = (
({"type": "thinking", "thinking": reply.reasoning, "signature": _signature(reply)},)
if reply.reasoning_tokens
else ()
)
if reply.tool_call is None:
return (*thinking, _block(type="text", text=reply.text))
tool: Final = _block(
type="tool_use",
id=ident(reply.seed, "toolu_", 24),
name=reply.tool_call.name,
input=reply.tool_call.parsed_arguments,
)
return (*thinking, tool)
def _block(**fields: object) -> dict[str, object]:
return dict(fields)
def _stop_reason(reply: Reply) -> str:
return "end_turn" if reply.tool_call is None else "tool_use"
def messages_json(reply: Reply) -> bytes:
body: Final = {
"model": reply.model,
"id": ident(reply.seed, "msg_", 24),
"type": "message",
"role": "assistant",
"content": list(_blocks(reply)),
"stop_reason": _stop_reason(reply),
"stop_sequence": None,
"usage": _usage(reply),
}
return json.dumps(body, separators=(",", ":")).encode()
def _block_events(
index: int, start: dict[str, object], deltas: tuple[dict[str, object], ...]
) -> tuple[dict[str, object], ...]:
return (
{"type": "content_block_start", "index": index, "content_block": start},
*({"type": "content_block_delta", "index": index, "delta": delta} for delta in deltas),
{"type": "content_block_stop", "index": index},
)
def messages_events(reply: Reply) -> tuple[dict[str, object], ...]:
"""The stream as event payloads, `message_start` to `message_stop`, each one's `type` its SSE event name."""
opening: Final[dict[str, object]] = {
"model": reply.model,
"id": ident(reply.seed, "msg_", 24),
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {**_usage(reply), "output_tokens": 1},
}
thinking: Final = (
_block_events(
0,
{"type": "thinking", "thinking": "", "signature": ""},
(
*({"type": "thinking_delta", "thinking": piece} for piece in pieces(reply.reasoning, DELTA_CHARS)),
{"type": "signature_delta", "signature": _signature(reply)},
),
)
if reply.reasoning_tokens
else ()
)
index: Final = 1 if reply.reasoning_tokens else 0
content: Final = (
_block_events(
index,
{"type": "text", "text": ""},
tuple({"type": "text_delta", "text": p} for p in pieces(reply.text, DELTA_CHARS)),
)
if reply.tool_call is None
else _block_events(
index,
{"type": "tool_use", "id": ident(reply.seed, "toolu_", 24), "name": reply.tool_call.name, "input": {}},
tuple(
{"type": "input_json_delta", "partial_json": p} for p in pieces(reply.tool_call.arguments, DELTA_CHARS)
),
)
)
closing: Final[tuple[dict[str, object], ...]] = (
{
"type": "message_delta",
"delta": {"stop_reason": _stop_reason(reply), "stop_sequence": None},
"usage": {**_usage(reply), "output_tokens": reply.output_tokens},
},
{"type": "message_stop"},
)
return ({"type": "message_start", "message": opening}, {"type": "ping"}, *thinking, *content, *closing)
def messages_chunks(reply: Reply) -> tuple[bytes, ...]:
return grouped(tuple(sse(event, str(event["type"])) for event in messages_events(reply)), CHUNK_BYTES)

View file

@ -0,0 +1,174 @@
"""Bedrock wire formats for a generated `Reply`: the `converse` JSON body and `converse-stream` frames, and
for Anthropic models the `invoke` body (Anthropic's own messages JSON) and `invoke-with-response-stream` frames
(each a `chunk` event carrying one base64 Anthropic stream event). Frames are `application/vnd.amazon.eventstream`
(prelude, string headers, JSON payload, CRC32s), what the gateway's botocore-based decoder reads, one per chunk."""
from __future__ import annotations
import base64
import json
import struct
import zlib
from collections.abc import Callable
from typing import Final
from reply_anthropic import messages_events
from reply_common import Reply, ident, pick, pieces
DELTA_CHARS: Final = 15
_PADDING: Final = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123"
def _usage(reply: Reply) -> dict[str, object]:
return {
"cacheReadInputTokenCount": reply.cached_tokens,
"cacheReadInputTokens": reply.cached_tokens,
"cacheWriteInputTokenCount": 0,
"cacheWriteInputTokens": 0,
"inputTokens": reply.uncached_tokens,
"outputTokens": reply.output_tokens,
"serverToolUsage": {},
"totalTokens": reply.total_tokens,
}
def _tool_use_id(reply: Reply) -> str:
return ident(reply.seed, "tooluse_", 22)
def _stop_reason(reply: Reply) -> str:
return "end_turn" if reply.tool_call is None else "tool_use"
def _text_block(text: str) -> dict[str, object]:
return {"text": text}
def _content(reply: Reply) -> tuple[dict[str, object], ...]:
reasoning: Final[tuple[dict[str, object], ...]] = (
(
{
"reasoningContent": {
"reasoningText": {"text": reply.reasoning, "signature": ident(reply.seed, "", 64, hex_digits=True)}
}
},
)
if reply.reasoning_tokens
else ()
)
if reply.tool_call is None:
return (*reasoning, _text_block(reply.text))
tool: Final[dict[str, object]] = {
"toolUse": {
"toolUseId": _tool_use_id(reply),
"name": reply.tool_call.name,
"input": reply.tool_call.parsed_arguments,
}
}
return (*reasoning, tool)
def converse_json(reply: Reply) -> bytes:
body: Final = {
"metrics": {"latencyMs": 40 * reply.output_tokens},
"output": {"message": {"role": "assistant", "content": list(_content(reply))}},
"stopReason": _stop_reason(reply),
"usage": _usage(reply),
}
return json.dumps(body, separators=(",", ":")).encode()
def _header(name: str, value: str) -> bytes:
encoded: Final = value.encode()
return struct.pack("!B", len(name)) + name.encode() + struct.pack("!BH", 7, len(encoded)) + encoded
def frame(event_type: str, payload: dict[str, object]) -> bytes:
"""One eventstream message: prelude (total length, headers length, CRC), the three string headers
bedrock-runtime sets, the JSON payload, and the CRC of everything before it."""
headers: Final = (
_header(":event-type", event_type)
+ _header(":content-type", "application/json")
+ _header(":message-type", "event")
)
body: Final = json.dumps(payload, separators=(",", ":")).encode()
total: Final = 4 + 4 + 4 + len(headers) + len(body) + 4
prelude: Final = struct.pack("!II", total, len(headers))
head: Final = prelude + struct.pack("!I", zlib.crc32(prelude))
message: Final = head + headers + body
return message + struct.pack("!I", zlib.crc32(message))
def _pad(seed: str, index: int) -> str:
return _PADDING[: 2 + pick(f"{seed}:pad", index, 50)]
def _text_events(reply: Reply, index: int, event: Callable[[int, str, dict[str, object]], bytes]) -> tuple[bytes, ...]:
return (
*(
event(i, "contentBlockDelta", {"contentBlockIndex": index, "delta": {"text": piece}})
for i, piece in enumerate(pieces(reply.text, DELTA_CHARS), start=1)
),
event(0, "contentBlockStop", {"contentBlockIndex": index}),
)
def _tool_events(reply: Reply, index: int, event: Callable[[int, str, dict[str, object]], bytes]) -> tuple[bytes, ...]:
assert reply.tool_call is not None
start: Final = {"toolUse": {"name": reply.tool_call.name, "toolUseId": _tool_use_id(reply)}}
return (
event(0, "contentBlockStart", {"contentBlockIndex": index, "start": start}),
*(
event(i, "contentBlockDelta", {"contentBlockIndex": index, "delta": {"toolUse": {"input": piece}}})
for i, piece in enumerate(pieces(reply.tool_call.arguments, DELTA_CHARS), start=1)
),
event(0, "contentBlockStop", {"contentBlockIndex": index}),
)
def converse_chunks(reply: Reply) -> tuple[bytes, ...]:
def event(index: int, event_type: str, payload: dict[str, object]) -> bytes:
return frame(event_type, {**payload, "p": _pad(reply.seed, index)})
reasoning: Final = (
(
*(
event(i, "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": piece}}})
for i, piece in enumerate(pieces(reply.reasoning, DELTA_CHARS), start=1)
),
event(
0,
"contentBlockDelta",
{
"contentBlockIndex": 0,
"delta": {"reasoningContent": {"signature": ident(reply.seed, "", 64, hex_digits=True)}},
},
),
event(0, "contentBlockStop", {"contentBlockIndex": 0}),
)
if reply.reasoning_tokens
else ()
)
index: Final = 1 if reply.reasoning_tokens else 0
content: Final = _text_events(reply, index, event) if reply.tool_call is None else _tool_events(reply, index, event)
return (
event(0, "messageStart", {"role": "assistant"}),
*reasoning,
*content,
event(0, "messageStop", {"stopReason": _stop_reason(reply)}),
event(0, "metadata", {"metrics": {"latencyMs": 40 * reply.output_tokens}, "usage": _usage(reply)}),
)
def invoke_chunks(reply: Reply) -> tuple[bytes, ...]:
"""The Anthropic stream events, each base64 in the `bytes` of one `chunk` frame."""
return tuple(
frame(
"chunk",
{
"bytes": base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode(),
"p": _pad(reply.seed, index),
},
)
for index, event in enumerate(messages_events(reply))
)

View file

@ -0,0 +1,182 @@
"""What every generated provider reply is made of: a `Reply` sized from the request (prompt tokens, the
output cap, whether tools were offered), deterministic ids and text drawn from the request marker, and the
chunking that makes a stream arrive the way the provider's TCP writes do. The provider modules render a
`Reply` in their own wire format; nothing here is provider specific."""
from __future__ import annotations
import asyncio
import hashlib
import json
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass
from itertools import accumulate, groupby
from typing import Final, Literal
from pydantic import TypeAdapter
Provider = Literal["openai", "anthropic", "bedrock"]
CHARS_PER_TOKEN: Final = 4
CACHED_PREFIX_TOKENS: Final = 1024
"""Prompt tokens the provider reports as a cache hit (system prompt and tool set) on the requests that hit."""
CACHE_HIT_SHARE: Final = 0.7
REASONING_SHARE: Final = 0.6
"""Share of the output spent on reasoning when the request offered tools, as the captured replies did."""
_BASE62: Final = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
_DIRS: Final = (
"management",
"guardrails",
"llm_translation",
"router",
"quota_management",
"mcp",
"logging",
"security",
)
_WORDS: Final = (
"spend", "budget", "team", "key", "fallback", "cooldown", "retry", "stream", "cache", "tag",
"alias", "quota", "webhook", "secret", "model", "org", "audit", "rerank", "batch", "realtime",
) # fmt: skip
_THOUGHT: Final = "Checking {path}: the except clause around the client call logs and returns None, so the caller never sees the error. "
def unit(seed: str) -> float:
"""A float in [0, 1) that is the same for the same seed on every mock pod."""
digest: Final = hashlib.sha1(seed.encode()).digest()
return int.from_bytes(digest[:8], "big") / 2**64
def pick(seed: str, index: int, size: int) -> int:
return int(hashlib.sha1(f"{seed}:{index}".encode()).hexdigest()[:8], 16) % size
def ident(seed: str, prefix: str, length: int, *, hex_digits: bool = False) -> str:
"""A provider-shaped id (`chatcmpl-...`, `msg_...`, `toolu_...`) that is a pure function of the seed."""
digest: Final = hashlib.sha256(f"{seed}:{prefix}".encode())
if hex_digits:
return prefix + digest.hexdigest()[:length]
return prefix + "".join(_BASE62[byte % len(_BASE62)] for byte in digest.digest()[:length])
def file_paths(seed: str, count: int) -> tuple[str, ...]:
return tuple(
f"tests/e2e/{_DIRS[pick(seed, i, len(_DIRS))]}/test_{_WORDS[pick(seed, i + 1000, len(_WORDS))]}"
f"_{_WORDS[pick(seed, i + 2000, len(_WORDS))]}_e2e.py"
for i in range(count)
)
def answer_text(seed: str, tokens: int) -> str:
"""The assistant's answer to the corpus question (which files swallow exceptions): paths, one per line."""
return "\n".join(file_paths(seed, max(1, tokens * CHARS_PER_TOKEN // 56)))
def thinking_text(seed: str, tokens: int) -> str:
if tokens == 0:
return ""
paths: Final = file_paths(f"{seed}:thinking", max(1, tokens * CHARS_PER_TOKEN // len(_THOUGHT)))
return "".join(_THOUGHT.format(path=path) for path in paths).rstrip()
def grep_arguments(seed: str, tokens: int) -> str:
"""A `Grep` call over the files the session read, its glob sized to the output budget."""
paths: Final = file_paths(f"{seed}:grep", max(1, tokens * CHARS_PER_TOKEN // 56))
arguments: Final = {
"pattern": r"\braise\s+\w+",
"path": "/workspace/litellm/tests/e2e",
"glob": "{" + ",".join(paths) + "}",
}
return json.dumps(arguments, separators=(",", ":"))
_ARGUMENTS: Final = TypeAdapter(dict[str, object])
@dataclass(frozen=True, slots=True)
class ToolCall:
name: str
arguments: str
@property
def parsed_arguments(self) -> dict[str, object]:
return _ARGUMENTS.validate_json(self.arguments)
@dataclass(frozen=True, slots=True)
class Reply:
"""One generated reply, every number the gateway will read already decided."""
provider: Provider
seed: str
model: str
prompt_tokens: int
cached_tokens: int
reasoning_tokens: int
output_tokens: int
reasoning: str
text: str
tool_call: ToolCall | None
@property
def visible_tokens(self) -> int:
return self.output_tokens - self.reasoning_tokens
@property
def uncached_tokens(self) -> int:
return self.prompt_tokens - self.cached_tokens
@property
def total_tokens(self) -> int:
return self.prompt_tokens + self.output_tokens
def plan_reply(
*, provider: Provider, seed: str, model: str, prompt_tokens: int, max_tokens: int, tool_names: Sequence[str]
) -> Reply:
"""Size the reply from the request: a quarter to three quarters of the output cap, a tool call (with
reasoning first) when tools were offered, otherwise a text answer, a prefix cache hit on most requests."""
output_tokens: Final = max(32, int(max_tokens * (0.25 + 0.5 * unit(f"{seed}:length"))))
tool: Final = "Grep" if "Grep" in tool_names else next(iter(tool_names), None)
reasoning_tokens: Final = int(output_tokens * REASONING_SHARE) if tool is not None else 0
visible: Final = output_tokens - reasoning_tokens
cached: Final = (
CACHED_PREFIX_TOKENS if unit(f"{seed}:cache") < CACHE_HIT_SHARE and prompt_tokens > CACHED_PREFIX_TOKENS else 0
)
return Reply(
provider=provider,
seed=seed,
model=model,
prompt_tokens=prompt_tokens,
cached_tokens=cached,
reasoning_tokens=reasoning_tokens,
output_tokens=output_tokens,
reasoning=thinking_text(seed, reasoning_tokens),
text="" if tool is not None else answer_text(seed, visible),
tool_call=ToolCall(tool, grep_arguments(seed, visible)) if tool is not None else None,
)
def pieces(text: str, size: int) -> tuple[str, ...]:
"""`text` cut into the deltas a provider streams it as, `size` characters each."""
return tuple(text[i : i + size] for i in range(0, len(text), size))
def sse(payload: dict[str, object], event: str | None = None) -> bytes:
data: Final = json.dumps(payload, separators=(",", ":"))
prefix: Final = f"event: {event}\n" if event is not None else ""
return f"{prefix}data: {data}\n\n".encode()
def grouped(events: Sequence[bytes], target_bytes: int) -> tuple[bytes, ...]:
"""Consecutive events joined until a chunk reaches `target_bytes`, the size the provider's writes arrive in."""
cumulative: Final = tuple(accumulate(len(event) for event in events))
keyed: Final = tuple(zip((total // target_bytes for total in cumulative), events, strict=True))
return tuple(b"".join(event for _, event in group) for _, group in groupby(keyed, key=lambda pair: pair[0]))
async def chunks_after(chunks: Sequence[bytes], delay_s: float) -> AsyncIterator[bytes]:
"""The chunks back to back once `delay_s` has passed."""
if delay_s > 0:
await asyncio.sleep(delay_s)
for chunk in chunks:
yield chunk

View file

@ -0,0 +1,340 @@
"""OpenAI wire formats for a generated `Reply`: `/v1/chat/completions` (JSON and SSE chunks) and
`/v1/responses` (JSON and typed SSE events), with the field set, ids, event order, delta sizes and chunk
sizes of what api.openai.com sends gpt-6-luna replies as."""
from __future__ import annotations
import json
import time
from collections.abc import Callable
from typing import Final
from reply_common import Reply, grouped, ident, pieces, sse
CHAT_DELTA_CHARS: Final = 4
CHAT_CHUNK_BYTES: Final = 560
RESPONSES_DELTA_CHARS: Final = 3
RESPONSES_CHUNK_BYTES: Final = 250
def _chat_usage(reply: Reply) -> dict[str, object]:
return {
"prompt_tokens": reply.prompt_tokens,
"completion_tokens": reply.output_tokens,
"total_tokens": reply.total_tokens,
"prompt_tokens_details": {"cached_tokens": reply.cached_tokens, "audio_tokens": 0},
"completion_tokens_details": {
"reasoning_tokens": reply.reasoning_tokens,
"audio_tokens": 0,
"accepted_prediction_tokens": 0,
"rejected_prediction_tokens": 0,
},
}
def _chat_message(reply: Reply) -> dict[str, object]:
if reply.tool_call is None:
return {"role": "assistant", "content": reply.text, "refusal": None, "annotations": []}
call: Final = {
"id": ident(reply.seed, "call_", 24),
"type": "function",
"function": {"name": reply.tool_call.name, "arguments": reply.tool_call.arguments},
}
return {"role": "assistant", "content": None, "refusal": None, "annotations": [], "tool_calls": [call]}
def chat_json(reply: Reply) -> bytes:
finish: Final = "stop" if reply.tool_call is None else "tool_calls"
body: Final = {
"id": ident(reply.seed, "chatcmpl-", 29),
"object": "chat.completion",
"created": int(time.time()),
"model": reply.model,
"choices": [{"index": 0, "message": _chat_message(reply), "logprobs": None, "finish_reason": finish}],
"usage": _chat_usage(reply),
"service_tier": "default",
"system_fingerprint": None,
}
return json.dumps(body, separators=(",", ":")).encode()
def _chat_deltas(reply: Reply, chunk: Callable[[dict[str, object]], bytes]) -> tuple[bytes, ...]:
if reply.tool_call is None:
return tuple(chunk({"content": piece}) for piece in pieces(reply.text, CHAT_DELTA_CHARS))
opening: Final = {
"index": 0,
"id": ident(reply.seed, "call_", 24),
"type": "function",
"function": {"name": reply.tool_call.name, "arguments": ""},
}
argument_pieces: Final = pieces(reply.tool_call.arguments, CHAT_DELTA_CHARS)
return (
chunk({"tool_calls": [opening]}),
*(chunk({"tool_calls": [{"index": 0, "function": {"arguments": piece}}]}) for piece in argument_pieces),
)
def chat_chunks(reply: Reply) -> tuple[bytes, ...]:
head: Final = {
"id": ident(reply.seed, "chatcmpl-", 29),
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": reply.model,
"service_tier": "default",
"system_fingerprint": None,
}
def chunk(delta: dict[str, object], finish: str | None = None) -> bytes:
choice: Final = {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}
return sse({**head, "choices": [choice], "usage": None})
deltas: Final = _chat_deltas(reply, chunk)
finish: Final = "stop" if reply.tool_call is None else "tool_calls"
events: Final = (
chunk({"role": "assistant", "content": "", "refusal": None}),
*deltas,
chunk({}, finish),
sse({**head, "choices": [], "usage": _chat_usage(reply)}),
b"data: [DONE]\n\n",
)
return grouped(events, CHAT_CHUNK_BYTES)
def _responses_usage(reply: Reply) -> dict[str, object]:
return {
"input_tokens": reply.prompt_tokens,
"input_tokens_details": {"cached_tokens": reply.cached_tokens},
"output_tokens": reply.output_tokens,
"output_tokens_details": {"reasoning_tokens": reply.reasoning_tokens},
"total_tokens": reply.total_tokens,
}
def _reasoning_item(reply: Reply) -> dict[str, object]:
return {"id": ident(reply.seed, "rs_", 50, hex_digits=True), "type": "reasoning", "summary": [], "content": []}
def _message_item(reply: Reply, text: str) -> dict[str, object]:
content: Final[dict[str, object]] = {"type": "output_text", "annotations": [], "logprobs": [], "text": text}
return {
"id": ident(reply.seed, "msg_", 50, hex_digits=True),
"type": "message",
"status": "completed",
"content": [content],
"role": "assistant",
}
def _function_item(reply: Reply, arguments: str, status: str) -> dict[str, object]:
assert reply.tool_call is not None
return {
"id": ident(reply.seed, "fc_", 50, hex_digits=True),
"type": "function_call",
"status": status,
"arguments": arguments,
"call_id": ident(reply.seed, "call_", 24),
"name": reply.tool_call.name,
}
def _output(reply: Reply) -> tuple[dict[str, object], ...]:
final: Final = (
_message_item(reply, reply.text)
if reply.tool_call is None
else _function_item(reply, reply.tool_call.arguments, "completed")
)
return (_reasoning_item(reply), final) if reply.reasoning_tokens else (final,)
def _response(reply: Reply, status: str, output: tuple[dict[str, object], ...]) -> dict[str, object]:
created: Final = int(time.time())
return {
"id": ident(reply.seed, "resp_", 50, hex_digits=True),
"object": "response",
"created_at": created,
"status": status,
"background": False,
"completed_at": created if status == "completed" else None,
"error": None,
"incomplete_details": None,
"instructions": None,
"max_output_tokens": None,
"model": reply.model,
"output": list(output),
"parallel_tool_calls": True,
"previous_response_id": None,
"reasoning": {"effort": "medium", "summary": None},
"service_tier": "default",
"store": True,
"temperature": 1.0,
"text": {"format": {"type": "text"}, "verbosity": "medium"},
"tool_choice": "auto",
"tools": [],
"top_p": 1.0,
"truncation": "disabled",
"usage": _responses_usage(reply) if status == "completed" else None,
"user": None,
"metadata": {},
}
def responses_json(reply: Reply) -> bytes:
return json.dumps(_response(reply, "completed", _output(reply)), separators=(",", ":")).encode()
def _message_events(reply: Reply, index: int, start: int) -> tuple[bytes, ...]:
item_id: Final = ident(reply.seed, "msg_", 50, hex_digits=True)
at: Final = {"item_id": item_id, "output_index": index, "content_index": 0}
deltas: Final = pieces(reply.text, RESPONSES_DELTA_CHARS)
part: Final[dict[str, object]] = {"type": "output_text", "annotations": [], "logprobs": [], "text": ""}
added: Final[dict[str, object]] = {
"id": item_id,
"type": "message",
"status": "in_progress",
"content": [],
"role": "assistant",
}
middle: Final = tuple(
sse(
{"type": "response.output_text.delta", **at, "delta": delta, "sequence_number": start + 2 + i},
"response.output_text.delta",
)
for i, delta in enumerate(deltas)
)
end: Final = start + 2 + len(deltas)
return (
sse(
{"type": "response.output_item.added", "output_index": index, "item": added, "sequence_number": start},
"response.output_item.added",
),
sse(
{"type": "response.content_part.added", **at, "part": part, "sequence_number": start + 1},
"response.content_part.added",
),
*middle,
sse(
{"type": "response.output_text.done", **at, "text": reply.text, "sequence_number": end},
"response.output_text.done",
),
sse(
{
"type": "response.content_part.done",
**at,
"part": {**part, "text": reply.text},
"sequence_number": end + 1,
},
"response.content_part.done",
),
sse(
{
"type": "response.output_item.done",
"output_index": index,
"item": _message_item(reply, reply.text),
"sequence_number": end + 2,
},
"response.output_item.done",
),
)
def _function_events(reply: Reply, index: int, start: int) -> tuple[bytes, ...]:
assert reply.tool_call is not None
item_id: Final = ident(reply.seed, "fc_", 50, hex_digits=True)
deltas: Final = pieces(reply.tool_call.arguments, RESPONSES_DELTA_CHARS)
middle: Final = tuple(
sse(
{
"type": "response.function_call_arguments.delta",
"item_id": item_id,
"output_index": index,
"delta": delta,
"sequence_number": start + 1 + i,
},
"response.function_call_arguments.delta",
)
for i, delta in enumerate(deltas)
)
end: Final = start + 1 + len(deltas)
return (
sse(
{
"type": "response.output_item.added",
"output_index": index,
"item": _function_item(reply, "", "in_progress"),
"sequence_number": start,
},
"response.output_item.added",
),
*middle,
sse(
{
"type": "response.function_call_arguments.done",
"item_id": item_id,
"output_index": index,
"arguments": reply.tool_call.arguments,
"sequence_number": end,
},
"response.function_call_arguments.done",
),
sse(
{
"type": "response.output_item.done",
"output_index": index,
"item": _function_item(reply, reply.tool_call.arguments, "completed"),
"sequence_number": end + 1,
},
"response.output_item.done",
),
)
def responses_chunks(reply: Reply) -> tuple[bytes, ...]:
opening: Final = (
sse(
{"type": "response.created", "response": _response(reply, "in_progress", ()), "sequence_number": 0},
"response.created",
),
sse(
{"type": "response.in_progress", "response": _response(reply, "in_progress", ()), "sequence_number": 1},
"response.in_progress",
),
)
reasoning: Final = (
(
sse(
{
"type": "response.output_item.added",
"output_index": 0,
"item": _reasoning_item(reply),
"sequence_number": 2,
},
"response.output_item.added",
),
sse(
{
"type": "response.output_item.done",
"output_index": 0,
"item": _reasoning_item(reply),
"sequence_number": 3,
},
"response.output_item.done",
),
)
if reply.reasoning_tokens
else ()
)
index: Final = len(reasoning) // 2
start: Final = 2 + len(reasoning)
body: Final = (
_message_events(reply, index, start) if reply.tool_call is None else _function_events(reply, index, start)
)
sequence: Final = start + len(body)
closing: Final = sse(
{
"type": "response.completed",
"response": _response(reply, "completed", _output(reply)),
"sequence_number": sequence,
},
"response.completed",
)
return grouped((*opening, *reasoning, *body, closing), RESPONSES_CHUNK_BYTES)

View file

@ -0,0 +1,43 @@
"""Run one perf mock out of process: `python server.py openai --port 8090`.
Each kind is its own pod (or local process) with its own request log; the gateway talks to it over
HTTP exactly as it would to the real provider, and the test reads `/mock/requests` (providers) or
`/mock/calls` (guardrail) afterwards to join what was received against what each client saw.
"""
from __future__ import annotations
import argparse
from collections.abc import Callable
from typing import Final
import uvicorn
from fastapi import FastAPI
import mock_anthropic
import mock_bedrock
import mock_guardrail
import mock_openai
KINDS: Final[dict[str, Callable[[], FastAPI]]] = {
"openai": mock_openai.build_app,
"anthropic": mock_anthropic.build_app,
"bedrock": mock_bedrock.build_app,
"guardrail": mock_guardrail.build_app,
}
def main() -> None:
parser: Final = argparse.ArgumentParser(description=__doc__)
_ = parser.add_argument("kind", choices=sorted(KINDS))
_ = parser.add_argument("--host", default="0.0.0.0")
_ = parser.add_argument("--port", type=int, default=8090)
args: Final = parser.parse_args()
kind: Final = str(args.kind) # pyright: ignore[reportAny]
port: Final = int(args.port) # pyright: ignore[reportAny]
host: Final = str(args.host) # pyright: ignore[reportAny]
uvicorn.run(KINDS[kind](), host=host, port=port, log_level="warning", access_log=False)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,298 @@
"""Turns what a run left behind (the client records, what each mock received, the guardrail's
calls, the workers' memory) into the case's metrics, and the case's limits into a verdict."""
from __future__ import annotations
import hashlib
import statistics
from collections import Counter
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Final
from e2e_http import NoBody, get_external, post_external, unwrap
from memory_readings import RssCapture, RssReading, WorkerKey
from pydantic import BaseModel
from setup.datastore import DbCounters, RedisCounters
from setup.mocks.reply_common import Provider
from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase
from setup.wire import WireRecord
class MockRequestSeen(BaseModel):
provider: Provider
marker: str | None
route: str
stream: bool
output_tokens: int
text_sha1: str
delay_s: float
class MockRequestsSeen(BaseModel):
requests: list[MockRequestSeen]
class GuardrailCallSeen(BaseModel):
marker: str | None
input_type: str
action: str
class GuardrailCallsSeen(BaseModel):
calls: list[GuardrailCallSeen]
class Cleared(BaseModel):
cleared: int
class LatencyBody(BaseModel):
seconds: float
jitter_seconds: float
def read_mock_requests(url: str) -> list[MockRequestSeen]:
return unwrap(get_external(f"{url}/mock/requests", response_type=MockRequestsSeen, timeout=60.0)).requests
def read_guardrail_calls(url: str) -> list[GuardrailCallSeen]:
return unwrap(get_external(f"{url}/mock/calls", response_type=GuardrailCallsSeen, timeout=60.0)).calls
def reset_mock(url: str) -> int:
return unwrap(post_external(f"{url}/mock/reset", json=NoBody(), response_type=Cleared)).cleared
def set_mock_latency(url: str, latency: MockLatency) -> LatencyBody:
body: Final = LatencyBody(seconds=latency.seconds, jitter_seconds=latency.jitter_seconds)
return unwrap(post_external(f"{url}/mock/latency", json=body, response_type=LatencyBody))
def percentile(values: Iterable[float], p: int) -> float | None:
ordered: Final = sorted(values)
if not ordered:
return None
index: Final = min(len(ordered) - 1, max(0, round(p / 100 * (len(ordered) - 1))))
return ordered[index]
def worst_worker_growth_mb(before: RssCapture, after: RssCapture) -> float | None:
"""Worst RSS growth across workers. A worker missing from `before` (a pod replaced during the
run) is measured against the median of the old workers."""
old: Final[dict[WorkerKey, RssReading]] = {reading.worker: reading for reading in before.readings}
if not old or not after.readings:
return None
median_before: Final = statistics.median(reading.ram_usage_mb for reading in old.values())
growth: list[float] = []
for reading in after.readings:
baseline = old[reading.worker].ram_usage_mb if reading.worker in old else median_before
growth.append(reading.ram_usage_mb - baseline)
return max(growth)
@dataclass(frozen=True, slots=True)
class Observation:
records: tuple[WireRecord, ...]
mock_requests: tuple[MockRequestSeen, ...]
guardrail_calls: tuple[GuardrailCallSeen, ...]
rss_before: RssCapture
rss_after: RssCapture
wall_seconds: float
spend_logged_share: float | None = None
db_before: DbCounters | None = None
db_after: DbCounters | None = None
redis_before: RedisCounters | None = None
redis_after: RedisCounters | None = None
gateway_cpu_seconds: float | None = None
def per_request(self, before: int | None, after: int | None) -> float | None:
if before is None or after is None or not self.records:
return None
return (after - before) / len(self.records)
@property
def retried_markers(self) -> frozenset[str]:
counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None)
return frozenset(marker for marker, count in counts.items() if count > 1)
@property
def served(self) -> tuple[WireRecord, ...]:
return tuple(record for record in self.records if record.ok and not record.blocked)
@property
def _delay_by_marker(self) -> dict[str, float]:
return {seen.marker: seen.delay_s for seen in self.mock_requests if seen.marker is not None}
def _clean(self, *, stream: bool) -> list[tuple[WireRecord, float]]:
"""Served requests of one kind that were answered on the first attempt, with the mock's delay."""
delays: Final = self._delay_by_marker
retried: Final = self.retried_markers
pairs: list[tuple[WireRecord, float]] = []
for record in self.served:
delay = record.mock_delay_s if record.mock_delay_s is not None else delays.get(record.marker)
if record.stream == stream and delay is not None and record.marker not in retried:
pairs.append((record, delay))
return pairs
def unexpected_status_count(self) -> int:
"""Requests the plan expected to succeed that came back with any other status (0 is a client timeout)."""
return sum(1 for record in self.records if not record.blocked and record.status != 200)
def gateway_added_seconds(self) -> list[float]:
"""Seconds the gateway added to each first-attempt request: total time minus the mock's own delay for a
non-streamed request, time to the first token minus the mock's delay to its first byte for a stream."""
non_stream: Final = [record.total_s - delay for record, delay in self._clean(stream=False)]
stream: Final = [
record.first_token_s - delay
for record, delay in self._clean(stream=True)
if record.first_token_s is not None
]
return non_stream + stream
def time_to_first_token_seconds(self) -> list[float]:
"""Seconds from the client sending each first-attempt streamed request to its first token, mock delay
included, comparable across versions because the mocks generate the same replies after the same wait."""
return [record.first_token_s for record, _ in self._clean(stream=True) if record.first_token_s is not None]
def bodies_off_reply(self) -> int:
"""Served requests whose text is not the reply the mock generated for them (plus the post-call rewrite
for requests in the intervene slice), a translation or streaming fault somewhere in the gateway."""
seen_by_marker: Final = {seen.marker: seen for seen in self.mock_requests if seen.marker is not None}
return sum(
1
for record in self.served
if (seen := seen_by_marker.get(record.marker)) is not None
and _text_sha1(record.text, intervened=record.marker.endswith("-i")) != seen.text_sha1
)
def guardrail_missed(self) -> int:
"""Served requests that did not get exactly one pre-call and one post-call guardrail call."""
pre: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "request")
post: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "response")
return sum(1 for record in self.served if pre[record.marker] != 1 or post[record.marker] != 1)
def _text_sha1(text: str, *, intervened: bool) -> str:
original: Final = text.removesuffix(" [post]") if intervened else text
return hashlib.sha1(original.encode()).hexdigest()
def observe(metric: Metric, percentile_rank: int | None, observation: Observation) -> float | None:
"""The observed value of one metric, at a percentile when the metric is per request."""
per_request: dict[str, list[float]] = {
"gateway_added_seconds": observation.gateway_added_seconds(),
"time_to_first_token_seconds": observation.time_to_first_token_seconds(),
}
if metric in per_request:
return percentile(per_request[metric], 50 if percentile_rank is None else percentile_rank)
if metric == "unexpected_status_count":
return float(observation.unexpected_status_count())
if metric == "postgres_calls_per_request":
db_before, db_after = observation.db_before, observation.db_after
return observation.per_request(db_before and db_before.commits, db_after and db_after.commits)
if metric == "redis_calls_per_request":
redis_before, redis_after = observation.redis_before, observation.redis_after
return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands)
if metric == "cpu_seconds_per_request":
cpu: Final = observation.gateway_cpu_seconds
return None if cpu is None or not observation.records else cpu / len(observation.records)
return worst_worker_growth_mb(observation.rss_before, observation.rss_after)
@dataclass(frozen=True, slots=True)
class Row:
label: str
observed: float | None
limit: float | None
@property
def over(self) -> bool:
return self.limit is not None and self.observed is not None and self.observed > self.limit
@property
def missing(self) -> bool:
return self.limit is not None and self.observed is None
REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = (
("unexpected_status_count", None),
("gateway_added_seconds", 50),
("gateway_added_seconds", 99),
("time_to_first_token_seconds", 50),
("time_to_first_token_seconds", 99),
("memory_growth_megabytes", None),
("postgres_calls_per_request", None),
("redis_calls_per_request", None),
("cpu_seconds_per_request", None),
)
def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[Row], list[str]]:
"""Every reported metric with its observed value and the case's limit if it has one, plus one line per
limit that was exceeded or that the run could not observe."""
limits: Final[dict[tuple[Metric, int | None], Limit]] = {
(limit.metric, limit.percentile): limit for limit in case.limits
}
keys: list[tuple[Metric, int | None]] = list(REPORTED)
keys.extend(key for key in limits if key not in keys)
rows: list[Row] = []
for metric, rank in keys:
limit = limits.get((metric, rank))
label = f"{metric} percentile {rank}" if rank is not None else metric
rows.append(Row(label=label, observed=observe(metric, rank, observation), limit=limit.max if limit else None))
exceeded: Final = [
f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}"
if row.over and row.observed is not None and row.limit is not None
else f"{row.label}: limit set but nothing observed"
for row in rows
if row.over or row.missing
]
return rows, exceeded
def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) -> str:
by_endpoint: Final = Counter((record.endpoint, record.stream) for record in observation.records)
failures: Final = [record for record in observation.records if not record.ok][:10]
lines: list[str] = [
f"perf case {case.name}: {len(observation.records)} requests, {len(observation.served)} served, "
f"{sum(1 for r in observation.records if r.blocked and r.ok)} blocked as planned, "
f"{sum(1 for r in observation.records if not r.ok)} wrong, wall {observation.wall_seconds:.1f}s",
" "
+ ", ".join(
f"{endpoint}{' stream' if stream else ''}={count}"
for (endpoint, stream), count in sorted(by_endpoint.items())
),
f" {'metric':32} {'observed':>12} {'limit':>12}",
]
for row in rows:
observed = "n/a" if row.observed is None else f"{row.observed:.4f}"
limit = "" if row.limit is None else f"{row.limit:.4f}"
flag = " OVER" if row.over else " MISSING" if row.missing else ""
lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}")
lines.extend(_diagnostics(observation))
for record in failures:
lines.append(
f" wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}"
)
return "\n".join(lines)
def _diagnostics(observation: Observation) -> list[str]:
"""Functional side-channels printed for whoever reads a red run, never limits of a case."""
count: Final = max(1, len(observation.records))
share: Final = observation.spend_logged_share
return [
" diagnostics (not limits):",
f" bodies off reply {observation.bodies_off_reply()}",
f" guardrail calls per request {len(observation.guardrail_calls) / count:.3f}",
f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}",
f" retried share {len(observation.retried_markers) / count:.4f}",
f" spend logged share {'n/a' if share is None else f'{share:.3f}'}",
f" redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}",
]
def _clients(counters: RedisCounters | None) -> str:
return "n/a" if counters is None else str(counters.connected_clients)

View file

@ -0,0 +1,191 @@
"""The schema of a perf test: a `PerformanceTestCase` names what the gateway is deployed with
(model groups and guardrails, both written out in the gateway's config) and lists what the test
itself creates and sends: teams, Keycloak users, virtual keys, the request shapes, the traffic plan
and the `Limit`s that decide pass or fail. Add a case by adding a module to `cases/`.
The same shape as `CostTrackingTestCase` in the cost suite: frozen dataclasses with plain defaults,
so a case is a literal a reviewer can read top to bottom.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Final, Literal
Endpoint = Literal["chat_completions", "responses", "messages"]
MockProvider = Literal["openai", "anthropic", "bedrock"]
GuardrailMode = Literal["pre_call", "post_call"]
Metric = Literal[
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
"postgres_calls_per_request",
"cpu_seconds_per_request",
"redis_calls_per_request",
]
@dataclass(frozen=True, slots=True)
class Team:
"""One team (`/team/new`). `models` names the model groups the team may call; empty means all."""
name: str
models: tuple[str, ...] = ()
max_budget: float | None = None
rpm_limit: int | None = None
tpm_limit: int | None = None
@dataclass(frozen=True, slots=True)
class User:
"""One Keycloak user who calls the gateway with a JWT. `team` puts them in that team's group."""
name: str
team: str | None = None
max_budget: float | None = None
@dataclass(frozen=True, slots=True)
class VirtualKey:
"""One virtual key (`/key/generate`), scoped to a team and/or a user when named."""
name: str
team: str | None = None
user: str | None = None
models: tuple[str, ...] = ()
rpm_limit: int | None = None
max_budget: float | None = None
@dataclass(frozen=True, slots=True)
class Guardrail:
"""A guardrail the gateway's config defines (a `generic_guardrail_api` pointed at the mock
guardrail, `default_on`); the case names it so the plan knows which hooks every request sees."""
name: str
mode: GuardrailMode
@dataclass(frozen=True, slots=True)
class RequestShape:
"""What one request looks like: a coding-agent session (`setup/corpus/session.py`, the system prompt, the
Read/Edit/Write/Bash/Grep/Glob tools when `tools` is set, then file-reading turns carrying real source files)
whose size in input tokens is drawn per request from the `input_tokens` range, and the output cap. The mocks
answer with a reply generated for the request, sized from that cap, a tool call when `tools` is set."""
name: str
input_tokens: tuple[int, int]
tools: bool = False
output_tokens: int = 1_024
SHORT: Final = RequestShape("short", input_tokens=(5_000, 100_000), tools=True)
LONG: Final = RequestShape("long", input_tokens=(300_000, 1_000_000), tools=True)
@dataclass(frozen=True, slots=True)
class Traffic:
"""The plan: `total_requests` drawn from `seed`, sent at `rps` by `concurrency` simulated
users. Shares are fractions of the requests: `endpoints` and `shapes` must each sum to 1,
`intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` picks
requests it blocks."""
seed: int
rps: float
total_requests: int
endpoints: Mapping[Endpoint, float] = field(default_factory=lambda: {"chat_completions": 1.0})
shapes: Mapping[RequestShape, float] = field(default_factory=lambda: {SHORT: 1.0})
stream_share: float = 0.0
intervene_share: float = 0.0
block_share: float = 0.0
concurrency: int = 20
@property
def expected_seconds(self) -> float:
return self.total_requests / self.rps
@dataclass(frozen=True, slots=True)
class Limit:
"""One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`.
Metrics without a Limit are still reported."""
metric: Metric
max: float
percentile: int | None = None
@property
def label(self) -> str:
return f"{self.metric} percentile {self.percentile}" if self.percentile is not None else self.metric
@dataclass(frozen=True, slots=True)
class MockLatency:
"""Seconds a mock provider waits before its first byte, plus or minus up to `jitter_seconds` per request
(drawn from the request marker, so a rerun waits the same). Zero means the mocks answer at once."""
seconds: float = 0.0
jitter_seconds: float = 0.0
@dataclass(frozen=True, slots=True)
class PerformanceTestCase:
"""`model_groups` and `guardrails` name what the gateway config deploys (setup checks they are
there); everything else the test creates or sends. `mock_latency` is how long every mock provider waits
before answering, set on the mock pods at setup, so the provider side of each request is a known
constant and the rest of the client's time is the gateway's."""
name: str
teams: tuple[Team, ...]
users: tuple[User, ...]
virtual_keys: tuple[VirtualKey, ...]
model_groups: tuple[str, ...]
guardrails: tuple[Guardrail, ...]
traffic: Traffic
limits: tuple[Limit, ...]
mock_latency: MockLatency = field(default_factory=lambda: MockLatency())
def team(self, name: str) -> Team:
return next(team for team in self.teams if team.name == name)
@property
def has_post_guardrail(self) -> bool:
return any(guardrail.mode == "post_call" for guardrail in self.guardrails)
def validate(self) -> None:
"""Fail fast on a case that cannot run: dangling names, shares that do not add up, a model
group no caller may use."""
team_names: Final = {team.name for team in self.teams}
user_names: Final = {user.name for user in self.users}
group_names: Final = set(self.model_groups)
for names, label in (
([team.name for team in self.teams], "team"),
([user.name for user in self.users], "user"),
([key.name for key in self.virtual_keys], "key"),
(list(self.model_groups), "model group"),
([guardrail.name for guardrail in self.guardrails], "guardrail"),
):
assert len(names) == len(set(names)), f"{self.name}: duplicate {label} names in {names}"
for team in self.teams:
assert set(team.models) <= group_names, f"{self.name}: team {team.name} names unknown models {team.models}"
for user in self.users:
assert user.team is None or user.team in team_names, (
f"{self.name}: user {user.name} in unknown team {user.team}"
)
for key in self.virtual_keys:
assert key.team is None or key.team in team_names, f"{self.name}: key {key.name} in unknown team {key.team}"
assert key.user is None or key.user in user_names, (
f"{self.name}: key {key.name} for unknown user {key.user}"
)
assert set(key.models) <= group_names, f"{self.name}: key {key.name} names unknown models {key.models}"
assert abs(sum(self.traffic.endpoints.values()) - 1.0) < 1e-9, f"{self.name}: endpoint shares must sum to 1"
assert abs(sum(self.traffic.shapes.values()) - 1.0) < 1e-9, f"{self.name}: shape shares must sum to 1"
assert 0 <= self.traffic.stream_share <= 1, f"{self.name}: stream_share out of range"
assert self.traffic.intervene_share + self.traffic.block_share <= 1, f"{self.name}: guardrail shares exceed 1"
assert self.traffic.total_requests > 0 and self.traffic.rps > 0, f"{self.name}: empty traffic"
assert self.model_groups, f"{self.name}: no model groups"
assert 0 <= self.mock_latency.jitter_seconds <= self.mock_latency.seconds or self.mock_latency.seconds == 0, (
f"{self.name}: mock latency jitter {self.mock_latency.jitter_seconds} s exceeds {self.mock_latency.seconds} s"
)

View file

@ -0,0 +1,122 @@
"""The deterministic request plan of a `PerformanceTestCase`: every request the run sends, in order,
drawn once from the case's seed so two runs of the same case send the same traffic. Each request
carries a marker (`perfreq-000123-n`) that the runner puts in the final user message and the mock
providers and the mock guardrail read back, so the test can join what the client saw against what
each mock received, and a `RequestShape` the runner hands the corpus to build the body."""
from __future__ import annotations
import random
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal, TypeVar
from setup.performance_test_case import Endpoint, PerformanceTestCase, RequestShape, Team
Flag = Literal["n", "i", "b"]
CallerKind = Literal["jwt", "key"]
T = TypeVar("T")
@dataclass(frozen=True, slots=True)
class Caller:
kind: CallerKind
name: str
@property
def label(self) -> str:
return f"{self.kind}:{self.name}"
@dataclass(frozen=True, slots=True)
class PlannedRequest:
seq: int
marker: str
caller: Caller
endpoint: Endpoint
model_group: str
stream: bool
flag: Flag
shape: RequestShape
input_tokens: int
@property
def blocked(self) -> bool:
return self.flag == "b"
def _allowed_models(case: PerformanceTestCase, team: Team | None, own: tuple[str, ...]) -> frozenset[str]:
if own:
return frozenset(own)
if team is not None and team.models:
return frozenset(team.models)
return frozenset(case.model_groups)
def _callers_by_group(case: PerformanceTestCase) -> dict[str, tuple[Caller, ...]]:
"""Which callers may use each model group, in case order, so a caller never asks for a model its
team or key does not allow and every denial in the run is a bug, not the plan."""
reach: list[tuple[Caller, frozenset[str]]] = []
for user in case.users:
user_team = case.team(user.team) if user.team else None
reach.append((Caller("jwt", user.name), _allowed_models(case, user_team, ())))
for key in case.virtual_keys:
key_team = case.team(key.team) if key.team else None
reach.append((Caller("key", key.name), _allowed_models(case, key_team, key.models)))
by_group: dict[str, tuple[Caller, ...]] = {}
for group in case.model_groups:
by_group[group] = tuple(caller for caller, models in reach if group in models)
assert by_group[group], f"{case.name}: no caller may use model group {group}"
return by_group
def _pick(rng: random.Random, shares: Mapping[T, float]) -> T:
draw: Final = rng.random()
cumulative = 0.0
chosen: T | None = None
for item, share in shares.items():
cumulative += share
chosen = item
if draw < cumulative:
break
assert chosen is not None, "empty shares"
return chosen
def _pick_flag(rng: random.Random, case: PerformanceTestCase, *, stream: bool) -> Flag:
if not case.guardrails:
return "n"
draw: Final = rng.random()
if draw < case.traffic.block_share:
return "b"
if draw < case.traffic.block_share + case.traffic.intervene_share and not stream:
return "i"
return "n"
def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]:
case.validate()
rng: Final = random.Random(case.traffic.seed)
callers: Final = _callers_by_group(case)
planned: list[PlannedRequest] = []
for seq in range(1, case.traffic.total_requests + 1):
endpoint = _pick(rng, case.traffic.endpoints)
group = rng.choice(case.model_groups)
stream = rng.random() < case.traffic.stream_share
flag = _pick_flag(rng, case, stream=stream)
shape = _pick(rng, case.traffic.shapes)
low, high = shape.input_tokens
planned.append(
PlannedRequest(
seq=seq,
marker=f"perfreq-{seq:06d}-{flag}",
caller=rng.choice(callers[group]),
endpoint=endpoint,
model_group=group,
stream=stream,
flag=flag,
shape=shape,
input_tokens=rng.randint(low, high),
)
)
return tuple(planned)

View file

@ -0,0 +1,126 @@
"""CPU seconds the gateway pods burned during the traffic phase, from metrics-server (`metrics.k8s.io`
PodMetrics), so a case can put a ceiling on the CPU the gateway costs per request: a slower gateway shows
up here at 5 rps, long before it saturates. The runner samples the pods matching
`E2E_PERF_GATEWAY_POD_SELECTOR` every few seconds with its own service account; `usage.cpu` is the
pods' average core use over metrics-server's window, integrated over the sampling interval. Without the
selector or the in-cluster token the reading is None and a Limit on it reports as missing."""
from __future__ import annotations
import os
import threading
import time
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from typing import Final
from urllib.parse import quote
from e2e_http import AuthHeaders, Success, get_external
from pydantic import BaseModel
SELECTOR_ENV: Final = "E2E_PERF_GATEWAY_POD_SELECTOR"
_SERVICE_ACCOUNT: Final = Path("/var/run/secrets/kubernetes.io/serviceaccount")
_INTERVAL_S: Final = 5.0
_SUFFIX_CORES: Final = {"n": 1e-9, "u": 1e-6, "m": 1e-3}
class _Usage(BaseModel):
cpu: str
class _Container(BaseModel):
usage: _Usage
class _PodMetrics(BaseModel):
containers: list[_Container]
class _PodMetricsList(BaseModel):
items: list[_PodMetrics]
def cores(quantity: str) -> float:
"""A Kubernetes CPU quantity (`1500m`, `250000000n`, `2`) as cores."""
if quantity and quantity[-1] in _SUFFIX_CORES:
return float(quantity[:-1]) * _SUFFIX_CORES[quantity[-1]]
return float(quantity)
class GatewayCpuSampler:
"""Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds."""
def __init__(self, selector: str, namespace: str, token: str, ca_bundle: str) -> None:
self._url: Final = (
f"https://kubernetes.default.svc/apis/metrics.k8s.io/v1beta1/namespaces/{namespace}/pods"
f"?labelSelector={quote(selector)}"
)
self._headers: Final = AuthHeaders(authorization=f"Bearer {token}")
self._ca_bundle: Final = ca_bundle
self._stop: Final = threading.Event()
self._thread: Final = threading.Thread(target=self._run, name="gateway-cpu-sampler", daemon=True)
self._cpu_seconds = 0.0
self._samples = 0
@property
def seconds(self) -> float | None:
return self._cpu_seconds if self._samples > 0 else None
def _read_cores(self) -> float | None:
result: Final = get_external(
self._url, response_type=_PodMetricsList, headers=self._headers, timeout=10.0, ca_bundle=self._ca_bundle
)
if not isinstance(result, Success):
return None
return sum(cores(container.usage.cpu) for pod in result.data.items for container in pod.containers)
def _run(self) -> None:
last: float = time.monotonic() # rebind-ok: the sampling clock advances every tick
while not self._stop.wait(_INTERVAL_S):
now: float = time.monotonic()
current: float | None = self._read_cores()
if current is not None:
self._cpu_seconds += current * (now - last)
self._samples += 1
last = now # rebind-ok: the sampling clock advances every tick
def start(self) -> None:
self._thread.start()
def stop(self) -> None:
self._stop.set()
self._thread.join(timeout=_INTERVAL_S * 3)
class NoSampler:
seconds: Final[None] = None
def start(self) -> None:
return None
def stop(self) -> None:
return None
def _in_cluster() -> tuple[str, str, str] | None:
token: Final = _SERVICE_ACCOUNT / "token"
ca: Final = _SERVICE_ACCOUNT / "ca.crt"
namespace: Final = _SERVICE_ACCOUNT / "namespace"
if not (token.is_file() and ca.is_file() and namespace.is_file()):
return None
return token.read_text().strip(), str(ca), namespace.read_text().strip()
@contextmanager
def gateway_cpu_sampler() -> Generator[GatewayCpuSampler | NoSampler]:
selector: Final = os.environ.get(SELECTOR_ENV, "")
cluster: Final = _in_cluster()
sampler: Final = (
GatewayCpuSampler(selector, cluster[2], cluster[0], cluster[1]) if selector and cluster else NoSampler()
)
sampler.start()
try:
yield sampler
finally:
sampler.stop()

View file

@ -0,0 +1,190 @@
"""Builds the caller side of a `PerformanceTestCase` on the gateway: the teams, the Keycloak users
(JWT callers) and the virtual keys, each suffixed with the run marker and deferred to the
`ResourceManager`, so two runs never collide and a failed run still tears down. Model groups and
guardrails are not created here: they are written out in the gateway's config (the helm values in
the lane, `gateway/perf_ci_config.yml` locally), and setup only checks the deployed gateway has what
the case names."""
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Final
from e2e_config import unique_marker
from e2e_http import NoBody, unwrap
from idp import TESTS_CLIENT_ID, Keycloak
from lifecycle import ResourceManager
from models import KeyGenerateBody, TeamNewBody, UserNewBody, UserNewResponse
from proxy_client import ProxyClient
from pydantic import BaseModel, ConfigDict
from setup.performance_test_case import MockProvider, PerformanceTestCase
from setup.plan import Caller
from setup.wire import TokenSource
MOCK_PROVIDERS: Final[tuple[MockProvider, ...]] = ("openai", "anthropic", "bedrock")
MOCK_URL_ENV: Final[dict[MockProvider, str]] = {
"openai": "E2E_PERF_MOCK_OPENAI_URL",
"anthropic": "E2E_PERF_MOCK_ANTHROPIC_URL",
"bedrock": "E2E_PERF_MOCK_BEDROCK_URL",
}
MOCK_URL_DEFAULT: Final[dict[MockProvider, str]] = {
"openai": "http://127.0.0.1:8090/v1",
"anthropic": "http://127.0.0.1:8091",
"bedrock": "http://127.0.0.1:8092",
}
GUARDRAIL_URL_ENV: Final = "E2E_PERF_MOCK_GUARDRAIL_URL"
GUARDRAIL_URL_DEFAULT: Final = "http://127.0.0.1:8093"
def mock_url(provider: MockProvider) -> str:
"""The mock's root, where its control routes live. The same variable feeds the gateway's
`api_base`, which for OpenAI ends in `/v1`, so that suffix is dropped here."""
url: Final = os.environ.get(MOCK_URL_ENV[provider], MOCK_URL_DEFAULT[provider]).rstrip("/")
return url.removesuffix("/v1")
def guardrail_url() -> str:
return os.environ.get(GUARDRAIL_URL_ENV, GUARDRAIL_URL_DEFAULT).rstrip("/")
class PerfTeamNewBody(TeamNewBody):
"""`/team/new` with the per-team settings a `Team` carries."""
max_budget: float | None = None
rpm_limit: int | None = None
tpm_limit: int | None = None
class ListedGuardrailParams(BaseModel):
model_config = ConfigDict(frozen=True)
mode: str | list[str]
class ListedGuardrail(BaseModel):
model_config = ConfigDict(frozen=True)
guardrail_name: str
litellm_params: ListedGuardrailParams
class GuardrailList(BaseModel):
model_config = ConfigDict(frozen=True)
guardrails: list[ListedGuardrail]
@dataclass(frozen=True, slots=True)
class CaseSetup:
marker: str
team_ids: dict[str, str]
key_bearers: dict[str, str]
token_sources: dict[str, TokenSource]
user_ids: dict[str, str]
def authorization(self, caller: Caller) -> str | None:
return f"Bearer {self.key_bearers[caller.name]}" if caller.kind == "key" else None
def token_source(self, caller: Caller) -> TokenSource | None:
return self.token_sources[caller.name] if caller.kind == "jwt" else None
def check_deployed(proxy: ProxyClient, case: PerformanceTestCase) -> None:
"""The gateway was deployed with every model group and guardrail the case names, with the
guardrail on the hook the case expects. A mismatch is a config drift, not a perf result."""
deployed: Final = sorted({entry.model_name for entry in proxy.model_info()})
missing: Final = [group for group in case.model_groups if group not in deployed]
assert not missing, f"{case.name}: model groups {missing} are not in the gateway config (deployed: {deployed})"
if not case.guardrails:
return
listed: Final = unwrap(
proxy.transport.get(
"/guardrails/list", headers=proxy.transport.master, params=NoBody(), response_type=GuardrailList
)
).guardrails
modes: Final = {guardrail.guardrail_name: guardrail.litellm_params.mode for guardrail in listed}
for guardrail in case.guardrails:
mode = modes.get(guardrail.name)
hooks = mode if isinstance(mode, list) else [mode]
assert guardrail.mode in hooks, (
f"{case.name}: guardrail {guardrail.name} should be on {guardrail.mode}, gateway has {modes}"
)
def build_setup(
case: PerformanceTestCase, *, proxy: ProxyClient, idp: Keycloak, resources: ResourceManager
) -> CaseSetup:
check_deployed(proxy, case)
marker: Final = unique_marker()
team_ids: dict[str, str] = {}
keycloak_groups: dict[str, str] = {}
for team in case.teams:
team_id = f"perf-{team.name}-{marker}"
team_ids[team.name] = proxy.create_team(
PerfTeamNewBody(
team_alias=team_id,
team_id=team_id,
models=list(team.models),
max_budget=team.max_budget,
rpm_limit=team.rpm_limit,
tpm_limit=team.tpm_limit,
)
)
resources.defer(lambda team_id=team_id: proxy.delete_team(team_id))
group_id = idp.create_group(team_id)
keycloak_groups[team_id] = group_id
resources.defer(lambda group_id=group_id: idp.delete_group(group_id))
user_ids: dict[str, str] = {}
token_sources: dict[str, TokenSource] = {}
for user in case.users:
groups = (team_ids[user.team],) if user.team else ()
identity = idp.provision_user(
marker=f"{marker}-{user.name}",
groups=groups,
group_ids=tuple(keycloak_groups[group] for group in groups),
defer=resources.defer,
)
user_ids[user.name] = identity.user_id
token_sources[user.name] = TokenSource(
token_url=idp.token_url(idp.realm),
client_id=TESTS_CLIENT_ID,
username=identity.username,
password=identity.password,
)
created = unwrap(
proxy.transport.post(
"/user/new",
headers=proxy.transport.master,
json=UserNewBody(
user_email=f"{identity.username}@example.com",
user_role="internal_user",
user_id=identity.user_id,
auto_create_key=False,
max_budget=user.max_budget,
),
response_type=UserNewResponse,
)
)
resources.defer(lambda user_id=created.user_id: proxy.delete_user(user_id))
key_bearers: dict[str, str] = {}
for key in case.virtual_keys:
generated = proxy.generate_key(
KeyGenerateBody(
models=list(key.models),
team_id=team_ids[key.team] if key.team else None,
user_id=user_ids[key.user] if key.user else None,
rpm_limit=key.rpm_limit,
max_budget=key.max_budget,
key_alias=f"perf-{key.name}-{marker}",
)
)
resources.defer(lambda generated=generated: proxy.delete_key(generated))
key_bearers[key.name] = generated
return CaseSetup(
marker=marker, team_ids=team_ids, key_bearers=key_bearers, token_sources=token_sources, user_ids=user_ids
)

129
tests/e2e/perf/setup/run.py Normal file
View file

@ -0,0 +1,129 @@
"""Sends a case's plan through the gateway from a separate process: the plan is written to disk as
wire requests, Locust (`locustfile.py`) replays it at the planned rate, and the records it
wrote come back as `WireRecord`s."""
from __future__ import annotations
import os
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from e2e_config import PROXY_BASE_URL
from pydantic import TypeAdapter
from setup.performance_test_case import PerformanceTestCase
from setup.plan import PlannedRequest
from setup.provision import CaseSetup
from setup.wire import WireRecord, WireRequest, WireShape
_LOCUSTFILE: Final = Path(__file__).with_name("locustfile.py")
_PERF_DIR: Final = Path(__file__).resolve().parents[1]
_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest])
_RECORD_ADAPTER: Final = TypeAdapter(WireRecord)
_PATHS: Final = {"chat_completions": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"}
@dataclass(frozen=True, slots=True)
class RunOutcome:
records: tuple[WireRecord, ...]
locust_exit: int
locust_output: str
@property
def wall_seconds(self) -> float:
if not self.records:
return 0.0
first: Final = min(record.started_at for record in self.records)
last: Final = max(record.started_at + record.total_s for record in self.records)
return last - first
def wire_request(item: PlannedRequest, setup: CaseSetup) -> WireRequest:
return WireRequest(
seq=item.seq,
marker=item.marker,
caller=item.caller.label,
endpoint=item.endpoint,
path=_PATHS[item.endpoint],
model=item.model_group,
shape=WireShape(
name=item.shape.name,
input_tokens=item.input_tokens,
tools=item.shape.tools,
output_tokens=item.shape.output_tokens,
),
authorization=setup.authorization(item.caller),
token_source=setup.token_source(item.caller),
stream=item.stream,
blocked=item.blocked,
)
def _keep_records_artifact(records_path: Path) -> None:
"""Copy the per-request records next to the JUnit file when the runner exposes an artifact
directory (`E2E_ARTIFACT_DIR`), so a red lane run can be broken down after the pod is gone."""
artifact_dir: Final = os.environ.get("E2E_ARTIFACT_DIR")
if not artifact_dir or not Path(artifact_dir).is_dir():
return
_ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl")
def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome:
"""Replay `plan` with Locust in a child process and collect its records. The run-time ceiling
is well past the planned duration so Locust's own clock never cuts a slow run short; the
locustfile quits itself once the last record is written."""
wire: Final = [wire_request(item, setup) for item in plan]
ceiling_s: Final = int(case.traffic.expected_seconds * 2 + 180)
with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory:
plan_path: Final = Path(directory) / "plan.json"
records_path: Final = Path(directory) / "records.jsonl"
_ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire))
records_path.touch()
env: Final = {
**os.environ,
"PYTHONPATH": os.pathsep.join(filter(None, (str(_PERF_DIR), os.environ.get("PYTHONPATH")))),
"PERF_PLAN_FILE": str(plan_path),
"PERF_RECORDS_FILE": str(records_path),
"PERF_RPS": str(case.traffic.rps),
"PERF_USERS": str(case.traffic.concurrency),
}
completed: Final = subprocess.run(
[
sys.executable,
"-P",
"-m",
"locust",
"--headless",
"--only-summary",
"--locustfile",
str(_LOCUSTFILE),
"--host",
PROXY_BASE_URL,
"--users",
str(case.traffic.concurrency),
"--spawn-rate",
str(case.traffic.concurrency),
"--run-time",
f"{ceiling_s}s",
"--exit-code-on-error",
"0",
],
env=env,
capture_output=True,
text=True,
timeout=ceiling_s + 120,
check=False,
)
records: Final = tuple(
_RECORD_ADAPTER.validate_json(line)
for line in records_path.read_text(encoding="utf-8").splitlines()
if line.strip()
)
_keep_records_artifact(records_path)
output: Final = (completed.stdout + completed.stderr)[-6000:]
return RunOutcome(records=records, locust_exit=completed.returncode, locust_output=output)

View file

@ -0,0 +1,91 @@
"""What crosses the process boundary between the test and the Locust runner: the wire plan (one
entry per planned request, with the shape the runner expands into the body right before sending,
so a 300k-token request costs a few hundred bytes in the plan file) and the record each request
leaves behind."""
from __future__ import annotations
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field
from setup.corpus.session import Endpoint as CorpusEndpoint
from setup.corpus.session import request_body as corpus_body
Endpoint = Literal["chat_completions", "responses", "messages"]
class TokenSource(BaseModel):
"""How the runner mints a JWT for a Keycloak user, so a run longer than the realm's access
token lifespan keeps authenticating as that user."""
token_url: str
client_id: str
username: str
password: str = Field(repr=False)
class WireShape(BaseModel):
"""`RequestShape` on the wire: what the runner hands the corpus to build the body."""
model_config = ConfigDict(frozen=True)
name: str
input_tokens: int
tools: bool
output_tokens: int
class WireRequest(BaseModel):
seq: int
marker: str
caller: str
endpoint: Endpoint
path: str
model: str
shape: WireShape
authorization: str | None = Field(default=None, repr=False)
token_source: TokenSource | None = None
stream: bool
blocked: bool
class WireRecord(BaseModel):
seq: int
marker: str
caller: str
endpoint: Endpoint
stream: bool
blocked: bool
started_at: float
status: int
header_s: float
first_byte_s: float | None
first_token_s: float | None
total_s: float
text: str
ok: bool
mock_delay_s: float | None
call_id: str | None
error: str | None
_CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType(
{"chat_completions": "chat", "responses": "responses", "messages": "messages"}
)
def request_body(item: WireRequest) -> dict[str, object]:
"""The request in the endpoint's own wire format: the planned shape's coding-agent session, the marker in
the final user message, built right before sending so the plan file stays small."""
return corpus_body(
endpoint=_CORPUS_ENDPOINT[item.endpoint],
model=item.model,
marker=item.marker,
input_tokens=item.shape.input_tokens,
tools=item.shape.tools,
output_tokens=item.shape.output_tokens,
stream=item.stream,
)

View file

@ -0,0 +1,106 @@
"""One perf run per `PerformanceTestCase` in `cases/`: check the gateway was deployed with the case's
model groups and guardrails, build its teams, users and keys, replay its deterministic plan from a separate Locust
process against the out-of-process mock providers, then check every `Limit` against what the
clients saw, what the mocks received and how the workers' memory moved."""
from __future__ import annotations
import time
from typing import Final
import pytest
from idp import Keycloak
from lifecycle import ResourceManager
from memory_readings import read_rss_everywhere
from models import SpendLogsParams
from cases import CASES
from setup.observe import (
Observation,
evaluate,
read_guardrail_calls,
read_mock_requests,
render,
reset_mock,
set_mock_latency,
)
from setup.performance_test_case import PerformanceTestCase
from setup.plan import build_plan
from setup.datastore import read_db_counters, read_redis_counters
from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url
from setup.run import run_plan
from proxy_client import ProxyClient
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
_SPEND_SAMPLE: Final = 40
def _spend_logged_share(proxy: ProxyClient, call_ids: tuple[str, ...]) -> float | None:
"""Share of a sample of served requests whose spend row the gateway wrote, read back by the
request id the gateway returned in `x-litellm-call-id`."""
if not call_ids:
return None
found = 0
for call_id in call_ids:
rows = proxy.spend_logs(SpendLogsParams(request_id=call_id))
found += 1 if rows else 0
return found / len(call_ids)
@pytest.mark.parametrize("case", CASES, ids=[case.name for case in CASES])
def test_perf_case(case: PerformanceTestCase, proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None:
plan: Final = build_plan(case)
mock_urls: Final = {provider: mock_url(provider) for provider in MOCK_PROVIDERS}
for url in (*mock_urls.values(), guardrail_url()):
_ = reset_mock(url)
for url in mock_urls.values():
_ = set_mock_latency(url, case.mock_latency)
setup_started: Final = time.monotonic()
setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources)
rss_before: Final = read_rss_everywhere(proxy)
db_before: Final = read_db_counters()
redis_before: Final = read_redis_counters()
run_started: Final = time.monotonic()
with gateway_cpu_sampler() as cpu:
outcome: Final = run_plan(case, plan, setup)
run_finished: Final = time.monotonic()
rss_after: Final = read_rss_everywhere(proxy)
rss_read: Final = time.monotonic()
assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}"
served_call_ids: Final = tuple(
record.call_id for record in outcome.records if record.ok and not record.blocked and record.call_id
)
spend_share: Final = _spend_logged_share(
proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE]
)
spend_read: Final = time.monotonic()
db_after: Final = read_db_counters()
redis_after: Final = read_redis_counters()
mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url))
guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else ()
mocks_read: Final = time.monotonic()
observation: Final = Observation(
records=outcome.records,
mock_requests=mock_requests,
guardrail_calls=guardrail_calls,
rss_before=rss_before,
rss_after=rss_after,
wall_seconds=outcome.wall_seconds,
spend_logged_share=spend_share,
db_before=db_before,
db_after=db_after,
redis_before=redis_before,
redis_after=redis_after,
gateway_cpu_seconds=cpu.seconds,
)
rows, exceeded_limits = evaluate(case, observation)
print(render(case, observation, rows))
print(
f" phases: setup {run_started - setup_started:.0f}s, locust {run_finished - run_started:.0f}s, "
f"observe {time.monotonic() - run_finished:.0f}s (rss {rss_read - run_finished:.0f}s, spend sample "
f"{spend_read - rss_read:.0f}s, mock logs {mocks_read - spend_read:.0f}s); locust exit {outcome.locust_exit}"
)
assert len(outcome.records) == len(plan), (
f"{len(plan)} planned, {len(outcome.records)} recorded; locust said:\n{outcome.locust_output}"
)
assert not exceeded_limits, "limits exceeded (observed > max):\n" + "\n".join(exceeded_limits)