Add in-memory heap storage implementation

- Implemented MemoryHeapStorage using Arc<RwLock<HashMap>>
- Added --memory CLI flag for in-memory storage mode
- Memory storage is ideal for testing and scenarios where persistence is not needed
- Data is stored in process memory and lost when the process terminates
- Updated AnyHeapStorage enum to include Memory variant
- Exported MemoryHeapStorage from mcp module

Author: Claude Bot

server/src/main.rs

@@ -12,7 +12,7 @@ use rmcp::transport::sse_server::{SseServer, SseServerConfig};
 use tokio_util::sync::CancellationToken;
 
 mod mcp;
-use mcp::{StatelessService, StatefulService, initialize_v8};
+use mcp::{StatelessService, StatefulService, initialize_v8, MemoryHeapStorage};
 use mcp::heap_storage::{AnyHeapStorage, S3HeapStorage, FileHeapStorage};
 
 /// Command line arguments for configuring heap storage
@@ -21,15 +21,19 @@ use mcp::heap_storage::{AnyHeapStorage, S3HeapStorage, FileHeapStorage};
 struct Cli {
 
     /// S3 bucket name (required if --use-s3)
-    #[arg(long, conflicts_with_all = ["directory_path", "stateless"])]
+    #[arg(long, conflicts_with_all = ["directory_path", "stateless", "memory"])]
     s3_bucket: Option<String>,
 
     /// Directory path for filesystem storage (required if --use-filesystem)
-    #[arg(long, conflicts_with_all = ["s3_bucket", "stateless"])]
+    #[arg(long, conflicts_with_all = ["s3_bucket", "stateless", "memory"])]
     directory_path: Option<String>,
 
+    /// Use in-memory storage for heap snapshots (data is lost when process terminates)
+    #[arg(long, conflicts_with_all = ["s3_bucket", "directory_path", "stateless"])]
+    memory: bool,
+
     /// Run in stateless mode - no heap snapshots are saved or loaded
-    #[arg(long, conflicts_with_all = ["s3_bucket", "directory_path"])]
+    #[arg(long, conflicts_with_all = ["s3_bucket", "directory_path", "memory"])]
     stateless: bool,
 
     /// HTTP port to listen on (if not specified, uses stdio transport)
@@ -81,6 +85,8 @@ async fn main() -> Result<()> {
             AnyHeapStorage::S3(S3HeapStorage::new(bucket).await)
         } else if let Some(dir) = cli.directory_path {
             AnyHeapStorage::File(FileHeapStorage::new(dir))
+        } else if cli.memory {
+            AnyHeapStorage::Memory(MemoryHeapStorage::new())
         } else {
             // default to file /tmp/mcp-v8-heaps
             AnyHeapStorage::File(FileHeapStorage::new("/tmp/mcp-v8-heaps"))

server/src/mcp.rs

@@ -11,6 +11,7 @@ use std::sync::Once;
 use v8::{self};
 
 pub(crate) mod heap_storage;
+pub use crate::mcp::heap_storage::MemoryHeapStorage;
 use crate::mcp::heap_storage::{HeapStorage, AnyHeapStorage};
 
 

server/src/mcp/heap_storage.rs

@@ -5,6 +5,8 @@ use aws_sdk_s3::ByteStream;
 use aws_config;
 use std::sync::Arc;
 use async_trait::async_trait;
+use std::collections::HashMap;
+use tokio::sync::RwLock;
 
 #[async_trait]
 pub trait HeapStorage: Send + Sync + 'static {
@@ -98,28 +100,60 @@ impl HeapStorage for S3HeapStorage {
     }
 }
 
+/// In-memory heap storage using a HashMap
+/// Data is stored in memory and will be lost when the process terminates
+#[derive(Clone)]
+pub struct MemoryHeapStorage {
+    store: Arc<RwLock<HashMap<String, Vec<u8>>>>,
+}
+
+impl MemoryHeapStorage {
+    pub fn new() -> Self {
+        Self {
+            store: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+}
+
+#[async_trait]
+impl HeapStorage for MemoryHeapStorage {
+    async fn put(&self, name: &str, data: &[u8]) -> Result<(), String> {
+        let mut store = self.store.write().await;
+        store.insert(name.to_string(), data.to_vec());
+        Ok(())
+    }
+
+    async fn get(&self, name: &str) -> Result<Vec<u8>, String> {
+        let store = self.store.read().await;
+        store
+            .get(name)
+            .cloned()
+            .ok_or_else(|| format!("Heap '{}' not found in memory storage", name))
+    }
+}
+
 #[derive(Clone)]
 pub enum AnyHeapStorage {
     #[allow(dead_code)]
     File(FileHeapStorage),
     S3(S3HeapStorage),
+    Memory(MemoryHeapStorage),
 }
 
-
-
-
 #[async_trait::async_trait]
 impl HeapStorage for AnyHeapStorage {
     async fn put(&self, name: &str, data: &[u8]) -> Result<(), String> {
         match self {
             AnyHeapStorage::File(inner) => inner.put(name, data).await,
             AnyHeapStorage::S3(inner) => inner.put(name, data).await,
+            AnyHeapStorage::Memory(inner) => inner.put(name, data).await,
         }
     }
     async fn get(&self, name: &str) -> Result<Vec<u8>, String> {
         match self {
             AnyHeapStorage::File(inner) => inner.get(name).await,
             AnyHeapStorage::S3(inner) => inner.get(name).await,
+            AnyHeapStorage::Memory(inner) => inner.get(name).await,
         }
     }
 }