perf: reduce gateway db pressure under load

crates/aether-testkit/src/bin/gateway_pressure_probe.rs

@@ -0,0 +1,353 @@
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use aether_testkit::{
+    fetch_prometheus_samples, find_metric_value_u64, run_http_load_probe, HttpLoadProbeConfig,
+    HttpLoadProbeResponseMode, HttpLoadProbeResult, PrometheusSample,
+};
+use reqwest::Method;
+use serde::Serialize;
+use tokio::sync::Mutex;
+
+#[derive(Debug, Clone)]
+struct Config {
+    load: HttpLoadProbeConfig,
+    metrics_url: String,
+    sample_interval: Duration,
+    output_path: Option<PathBuf>,
+}
+
+#[derive(Debug, Serialize)]
+struct GatewayPressureReport {
+    suite: &'static str,
+    target_url: String,
+    metrics_url: String,
+    sample_interval_ms: u64,
+    load: HttpLoadProbeResult,
+    metrics: GatewayPressureMetricsSummary,
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+struct GatewayPressureMetricsSummary {
+    samples: usize,
+    db_pool_max_checked_out: u64,
+    db_pool_min_idle: Option<u64>,
+    db_pool_max_size: u64,
+    db_pool_max_connections: u64,
+    db_pool_max_usage_basis_points: u64,
+    db_pool_max_idle_reserve: u64,
+    db_pool_pressure_samples: usize,
+    gateway_requests_max_in_flight: u64,
+    gateway_requests_max_rejected_total: u64,
+    gateway_requests_distributed_max_in_flight: u64,
+    gateway_requests_distributed_max_rejected_total: u64,
+}
+
+impl GatewayPressureMetricsSummary {
+    fn observe(&mut self, samples: &[PrometheusSample]) {
+        self.samples += 1;
+        self.db_pool_max_checked_out = self
+            .db_pool_max_checked_out
+            .max(metric_max(samples, "database_pool_checked_out_connections"));
+        let idle = metric_min(samples, "database_pool_idle_connections");
+        self.db_pool_min_idle = match (self.db_pool_min_idle, idle) {
+            (Some(current), Some(next)) => Some(current.min(next)),
+            (None, Some(next)) => Some(next),
+            (current, None) => current,
+        };
+        self.db_pool_max_size = self
+            .db_pool_max_size
+            .max(metric_max(samples, "database_pool_size_connections"));
+        self.db_pool_max_connections = self
+            .db_pool_max_connections
+            .max(metric_max(samples, "database_pool_max_connections"));
+        self.db_pool_max_usage_basis_points = self
+            .db_pool_max_usage_basis_points
+            .max(metric_max(samples, "database_pool_usage_basis_points"));
+        self.db_pool_max_idle_reserve = self.db_pool_max_idle_reserve.max(metric_max(
+            samples,
+            "database_pool_idle_reserve_connections",
+        ));
+        if metric_max(samples, "database_pool_under_maintenance_pressure") > 0 {
+            self.db_pool_pressure_samples += 1;
+        }
+        self.gateway_requests_max_in_flight = self.gateway_requests_max_in_flight.max(
+            find_metric_value_u64(
+                samples,
+                "concurrency_in_flight",
+                &[("gate", "gateway_requests")],
+            )
+            .unwrap_or_default(),
+        );
+        self.gateway_requests_max_rejected_total = self.gateway_requests_max_rejected_total.max(
+            find_metric_value_u64(
+                samples,
+                "concurrency_rejected_total",
+                &[("gate", "gateway_requests")],
+            )
+            .unwrap_or_default(),
+        );
+        self.gateway_requests_distributed_max_in_flight =
+            self.gateway_requests_distributed_max_in_flight.max(
+                find_metric_value_u64(
+                    samples,
+                    "concurrency_in_flight",
+                    &[("gate", "gateway_requests_distributed")],
+                )
+                .unwrap_or_default(),
+            );
+        self.gateway_requests_distributed_max_rejected_total =
+            self.gateway_requests_distributed_max_rejected_total.max(
+                find_metric_value_u64(
+                    samples,
+                    "concurrency_rejected_total",
+                    &[("gate", "gateway_requests_distributed")],
+                )
+                .unwrap_or_default(),
+            );
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let config = parse_args(std::env::args().skip(1).collect())?;
+    let stop = Arc::new(AtomicBool::new(false));
+    let summary = Arc::new(Mutex::new(GatewayPressureMetricsSummary::default()));
+    let sampler = spawn_metrics_sampler(
+        config.metrics_url.clone(),
+        config.sample_interval,
+        Arc::clone(&stop),
+        Arc::clone(&summary),
+    );
+
+    let load = run_http_load_probe(&config.load)
+        .await
+        .map_err(std::io::Error::other)?;
+    stop.store(true, Ordering::Release);
+    sampler.await??;
+
+    if let Ok(samples) = fetch_prometheus_samples(&config.metrics_url).await {
+        summary.lock().await.observe(&samples);
+    }
+
+    let report = GatewayPressureReport {
+        suite: "gateway_pressure_probe",
+        target_url: config.load.url,
+        metrics_url: config.metrics_url,
+        sample_interval_ms: config.sample_interval.as_millis() as u64,
+        load,
+        metrics: Arc::try_unwrap(summary)
+            .unwrap_or_else(|_| panic!("metrics summary still referenced"))
+            .into_inner(),
+    };
+    let raw = serde_json::to_string_pretty(&report)?;
+    println!("{raw}");
+    if let Some(path) = config.output_path.as_ref() {
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent)?;
+        }
+        std::fs::write(path, format!("{raw}\n"))?;
+    }
+    Ok(())
+}
+
+fn spawn_metrics_sampler(
+    metrics_url: String,
+    interval: Duration,
+    stop: Arc<AtomicBool>,
+    summary: Arc<Mutex<GatewayPressureMetricsSummary>>,
+) -> tokio::task::JoinHandle<Result<(), std::io::Error>> {
+    tokio::spawn(async move {
+        while !stop.load(Ordering::Acquire) {
+            match fetch_prometheus_samples(&metrics_url).await {
+                Ok(samples) => summary.lock().await.observe(&samples),
+                Err(err) => {
+                    eprintln!("gateway pressure probe metrics sample failed: {err}");
+                }
+            }
+            tokio::time::sleep(interval).await;
+        }
+        Ok(())
+    })
+}
+
+fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
+    let mut target_url: Option<String> = None;
+    let mut metrics_url: Option<String> = None;
+    let mut total_requests: Option<usize> = None;
+    let mut concurrency: Option<usize> = None;
+    let mut timeout_ms: Option<u64> = None;
+    let mut sample_interval_ms: u64 = 500;
+    let mut method = Method::GET;
+    let mut headers = BTreeMap::new();
+    let mut body: Option<Vec<u8>> = None;
+    let mut response_mode = HttpLoadProbeResponseMode::HeadersOnly;
+    let mut output_path = None;
+
+    let mut iter = args.into_iter();
+    while let Some(arg) = iter.next() {
+        match arg.as_str() {
+            "--url" => target_url = Some(next_value(&mut iter, "--url")?),
+            "--metrics-url" => metrics_url = Some(next_value(&mut iter, "--metrics-url")?),
+            "--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?),
+            "--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?),
+            "--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?),
+            "--sample-interval-ms" => {
+                sample_interval_ms = next_value(&mut iter, "--sample-interval-ms")?.parse()?
+            }
+            "--method" => {
+                method = Method::from_bytes(next_value(&mut iter, "--method")?.as_bytes())?
+            }
+            "--header" | "-H" => {
+                let (name, value) = parse_header_arg(&next_value(&mut iter, "--header")?)?;
+                headers.insert(name, value);
+            }
+            "--body" => body = Some(next_value(&mut iter, "--body")?.into_bytes()),
+            "--body-file" => body = Some(std::fs::read(next_value(&mut iter, "--body-file")?)?),
+            "--response-mode" => {
+                response_mode = parse_response_mode(&next_value(&mut iter, "--response-mode")?)?
+            }
+            "--output" => output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?)),
+            "--help" | "-h" => {
+                print_usage();
+                std::process::exit(0);
+            }
+            other => {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidInput,
+                    format!("unknown argument: {other}"),
+                )
+                .into());
+            }
+        }
+    }
+
+    let mut load = HttpLoadProbeConfig {
+        url: target_url.ok_or_else(|| {
+            std::io::Error::new(std::io::ErrorKind::InvalidInput, "missing required --url")
+        })?,
+        total_requests: total_requests.ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "missing required --requests",
+            )
+        })?,
+        concurrency: concurrency.ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "missing required --concurrency",
+            )
+        })?,
+        method,
+        headers,
+        body,
+        response_mode,
+        ..HttpLoadProbeConfig::default()
+    };
+    if let Some(timeout_ms) = timeout_ms {
+        load.timeout = Duration::from_millis(timeout_ms);
+    }
+    load.validate()
+        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
+    if sample_interval_ms == 0 {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "--sample-interval-ms must be positive",
+        )
+        .into());
+    }
+    Ok(Config {
+        load,
+        metrics_url: metrics_url.ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "missing required --metrics-url",
+            )
+        })?,
+        sample_interval: Duration::from_millis(sample_interval_ms),
+        output_path,
+    })
+}
+
+fn parse_header_arg(value: &str) -> Result<(String, String), Box<dyn std::error::Error>> {
+    let (name, value) = value
+        .split_once(':')
+        .or_else(|| value.split_once('='))
+        .ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "--header expects `Name: value` or `Name=value`",
+            )
+        })?;
+    let name = name.trim();
+    if name.is_empty() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "--header name cannot be empty",
+        )
+        .into());
+    }
+    Ok((name.to_string(), value.trim().to_string()))
+}
+
+fn parse_response_mode(
+    value: &str,
+) -> Result<HttpLoadProbeResponseMode, Box<dyn std::error::Error>> {
+    match value.trim().to_ascii_lowercase().as_str() {
+        "headers" | "headers-only" | "header" => Ok(HttpLoadProbeResponseMode::HeadersOnly),
+        "full" | "full-body" | "body" => Ok(HttpLoadProbeResponseMode::FullBody),
+        other => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("unsupported --response-mode {other}; expected headers or full"),
+        )
+        .into()),
+    }
+}
+
+fn next_value(
+    iter: &mut impl Iterator<Item = String>,
+    flag: &str,
+) -> Result<String, Box<dyn std::error::Error>> {
+    iter.next().ok_or_else(|| {
+        std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("missing value for {flag}"),
+        )
+        .into()
+    })
+}
+
+fn metric_max(samples: &[PrometheusSample], metric_name: &str) -> u64 {
+    samples
+        .iter()
+        .filter(|sample| metric_name_matches(&sample.name, metric_name))
+        .filter_map(|sample| sample.value.parse::<u64>().ok())
+        .max()
+        .unwrap_or_default()
+}
+
+fn metric_min(samples: &[PrometheusSample], metric_name: &str) -> Option<u64> {
+    samples
+        .iter()
+        .filter(|sample| metric_name_matches(&sample.name, metric_name))
+        .filter_map(|sample| sample.value.parse::<u64>().ok())
+        .min()
+}
+
+fn metric_name_matches(actual: &str, expected: &str) -> bool {
+    actual == expected
+        || actual
+            .rsplit_once('_')
+            .map(|(_, suffix)| suffix == expected)
+            .unwrap_or(false)
+        || actual.ends_with(&format!("_{expected}"))
+}
+
+fn print_usage() {
+    eprintln!(
+        "usage: cargo run -p aether-testkit --bin gateway_pressure_probe -- --url <URL> --metrics-url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000] [--sample-interval-ms 500] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full] [--output /tmp/gateway_pressure.json]"
+    );
+}

