af_web/

lib.rs

//! Reusable HTTP/gRPC host infrastructure for Agent Factory services.
//!
//! - [`ApiError`] — JSON error envelope carrying the request id.
//! - [`middleware::request_id`] / [`middleware::latency_slo`] — request-id
//!   propagation and the API latency SLO (warn past [`middleware::SLO_BUDGET_MS`]).
//! - [`with_standard_middleware`] — wrap any `Router` with the standard stack.
//! - [`health`] — a ready-made health handler.
//! - [`metrics`] — a minimal Prometheus-compatible liveness metric.
//!
//! A service composes its product routes and wraps them:
//!
//! ```ignore
//! let app = with_standard_middleware(
//!     Router::new().route("/health", get(health)).merge(product_routes),
//! );
//! ```

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

pub mod error;
pub mod middleware;

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use af_context::{PlatformContextProvider, RequestContext, RequestMetadata};
use axum::extract::{DefaultBodyLimit, State};
use axum::http::HeaderValue;
use axum::routing::get;
use axum::{Json, Router};
use serde_json::{json, Value};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tonic::{Request, Status};
use tower::limit::ConcurrencyLimitLayer;
use tower_http::cors::{AllowOrigin, CorsLayer};
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::TraceLayer;

pub use af_context::RequestId;
pub use error::ApiError;
pub use middleware::{latency_slo, rate_limit, request_id, RateLimiter, SLO_BUDGET_MS};

/// Production edge limits shared by reference hosts.
#[derive(Clone)]
pub struct EdgeConfig {
    /// Explicit allowed origins. Empty denies cross-origin requests.
    pub allowed_origins: Vec<HeaderValue>,
    /// Maximum request body bytes.
    pub max_body_bytes: usize,
    /// Unary request timeout.
    pub request_timeout: Duration,
    /// Maximum in-flight requests.
    pub max_concurrency: usize,
    /// Maximum requests in each rate window.
    pub requests_per_window: u64,
    /// Rate-limit window.
    pub rate_window: Duration,
}

impl Default for EdgeConfig {
    fn default() -> Self {
        Self {
            allowed_origins: Vec::new(),
            max_body_bytes: 1024 * 1024,
            request_timeout: Duration::from_secs(30),
            max_concurrency: 128,
            requests_per_window: 120,
            rate_window: Duration::from_secs(1),
        }
    }
}

/// Mutable readiness signal; liveness remains independent.
#[derive(Clone, Default)]
pub struct Readiness(Arc<AtomicBool>);

impl Readiness {
    /// Changes whether the host can accept traffic.
    pub fn set(&self, ready: bool) {
        self.0.store(ready, Ordering::Release);
    }

    /// Returns the current readiness state.
    pub fn is_ready(&self) -> bool {
        self.0.load(Ordering::Acquire)
    }
}

/// Standard health endpoint payload.
pub async fn health() -> Json<Value> {
    Json(json!({ "status": "ok" }))
}

/// Liveness endpoint; success means the process event loop is responsive.
pub async fn liveness() -> Json<Value> {
    health().await
}

/// Readiness endpoint backed by the host's dependency state.
pub async fn readiness(State(readiness): State<Readiness>) -> impl axum::response::IntoResponse {
    let status = if readiness.is_ready() {
        axum::http::StatusCode::OK
    } else {
        axum::http::StatusCode::SERVICE_UNAVAILABLE
    };
    (
        status,
        Json(json!({ "status": if readiness.is_ready() { "ready" } else { "not_ready" } })),
    )
}

/// Baseline metric shared by reference hosts. Products merge their own metric
/// families into the same endpoint when they need richer observability.
pub async fn metrics() -> ([(&'static str, &'static str); 1], &'static str) {
    (
        [("content-type", "text/plain; version=0.0.4")],
        "agent_factory_up 1\n",
    )
}

/// Wrap a router with the standard middleware stack: latency SLO (outermost,
/// so it times everything) over request-id (so the id exists for SLO logs).
pub fn with_standard_middleware<S>(router: Router<S>) -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    router
        .layer(axum::middleware::from_fn(latency_slo))
        .layer(axum::middleware::from_fn(request_id))
}

