af-web 0.5.0

//! Cross-cutting middleware. Ports the reusable concerns of `agent_core/api/`:
//! request-id propagation ([`request_id`]) and the API latency SLO
//! ([`latency_slo`]).

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, Instant};

use af_context::RequestId;
use axum::extract::Request;
use axum::http::{HeaderName, HeaderValue, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};

static REQ_SEQ: AtomicU64 = AtomicU64::new(1);

/// The API latency budget (p95 ≤ 1s is the product line). Requests over this
/// are logged at `warn` so a regression surfaces without a profiler.
pub const SLO_BUDGET_MS: u128 = 1_000;

const REQUEST_ID_HEADER: &str = "x-request-id";
const RESPONSE_TIME_HEADER: &str = "x-response-time-ms";

/// Mint (or honor a client-supplied) request id, stash it in extensions for
/// handlers, and echo it on the response.
pub async fn request_id(mut req: Request, next: Next) -> Response {
    let id = match req
        .headers()
        .get(REQUEST_ID_HEADER)
        .and_then(|v| v.to_str().ok())
        .and_then(|value| RequestId::parse(value).ok())
        .or_else(mint_request_id)
    {
        Some(id) => id,
        None => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
    };

    req.extensions_mut().insert(id.clone());
    let mut resp = next.run(req).await;

    if let Ok(value) = HeaderValue::from_str(id.as_str()) {
        resp.headers_mut()
            .insert(HeaderName::from_static(REQUEST_ID_HEADER), value);
    }
    resp
}

/// Time each request, attach `x-response-time-ms`, and warn past the SLO budget.
pub async fn latency_slo(req: Request, next: Next) -> Response {
    let method = req.method().clone();
    let path = req.uri().path().to_string();
    let request_id = req
        .extensions()
        .get::<RequestId>()
        .map(ToString::to_string)
        .unwrap_or_default();

    let started = Instant::now();
    let mut resp = next.run(req).await;
    let elapsed = started.elapsed();
    let ms = elapsed.as_millis();

    if let Ok(value) = HeaderValue::from_str(&ms.to_string()) {
        resp.headers_mut()
            .insert(HeaderName::from_static(RESPONSE_TIME_HEADER), value);
    }

    if ms > SLO_BUDGET_MS {
        tracing::warn!(
            target: "web.slo",
            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
            "api_slo_violation"
        );
    } else {
        tracing::debug!(
            target: "web.request",
            %method, path, request_id, elapsed_ms = ms, status = resp.status().as_u16(),
            "request"
        );
    }
    resp
}

fn mint_request_id() -> Option<RequestId> {
    RequestId::parse(format!(
        "req_{:016x}",
        REQ_SEQ.fetch_add(1, Ordering::Relaxed)
    ))
    .ok()
}

/// Shared fixed-window request limiter used by the host middleware stack.
#[derive(Clone)]
pub struct RateLimiter {
    state: Arc<Mutex<RateWindow>>,
    limit: u64,
    window: Duration,
}

struct RateWindow {
    started: Instant,
    count: u64,
}

impl RateLimiter {
    /// Creates a limiter allowing `limit` requests per `window`.
    pub fn new(limit: u64, window: Duration) -> Self {
        Self {
            state: Arc::new(Mutex::new(RateWindow {
                started: Instant::now(),
                count: 0,
            })),
            limit: limit.max(1),
            window,
        }
    }
}

/// Rejects requests beyond the configured fixed-window rate.
pub async fn rate_limit(
    axum::extract::State(limiter): axum::extract::State<RateLimiter>,
    request: Request,
    next: Next,
) -> Response {
    if matches!(
        request.uri().path(),
        "/health" | "/live" | "/ready" | "/metrics"
    ) {
        return next.run(request).await;
    }
    let allowed = {
        let mut state = limiter.state.lock().unwrap_or_else(PoisonError::into_inner);
        if state.started.elapsed() >= limiter.window {
            state.started = Instant::now();
            state.count = 0;
        }
        state.count += 1;
        state.count <= limiter.limit
    };
    if allowed {
        next.run(request).await
    } else {
        (StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded").into_response()
    }
}