crates/aether-testkit/src/bin/http_load_probe.rs

@@ -19,6 +19,9 @@ fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::err
     let mut concurrency: Option<usize> = None;
     let mut timeout_ms: Option<u64> = None;
     let mut method = Method::GET;
+    let mut headers = std::collections::BTreeMap::new();
+    let mut body: Option<Vec<u8>> = None;
+    let mut response_mode = aether_testkit::HttpLoadProbeResponseMode::HeadersOnly;
 
     let mut iter = args.into_iter();
     while let Some(arg) = iter.next() {
@@ -30,6 +33,15 @@ fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::err
             "--method" => {
                 method = Method::from_bytes(next_value(&mut iter, "--method")?.as_bytes())?
             }
+            "--header" | "-H" => {
+                let (name, value) = parse_header_arg(&next_value(&mut iter, "--header")?)?;
+                headers.insert(name, value);
+            }
+            "--body" => body = Some(next_value(&mut iter, "--body")?.into_bytes()),
+            "--body-file" => body = Some(std::fs::read(next_value(&mut iter, "--body-file")?)?),
+            "--response-mode" => {
+                response_mode = parse_response_mode(&next_value(&mut iter, "--response-mode")?)?
+            }
             "--help" | "-h" => {
                 print_usage();
                 std::process::exit(0);
@@ -61,6 +73,9 @@ fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::err
             )
         })?,
         method,
