Initial public release

src/lib.rs

@@ -0,0 +1,150 @@
+//! 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)
+        .nest("/mcp", 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(())
+}