Skip to main content

af_web/
middleware.rs

1//! Cross-cutting middleware. Ports the reusable concerns of `agent_core/api/`:
2//! request-id propagation ([`request_id`]) and the API latency SLO
3//! ([`latency_slo`]).
4
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Instant;
7
8use axum::extract::Request;
9use axum::http::{HeaderName, HeaderValue};
10use axum::middleware::Next;
11use axum::response::Response;
12
13static REQ_SEQ: AtomicU64 = AtomicU64::new(1);
14
15/// The API latency budget (p95 ≤ 1s is the product line). Requests over this
16/// are logged at `warn` so a regression surfaces without a profiler.
17pub const SLO_BUDGET_MS: u128 = 1_000;
18
19const REQUEST_ID_HEADER: &str = "x-request-id";
20const RESPONSE_TIME_HEADER: &str = "x-response-time-ms";
21
22/// A request id, available to handlers via `Extension<RequestId>`.
23#[derive(Debug, Clone)]
24pub struct RequestId(pub String);
25
26/// Mint (or honor a client-supplied) request id, stash it in extensions for
27/// handlers, and echo it on the response.
28pub async fn request_id(mut req: Request, next: Next) -> Response {
29    let id = req
30        .headers()
31        .get(REQUEST_ID_HEADER)
32        .and_then(|v| v.to_str().ok())
33        .map(|s| s.to_string())
34        .unwrap_or_else(|| format!("req_{:016x}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)));
35
36    req.extensions_mut().insert(RequestId(id.clone()));
37    let mut resp = next.run(req).await;
38
39    if let Ok(value) = HeaderValue::from_str(&id) {
40        resp.headers_mut()
41            .insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
42    }
43    resp
44}
45
46/// Time each request, attach `x-response-time-ms`, and warn past the SLO budget.
47pub async fn latency_slo(req: Request, next: Next) -> Response {
48    let method = req.method().clone();
49    let path = req.uri().path().to_string();
50    let request_id = req
51        .extensions()
52        .get::<RequestId>()
53        .map(|r| r.0.clone())
54        .unwrap_or_default();
55
56    let started = Instant::now();
57    let mut resp = next.run(req).await;
58    let elapsed = started.elapsed();
59    let ms = elapsed.as_millis();
60
61    if let Ok(value) = HeaderValue::from_str(&ms.to_string()) {
62        resp.headers_mut()
63            .insert(HeaderName::from_static(RESPONSE_TIME_HEADER), value);
64    }
65
66    if ms > SLO_BUDGET_MS {
67        tracing::warn!(
68            target: "web.slo",
69            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
70            "api_slo_violation"
71        );
72    } else {
73        tracing::debug!(
74            target: "web.request",
75            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
76            "request"
77        );
78    }
79    resp
80}