+        headers,
+        body,
+        response_mode,
         ..HttpLoadProbeConfig::default()
     };
     if let Some(timeout_ms) = timeout_ms {
@@ -72,6 +87,43 @@ fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::err
     Ok(config)
 }
 
+fn parse_header_arg(value: &str) -> Result<(String, String), Box<dyn std::error::Error>> {
+    let (name, value) = value
+        .split_once(':')
+        .or_else(|| value.split_once('='))
+        .ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "--header expects `Name: value` or `Name=value`",
+            )
+        })?;
+    let name = name.trim();
+    if name.is_empty() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "--header name cannot be empty",
+        )
+        .into());
+    }
+    Ok((name.to_string(), value.trim().to_string()))
+}
+
+fn parse_response_mode(
+    value: &str,
+) -> Result<aether_testkit::HttpLoadProbeResponseMode, Box<dyn std::error::Error>> {
+    match value.trim().to_ascii_lowercase().as_str() {
+        "headers" | "headers-only" | "header" => {
+            Ok(aether_testkit::HttpLoadProbeResponseMode::HeadersOnly)
+        }
+        "full" | "full-body" | "body" => Ok(aether_testkit::HttpLoadProbeResponseMode::FullBody),
+        other => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("unsupported --response-mode {other}; expected headers or full"),
+        )
+        .into()),
+    }
+}
+
 fn next_value(
     iter: &mut impl Iterator<Item = String>,
     flag: &str,
@@ -87,6 +139,6 @@ fn next_value(
 
 fn print_usage() {
     eprintln!(
-        "usage: cargo run -p aether-testkit --bin http_load_probe -- --url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000]"
+        "usage: cargo run -p aether-testkit --bin http_load_probe -- --url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full]"
     );
 }

crates/aether-testkit/src/bin/mock_openai_upstream.rs

@@ -0,0 +1,348 @@
+use std::convert::Infallible;
+use std::net::SocketAddr;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use axum::body::{Body, Bytes};
+use axum::extract::State;
+use axum::http::{header, HeaderValue, StatusCode};
+use axum::response::{IntoResponse, Response};
+use axum::routing::{get, post};
+use axum::Router;
+use serde_json::json;
+
+#[derive(Debug, Clone)]
+struct Config {
+    bind: SocketAddr,
+    chunks: u64,
+    first_byte_delay: Duration,
+    chunk_delay: Duration,
+    payload_bytes: usize,
+    status: StatusCode,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            bind: "127.0.0.1:18181"
+                .parse()
+                .expect("default bind address should parse"),
+            chunks: 8,
+            first_byte_delay: Duration::from_millis(0),
+            chunk_delay: Duration::from_millis(20),
+            payload_bytes: 32,
+            status: StatusCode::OK,
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+struct Metrics {
+    requests_total: AtomicU64,
+    completed_total: AtomicU64,
+    in_flight: AtomicU64,
+    max_in_flight: AtomicU64,
+}
+
+#[derive(Debug, Clone)]
+struct App {
+    config: Config,
+    metrics: Arc<Metrics>,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let config = parse_args(std::env::args().skip(1).collect())?;
+    let app_state = App {
+        config: config.clone(),
+        metrics: Arc::new(Metrics::default()),
+    };
+
+    let app = Router::new()
+        .route("/health", get(health))
+        .route("/metrics", get(metrics))
+        .route("/v1/chat/completions", post(chat_completions))
+        .route("/chat/completions", post(chat_completions))
+        .route("/v1/responses", post(responses))
+        .route("/responses", post(responses))
+        .with_state(app_state);
+
+    let listener = tokio::net::TcpListener::bind(config.bind).await?;
+    eprintln!("mock OpenAI upstream listening on http://{}", config.bind);
+    axum::serve(listener, app).await?;
+    Ok(())
+}
+
+async fn health() -> impl IntoResponse {
+    (StatusCode::OK, "ok\n")
+}
+
+async fn metrics(State(app): State<App>) -> Response {
+    let body = format!(
+        concat!(
+            "# HELP mock_upstream_requests_total Total requests accepted by the mock upstream.\n",
+            "# TYPE mock_upstream_requests_total counter\n",
+            "mock_upstream_requests_total {}\n",
+            "# HELP mock_upstream_completed_total Total requests completed by the mock upstream.\n",
+            "# TYPE mock_upstream_completed_total counter\n",
+            "mock_upstream_completed_total {}\n",
+            "# HELP mock_upstream_in_flight Current in-flight requests/streams.\n",
+            "# TYPE mock_upstream_in_flight gauge\n",
+            "mock_upstream_in_flight {}\n",
+            "# HELP mock_upstream_max_in_flight Maximum in-flight requests/streams observed.\n",
+            "# TYPE mock_upstream_max_in_flight gauge\n",
+            "mock_upstream_max_in_flight {}\n"
+        ),
+        app.metrics.requests_total.load(Ordering::Acquire),
+        app.metrics.completed_total.load(Ordering::Acquire),
+        app.metrics.in_flight.load(Ordering::Acquire),
+        app.metrics.max_in_flight.load(Ordering::Acquire),
+    );
+    (StatusCode::OK, body).into_response()
+}
+
+async fn chat_completions(State(app): State<App>, body: Bytes) -> Response {
+    let stream = request_wants_stream(&body);
+    record_request_started(&app.metrics);
+    if app.config.status != StatusCode::OK {
+        record_request_completed(&app.metrics);
+        return (app.config.status, "mock upstream error\n").into_response();
+    }
+    if stream {
+        return build_chat_sse_response(app);
+    }
+    let payload = json!({
+        "id": "chatcmpl-mock",
+        "object": "chat.completion",
+        "created": current_unix_secs(),
+        "model": request_model(&body).unwrap_or_else(|| "mock-model".to_string()),
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": mock_payload(app.config.payload_bytes)
+            },
+            "finish_reason": "stop"
+        }],
+        "usage": {
+            "prompt_tokens": 1,
+            "completion_tokens": app.config.chunks.max(1),
+            "total_tokens": app.config.chunks.max(1) + 1
+        }
+    });
+    record_request_completed(&app.metrics);
+    axum::Json(payload).into_response()
+}
+
+async fn responses(State(app): State<App>, body: Bytes) -> Response {
+    let stream = request_wants_stream(&body);
+    record_request_started(&app.metrics);
+    if app.config.status != StatusCode::OK {
+        record_request_completed(&app.metrics);
+        return (app.config.status, "mock upstream error\n").into_response();
+    }
+    if stream {
+        return build_responses_sse_response(app);
+    }
+    let payload = json!({
+        "id": "resp_mock",
+        "object": "response",
+        "created_at": current_unix_secs(),
+        "model": request_model(&body).unwrap_or_else(|| "mock-model".to_string()),
+        "output": [{
+            "type": "message",
+            "id": "msg_mock",
+            "role": "assistant",
+            "content": [{
+                "type": "output_text",
+                "text": mock_payload(app.config.payload_bytes)
+            }]
+        }],
+        "usage": {
+            "input_tokens": 1,
+            "output_tokens": app.config.chunks.max(1),
+            "total_tokens": app.config.chunks.max(1) + 1
+        }
+    });
+    record_request_completed(&app.metrics);
+    axum::Json(payload).into_response()
+}
+
+fn build_chat_sse_response(app: App) -> Response {
+    let metrics = Arc::clone(&app.metrics);
+    let config = app.config.clone();
+    let stream = async_stream::stream! {
+        if !config.first_byte_delay.is_zero() {
+            tokio::time::sleep(config.first_byte_delay).await;
+        }
+        for index in 0..config.chunks {
+            if index > 0 && !config.chunk_delay.is_zero() {
+                tokio::time::sleep(config.chunk_delay).await;
+            }
+            let payload = json!({
+                "id": "chatcmpl-mock",
+                "object": "chat.completion.chunk",
+                "created": current_unix_secs(),
+                "model": "mock-model",
+                "choices": [{
+                    "index": 0,
+                    "delta": {
+                        "content": mock_payload(config.payload_bytes)
+                    },
+                    "finish_reason": serde_json::Value::Null
+                }]
+            });
+            yield Ok::<Bytes, Infallible>(Bytes::from(format!("data: {payload}\n\n")));
+        }
+        yield Ok::<Bytes, Infallible>(Bytes::from(
+            "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
+        ));
+        yield Ok::<Bytes, Infallible>(Bytes::from("data: [DONE]\n\n"));
+        record_request_completed(&metrics);
+    };
+    sse_response(Body::from_stream(stream))
+}
+
+fn build_responses_sse_response(app: App) -> Response {
+    let metrics = Arc::clone(&app.metrics);
+    let config = app.config.clone();
+    let stream = async_stream::stream! {
+        if !config.first_byte_delay.is_zero() {
+            tokio::time::sleep(config.first_byte_delay).await;
+        }
+        yield Ok::<Bytes, Infallible>(Bytes::from(
+            "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_mock\",\"status\":\"in_progress\"}}\n\n",
+        ));
+        for index in 0..config.chunks {
+            if index > 0 && !config.chunk_delay.is_zero() {
+                tokio::time::sleep(config.chunk_delay).await;
+            }
+            let payload = json!({
+                "type": "response.output_text.delta",
+                "item_id": "msg_mock",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": mock_payload(config.payload_bytes)
+            });
+            yield Ok::<Bytes, Infallible>(Bytes::from(format!("event: response.output_text.delta\ndata: {payload}\n\n")));
+        }
+        yield Ok::<Bytes, Infallible>(Bytes::from(
+            "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_mock\",\"status\":\"completed\"}}\n\n",
+        ));
+        yield Ok::<Bytes, Infallible>(Bytes::from("data: [DONE]\n\n"));
+        record_request_completed(&metrics);
+    };
+    sse_response(Body::from_stream(stream))
+}
+
+fn sse_response(body: Body) -> Response {
+    let mut response = Response::new(body);
+    let headers = response.headers_mut();
+    headers.insert(
+        header::CONTENT_TYPE,
+        HeaderValue::from_static("text/event-stream; charset=utf-8"),
+    );
+    headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
+    response
+}
+
+fn request_wants_stream(body: &[u8]) -> bool {
+    serde_json::from_slice::<serde_json::Value>(body)
+        .ok()
+        .and_then(|value| value.get("stream").and_then(serde_json::Value::as_bool))
+        .unwrap_or(false)
+}
+
+fn request_model(body: &[u8]) -> Option<String> {
+    serde_json::from_slice::<serde_json::Value>(body)
+        .ok()
+        .and_then(|value| {
+            value
+                .get("model")
+                .and_then(serde_json::Value::as_str)
+                .map(str::to_string)
+        })
+}
+
+fn mock_payload(bytes: usize) -> String {
+    if bytes == 0 {
+        return String::new();
+    }
+    "x".repeat(bytes)
+}
+
+fn record_request_started(metrics: &Metrics) {
+    metrics.requests_total.fetch_add(1, Ordering::AcqRel);
+    let in_flight = metrics.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
+    metrics.max_in_flight.fetch_max(in_flight, Ordering::AcqRel);
+}
+
+fn record_request_completed(metrics: &Metrics) {
+    metrics.completed_total.fetch_add(1, Ordering::AcqRel);
+    metrics.in_flight.fetch_sub(1, Ordering::AcqRel);
+}
+
+fn current_unix_secs() -> u64 {
+    std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .map(|duration| duration.as_secs())
+        .unwrap_or(0)
+}
+
+fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
+    let mut config = Config::default();
+    let mut iter = args.into_iter();
+    while let Some(arg) = iter.next() {
+        match arg.as_str() {
+            "--bind" => config.bind = next_value(&mut iter, "--bind")?.parse()?,
+            "--chunks" => config.chunks = next_value(&mut iter, "--chunks")?.parse()?,
+            "--first-byte-delay-ms" => {
+                config.first_byte_delay =
+                    Duration::from_millis(next_value(&mut iter, "--first-byte-delay-ms")?.parse()?)
+            }
+            "--chunk-delay-ms" => {
+                config.chunk_delay =
+                    Duration::from_millis(next_value(&mut iter, "--chunk-delay-ms")?.parse()?)
+            }
+            "--payload-bytes" => {
+                config.payload_bytes = next_value(&mut iter, "--payload-bytes")?.parse()?
+            }
+            "--status" => {
+                let status = next_value(&mut iter, "--status")?.parse::<u16>()?;
+                config.status = StatusCode::from_u16(status)?;
+            }
+            "--help" | "-h" => {
+                print_usage();
+                std::process::exit(0);
+            }
+            other => {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidInput,
+                    format!("unknown argument: {other}"),
+                )
+                .into());
+            }
+        }
+    }
+    Ok(config)
+}
+
+fn next_value(
+    iter: &mut impl Iterator<Item = String>,
+    flag: &str,
+) -> Result<String, Box<dyn std::error::Error>> {
+    iter.next().ok_or_else(|| {
+        std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!("missing value for {flag}"),
+        )
+        .into()
+    })
+}
+
+fn print_usage() {
+    eprintln!(
+        "usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181] [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]"
+    );
+}

