feat: implement dual-transport WAL architecture with permanent background leader and lightweight stubs
This commit is contained in:
1 parent
a2febd1b00
commit
e4ff476b6d
16 files changed
+1910
-52
No files matched your search
Generated
+1355
-10
File diff suppressed because it is too large.
Load diff
+6
-16
@@ -1,16 +1,6 @@
|
||||
[package]
|
||||
name = "mcp-memory"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1.92"
|
||||
clap = { version = "4.6.6", features = ["derive"] }
|
||||
dirs = "6.0.0"
|
||||
glob = "0.3.4"
|
||||
rust-mcp-sdk = { version = "2.0.0", default-features = false, features = ["server", "macros", "stdio"] }
|
||||
serde = { version = "1.0.229", features = ["derive"] }
|
||||
serde_json = "1.0.151"
|
||||
strsim = "0.11.1"
|
||||
tokio = { version = "1.53.1", features = ["full"] }
|
||||
uuid = { version = "1.26.0", features = ["v4"] }
|
||||
[workspace]
|
||||
members = [
|
||||
"server",
|
||||
"stub"
|
||||
]
|
||||
resolver = "2"
|
||||
@@ -0,0 +1,70 @@
|
||||
# mcp-memory
|
||||
A high-performance, persistent Knowledge Graph and Context daemon for Antigravity, implementing the Model Context Protocol (MCP).
|
||||
|
||||
## Overview
|
||||
mcp-memory acts as the persistent "brain" for the agy CLI agents. It tracks entities, relations, background tasks, engineering debt, and architectural decisions across sessions.
|
||||
|
||||
To eliminate heavy Cross-OS I/O penalties when using WSL and Windows simultaneously, mcp-memory operates using a **Dual-Transport Leader/Stub Architecture**:
|
||||
* **The Server (mcp-memory-server)**: Runs natively on the Windows host. It binds to .0.0.0:3000, serving standard stdio to the primary Windows agy instance while simultaneously hosting an Axum HTTP server for secondary clients.
|
||||
* **The Stub (mcp-memory-stub)**: An ultra-lightweight proxy binary. WSL agy instances run this native Linux stub, which transparently pipes stdio JSON-RPC traffic over the network to the Windows HTTP server (http://127.0.0.1:3000), completely bypassing WSL NTFS mounts.
|
||||
|
||||
## Quick Start & Usage
|
||||
|
||||
### 1. Windows Installation (The Server & Stub)
|
||||
Compile the main daemon and lightweight stub natively for Windows:
|
||||
`powershell
|
||||
cargo build --release
|
||||
Copy-Item target\release\mcp-memory-server.exe C:\Users\reazul.ashraf\.local\bin\mcp-memory.exe
|
||||
Copy-Item target\release\mcp-memory-stub.exe C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe
|
||||
`
|
||||
|
||||
**Step 1:** To bypass Antigravity's lazy-loading and ensure the server is instantly available for WSL, configure your PowerShell profile to auto-start the background server when you open a terminal:
|
||||
`powershell
|
||||
# Add this to your PowerShell profile:
|
||||
if (-not (Get-Process mcp-memory -ErrorAction SilentlyContinue)) { Start-Process -FilePath "C:\Users\reazul.ashraf\.local\bin\mcp-memory.exe" -WindowStyle Hidden -ErrorAction SilentlyContinue }
|
||||
`
|
||||
|
||||
**Step 2:** Update your Windows ~/.gemini/config/mcp_config.json to point the CLI to the ultra-lightweight stub (since the server is already running in the background):
|
||||
`json
|
||||
{
|
||||
"mcpServers": {
|
||||
"memory": {
|
||||
"command": "C:\\Users\\reazul.ashraf\\.local\\bin\\mcp-memory-stub.exe",
|
||||
"args": []
|
||||
}
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
### 2. WSL / Linux Installation (The Stub)
|
||||
Compile the ultra-lightweight stub as a static Linux binary (from the Windows host):
|
||||
`powershell
|
||||
cargo zigbuild --target x86_64-unknown-linux-musl --release -p mcp-memory-stub
|
||||
wsl.exe -d Ubuntu -e bash -c "cp /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory/target/x86_64-unknown-linux-musl/release/mcp-memory-stub ~/.local/bin/mcp-memory-stub && chmod +x ~/.local/bin/mcp-memory-stub"
|
||||
`
|
||||
|
||||
Update your WSL ~/.gemini/config/mcp_config.json:
|
||||
`json
|
||||
{
|
||||
"mcpServers": {
|
||||
"memory": {
|
||||
"command": "/home/riz/.local/bin/mcp-memory-stub",
|
||||
"args": [
|
||||
"--target", "http://127.0.0.1:3000",
|
||||
"--wake-cmd", "/mnt/c/Windows/System32/cmd.exe /c start /B C:\\Users\\reazul.ashraf\\.local\\bin\\mcp-memory.exe"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
`
|
||||
*Note: The --wake-cmd ensures that if you start WSL while Windows is completely asleep, the Linux stub will use WSL interop to silently spin up the Windows daemon in the background before connecting.*
|
||||
|
||||
## Push Safety Gates
|
||||
The daemon also operates as a global safety gate for Git. Before pushing code, run:
|
||||
`ash
|
||||
mcp-memory gate verify
|
||||
`
|
||||
This queries the daemon (via HTTP) to confirm if pre-push validation (like running tests via PrePushAuditor) has been cleared by the agent.
|
||||
|
||||
## Further Reading
|
||||
For a deep dive into the architecture, Write-Ahead Logging (WAL), locking mechanisms, and the HTTP SSE event loop, consult the design.md file in this repository.
|
||||
@@ -0,0 +1,12 @@
|
||||
@echo off
|
||||
echo Building mcp-memory in release mode...
|
||||
cargo build --release
|
||||
|
||||
if %ERRORLEVEL% EQU 0 (
|
||||
echo.
|
||||
echo Build successful!
|
||||
echo Executable is located at: target\release\mcp-memory.exe
|
||||
) else (
|
||||
echo.
|
||||
echo Build failed with error code %ERRORLEVEL%.
|
||||
)
|
||||
@@ -0,0 +1,121 @@
|
||||
# MCP Memory Server Architecture & Workflow Design
|
||||
|
||||
## 1. Core Architecture
|
||||
The `mcp-memory` system runs as a **single, continuously running background daemon natively on Windows**. It manages the state of the Antigravity knowledge graph and serves as a universal backend for both Windows and WSL environments.
|
||||
|
||||
### Why this design?
|
||||
* **Cross-OS I/O Optimization:** Prevents the WSL agent from performing slow, heavy filesystem writes against the mounted Windows `C:\` drive.
|
||||
* **Concurrency & Locking:** A single daemon holds the lock on the `mcp_memory` JSON files, preventing data corruption and eliminating complex delta-reconciliation between parallel WSL and Windows processes.
|
||||
* **Dual Transport System:** Utilizes both `stdio` and HTTP transports concurrently. The local Windows `agy` instance connects natively via standard `stdio`, while the Axum HTTP server provides synchronous, non-blocking access for WSL clients and external scripts.
|
||||
|
||||
## 2. Server Transport & Endpoints (Axum)
|
||||
The daemon uses the `axum` and `rust-mcp-axum` crates, binding to `0.0.0.0:3000` to serve the network.
|
||||
|
||||
### Standard MCP Endpoints
|
||||
* `GET /sse`: The Server-Sent Events (SSE) endpoint. Antigravity clients connect here to keep a one-way pipe open for receiving pushed responses and event notifications from the server.
|
||||
* `POST /messages`: The JSON-RPC endpoint. Clients use this to send tool calls and resource reads up to the server.
|
||||
|
||||
### Custom Integration Endpoints
|
||||
* `GET /gate/verify`: A lightweight, deterministic endpoint used by external scripts to verify if an action (like a `git push`) is authorized based on the current state.
|
||||
* *(Future)* `GET /health` or `GET /graph` for immediate system inspection.
|
||||
|
||||
## 3. Client Connections
|
||||
The Antigravity configurations (`mcp_config.json`) utilize the dual-transport system:
|
||||
* **Windows `agy`:** Spawns and connects to the local daemon natively via `stdio` subprocess execution.
|
||||
* **WSL (Linux) `agy`:** Connects to the running Windows HTTP daemon via the host network proxy, e.g., `http://127.0.0.1:3000/sse`.
|
||||
|
||||
## 4. Git Integration (Global Wrapper)
|
||||
Instead of relying on localized per-repository Git hooks (like `.git/hooks/pre-push`), the system leverages a **global bash/PowerShell alias wrapper** for the `git` command. This intercepts `git` commands universally across the OS.
|
||||
|
||||
### Workflow Example (Push Safety Gate)
|
||||
1. The user types `git push`.
|
||||
2. The global wrapper intercepts the command.
|
||||
3. It makes a synchronous HTTP request to the local daemon: `curl -s http://localhost:3000/gate/verify`.
|
||||
4. If the endpoint returns `200 OK` (indicating the `PrePushAuditor` subagent has verified that unit tests pass and history is squashed), the push proceeds.
|
||||
5. If not `200 OK`, the wrapper blocks the push and alerts the user to fix tests or run `gsquash`.
|
||||
|
||||
### Benefits of the Wrapper Approach
|
||||
* **Universal Enforcement:** The push safety gate is protected across all repositories automatically, without copying hook scripts.
|
||||
* **No File Locks:** Uses safe, lightweight, parallelizable HTTP requests rather than executing the Rust binary directly.
|
||||
* **Action Logging:** The wrapper can be seamlessly extended to log actions (like `checkout` or `commit`) directly into the knowledge graph in real-time.
|
||||
## 5. Storage Architecture & Persistence (WAL)
|
||||
The daemon utilizes a Write-Ahead Logging (WAL) architecture for robust, high-performance state management, replacing fragmented delta-file reconciliation and massive synchronous JSON rewrites.
|
||||
|
||||
### Key Principles:
|
||||
* **Append-Only Log:** Every state change (e.g., adding an entity, updating a task) is instantly appended to a sequential, append-only log file (wal.log) *before* in-memory state is altered.
|
||||
* **Crash Resilience:** Eliminates data corruption. If the daemon is forcefully terminated, the system guarantees zero data loss by replaying the WAL against the last known valid checkpoint upon restart.
|
||||
* **Asynchronous Checkpointing:** Master JSON store files are no longer rewritten on every tool call. Checkpointing (flushing memory to the master JSON files) is pushed to a background thread to run periodically or upon graceful shutdown, drastically reducing disk I/O.
|
||||
|
||||
## 6. Domain Models & Component Stores
|
||||
While the core architecture relies on a unified Knowledge Graph (Entities, Relations, Observations), the daemon leverages a modular, thread-safe generic Store<T> pattern for specialized domains.
|
||||
Currently implemented persistent stores include:
|
||||
* **Audit Ledger & Tasks:** Tracks agent actions and active background tasks.
|
||||
* **Context & Handoffs:** Sticky notes, Session Summaries, Handoff Memos, and Context Workspaces.
|
||||
* **Engineering Tracking:** ADRs (Architecture Decision Records), Snippets, Error Fixes, Tech Debt, and PR Checklists.
|
||||
* **Environment State:** Pinned Files, Env Fingerprints, Milestones, and Environments.
|
||||
|
||||
## 7. Background Maintenance (Reconcile Worker)
|
||||
The daemon runs a continuous asynchronous background task (
|
||||
econcile_worker) on a 5-second polling loop responsible for:
|
||||
* **State Reconciliation:** Reading the append-only wal.jsonl file, squashing the mutations into the master knowledge graph, and cleanly truncating the WAL.
|
||||
* **Ledger Pruning:** Automatically truncating the audit_ledger.json to keep only the last 7 days of activity, with a hard cap of 1,000 records to prevent infinite bloat.
|
||||
* **Ephemeral Data Cleanup:** Automatically expiring and purging sticky_notes.json that are older than 24 hours.
|
||||
|
||||
## 8. The Gate System (Push Safety Verification)
|
||||
The binary includes dedicated CLI subcommands (gate set and gate verify) that interact with a persistent gates.json store to enforce safety policies (like ensuring tests pass before a git push).
|
||||
* **mcp-memory gate set**: Records an authorization status (authorized, blocked, or pending) for a specific target and namespace, alongside optional failure reasons and parameters.
|
||||
* **mcp-memory gate verify**: Evaluates a pending action against the gate store. It returns exit code if authorized, 1 if explicitly blocked, and 2 if no gate record exists. It also supports a --consume flag to immediately revoke the authorization after a successful check.
|
||||
|
||||
## 9. Operational Configuration & Paths
|
||||
The physical storage location of the knowledge graph and all persistent stores is strictly controlled by the MCP_MEMORY_STORE_DIR environment variable.
|
||||
* **Default Path:** If not set, the daemon defaults to ~/.gemini/mcp_memory.
|
||||
* **Port Binding:** The Axum HTTP server strictly binds to .0.0.0:3000.
|
||||
|
||||
## 10. Exposed MCP Capabilities (Tools)
|
||||
The server implements the Model Context Protocol (MCP) by exposing a vast suite of tools via the JSON-RPC interface, categorized broadly into:
|
||||
* **Graph Management:** create_entities, create_relations, merge_entities,
|
||||
ead_graph, etc.
|
||||
* **Task & Context Tracking:** add_task, add_sticky_note, add_session_summary, etc.
|
||||
* **Engineering & DevOps:** log_code_change, log_error_fix, log_tech_debt, add_pr_checklist_item.
|
||||
* **Environment & Workspaces:**
|
||||
egister_environment, save_context_workspace, pin_file.
|
||||
|
||||
## 11. Concurrency & Thread Safety
|
||||
With the introduction of the Dual Transport System, the daemon must safely handle simultaneous read/write requests from both Stdio (Windows agy) and HTTP (WSL agy) clients.
|
||||
* **Shared State:** The entire server operates on a cloned Arc<MemoryState>.
|
||||
* **Locking Mechanism:** The unified Knowledge Graph and modular Store<T> components are protected by RwLock primitives.
|
||||
* **Safe Mutation:** Store modifications utilize closure-based modify(|store| { ... }) methods to ensure locks are safely acquired, mutations applied, and file writes executed sequentially without deadlocking the asynchronous Tokio runtimes.
|
||||
|
||||
## 12. Lifecycle & Startup Management (Leader/Stub Architecture)
|
||||
The mcp-memory daemon employs a self-healing, morphing Server/Client paradigm to allow multiple concurrent agy sessions without port collisions or file lock contention.
|
||||
|
||||
### The Windows Morphing Daemon (Leader vs. Stub)
|
||||
Every time a Windows agy session starts, it blindly spawns mcp-memory.exe as a stdio subprocess. On startup, the binary evaluates its environment:
|
||||
* **The Leader (Server):** The binary attempts to bind to .0.0.0:3000. If successful, it becomes the Leader. It initializes the file locks, starts the Axum HTTP server, and natively processes the stdio input from its parent agy session.
|
||||
* **The Stub (Proxy Client):** If port 3000 is already in use, the binary instantly morphs into a Stub. Instead of crashing, it acts as a proxy: reading JSON-RPC requests from its parent agy's stdin, forwarding them via HTTP POST to the Leader at http://localhost:3000/messages, and writing the HTTP responses back to stdout. It also pipes the Leader's /sse stream directly to stdout. **Graceful Shutdown:** The stub detects EOF on its stdin pipe when the CLI exits, safely closing all HTTP streams and terminating cleanly.
|
||||
|
||||
### Automatic Failover & Self-Healing
|
||||
If the terminal hosting the Leader is closed, the Leader process dies and releases port 3000.
|
||||
* **The Race:** The active Stubs immediately detect the dropped SSE HTTP connection. Instead of crashing, they race to bind to port 3000.
|
||||
* **Resolution:** The first Stub to bind successfully promotes itself to the new Leader (taking over file locks, WAL management, and the Axum server). The losing Stubs recognize the new Leader, reconnect their HTTP streams, and resume proxying. This results in near-zero downtime for the user's agy clients.
|
||||
|
||||
### WSL (Linux) Client Lifecycle (Permanent Stub)
|
||||
To provide a seamless experience without complex configuration drift, the WSL environment utilizes a native Linux binary (`/home/riz/.local/bin/mcp-memory`) that acts as a **Permanent Stub**.
|
||||
* **Transparent Proxying:** The WSL `agy` CLI spawns this Linux binary via standard `stdio`. The binary immediately proxies all `stdio` JSON-RPC requests over HTTP to the Windows Leader at `http://127.0.0.1:3000/messages`, handling SSE streams transparently.
|
||||
* **Startup via Interop:** Upon launch, the Linux Stub pings the Windows host. If port 3000 is dead, the Linux binary automatically executes WSL interop (`cmd.exe /c start /B C:\Users\reazul.ashraf\.local\bin\mcp-memory.exe`) to silently wake up the Windows Leader before commencing the proxy loop.
|
||||
* **Zero I/O Penalty:** This ensures the Linux binary never directly touches the Windows NTFS files, reserving all heavy disk operations for the native Windows host.
|
||||
## 13. Cargo Workspace & Binary Artifacts
|
||||
To optimize for different environments, the codebase is structured as a Cargo Workspace containing two distinct crates:
|
||||
|
||||
### 1. mcp-memory-server (The "Full-Fat" Daemon)
|
||||
* **Path:** server/
|
||||
* **Size/Complexity:** Heavy (contains Axum, MCP SDK, JSON parsing, Tokio runtime).
|
||||
* **Role:** This is the primary background daemon. It binds to .0.0.0:3000, holds file locks, and manages the graph.
|
||||
* **Windows Behavior:** It employs the Morphing logic. The first instance becomes the Leader (Server). Subsequent instances spawned by agy detect the port is in use and seamlessly morph into Stubs, proxying stdio to HTTP.
|
||||
|
||||
### 2. mcp-memory-stub (The Ultra-Lightweight Proxy)
|
||||
* **Path:** stub/
|
||||
* **Size/Complexity:** Extremely light (only relies on
|
||||
eqwest and okio).
|
||||
* **Role:** A dedicated, OS-agnostic proxy binary used strictly for routing stdio JSON-RPC traffic over HTTP to a remote Leader. Windows clients can point directly to this binary to bypass loading the heavy Server daemon into memory.
|
||||
* **WSL Behavior:** Compiled as a Linux native binary (x86_64-unknown-linux-musl). When executed by WSL agy, it acts as a transparent proxy to http://127.0.0.1:3000. It can also execute wake_cmd (e.g., WSL interop) to silently wake the Windows host if the Leader is offline.
|
||||
+1
-1
@@ -3,7 +3,7 @@
|
||||
This server provides an OS-agnostic, persistent Knowledge Graph and working memory. The following guidelines dictate **when** and **how** to best utilize these tools.
|
||||
|
||||
## 1. Ephemeral Working Memory (Sticky Notes)
|
||||
- **Tools:** dd_sticky_note, ␍ead_sticky_notes
|
||||
- **Tools:** add_sticky_note, ␍ead_sticky_notes
|
||||
- **When to use:**
|
||||
- When you need to remember a specific context, broken build state, or pending task across conversation boundaries (e.g., "We are in the middle of refactoring main.rs, next step is testing").
|
||||
- When saving short-term scratchpad thoughts that don't belong in the permanent knowledge graph.
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "mcp-memory-server"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1.92"
|
||||
axum = "0.8"
|
||||
clap = { version = "4.6.6", features = ["derive"] }
|
||||
dirs = "6.0.0"
|
||||
futures-util = "0.3.34"
|
||||
glob = "0.3.4"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls"] }
|
||||
rust-mcp-axum = "2.0.0"
|
||||
rust-mcp-sdk = "2.0.0"
|
||||
serde = { version = "1.0.229", features = ["derive"] }
|
||||
serde_json = "1.0.151"
|
||||
strsim = "0.11.1"
|
||||
tokio = { version = "1.53.1", features = ["full"] }
|
||||
tokio-util = { version = "0.7.19", features = ["io"] }
|
||||
uuid = { version = "1.26.0", features = ["v4"] }
|
||||
File renamed without changes.
@@ -9,9 +9,9 @@ use crate::models::*;
|
||||
use crate::state::MemoryState;
|
||||
use crate::store::Store;
|
||||
use rust_mcp_sdk::{
|
||||
McpServer, ServerDetails, StdioTransport, ToMcpServerHandler, TransportOptions,
|
||||
ServerDetails, ToMcpServerHandler,
|
||||
error::SdkResult,
|
||||
mcp_server::{McpServerOptions, server_runtime},
|
||||
|
||||
schema::{Implementation, ServerCapabilities, ServerCapabilitiesTools},
|
||||
};
|
||||
use std::fs;
|
||||
@@ -28,6 +28,8 @@ use std::collections::HashMap;
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Option<Commands>,
|
||||
#[arg(long)]
|
||||
target: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
@@ -103,6 +105,8 @@ async fn reconcile_worker(state: Arc<MemoryState>) {
|
||||
}
|
||||
|
||||
|
||||
use rust_mcp_axum::{create_axum_server, AxumServerOptions};
|
||||
|
||||
fn run_server(state: Arc<MemoryState>) -> SdkResult<()> {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
rt.block_on(async {
|
||||
@@ -129,19 +133,68 @@ fn run_server(state: Arc<MemoryState>) -> SdkResult<()> {
|
||||
instructions: None,
|
||||
meta: None,
|
||||
};
|
||||
let server = server_runtime::create_server(McpServerOptions {
|
||||
transport: StdioTransport::new(TransportOptions::default())?,
|
||||
handler: MemoryHandler { state }.to_mcp_server_handler(),
|
||||
|
||||
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,
|
||||
message_observer: None,
|
||||
});
|
||||
server.start().await
|
||||
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
|
||||
}),
|
||||
);
|
||||
|
||||
println!("MCP Memory Server running on http://0.0.0.0:3000/sse");
|
||||
server.start().await.unwrap();
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
mod proxy;
|
||||
|
||||
fn main() -> SdkResult<()> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
loop {
|
||||
if let Err(_) = std::net::TcpListener::bind("0.0.0.0:3000") {
|
||||
// Port in use, become a stub proxy!
|
||||
let target_url = cli.target.as_deref().unwrap_or("http://127.0.0.1:3000");
|
||||
match proxy::run_proxy(target_url) {
|
||||
Ok(true) => {
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
continue; // Leader died, race to bind 3000
|
||||
}
|
||||
Ok(false) => return Ok(()), // Stdin closed, user exited
|
||||
Err(_) => std::thread::sleep(std::time::Duration::from_millis(1000)),
|
||||
}
|
||||
} else {
|
||||
break; // Proceed as Leader
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
// This shouldn't be executed on linux natively anymore due to workspace split,
|
||||
// but keeping it as a fallback.
|
||||
let target_url = cli.target.as_deref().unwrap_or("http://host.docker.internal:3000");
|
||||
let _ = proxy::run_proxy(target_url);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
@@ -155,7 +208,6 @@ fn main() -> SdkResult<()> {
|
||||
|
||||
let state = Arc::new(MemoryState {
|
||||
master_path: base.join("knowledge_graph_master.json"),
|
||||
delta_path: base.join(format!("delta_{}.json", uuid::Uuid::new_v4())),
|
||||
session_graph: RwLock::new(KnowledgeGraph::default()),
|
||||
base_dir: base.clone(),
|
||||
master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)),
|
||||
File renamed without changes.
@@ -0,0 +1,93 @@
|
||||
use rust_mcp_sdk::error::SdkResult;
|
||||
use tokio_util::io::StreamReader;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use futures_util::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
pub fn run_proxy(target_url: &str) -> SdkResult<bool> {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
rt.block_on(async {
|
||||
let client = reqwest::Client::builder().build().unwrap();
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
|
||||
let resp = match client.get(&sse_url).send().await {
|
||||
Ok(r) => r,
|
||||
Err(_) => return Ok(true), // Connection failed (Leader is dead)
|
||||
};
|
||||
|
||||
let post_url = Arc::new(RwLock::new(format!("{}/messages", target_url)));
|
||||
let post_url_clone = Arc::clone(&post_url);
|
||||
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
|
||||
|
||||
let tx_clone = tx.clone();
|
||||
let client_clone = client.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let stdin = std::io::stdin();
|
||||
let mut handle = stdin.lock();
|
||||
let mut buffer = String::new();
|
||||
while let Ok(bytes) = std::io::BufRead::read_line(&mut handle, &mut buffer) {
|
||||
if bytes == 0 { break; }
|
||||
let body = buffer.clone();
|
||||
buffer.clear();
|
||||
let client = client_clone.clone();
|
||||
let url_arc = Arc::clone(&post_url_clone);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut url = "".to_string();
|
||||
for _ in 0..50 {
|
||||
let u = url_arc.read().await.clone();
|
||||
if u.contains("sessionId") {
|
||||
url = u;
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||
}
|
||||
if url.is_empty() {
|
||||
url = url_arc.read().await.clone();
|
||||
}
|
||||
let _ = client.post(&url).header("Content-Type", "application/json").body(body).send().await;
|
||||
});
|
||||
}
|
||||
let _ = tx_clone.blocking_send(false); // Stdin EOF
|
||||
});
|
||||
|
||||
let target_url = target_url.to_string();
|
||||
let tx_clone2 = tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
|
||||
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
|
||||
let mut line = String::new();
|
||||
let mut is_message = false;
|
||||
let mut is_endpoint = false;
|
||||
|
||||
while let Ok(bytes) = reader.read_line(&mut line).await {
|
||||
if bytes == 0 { break; }
|
||||
let trimmed = line.trim();
|
||||
if trimmed.starts_with("event: message") {
|
||||
is_message = true;
|
||||
is_endpoint = false;
|
||||
} else if trimmed.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
is_message = false;
|
||||
} else if trimmed.starts_with("data: ") {
|
||||
if is_message {
|
||||
println!("{}", &trimmed[6..]);
|
||||
is_message = false;
|
||||
} else if is_endpoint {
|
||||
let ep = &trimmed[6..];
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, ep);
|
||||
is_endpoint = false;
|
||||
}
|
||||
}
|
||||
line.clear();
|
||||
}
|
||||
let _ = tx_clone2.send(true).await; // Stream dropped (Leader dead)
|
||||
});
|
||||
|
||||
let dropped = rx.recv().await.unwrap_or(true);
|
||||
Ok(dropped)
|
||||
})
|
||||
}
|
||||
@@ -9,7 +9,6 @@ use std::time::{Duration, SystemTime};
|
||||
pub struct MemoryState {
|
||||
pub base_dir: PathBuf,
|
||||
pub master_path: PathBuf,
|
||||
pub delta_path: PathBuf,
|
||||
pub session_graph: RwLock<KnowledgeGraph>,
|
||||
pub master_cache: RwLock<(KnowledgeGraph, SystemTime)>,
|
||||
pub ledger: Store<Vec<CodeChange>>,
|
||||
@@ -83,24 +82,28 @@ impl MemoryState {
|
||||
|
||||
pub fn get_full_graph(&self) -> KnowledgeGraph {
|
||||
let mut master = self.read_master_cached();
|
||||
let pattern = format!("{}/delta_*.json", self.base_dir.display());
|
||||
if let Ok(paths) = glob::glob(&pattern) {
|
||||
for path in paths.flatten() {
|
||||
if path == self.delta_path {
|
||||
let session_graph = self.session_graph.read().unwrap();
|
||||
Self::merge_graphs(&mut master, &session_graph);
|
||||
} else {
|
||||
let d: KnowledgeGraph = read_json_file(&path);
|
||||
let wal_path = self.base_dir.join("wal.jsonl");
|
||||
if let Ok(content) = std::fs::read_to_string(&wal_path) {
|
||||
for line in content.lines() {
|
||||
if let Ok(d) = serde_json::from_str::<KnowledgeGraph>(line) {
|
||||
Self::merge_graphs(&mut master, &d);
|
||||
}
|
||||
}
|
||||
}
|
||||
let session_graph = self.session_graph.read().unwrap();
|
||||
Self::merge_graphs(&mut master, &session_graph);
|
||||
master
|
||||
}
|
||||
pub fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
|
||||
let mut session_graph = self.session_graph.write().unwrap();
|
||||
update_fn(&mut session_graph);
|
||||
let _ = write_json_atomic(&self.delta_path, &*session_graph);
|
||||
let wal_path = self.base_dir.join("wal.jsonl");
|
||||
if let Ok(payload) = serde_json::to_string(&*session_graph) {
|
||||
if let Ok(mut file) = std::fs::OpenOptions::new().create(true).append(true).open(&wal_path) {
|
||||
use std::io::Write;
|
||||
let _ = writeln!(file, "{}", payload);
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn apply_sync_write<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
|
||||
let lock_path = self.base_dir.join("master.lock");
|
||||
@@ -121,12 +124,8 @@ impl MemoryState {
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
let mut master = self.get_full_graph();
|
||||
let pattern = format!("{}/delta_*.json", self.base_dir.display());
|
||||
if let Ok(paths) = glob::glob(&pattern) {
|
||||
for path in paths.flatten() {
|
||||
let _ = fs::remove_file(&path);
|
||||
}
|
||||
}
|
||||
let wal_path = self.base_dir.join("wal.jsonl");
|
||||
let _ = fs::remove_file(&wal_path);
|
||||
*self.session_graph.write().unwrap() = KnowledgeGraph::default();
|
||||
update_fn(&mut master);
|
||||
let _ = write_json_atomic(&self.master_path, &master);
|
||||
File renamed without changes.
File renamed without changes.
@@ -0,0 +1,11 @@
|
||||
[package]
|
||||
name = "mcp-memory-stub"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
clap = { version = "4.6.6", features = ["derive"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls"] }
|
||||
tokio = { version = "1.53.1", features = ["full"] }
|
||||
tokio-util = { version = "0.7.19", features = ["io"] }
|
||||
futures-util = "0.3.34"
|
||||
@@ -0,0 +1,144 @@
|
||||
use clap::Parser;
|
||||
use futures_util::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_util::io::StreamReader;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "mcp-memory-stub")]
|
||||
struct Cli {
|
||||
#[arg(long, default_value = "http://localhost:3000")]
|
||||
target: String,
|
||||
#[arg(long)]
|
||||
wake_cmd: Option<String>,
|
||||
}
|
||||
|
||||
fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error>> {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
rt.block_on(async {
|
||||
let client = reqwest::Client::builder().build()?;
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
|
||||
eprintln!("[PROXY] Connecting to SSE: {}", sse_url);
|
||||
let resp = match client.get(&sse_url).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
eprintln!("[PROXY] SSE connection failed: {}", e);
|
||||
return Ok(true);
|
||||
}
|
||||
};
|
||||
|
||||
let post_url = Arc::new(RwLock::new(format!("{}/messages", target_url)));
|
||||
let post_url_clone = Arc::clone(&post_url);
|
||||
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
|
||||
|
||||
let tx_clone = tx.clone();
|
||||
let client_clone = client.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let stdin = std::io::stdin();
|
||||
let mut handle = stdin.lock();
|
||||
let mut buffer = String::new();
|
||||
while let Ok(bytes) = std::io::BufRead::read_line(&mut handle, &mut buffer) {
|
||||
if bytes == 0 { break; }
|
||||
let body = buffer.clone();
|
||||
buffer.clear();
|
||||
eprintln!("[PROXY] Stdin read: {}", body.trim());
|
||||
let client = client_clone.clone();
|
||||
let url_arc = Arc::clone(&post_url_clone);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut url = "".to_string();
|
||||
for _ in 0..50 {
|
||||
let u = url_arc.read().await.clone();
|
||||
if u.contains("sessionId") {
|
||||
url = u;
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||
}
|
||||
if url.is_empty() {
|
||||
url = url_arc.read().await.clone();
|
||||
}
|
||||
eprintln!("[PROXY] POSTing to {}", url);
|
||||
let res = client.post(&url).header("Content-Type", "application/json").body(body).send().await;
|
||||
eprintln!("[PROXY] POST response: {:?}", res.map(|r| r.status()));
|
||||
});
|
||||
}
|
||||
eprintln!("[PROXY] Stdin closed");
|
||||
let _ = tx_clone.blocking_send(false);
|
||||
});
|
||||
|
||||
let target_url = target_url.to_string();
|
||||
let tx_clone2 = tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
|
||||
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
|
||||
let mut line = String::new();
|
||||
let mut is_message = false;
|
||||
let mut is_endpoint = false;
|
||||
while let Ok(bytes) = reader.read_line(&mut line).await {
|
||||
if bytes == 0 { break; }
|
||||
let trimmed = line.trim();
|
||||
eprintln!("[PROXY] SSE read: {}", trimmed);
|
||||
if trimmed.starts_with("event: message") {
|
||||
is_message = true;
|
||||
is_endpoint = false;
|
||||
} else if trimmed.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
is_message = false;
|
||||
} else if trimmed.starts_with("data: ") {
|
||||
if is_message {
|
||||
println!("{}", &trimmed[6..]);
|
||||
is_message = false;
|
||||
} else if is_endpoint {
|
||||
let ep = &trimmed[6..];
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, ep);
|
||||
eprintln!("[PROXY] Endpoint updated: {}", *p);
|
||||
is_endpoint = false;
|
||||
}
|
||||
}
|
||||
line.clear();
|
||||
}
|
||||
eprintln!("[PROXY] SSE stream closed");
|
||||
let _ = tx_clone2.send(true).await;
|
||||
});
|
||||
|
||||
let dropped = rx.recv().await.unwrap_or(true);
|
||||
Ok(dropped)
|
||||
})
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let cli = Cli::parse();
|
||||
loop {
|
||||
if let Err(_) = std::net::TcpStream::connect(cli.target.replace("http://", "").replace("https://", "")) {
|
||||
eprintln!("[PROXY] Target {} offline, sleeping", cli.target);
|
||||
|
||||
if let Some(ref cmd) = cli.wake_cmd {
|
||||
eprintln!("[PROXY] Executing wake command...");
|
||||
let parts: Vec<&str> = cmd.split_whitespace().collect();
|
||||
if !parts.is_empty() {
|
||||
let _ = std::process::Command::new(parts[0])
|
||||
.args(&parts[1..])
|
||||
.spawn();
|
||||
}
|
||||
}
|
||||
|
||||
std::thread::sleep(std::time::Duration::from_secs(1));
|
||||
}
|
||||
match run_proxy(&cli.target) {
|
||||
Ok(true) => {
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
Ok(false) => {
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
std::thread::sleep(std::time::Duration::from_millis(1000));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user