systemprompt-api 0.70.0

//! Liveness and readiness probes for balancers and orchestrators.
//!
//! `/livez` answers as soon as the port is bound; `/readyz` is the admission
//! signal and answers 503 before boot completes, after the drain signal, while
//! the in-flight ceiling is saturated, and whenever the database probe fails.
//! A ready answer carries a `config` block: per declared config plane, the
//! authority, declaration hash, remaining drift and outcome of the last
//! reconcile pass. Drift never makes a replica unready.
//! It also carries `fingerprint`, this replica's configuration fingerprint,
//! so replicas running different configuration are visible.
//! A node whose database a newer release has migrated stays ready and says so
//! with `"schema": "ahead"` and the migrations it does not know, unless one of
//! them is `@compat: breaking`: then it answers 503 with `"schema":
//! "breaking"`. The ahead set is the one the node's schema watch last read,
//! not the boot snapshot.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use std::sync::Arc;

use axum::Json;
use serde_json::json;
use systemprompt_runtime::AppContext;
use systemprompt_scheduler::repository::ConfigReconcileRepository;

use super::health::{HEALTH_CHECK_QUERY, HEALTH_PROBE_TIMEOUT};
use crate::services::middleware::LoadShed;

pub(crate) async fn handle_livez(
    axum::extract::State(ctx): axum::extract::State<AppContext>,
) -> impl axum::response::IntoResponse {
    Json(json!({
        "status": "alive",
        "instance": ctx.config().instance_id,
        "version": env!("CARGO_PKG_VERSION")
    }))
}

pub(crate) async fn handle_readyz(
    axum::extract::State(ctx): axum::extract::State<AppContext>,
    shed: Option<axum::Extension<Arc<LoadShed>>>,
    planes: Option<axum::Extension<Arc<ConfigReconcileRepository>>>,
) -> impl axum::response::IntoResponse {
    use axum::http::StatusCode;
    use systemprompt_database::DatabaseProvider;

    let instance = ctx.config().instance_id.clone();
    let fingerprint = super::config_fingerprint::current();
    let respond = |status: StatusCode, mut body: serde_json::Value, config: serde_json::Value| {
        if let Some(fields) = body.as_object_mut() {
            fields.insert("instance".to_owned(), json!(instance));
            fields.insert("version".to_owned(), json!(env!("CARGO_PKG_VERSION")));
            fields.insert("config".to_owned(), config);
        }
        (status, Json(body))
    };
    let bare = json!({ "fingerprint": fingerprint });

    if !super::readiness::is_ready() {
        return respond(
            StatusCode::SERVICE_UNAVAILABLE,
            json!({ "status": "draining" }),
            bare,
        );
    }

    if let Some(axum::Extension(shed)) = shed
        && shed.saturated()
    {
        let body = json!({
            "status": "saturated",
            "in_flight": shed.in_flight(),
            "limit": shed.limit(),
        });
        return respond(StatusCode::SERVICE_UNAVAILABLE, body, bare);
    }

    let probe = ctx.db_pool().fetch_optional(&HEALTH_CHECK_QUERY, &[]);
    let db_ready = matches!(
        tokio::time::timeout(HEALTH_PROBE_TIMEOUT, probe).await,
        Ok(Ok(_))
    );
    if !db_ready {
        let body = json!({ "status": "unready", "database": "unreachable" });
        return respond(StatusCode::SERVICE_UNAVAILABLE, body, bare);
    }

    let config = config_block(fingerprint, planes.as_deref().map(AsRef::as_ref)).await;
    let schema = ctx.schema_state();
    let ahead = schema.ahead();
    if schema.breaks_this_release() {
        let body = json!({ "status": "unready", "schema": "breaking", "schema_ahead": ahead });
        return respond(StatusCode::SERVICE_UNAVAILABLE, body, config);
    }
    let body = if ahead.is_empty() {
        json!({ "status": "ready" })
    } else {
        json!({ "status": "ready", "schema": "ahead", "schema_ahead": ahead })
    };
    respond(StatusCode::OK, body, config)
}

// JSON: the `/readyz` body is assembled as an open JSON object.
async fn config_block(
    fingerprint: Option<systemprompt_config::ConfigFingerprint>,
    repository: Option<&ConfigReconcileRepository>,
) -> serde_json::Value {
    let Some(repository) = repository else {
        return json!({ "fingerprint": fingerprint, "planes": [] });
    };
    match tokio::time::timeout(HEALTH_PROBE_TIMEOUT, repository.list_states()).await {
        Ok(Ok(rows)) => json!({
            "fingerprint": fingerprint,
            "planes": rows.iter().map(|row| json!({
                "plane": row.plane,
                "authority": row.authority,
                "hash": row.declared_hash,
                "drift_rows": row.drift_rows,
                "outcome": row.outcome,
                "reconciled_at": row.reconciled_at,
            })).collect::<Vec<_>>()
        }),
        Ok(Err(error)) => {
            tracing::warn!(error = %error, "readyz could not read config_plane_state");
            json!({ "fingerprint": fingerprint, "planes": [], "error": "unavailable" })
        },
        Err(_) => json!({ "fingerprint": fingerprint, "planes": [], "error": "timeout" }),
    }
}