/// Applies bounded production edge controls. CORS is deny-by-default.
pub fn with_edge_middleware(router: Router, config: EdgeConfig) -> Router {
    let cors = if config.allowed_origins.is_empty() {
        CorsLayer::new()
    } else {
        CorsLayer::new().allow_origin(AllowOrigin::list(config.allowed_origins))
    };
    let limiter = RateLimiter::new(config.requests_per_window, config.rate_window);
    with_standard_middleware(router)
        .layer(TraceLayer::new_for_http())
        .layer(ConcurrencyLimitLayer::new(config.max_concurrency.max(1)))
        .layer(TimeoutLayer::with_status_code(
            axum::http::StatusCode::REQUEST_TIMEOUT,
            config.request_timeout,
        ))
        .layer(DefaultBodyLimit::max(config.max_body_bytes))
        .layer(axum::middleware::from_fn_with_state(limiter, rate_limit))
        .layer(cors)
}

/// A `Router` preseeded with `GET /health`. Services merge their routes onto it.
pub fn base_router<S>() -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    Router::new()
        .route("/health", get(health))
        .route("/metrics", get(metrics))
}

/// Router exposing separate process liveness and dependency readiness.
pub fn operations_router(state: Readiness) -> Router {
    Router::new()
        .route("/health", get(health))
        .route("/live", get(liveness))
        .route("/ready", get(readiness))
        .route("/metrics", get(metrics))
        .with_state(state)
}

/// Resolves verified caller context from tonic metadata for every gRPC service.
pub async fn resolve_tonic_context<T>(
    provider: &dyn PlatformContextProvider,
    request: &Request<T>,
) -> Result<RequestContext, Status> {
    let metadata = request
        .metadata()
        .iter()
        .filter_map(|entry| match entry {
            tonic::metadata::KeyAndValueRef::Ascii(key, value) => value
                .to_str()
                .ok()
                .map(|value| (key.as_str().to_string(), value.to_string())),
            tonic::metadata::KeyAndValueRef::Binary(_, _) => None,
        })
        .collect::<BTreeMap<_, _>>();
    provider
        .resolve(&RequestMetadata(metadata))
        .await
        .map_err(|error| Status::unauthenticated(error.to_string()))
}

/// Waits for SIGINT/SIGTERM and cancels all host tasks.
pub async fn shutdown_signal(cancel: CancellationToken) {
    #[cfg(unix)]
    {
        if let Ok(mut terminate) =
            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
        {
            tokio::select! {
                _ = tokio::signal::ctrl_c() => {}
                _ = terminate.recv() => {}
            }
        } else {
            let _ = tokio::signal::ctrl_c().await;
        }
    }
    #[cfg(not(unix))]
    let _ = tokio::signal::ctrl_c().await;
    cancel.cancel();
}

/// Joins background tasks within a bounded shutdown grace period.
pub async fn join_tasks(tasks: Vec<JoinHandle<()>>, grace: Duration) -> bool {
    let aborts: Vec<_> = tasks.iter().map(JoinHandle::abort_handle).collect();
    let joined = tokio::time::timeout(grace, futures_join(tasks))
        .await
        .unwrap_or(false);
    if !joined {
        for task in aborts {
            task.abort();
        }
    }
    joined
}

/// Fails readiness and cancels sibling workers if any host task exits early.
/// Returns true only for a requested shutdown with all workers drained.
pub async fn supervise_tasks(
    mut tasks: Vec<JoinHandle<()>>,
    readiness: Readiness,
    cancel: CancellationToken,
    grace: Duration,
) -> bool {
    use std::future::Future;
    let requested = tokio::select! {
        biased;
        _ = cancel.cancelled() => true,
        (index, outcome) = std::future::poll_fn(|cx| {
            for (index, task) in tasks.iter_mut().enumerate() {
                if let std::task::Poll::Ready(outcome) = std::pin::Pin::new(task).poll(cx) {
                    return std::task::Poll::Ready((index, outcome));
                }
            }
            std::task::Poll::Pending
        }) => {
            drop(tasks.swap_remove(index));
            tracing::error!(exit_reason = if outcome.is_err() { "task_panicked" } else { "task_exited" }, "host worker stopped unexpectedly");
            false
        }
    };
    readiness.set(false);
    cancel.cancel();
    let joined = join_tasks(tasks, grace).await;
    tracing::info!(
        exit_reason = if requested && joined {
            "shutdown"
        } else {
            "worker_failure"
        },
        "host stopped"
    );
    requested && joined
}

async fn futures_join(tasks: Vec<JoinHandle<()>>) -> bool {
    for task in tasks {
        if task.await.is_err() {
            return false;
        }
    }
    true
}