crates/aether-testkit/src/load.rs

@@ -79,6 +79,7 @@ pub struct HttpLoadProbeResult {
     pub mean_ms: u64,
     pub runtime: BenchmarkRuntimeSnapshot,
     pub status_counts: BTreeMap<u16, usize>,
+    pub error_counts: BTreeMap<String, usize>,
 }
 
 #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
@@ -100,6 +101,7 @@ pub struct MultiUrlHttpLoadProbeResult {
     pub mean_ms: u64,
     pub runtime: BenchmarkRuntimeSnapshot,
     pub status_counts: BTreeMap<u16, usize>,
+    pub error_counts: BTreeMap<String, usize>,
 }
 
 pub async fn run_http_load_probe(
@@ -129,6 +131,7 @@ pub async fn run_http_load_probe(
             mean_ms: result.mean_ms,
             runtime: result.runtime,
             status_counts: result.status_counts,
+            error_counts: result.error_counts,
         })
 }
 
@@ -161,6 +164,7 @@ async fn run_http_load_probe_against_urls(
     let next_request = Arc::new(AtomicUsize::new(0));
     let latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests)));
     let status_counts = Arc::new(Mutex::new(BTreeMap::<u16, usize>::new()));
+    let error_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
     let target_request_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
     let failed_requests = Arc::new(AtomicUsize::new(0));
     let completed_requests = Arc::new(AtomicUsize::new(0));
