path: root/makima/src/server/mod.rs

//! Web server module for the makima audio API.

pub mod auth;
pub mod handlers;
pub mod messages;
pub mod openapi;
pub mod state;

use axum::{
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use serde::Serialize;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi;

use crate::server::handlers::{api_keys, chat, file_ws, files, listen, mesh, mesh_chat, mesh_daemon, mesh_merge, mesh_ws, users, versions};
use crate::server::openapi::ApiDoc;
use crate::server::state::SharedState;

#[derive(Serialize)]
struct HealthResponse {
    status: &'static str,
    version: &'static str,
}

/// Health check endpoint for load balancers and orchestrators.
async fn health_check() -> impl IntoResponse {
    (
        StatusCode::OK,
        Json(HealthResponse {
            status: "healthy",
            version: env!("CARGO_PKG_VERSION"),
        }),
    )
}

/// Create the axum Router with all routes configured.
pub fn make_router(state: SharedState) -> Router {
    // API v1 routes
    let api_v1 = Router::new()
        .route("/listen", get(listen::websocket_handler))
        .route("/files/subscribe", get(file_ws::file_subscription_handler))
        .route("/files", get(files::list_files).post(files::create_file))
        .route(
            "/files/{id}",
            get(files::get_file)
                .put(files::update_file)
                .delete(files::delete_file),
        )
        .route("/files/{id}/chat", post(chat::chat_handler))
        // Version history endpoints
        .route("/files/{id}/versions", get(versions::list_versions))
        .route("/files/{id}/versions/{version}", get(versions::get_version))
        .route("/files/{id}/versions/restore", post(versions::restore_version))
        // Mesh/task orchestration endpoints
        .route(
            "/mesh/tasks",
            get(mesh::list_tasks).post(mesh::create_task),
        )
        .route(
            "/mesh/tasks/{id}",
            get(mesh::get_task)
                .put(mesh::update_task)
                .delete(mesh::delete_task),
        )
        .route("/mesh/tasks/{id}/subtasks", get(mesh::list_subtasks))
        .route("/mesh/tasks/{id}/events", get(mesh::list_task_events))
        .route("/mesh/tasks/{id}/output", get(mesh::get_task_output))
        .route("/mesh/tasks/{id}/start", post(mesh::start_task))
        .route("/mesh/tasks/{id}/stop", post(mesh::stop_task))
        .route("/mesh/tasks/{id}/message", post(mesh::send_message))
        .route("/mesh/tasks/{id}/retry-completion", post(mesh::retry_completion_action))
        .route("/mesh/tasks/{id}/clone", post(mesh::clone_worktree))
        .route("/mesh/tasks/{id}/check-target", post(mesh::check_target_exists))
        .route("/mesh/chat", post(mesh_chat::mesh_toplevel_chat_handler))
        .route(
            "/mesh/chat/history",
            get(mesh_chat::get_chat_history).delete(mesh_chat::clear_chat_history),
        )
        .route("/mesh/tasks/{id}/chat", post(mesh_chat::mesh_chat_handler))
        .route("/mesh/daemons", get(mesh::list_daemons))
        .route("/mesh/daemons/directories", get(mesh::get_daemon_directories))
        .route("/mesh/daemons/{id}", get(mesh::get_daemon))
        // Merge endpoints for orchestrators
        .route("/mesh/tasks/{id}/branches", get(mesh_merge::list_branches))
        .route("/mesh/tasks/{id}/merge/start", post(mesh_merge::merge_start))
        .route("/mesh/tasks/{id}/merge/status", get(mesh_merge::merge_status))
        .route("/mesh/tasks/{id}/merge/resolve", post(mesh_merge::merge_resolve))
        .route("/mesh/tasks/{id}/merge/commit", post(mesh_merge::merge_commit))
        .route("/mesh/tasks/{id}/merge/abort", post(mesh_merge::merge_abort))
        .route("/mesh/tasks/{id}/merge/skip", post(mesh_merge::merge_skip))
        .route("/mesh/tasks/{id}/merge/check", get(mesh_merge::merge_check))
        // Mesh WebSocket endpoints
        .route("/mesh/tasks/subscribe", get(mesh_ws::task_subscription_handler))
        .route("/mesh/daemons/connect", get(mesh_daemon::daemon_handler))
        // API key management endpoints
        .route(
            "/auth/api-keys",
            post(api_keys::create_api_key_handler)
                .get(api_keys::get_api_key_handler)
                .delete(api_keys::revoke_api_key_handler),
        )
        .route("/auth/api-keys/refresh", post(api_keys::refresh_api_key_handler))
        // User account management endpoints
        .route(
            "/users/me",
            axum::routing::delete(users::delete_account_handler),
        )
        .route("/users/me/password", axum::routing::put(users::change_password_handler))
        .route("/users/me/email", axum::routing::put(users::change_email_handler))
        .with_state(state);

    let swagger = SwaggerUi::new("/swagger-ui")
        .url("/api-docs/openapi.json", ApiDoc::openapi());

    Router::new()
        .route("/api/v1/healthcheck", get(health_check))
        .nest("/api/v1", api_v1)
        .merge(swagger)
        .layer(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                .allow_headers(Any),
        )
        .layer(TraceLayer::new_for_http())
}

/// Run the HTTP server with graceful shutdown support.
///
/// # Arguments
/// * `state` - Shared application state containing ML models
/// * `addr` - Address to bind to (e.g., "0.0.0.0:8080")
pub async fn run_server(state: SharedState, addr: &str) -> anyhow::Result<()> {
    let app = make_router(state);
    let listener = tokio::net::TcpListener::bind(addr).await?;

    tracing::info!("Server listening on {}", addr);
    tracing::info!("Swagger UI available at http://{}/swagger-ui", addr);

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    Ok(())
}

/// Wait for shutdown signals (Ctrl+C or SIGTERM).
async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("Failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    tracing::info!("Shutdown signal received, starting graceful shutdown");
}