From 3b3bb36b323e52e21640823bd8720e175ab09509 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Mon, 12 Oct 2026 06:11:36 +0100 Subject: [PATCH] refactor(server): strictly enforce CLI port argument for bind, dropping MCP_PORT env var fallback --- justfile | 12 +++++----- server/src/bin/test_server_shutdown.rs | 4 +++- server/src/lib.rs | 33 ++++++++++++++++++-------- server/src/router.rs | 23 ++++++++++++++++++ stub/tests/e2e.rs | 2 ++ 5 files changed, 57 insertions(+), 17 deletions(-) diff --git a/justfile b/justfile index 132c253..9fb97c6 100644 --- a/justfile +++ b/justfile @@ -5,10 +5,10 @@ default: @just --list # Build and deploy everything across Windows and WSL -all port: build-all (deploy-all port) +all port: (build-all port) (deploy-all port) # Build all binaries across Windows and WSL -build-all: build-win build-wsl +build-all port: (build-win port) build-wsl @Write-Host "Build complete across both OS boundaries." -ForegroundColor Green # Deploy all binaries and agent rules across Windows and WSL @@ -16,7 +16,7 @@ deploy-all port: (deploy-win port) deploy-wsl @Write-Host "Deployment complete across both OS boundaries." -ForegroundColor Green # Build Windows-native binaries -build-win: build-ui build-server build-stub-win build-nvim-win +build-win port: build-ui (build-server port) build-stub-win build-nvim-win # Deploy Windows-native binaries and rules deploy-win port: (deploy-server port) deploy-stub-win deploy-nvim-win deploy-rules-win @@ -68,7 +68,7 @@ check-ui: deno check server/src/dashboard.ts # Build Windows-native server binary in release mode (ensuring fresh UI build) -build-server: build-ui +build-server port: build-ui @Write-Host "Building Windows Server..." -ForegroundColor Cyan cargo build --release -p mcp-memory-server @@ -179,7 +179,7 @@ improve-tests file="server/src/ollama.rs": # ========================================================= # Complete lifecycle: stop, build, deploy, start, and verify Windows server -all-server-win port: (stop port) build-server (deploy-server port) (start port) (verify port) +all-server-win port: (stop port) (build-server port) (deploy-server port) (start port) (verify port) # Complete lifecycle: build and deploy Windows stub all-stub-win: build-stub-win deploy-stub-win @@ -194,7 +194,7 @@ all-stub-wsl: build-stub deploy-stub all-nvim-wsl: build-nvim deploy-nvim # Complete lifecycle: stop, build, deploy, start, and verify WSL server -all-server-wsl port: (stop port) build-server (deploy-server port) (start port) (verify port) +all-server-wsl port: (stop port) (build-server port) (deploy-server port) (start port) (verify port) # Install git hooks (post-commit reconciliation and pre-push gatekeeper) install-git-hooks: diff --git a/server/src/bin/test_server_shutdown.rs b/server/src/bin/test_server_shutdown.rs index 79773d1..564e44c 100644 --- a/server/src/bin/test_server_shutdown.rs +++ b/server/src/bin/test_server_shutdown.rs @@ -36,7 +36,9 @@ async fn main() { std::fs::create_dir_all(&temp_dir).unwrap(); let mut child = Command::new(exe_path) - .env("MCP_PORT", port.to_string()) + .arg("--port") + .arg(port.to_string()) + .env("MCP_PORT", port.to_string()) // Keep for backward compatibility of clients .env("MCP_UDP_PORT1", udp1.to_string()) .env("MCP_UDP_PORT2", udp2.to_string()) .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap()) diff --git a/server/src/lib.rs b/server/src/lib.rs index 3193658..4695296 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -45,6 +45,9 @@ pub struct Cli { /// Send a shutdown request to the existing server and wait for it to exit #[arg(long)] pub restart: bool, + /// Port to bind the server to (or to connect to for exit/gate commands) + #[arg(short, long)] + pub port: u16, } #[derive(Subcommand)] @@ -343,7 +346,7 @@ pub async fn memory_consolidation_worker(state: Arc) { } } -pub async fn run_server(state: Arc) -> Result<(), Box> { +pub async fn run_server(state: Arc, port: u16) -> Result<(), Box> { let state_for_index = Arc::clone(&state); tokio::spawn(async move { state_for_index.rebuild_index().await; @@ -586,12 +589,11 @@ pub async fn run_server(state: Arc) -> Result<(), Box Result<(), Box> { let base = PathBuf::from(base_dir); if cli.exit { - let port = crate::config::get_env("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); + let port = cli.port; let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default(); let rt = tokio::runtime::Runtime::new()?; rt.block_on(async { @@ -684,7 +686,7 @@ pub fn run_cli() -> Result<(), Box> { } if let Some(Commands::Gate { subcmd }) = cli.command { - let port = crate::config::get_env("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); + let port = cli.port; let rt = tokio::runtime::Runtime::new()?; match subcmd { GateCommands::Set { @@ -777,8 +779,9 @@ pub fn run_cli() -> Result<(), Box> { rt.block_on(async { let state = Arc::new(MemoryState::new(&base.to_string_lossy())); - - if let Err(e) = run_server(state).await { + + let port = cli.port; + if let Err(e) = run_server(state, port).await { tracing::error!("Server error: {}", e); } }); @@ -793,16 +796,19 @@ mod tests { #[test] fn test_cli_parsing_default() { - let cli = Cli::try_parse_from(["mcp-memory-server"]).unwrap(); + let cli = Cli::try_parse_from(["mcp-memory-server", "--port", "3000"]).unwrap(); assert!(cli.command.is_none()); assert!(!cli.exit); assert!(!cli.restart); + assert_eq!(cli.port, 3000); } #[test] fn test_cli_parsing_exit_and_target() { let cli = Cli::try_parse_from([ "mcp-memory-server", + "--port", + "3000", "--exit", "--target", "http://localhost:3000", @@ -810,12 +816,15 @@ mod tests { .unwrap(); assert!(cli.exit); assert_eq!(cli.target.as_deref(), Some("http://localhost:3000")); + assert_eq!(cli.port, 3000); } #[test] fn test_cli_parsing_gate_set() { let cli = Cli::try_parse_from([ "mcp-memory-server", + "--port", + "3000", "gate", "set", "--action", @@ -828,6 +837,7 @@ mod tests { ]) .unwrap(); + assert_eq!(cli.port, 3000); if let Some(Commands::Gate { subcmd: GateCommands::Set { @@ -852,6 +862,8 @@ mod tests { fn test_cli_parsing_gate_verify() { let cli = Cli::try_parse_from([ "mcp-memory-server", + "--port", + "3000", "gate", "verify", "--action", @@ -862,6 +874,7 @@ mod tests { ]) .unwrap(); + assert_eq!(cli.port, 3000); if let Some(Commands::Gate { subcmd: GateCommands::Verify { @@ -933,7 +946,7 @@ mod tests { } let server_handle = tokio::spawn(async move { - let _ = run_server(state).await; + let _ = run_server(state, port).await; }); tokio::time::sleep(Duration::from_millis(300)).await; diff --git a/server/src/router.rs b/server/src/router.rs index 1e10097..58df78b 100644 --- a/server/src/router.rs +++ b/server/src/router.rs @@ -1180,6 +1180,23 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) -> .unwrap_or_default(); ("Open Nodes", names) } + "search_nodes" => { + let query = args.get("query").and_then(|v| v.as_str()).unwrap_or(""); + ("Search Nodes", query.to_string()) + } + "get_subgraph" | "summarize_subgraph" => { + let name = args.get("entity_name").and_then(|v| v.as_str()).unwrap_or(""); + ("Subgraph", name.to_string()) + } + "query_graph_path" => { + let source = args.get("source").and_then(|v| v.as_str()).unwrap_or("?"); + let target = args.get("target").and_then(|v| v.as_str()).unwrap_or("?"); + ("Graph Path", format!("{} -> {}", source, target)) + } + "condense_entity" => { + let name = args.get("entity_name").and_then(|v| v.as_str()).unwrap_or(""); + ("Condense Entity", name.to_string()) + } _ => return name.to_string(), }; if detail.is_empty() { @@ -1221,6 +1238,12 @@ mod tests { let msg3 = format_tool_activity_description("read_graph", &json!({})); assert_eq!(msg3, "read_graph"); + + let msg4 = format_tool_activity_description( + "search_nodes", + &serde_json::json!({ "query": "database connection" }), + ); + assert_eq!(msg4, "Search Nodes: database connection"); } #[tokio::test] diff --git a/stub/tests/e2e.rs b/stub/tests/e2e.rs index 771c493..64bf7e4 100644 --- a/stub/tests/e2e.rs +++ b/stub/tests/e2e.rs @@ -59,6 +59,8 @@ async fn test_full_system_e2e_performance() { // 1. Start Server let _server = ChildGuard( Command::new(&server_exe) + .arg("--port") + .arg(test_port) .env("MCP_PORT", test_port) .env("RUST_LOG", "debug") .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())