Create mcp.rs

src/mcp.rs

@@ -0,0 +1,361 @@
+use crate::indexer::TantivyEngine;
+use crate::parser::{DocumentParser, ParsedDocument};
+use anyhow::Result;
+use axum::{
+    extract::State,
+    response::sse::{Event, KeepAlive, Sse},
+    response::IntoResponse,
+    routing::{get, post},
+    Json, Router,
+};
+use futures::stream::Stream;
+use rayon::prelude::*;
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+use std::convert::Infallible;
+use std::path::Path;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tracing::{error, info};
+use walkdir::WalkDir;
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct JsonRpcRequest {
+    pub jsonrpc: String,
+    pub id: Option<Value>,
+    pub method: String,
+    pub params: Option<Value>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct JsonRpcResponse {
+    pub jsonrpc: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub id: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub result: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<JsonRpcError>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct JsonRpcError {
+    pub code: i64,
+    pub message: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub data: Option<Value>,
+}
+
+pub struct McpServer {
+    engine: Arc<TantivyEngine>,
+}
+
+impl McpServer {
+    pub fn new(engine: Arc<TantivyEngine>) -> Self {
+        Self { engine }
+    }
+
+    pub async fn handle_request(&self, req: JsonRpcRequest) -> JsonRpcResponse {
+        let req_id = req.id.clone();
+
+        match req.method.as_str() {
+            "initialize" => {
+                let init_result = json!({
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {
+                        "tools": {}
+                    },
+                    "serverInfo": {
+                        "name": "rust-mcp-search",
+                        "version": "0.1.0"
+                    }
+                });
+                Self::success_response(req_id, init_result)
+            }
+            "notifications/initialized" | "ping" => {
+                Self::success_response(req_id, json!({}))
+            }
+            "tools/list" => {
+                let tools = json!({
+                    "tools": [
+                        {
+                            "name": "parse_and_index",
+                            "description": "Recursively scans a directory and indexes plain text, Markdown, JSON, PDF, CSV, and DOCX files in parallel using Rayon and Tantivy.",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {
+                                    "directory_path": {
+                                        "type": "string",
+                                        "description": "Absolute or relative directory path to index"
+                                    },
+                                    "collection_name": {
+                                        "type": "string",
+                                        "description": "Collection namespace identifier (default: 'default')"
+                                    }
+                                },
+                                "required": ["directory_path"]
+                            }
+                        },
+                        {
+                            "name": "search_documents",
+                            "description": "Full-text BM25 search over indexed documents with highlighted snippets and relevance scores.",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {
+                                    "query": {
+                                        "type": "string",
+                                        "description": "Tantivy BM25 search query"
+                                    },
+                                    "collection_name": {
+                                        "type": "string",
+                                        "description": "Filter results by collection"
+                                    },
+                                    "limit": {
+                                        "type": "integer",
+                                        "description": "Max hits to return (default: 10)"
+                                    }
+                                },
+                                "required": ["query"]
+                            }
+                        },
+                        {
+                            "name": "extract_document_text",
+                            "description": "High-speed raw text extraction from a single file (TXT, MD, PDF, DOCX, CSV, JSON).",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {
+                                    "file_path": {
+                                        "type": "string",
+                                        "description": "Path to document"
+                                    }
+                                },
+                                "required": ["file_path"]
+                            }
+                        },
+                        {
+                            "name": "get_index_stats",
+                            "description": "Returns indexed document counts, segment counts, active collections, and memory/storage status.",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {}
+                            }
+                        }
+                    ]
+                });
+                Self::success_response(req_id, tools)
+            }
+            "tools/call" => {
+                let params = req.params.unwrap_or_default();
+                let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
+                let args = params.get("arguments").cloned().unwrap_or(json!({}));
+
+                match self.dispatch_tool(tool_name, args).await {
+                    Ok(tool_output) => Self::success_response(
+                        req_id,
+                        json!({
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": tool_output
+                                }
+                            ],
+                            "isError": false
+                        }),
+                    ),
+                    Err(e) => Self::success_response(
+                        req_id,
+                        json!({
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": format!("Error executing tool '{}': {}", tool_name, e)
+                                }
+                            ],
+                            "isError": true
+                        }),
+                    ),
+                }
+            }
+            _ => Self::error_response(req_id, -32601, format!("Method '{}' not found", req.method)),
+        }
+    }
+
+    async fn dispatch_tool(&self, name: &str, args: Value) -> Result<String> {
+        match name {
+            "parse_and_index" => {
+                let dir = args
+                    .get("directory_path")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| anyhow::anyhow!("Missing 'directory_path' argument"))?;
+                let collection = args
+                    .get("collection_name")
+                    .and_then(|v| v.as_str())
+                    .unwrap_or("default");
+
+                let dir_path = Path::new(dir).to_path_buf();
+                if !dir_path.exists() {
+                    return Err(anyhow::anyhow!("Directory does not exist: {:?}", dir_path));
+                }
+
+                // 1. Collect all valid candidate files
+                let paths: Vec<_> = WalkDir::new(&dir_path)
+                    .into_iter()
+                    .filter_map(|e| e.ok())
+                    .filter(|e| e.file_type().is_file())
+                    .map(|e| e.into_path())
+                    .collect();
+
+                let total_found = paths.len();
+
+                // 2. Parallel ingestion with Rayon
+                let parsed_docs: Vec<ParsedDocument> = paths
+                    .par_iter()
+                    .filter_map(|path| match DocumentParser::parse_file(path) {
+                        Ok(doc) => Some(doc),
+                        Err(err) => {
+                            tracing::warn!("Skipping {:?}: {}", path, err);
+                            None
+                        }
+                    })
+                    .collect();
+
+                let parsed_count = parsed_docs.len();
+
+                // 3. Batch commit to Tantivy engine
+                let indexed_count = self.engine.add_documents(collection, parsed_docs).await?;
+
+                Ok(json!({
+                    "status": "success",
+                    "collection": collection,
+                    "scanned_files": total_found,
+                    "parsed_documents": parsed_count,
+                    "indexed_documents": indexed_count,
+                    "source_path": dir
+                })
+                .to_string())
+            }
+            "search_documents" => {
+                let query = args
+                    .get("query")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| anyhow::anyhow!("Missing 'query' argument"))?;
+                let collection = args.get("collection_name").and_then(|v| v.as_str());
+                let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
+
+                let hits = self.engine.search(query, collection, limit)?;
+                Ok(serde_json::to_string_pretty(&hits)?)
+            }
+            "extract_document_text" => {
+                let path_str = args
+                    .get("file_path")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| anyhow::anyhow!("Missing 'file_path' argument"))?;
+
+                let doc = DocumentParser::parse_file(path_str)?;
+                Ok(json!({
+                    "path": doc.path,
+                    "title": doc.title,
+                    "extension": doc.extension,
+                    "size_bytes": doc.size_bytes,
+                    "content": doc.content
+                })
+                .to_string())
+            }
+            "get_index_stats" => {
+                let stats = self.engine.get_stats()?;
+                Ok(serde_json::to_string_pretty(&stats)?)
+            }
+            unknown => Err(anyhow::anyhow!("Unknown tool name: {}", unknown)),
+        }
+    }
+
+    fn success_response(id: Option<Value>, result: Value) -> JsonRpcResponse {
+        JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id,
+            result: Some(result),
+            error: None,
+        }
+    }
+
+    fn error_response(id: Option<Value>, code: i64, message: String) -> JsonRpcResponse {
+        JsonRpcResponse {
+            jsonrpc: "2.0".to_string(),
+            id,
+            result: None,
+            error: Some(JsonRpcError {
+                code,
+                message,
+                data: None,
+            }),
+        }
+    }
+
+    /// Runs the server loop over standard I/O (STDIO) transport
+    pub async fn run_stdio(self: Arc<Self>) -> Result<()> {
+        let stdin = tokio::io::stdin();
+        let mut reader = BufReader::new(stdin).lines();
+        let mut stdout = tokio::io::stdout();
+
+        while let Some(line) = reader.next_line().await? {
+            if line.trim().is_empty() {
+                continue;
+            }
+
+            match serde_json::from_str::<JsonRpcRequest>(&line) {
+                Ok(request) => {
+                    let response = self.handle_request(request).await;
+                    let out_json = serde_json::to_string(&response)?;
+                    stdout.write_all(out_json.as_bytes()).await?;
+                    stdout.write_all(b"\n").await?;
+                    stdout.flush().await?;
+                }
+                Err(e) => {
+                    let err_resp = Self::error_response(None, -32700, format!("Parse error: {}", e));
+                    let out_json = serde_json::to_string(&err_resp)?;
+                    stdout.write_all(out_json.as_bytes()).await?;
+                    stdout.write_all(b"\n").await?;
+                    stdout.flush().await?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Runs the server over Axum SSE/HTTP transport
+    pub async fn run_http(self: Arc<Self>, host: &str, port: u16) -> Result<()> {
+        let app = Router::new()
+            .route("/rpc", post(handle_http_rpc))
+            .route("/sse", get(handle_sse))
+            .route("/health", get(|| async { "healthy" }))
+            .layer(tower_http::cors::CorsLayer::permissive())
+            .with_state(self);
+
+        let addr = format!("{}:{}", host, port);
+        info!("MCP HTTP/SSE server listening on http://{}", addr);
+        let listener = tokio::net::TcpListener::bind(&addr).await?;
+        axum::serve(listener, app).await?;
+        Ok(())
+    }
+}
+
+async fn handle_http_rpc(
+    State(server): State<Arc<McpServer>>,
+    Json(payload): Json<JsonRpcRequest>,
+) -> Json<JsonRpcResponse> {
+    let resp = server.handle_request(payload).await;
+    Json(resp)
+}
+
+async fn handle_sse(
+    State(_server): State<Arc<McpServer>>,
+) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
+    let stream = tokio_stream::iter(vec![
+        Ok(Event::default().event("endpoint").data("/rpc"))
+    ])
+    .chain(tokio_stream::wrappers::IntervalStream::new(tokio::time::interval(Duration::from_secs(15)))
+        .map(|_| Ok(Event::default().comment("keep-alive"))));
+
+    Sse::new(stream).keep_alive(KeepAlive::default())
+}