@@ -171,6 +175,7 @@ async fn run_http_load_probe_against_urls(
         let next_request = Arc::clone(&next_request);
         let latencies_ms = Arc::clone(&latencies_ms);
         let status_counts = Arc::clone(&status_counts);
+        let error_counts = Arc::clone(&error_counts);
         let target_request_counts = Arc::clone(&target_request_counts);
         let failed_requests = Arc::clone(&failed_requests);
         let completed_requests = Arc::clone(&completed_requests);
@@ -200,27 +205,33 @@ async fn run_http_load_probe_against_urls(
                         let status = response.status().as_u16();
                         let body_result = match response_mode {
                             HttpLoadProbeResponseMode::HeadersOnly => Ok(()),
-                            HttpLoadProbeResponseMode::FullBody => {
-                                response.bytes().await.map(|_| ()).map_err(|_| ())
-                            }
+                            HttpLoadProbeResponseMode::FullBody => response
+                                .bytes()
+                                .await
+                                .map(|_| ())
+                                .map_err(|err| classify_reqwest_error(&err)),
                         };
-                        if body_result.is_ok() {
+                        if let Err(error_kind) = body_result {
+                            failed_requests.fetch_add(1, Ordering::AcqRel);
+                            let mut counts = error_counts.lock().await;
+                            *counts.entry(error_kind).or_insert(0) += 1;
+                        } else {
                             let mut counts = status_counts.lock().await;
                             *counts.entry(status).or_insert(0) += 1;
                             drop(counts);
                             let mut target_counts = target_request_counts.lock().await;
                             *target_counts.entry(url).or_insert(0) += 1;
-                        } else {
-                            failed_requests.fetch_add(1, Ordering::AcqRel);
                         }
                         let latency_ms = started_at.elapsed().as_millis() as u64;
                         latencies_ms.lock().await.push(latency_ms);
                         completed_requests.fetch_add(1, Ordering::AcqRel);
                     }
-                    Err(_) => {
+                    Err(err) => {
                         let latency_ms = started_at.elapsed().as_millis() as u64;
                         latencies_ms.lock().await.push(latency_ms);
                         failed_requests.fetch_add(1, Ordering::AcqRel);
+                        let mut counts = error_counts.lock().await;
+                        *counts.entry(classify_reqwest_error(&err)).or_insert(0) += 1;
                         completed_requests.fetch_add(1, Ordering::AcqRel);
                     }
                 }
@@ -233,6 +244,7 @@ async fn run_http_load_probe_against_urls(
     }
 
     let status_counts = status_counts.lock().await.clone();
+    let error_counts = error_counts.lock().await.clone();
     let target_request_counts = target_request_counts.lock().await.clone();
     let mut latencies = latencies_ms.lock().await.clone();
     latencies.sort_unstable();
@@ -262,9 +274,32 @@ async fn run_http_load_probe_against_urls(
         mean_ms,
         runtime: runtime_sampler.snapshot(),
         status_counts,
+        error_counts,
     })
 }
 
+fn classify_reqwest_error(err: &reqwest::Error) -> String {
+    if err.is_timeout() {
+        return "timeout".to_string();
+    }
+    if err.is_connect() {
+        return "connect".to_string();
+    }
+    if err.is_body() {
+        return "body".to_string();
+    }
+    if err.is_request() {
+        return "request".to_string();
+    }
+    if err.is_decode() {
+        return "decode".to_string();
+    }
+    if err.is_redirect() {
+        return "redirect".to_string();
+    }
+    "other".to_string()
+}
+
 fn build_headers(headers: &BTreeMap<String, String>) -> Result<HeaderMap, String> {
     let mut result = HeaderMap::new();
     for (name, value) in headers {