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::sync::{Arc, Mutex, PoisonError};
7use std::time::{Duration, Instant};
8
9use af_context::RequestId;
10use axum::extract::Request;
11use axum::http::{HeaderName, HeaderValue, StatusCode};
12use axum::middleware::Next;
13use axum::response::{IntoResponse, Response};
14
15static REQ_SEQ: AtomicU64 = AtomicU64::new(1);
16
17/// The API latency budget (p95 ≤ 1s is the product line). Requests over this
18/// are logged at `warn` so a regression surfaces without a profiler.
19pub const SLO_BUDGET_MS: u128 = 1_000;
20
21const REQUEST_ID_HEADER: &str = "x-request-id";
22const RESPONSE_TIME_HEADER: &str = "x-response-time-ms";
23
24/// Mint (or honor a client-supplied) request id, stash it in extensions for
25/// handlers, and echo it on the response.
26pub async fn request_id(mut req: Request, next: Next) -> Response {
27    let id = match req
28        .headers()
29        .get(REQUEST_ID_HEADER)
30        .and_then(|v| v.to_str().ok())
31        .and_then(|value| RequestId::parse(value).ok())
32        .or_else(mint_request_id)
33    {
34        Some(id) => id,
35        None => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
36    };
37
38    req.extensions_mut().insert(id.clone());
39    let mut resp = next.run(req).await;
40
41    if let Ok(value) = HeaderValue::from_str(id.as_str()) {
42        resp.headers_mut()
43            .insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
44    }
45    resp
46}
47
48/// Time each request, attach `x-response-time-ms`, and warn past the SLO budget.
49pub async fn latency_slo(req: Request, next: Next) -> Response {
50    let method = req.method().clone();
51    let path = req.uri().path().to_string();
52    let request_id = req
53        .extensions()
54        .get::<RequestId>()
55        .map(ToString::to_string)
56        .unwrap_or_default();
57
58    let started = Instant::now();
59    let mut resp = next.run(req).await;
60    let elapsed = started.elapsed();
61    let ms = elapsed.as_millis();
62
63    if let Ok(value) = HeaderValue::from_str(&ms.to_string()) {
64        resp.headers_mut()
65            .insert(HeaderName::from_static(RESPONSE_TIME_HEADER), value);
66    }
67
68    if ms > SLO_BUDGET_MS {
69        tracing::warn!(
70            target: "web.slo",
71            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
72            "api_slo_violation"
73        );
74    } else {
75        tracing::debug!(
76            target: "web.request",
77            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
78            "request"
79        );
80    }
81    resp
82}
83
84fn mint_request_id() -> Option<RequestId> {
85    RequestId::parse(format!(
86        "req_{:016x}",
87        REQ_SEQ.fetch_add(1, Ordering::Relaxed)
88    ))
89    .ok()
90}
91
92/// Shared fixed-window request limiter used by the host middleware stack.
93#[derive(Clone)]
94pub struct RateLimiter {
95    state: Arc<Mutex<RateWindow>>,
96    limit: u64,
97    window: Duration,
98}
99
100struct RateWindow {
101    started: Instant,
102    count: u64,
103}
104
105impl RateLimiter {
106    /// Creates a limiter allowing `limit` requests per `window`.
107    pub fn new(limit: u64, window: Duration) -> Self {
108        Self {
109            state: Arc::new(Mutex::new(RateWindow {
110                started: Instant::now(),
111                count: 0,
112            })),
113            limit: limit.max(1),
114            window,
115        }
116    }
117}
118
119/// Rejects requests beyond the configured fixed-window rate.
120pub async fn rate_limit(
121    axum::extract::State(limiter): axum::extract::State<RateLimiter>,
122    request: Request,
123    next: Next,
124) -> Response {
125    if matches!(
126        request.uri().path(),
127        "/health" | "/live" | "/ready" | "/metrics"
128    ) {
129        return next.run(request).await;
130    }
131    let allowed = {
132        let mut state = limiter.state.lock().unwrap_or_else(PoisonError::into_inner);
133        if state.started.elapsed() >= limiter.window {
134            state.started = Instant::now();
135            state.count = 0;
136        }
137        state.count += 1;
138        state.count <= limiter.limit
139    };
140    if allowed {
141        next.run(request).await
142    } else {
143        (StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded").into_response()
144    }
145}
af_web/middleware.rs

af_web/
middleware.rs
