systemprompt-api 0.63.1

//! Health probe and system-statistics endpoints.
//!
//! [`handle_health`] is the lightweight liveness probe (a single `SELECT 1`).
//! `get_system_stats` gathers database, disk, and audit-log statistics for
//! the detailed health surface, with platform-specific process-memory and
//! disk-usage helpers.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use std::time::Duration;

use axum::Json;
use serde_json::json;
use systemprompt_database::{DatabaseQuery, JsonRow};
use systemprompt_runtime::AppContext;

pub(super) const HEALTH_CHECK_QUERY: DatabaseQuery = DatabaseQuery::new("SELECT 1");

const DB_SIZE_QUERY: DatabaseQuery = DatabaseQuery::new(
    "SELECT pg_database_size(current_database()) as size_bytes, current_database() as db_name",
);

const TABLE_SIZES_QUERY: DatabaseQuery = DatabaseQuery::new(
    "SELECT relname as table_name, pg_total_relation_size(relid) as total_bytes, n_live_tup as \
     row_estimate FROM pg_stat_user_tables ORDER BY pg_total_relation_size(relid) DESC LIMIT 15",
);

const TABLE_COUNT_QUERY: DatabaseQuery =
    DatabaseQuery::new("SELECT COUNT(*) as count FROM pg_stat_user_tables");

const AUDIT_LOG_QUERY: DatabaseQuery = DatabaseQuery::new(
    "SELECT COUNT(*) as row_count, pg_total_relation_size('audit_log') as size_bytes, \
     MIN(created_at) as oldest, MAX(created_at) as newest FROM audit_log",
);

pub use super::health_stats::{
    AuditLogStats, DatabaseStats, DiskUsage, ProcessMemory, SystemStats, TableStats,
};

#[cfg(target_os = "linux")]
pub fn parse_proc_status_kb(content: &str, key: &str) -> Option<u64> {
    content
        .lines()
        .find(|line| line.starts_with(key))
        .and_then(|line| {
            line.split_whitespace()
                .nth(1)
                .and_then(|v| v.parse::<u64>().ok())
        })
}

#[cfg(target_os = "linux")]
pub(super) fn get_process_memory() -> Option<ProcessMemory> {
    let content = std::fs::read_to_string("/proc/self/status").ok()?;

    let rss_kb = parse_proc_status_kb(&content, "VmRSS:");
    let virt_kb = parse_proc_status_kb(&content, "VmSize:");
    let peak_kb = parse_proc_status_kb(&content, "VmPeak:");

    Some(ProcessMemory {
        rss: rss_kb.map(|kb| kb / 1024),
        virtual_size: virt_kb.map(|kb| kb / 1024),
        peak: peak_kb.map(|kb| kb / 1024),
    })
}

#[cfg(not(target_os = "linux"))]
pub(super) const fn get_process_memory() -> Option<ProcessMemory> {
    None
}

pub fn human_bytes(bytes: i64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut size: f64 = bytes as f64;
    let mut idx = 0;
    while size >= 1024.0 && idx < UNITS.len() - 1 {
        size /= 1024.0;
        idx += 1;
    }
    format!("{size:.1} {}", UNITS[idx])
}

// Why: statvfs field widths vary by target; generic widening avoids
// useless_conversion on u64 targets.
fn widen(value: impl Into<u64>) -> u64 {
    value.into()
}

fn get_disk_usage() -> Option<DiskUsage> {
    let stat = nix::sys::statvfs::statvfs(".").ok()?;

    let block_size = widen(stat.fragment_size());
    let total = widen(stat.blocks()).saturating_mul(block_size);
    let available = widen(stat.blocks_available()).saturating_mul(block_size);
    let free = widen(stat.blocks_free()).saturating_mul(block_size);
    let used = total.saturating_sub(free);

    let usage_pct = if total > 0 {
        (used as f64 / total as f64) * 100.0
    } else {
        0.0
    };

    Some(DiskUsage {
        total: human_bytes(total as i64),
        used: human_bytes(used as i64),
        available: human_bytes(available as i64),
        usage_percent: (usage_pct * 10.0).round() / 10.0,
    })
}

