openbrain-mcp/src/lib.rs

//! OpenBrain MCP - High-performance vector memory for AI agents

pub mod auth;
pub mod config;
pub mod db;
pub mod embedding;
pub mod migrations;
pub mod tools;
pub mod transport;

use anyhow::Result;
use axum::{Router, Json, http::StatusCode, middleware};
use serde_json::json;
use std::sync::Arc;
use tokio::net::TcpListener;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
use tracing::{info, error};

use crate::auth::auth_middleware;
use crate::config::Config;
use crate::db::Database;
use crate::embedding::EmbeddingEngine;
use crate::transport::McpState;

/// Service readiness state
#[derive(Clone, Debug, PartialEq)]
pub enum ReadinessState {
    Initializing,
    Ready,
    Failed(String),
}

/// Shared application state
pub struct AppState {
    pub db: Database,
    pub embedding: tokio::sync::RwLock<Option<Arc<EmbeddingEngine>>>,
    pub config: Config,
    pub readiness: tokio::sync::RwLock<ReadinessState>,
}

impl AppState {
    /// Get embedding engine, returns None if not ready
    pub async fn get_embedding(&self) -> Option<Arc<EmbeddingEngine>> {
        self.embedding.read().await.clone()
    }
}

/// Health check endpoint - always returns OK if server is running
async fn health_handler() -> Json<serde_json::Value> {
    Json(json!({"status": "ok"}))
}

/// Readiness endpoint - returns 503 if embedding not ready
async fn readiness_handler(
    state: axum::extract::State<Arc<AppState>>,
) -> (StatusCode, Json<serde_json::Value>) {
    let readiness = state.readiness.read().await.clone();
    match readiness {
        ReadinessState::Ready => (
            StatusCode::OK,
            Json(json!({"status": "ready", "embedding": true}))
        ),
        ReadinessState::Initializing => (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({"status": "initializing", "embedding": false}))
        ),
        ReadinessState::Failed(err) => (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({"status": "failed", "error": err}))
        ),
    }
}

/// Run the MCP server
pub async fn run_server(config: Config, db: Database) -> Result<()> {
    // Create state with None embedding (will init in background)
    let state = Arc::new(AppState {
        db,
        embedding: tokio::sync::RwLock::new(None),
        config: config.clone(),
        readiness: tokio::sync::RwLock::new(ReadinessState::Initializing),
    });

    // Spawn background task to initialize embedding with retry
    let state_clone = state.clone();
    let embedding_config = config.embedding.clone();
    tokio::spawn(async move {
        let max_retries = 3;
        let mut attempt = 0;
        
        loop {
            attempt += 1;
            info!("Initializing embedding engine (attempt {}/{})", attempt, max_retries);
            
            match EmbeddingEngine::new(&embedding_config).await {
                Ok(engine) => {
                    let engine = Arc::new(engine);
                    *state_clone.embedding.write().await = Some(engine);
                    *state_clone.readiness.write().await = ReadinessState::Ready;
                    info!("Embedding engine initialized successfully");
                    break;
                }
                Err(e) => {
                    error!("Failed to init embedding (attempt {}): {:?}", attempt, e);
                    if attempt >= max_retries {
                        let err_msg = format!("Failed after {} attempts: {:?}", max_retries, e);
                        *state_clone.readiness.write().await = ReadinessState::Failed(err_msg);
                        break;
                    }
                    // Exponential backoff: 2s, 4s, 8s...
                    tokio::time::sleep(tokio::time::Duration::from_secs(2u64.pow(attempt))).await;
                }
            }
        }
    });

    // Create MCP state for SSE transport
    let mcp_state = McpState::new(state.clone());

    // Build router with health/readiness endpoints (no auth required)
    let health_router = Router::new()
        .route("/health", axum::routing::get(health_handler))
        .route("/ready", axum::routing::get(readiness_handler))
        .with_state(state.clone());
    
    // Build MCP router with auth middleware
    let mcp_router = transport::mcp_router(mcp_state)
        .layer(middleware::from_fn_with_state(state.clone(), auth_middleware));

    let app = Router::new()
        .merge(health_router)
        .merge(mcp_router)
        .layer(TraceLayer::new_for_http())
        .layer(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                .allow_headers(Any),
        );

    // Start server immediately
    let bind_addr = format!("{}:{}", config.server.host, config.server.port);
    let listener = TcpListener::bind(&bind_addr).await?;
    info!("Server listening on {}", bind_addr);

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

    Ok(())
}