Skip to main content

rust_memex/handlers/
mod.rs

1use anyhow::Result;
2use serde_json::Value;
3use std::sync::Arc;
4use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
5
6use crate::{
7    ServerConfig,
8    mcp_core::{
9        McpCore, McpDispatch, McpTransport, dispatch_mcp_jsonrpc_request, dispatch_mcp_payload,
10    },
11    mcp_runtime::build_mcp_core,
12};
13
14pub struct MCPServer {
15    mcp_core: Arc<McpCore>,
16}
17
18impl MCPServer {
19    /// Get the shared MCP core for reuse across transports.
20    pub fn mcp_core(&self) -> Arc<McpCore> {
21        self.mcp_core.clone()
22    }
23
24    pub async fn run_stdio(self) -> Result<()> {
25        let stdin = tokio::io::stdin();
26        let mut stdout = tokio::io::stdout();
27        let mut reader = BufReader::new(stdin);
28        let mut line = String::new();
29
30        loop {
31            line.clear();
32            let read = reader.read_line(&mut line).await?;
33            if read == 0 {
34                break;
35            }
36
37            let trimmed = line.trim();
38            if trimmed.is_empty() {
39                continue;
40            }
41
42            if let Some(response) =
43                dispatch_mcp_payload(self.mcp_core.as_ref(), trimmed, McpTransport::Stdio).await
44            {
45                write_json_line(&mut stdout, &response).await?;
46            }
47        }
48
49        Ok(())
50    }
51
52    pub async fn run(self) -> Result<()> {
53        self.run_stdio().await
54    }
55
56    /// Dispatch a parsed JSON-RPC request through the shared stdio MCP core.
57    ///
58    /// Notifications return `None` so library callers see the same semantics as
59    /// the real stdio transport instead of a synthetic success response.
60    pub async fn dispatch_request(&self, request: Value) -> Option<Value> {
61        self.dispatch_jsonrpc_request(request).await.into_option()
62    }
63
64    /// Dispatch a parsed JSON-RPC request and preserve notification semantics.
65    pub async fn dispatch_jsonrpc_request(&self, request: Value) -> McpDispatch {
66        dispatch_mcp_jsonrpc_request(self.mcp_core.as_ref(), request, McpTransport::Stdio).await
67    }
68
69    #[cfg(test)]
70    pub(crate) fn from_mcp_core(mcp_core: Arc<McpCore>) -> Self {
71        Self { mcp_core }
72    }
73}
74
75pub async fn create_server(config: ServerConfig) -> Result<MCPServer> {
76    let mcp_core = build_mcp_core(config).await?;
77
78    Ok(MCPServer { mcp_core })
79}
80
81async fn write_json_line(stdout: &mut tokio::io::Stdout, response: &Value) -> anyhow::Result<()> {
82    let payload = serde_json::to_string(response)?;
83    stdout.write_all(payload.as_bytes()).await?;
84    stdout.write_all(b"\n").await?;
85    stdout.flush().await?;
86    Ok(())
87}