pub(super) async fn get_system_stats(
    db: &dyn systemprompt_database::DatabaseProvider,
) -> Option<SystemStats> {
    let db_size_fut = db.fetch_one(&DB_SIZE_QUERY, &[]);
    let table_sizes_fut = db.fetch_all(&TABLE_SIZES_QUERY, &[]);
    let table_count_fut = db.fetch_one(&TABLE_COUNT_QUERY, &[]);
    let audit_fut = db.fetch_optional(&AUDIT_LOG_QUERY, &[]);
    let disk = get_disk_usage();

    let (db_size, table_sizes, table_count, audit) =
        tokio::join!(db_size_fut, table_sizes_fut, table_count_fut, audit_fut);

    let database =
        if let (Ok(size_row), Ok(tables), Ok(count_row)) = (&db_size, &table_sizes, &table_count) {
            Some(database_stats(size_row, tables, count_row))
        } else {
            None
        };

    let logs = audit
        .inspect_err(|error| tracing::warn!(%error, "Audit log stats query failed"))
        .ok()
        .flatten()
        .map(|row| audit_log_stats(&row));

    Some(SystemStats {
        database,
        disk,
        logs,
    })
}

pub fn database_stats(
    size_row: &JsonRow,
    tables: &[JsonRow],
    count_row: &JsonRow,
) -> DatabaseStats {
    let size_bytes = size_row
        .get("size_bytes")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);
    let db_name = size_row
        .get("db_name")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("unknown");
    let tbl_count = count_row
        .get("count")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);

    DatabaseStats {
        name: db_name.to_owned(),
        total_size: human_bytes(size_bytes),
        total_size_bytes: size_bytes,
        table_count: tbl_count,
        top_tables: tables.iter().map(table_stats).collect(),
    }
}

pub fn table_stats(row: &JsonRow) -> TableStats {
    let name = row
        .get("table_name")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("?");
    let total = row
        .get("total_bytes")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);
    let rows = row
        .get("row_estimate")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);
    TableStats {
        table_name: name.to_owned(),
        total_size: human_bytes(total),
        total_size_bytes: total,
        row_estimate: rows,
    }
}

pub fn audit_log_stats(row: &JsonRow) -> AuditLogStats {
    let row_count = row
        .get("row_count")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);
    let size_bytes = row
        .get("size_bytes")
        .and_then(serde_json::Value::as_i64)
        .unwrap_or(0);
    AuditLogStats {
        audit_rows: row_count,
        audit_size: human_bytes(size_bytes),
        audit_size_bytes: size_bytes,
        oldest: row.get("oldest").cloned(),
        newest: row.get("newest").cloned(),
    }
}

pub(super) const HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(2);

pub async fn handle_health(
    axum::extract::State(ctx): axum::extract::State<AppContext>,
) -> impl axum::response::IntoResponse {
    use axum::http::StatusCode;
    use systemprompt_database::DatabaseProvider;

    use super::scheduler_health;

    let probe = ctx.db_pool().fetch_optional(&HEALTH_CHECK_QUERY, &[]);
    let db_healthy = matches!(
        tokio::time::timeout(HEALTH_PROBE_TIMEOUT, probe).await,
        Ok(Ok(_))
    );

    let degraded_jobs = scheduler_health::degraded();
    let relay_status = ctx
        .event_bridge()
        .get()
        .map_or(systemprompt_events::RelayStatus::NotStarted, |handle| {
            handle.status()
        });
    let relay_listening = relay_status.is_listening();

    let (status, http_status) = if !db_healthy {
        ("unhealthy", StatusCode::SERVICE_UNAVAILABLE)
    } else if degraded_jobs.is_empty() && relay_listening {
        ("healthy", StatusCode::OK)
    } else {
        ("degraded", StatusCode::OK)
    };

    let mut body = json!({ "status": status, "version": env!("CARGO_PKG_VERSION") });
    if !degraded_jobs.is_empty() {
        body["scheduler"] = json!({ "degraded_jobs": degraded_jobs });
    }
    if !relay_listening {
        body["events"] = json!({ "relay": relay_label(relay_status) });
    }

    (http_status, Json(body))
}

const fn relay_label(status: systemprompt_events::RelayStatus) -> &'static str {
    match status {
        systemprompt_events::RelayStatus::NotStarted => "not_started",
        systemprompt_events::RelayStatus::Listening => "listening",
        systemprompt_events::RelayStatus::Reconnecting => "reconnecting",
        systemprompt_events::RelayStatus::Stopped => "stopped",
    }
}