Genarrative/server-rs/crates/api-server/src/health.rs

use axum::{
    Json,
    extract::{Extension, State},
    http::StatusCode,
    response::{IntoResponse, Response},
};
use serde_json::{Value, json};

use crate::{
    api_response::json_success_body, http_error::AppError, request_context::RequestContext,
    state::AppState,
};
use spacetime_client::SpacetimeClientHealthSnapshot;

pub async fn health_check(Extension(request_context): Extension<RequestContext>) -> Json<Value> {
    json_success_body(
        Some(&request_context),
        json!({
            "ok": true,
            "service": "genarrative-api-server",
        }),
    )
}

pub async fn readiness_check(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
) -> Response {
    if !state.is_ready() {
        return AppError::from_status(StatusCode::SERVICE_UNAVAILABLE)
            .with_message("api-server 正在退出，不再接收新流量")
            .with_details(json!({
                "reason": "api_server_draining",
                "ready": false,
            }))
            .into_response_with_context(Some(&request_context));
    }

    let spacetime_health = state.spacetime_health_check().await;
    if spacetime_health.ok {
        return json_success_body(
            Some(&request_context),
            json!({
                "ok": true,
                "ready": true,
                "service": "genarrative-api-server",
                "spacetime": spacetime_health_to_json(&spacetime_health),
            }),
        )
        .into_response();
    }

    AppError::from_status(StatusCode::SERVICE_UNAVAILABLE)
        .with_message("SpacetimeDB 连接健康检查失败，api-server 暂不接收新流量")
        .with_details(json!({
            "reason": "spacetime_unhealthy",
            "ready": false,
            "spacetime": spacetime_health_to_json(&spacetime_health),
        }))
        .into_response_with_context(Some(&request_context))
}

fn spacetime_health_to_json(snapshot: &SpacetimeClientHealthSnapshot) -> Value {
    json!({
        "ok": snapshot.ok,
        "stage": snapshot.stage.as_str(),
        "checkedAtMicros": snapshot.checked_at_micros,
        "elapsedMs": snapshot.elapsed_ms,
        "timeoutMs": snapshot.timeout_ms,
        "error": snapshot.error,
        "lastSuccessAtMicros": snapshot.last_success_at_micros,
        "lastError": snapshot.last_error,
    })
}