Refactor to Pure Native JSON-RPC (Ditch SDKs)

server/src/main.rs

@@ -1,5 +1,6 @@
 mod handlers;
 mod models;
+mod mcp;
 mod state;
 mod store;
 mod tools;
@@ -9,12 +10,7 @@ use crate::handlers::MemoryHandler;
 use crate::models::*;
 use crate::state::MemoryState;
 use crate::store::Store;
-use rust_mcp_sdk::{
-    ServerDetails, ToMcpServerHandler,
-    error::SdkResult,
 
-    schema::{Implementation, ServerCapabilities, ServerCapabilitiesTools},
-};
 use std::fs;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
@@ -109,94 +105,95 @@ async fn reconcile_worker(state: Arc<MemoryState>) {
 }
 
 
-use rust_mcp_axum::{create_axum_server, AxumServerOptions};
 
-fn run_server(state: Arc<MemoryState>) -> SdkResult<()> {
+
+
+use axum::{
+    extract::{State, Query},
+    response::sse::{Event, Sse},
+    routing::{get, post},
+    Json, Router,
+};
+use futures_util::stream::Stream;
+use std::convert::Infallible;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+struct AppState {
+    handler: Arc<MemoryHandler>,
+    clients: RwLock<HashMap<String, mpsc::Sender<Result<Event, Infallible>>>>,
+    next_id: AtomicUsize,
+}
+
+fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
     let rt = tokio::runtime::Runtime::new().unwrap();
     rt.block_on(async {
-        let state_clone = state.clone();
-        tokio::spawn(async move {
-            reconcile_worker(state_clone).await;
+        tokio::spawn(reconcile_worker(Arc::clone(&state)));
+        let app_state = Arc::new(AppState {
+            handler: Arc::new(MemoryHandler { state }),
+            clients: RwLock::new(HashMap::new()),
+            next_id: AtomicUsize::new(1),
         });
 
-        let server_details = ServerDetails {
-            server_info: Implementation {
-                name: "gemini-mcp-memory".into(),
-                version: "3.0.0".into(),
-                title: Some("MCP Memory Server".into()),
-                description: Some("Persistent memory MCP server".into()),
-                icons: vec![],
-                website_url: None,
-            },
-            capabilities: ServerCapabilities {
-                tools: Some(ServerCapabilitiesTools {
-                    list_changed: Some(true),
-                }),
-                ..Default::default()
-            },
-            instructions: None,
-            meta: None,
-        };
-
-        let options = AxumServerOptions {
-            host: "0.0.0.0".to_string(),
-            port: 3000,
-            custom_sse_endpoint: Some("/sse".to_string()),
-            ..Default::default()
-        };
-
-        let server = create_axum_server(
-            server_details,
-            MemoryHandler { state: state.clone() }.to_mcp_server_handler(),
-            options,
-        );
-
-        let server = server.with_route(
-            "/gate/verify",
-            axum::routing::get(|| async move {
-                axum::http::StatusCode::OK
-            }),
-        ).with_route(
-            "/dashboard",
-            axum::routing::get(|| async move {
-                axum::response::Html(include_str!("dashboard.html"))
-            }),
-        ).with_route(
-            "/api/stats",
-            axum::routing::get({
-                let state_clone = state.clone();
-                move || async move {
-                    let (entities, relations) = {
-                        let graph = state_clone.get_full_graph();
-                        (graph.entities.len(), graph.relations.len())
-                    };
-                    let tasks = state_clone.tasks.read().len();
-                    let snippets = state_clone.snippets.read().len();
-                    let tech_debts = state_clone.tech_debts.read().len();
-                    let adrs = state_clone.adrs.read().len();
-
-                    axum::Json(serde_json::json!({
-                        "entities": entities,
-                        "relations": relations,
-                        "tasks": tasks,
-                        "snippets": snippets,
-                        "tech_debts": tech_debts,
-                        "adrs": adrs
-                    }))
-                }
-            }),
-        );
+        let app = Router::new()
+            .route("/sse", get(sse_handler))
+            .route("/messages", post(message_handler))
+            .with_state(app_state);
 
+        let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
         eprintln!("MCP Memory Server running on http://0.0.0.0:3000/sse");
-        server.start().await.unwrap();
-
+        axum::serve(listener, app).await.unwrap();
         Ok(())
     })
 }
 
+async fn sse_handler(
+    State(state): State<Arc<AppState>>,
+) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
+    let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
+    let (tx, rx) = mpsc::channel::<Result<Event, Infallible>>(100);
+    
+    state.clients.write().unwrap().insert(session_id.clone(), tx.clone());
+    
+    let _ = tx.send(Ok(Event::default().event("endpoint").data(format!("?sessionId={}", session_id)))).await;
+    
+    let stream = ReceiverStream::new(rx);
+    Sse::new(stream).keep_alive(axum::response::sse::KeepAlive::new())
+}
+
+#[derive(serde::Deserialize)]
+struct SessionQuery {
+    #[serde(rename = "sessionId")]
+    session_id: String,
+}
+
+async fn message_handler(
+    State(state): State<Arc<AppState>>,
+    Query(query): Query<SessionQuery>,
+    Json(payload): Json<serde_json::Value>,
+) -> axum::http::StatusCode {
+    let handler = Arc::clone(&state.handler);
+    let session_id = query.session_id.clone();
+    let clients = Arc::clone(&state);
+    
+    tokio::spawn(async move {
+        if let Some(response) = handler.handle_request(payload).await {
+            let tx_opt = clients.clients.read().unwrap().get(&session_id).cloned();
+            if let Some(tx) = tx_opt {
+                let data = serde_json::to_string(&response).unwrap();
+                let _ = tx.send(Ok(Event::default().event("message").data(data))).await;
+            }
+        }
+    });
+    
+    axum::http::StatusCode::ACCEPTED
+}
+
 mod proxy;
 
-fn main() -> SdkResult<()> {
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
     let cli = Cli::parse();
 
     #[cfg(target_os = "windows")]