AkurAI-Notes / crates/notes-mcp/src/server.rs

Latest change 705e1ce035ea4b9fb4c413dd38f550650185b571 - release: v0.0.5 by Claude Code

use crate::tools::ToolRegistry;
use akurai_http::{Method, Request, Response};
use akurai_storage::BTree;
use std::sync::{Arc, Mutex};

pub struct McpConfig {
    pub host: String,
    pub port: u16,
    pub db: String,
    pub token: String,
    pub embed_url: Option<String>,
    pub embed_model: Option<String>,
    pub embed_api_key: Option<String>,
    pub owner_sub: String,
}

pub struct McpServer {
    config: McpConfig,
    tools: Arc<ToolRegistry>,
}

impl McpServer {
    pub fn new(config: McpConfig) -> Result<Self, String> {
        let btree =
            BTree::open(&config.db).map_err(|e| format!("failed to open MCP database: {e}"))?;
        let tools = ToolRegistry::new(
            Arc::new(Mutex::new(btree)),
            config.embed_url.clone(),
            config.embed_model.clone(),
            config.embed_api_key.clone(),
            config.owner_sub.clone(),
        );
        Ok(Self {
            config,
            tools: Arc::new(tools),
        })
    }

    pub fn handle_request(&self, body: &str) -> String {
        self.tools.handle_request(body)
    }

    pub fn serve(&self) -> Result<(), String> {
        let addr = format!("{}:{}", self.config.host, self.config.port);
        println!("akurai-notes MCP server listening on http://{addr}/mcp");

        let token = self.config.token.clone();
        let tools = Arc::clone(&self.tools);
        let handler = move |req: &Request| -> Response {
            if req.path != "/mcp" {
                return Response::new(404).with_text("Not found");
            }
            if req.method != Method::Post {
                return Response::new(405).with_text("Method not allowed");
            }
            if !token.is_empty() {
                let authorized = req
                    .header("Authorization")
                    .and_then(|s| s.strip_prefix("Bearer "))
                    .map(|t| t.trim() == token)
                    .unwrap_or(false);
                if !authorized {
                    return Response::new(401).with_text("unauthorized");
                }
            }

            let body = req.body_str();
            let result = tools.handle_request(body.as_ref());

            // Empty result = JSON-RPC notification: acknowledge with 202 Accepted and an
            // empty body (MCP Streamable-HTTP spec). Strict clients (codex/rmcp) break if
            // a notification receives a JSON-RPC response body.
            if result.is_empty() {
                return Response::new(202).with_header("Access-Control-Allow-Origin", "*");
            }

            Response::ok()
                .with_header("Access-Control-Allow-Origin", "*")
                .with_header("Content-Type", "application/json")
                .with_body("application/json", result.into_bytes())
        };

        let server = akurai_http::Server::bind(&addr)
            .map_err(|e| format!("failed to bind MCP server: {e}"))?;
        server
            .run(handler)
            .map_err(|e| format!("server error: {e}"))
    }
}