diff --git a/CHANGELOG.md b/CHANGELOG.md index a45aa574b..a7f439bc6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). (`zeph-a2a/src/server/router.rs`) so `rate_limit_middleware` wraps `auth_middleware`, guaranteeing every request — including failed-auth ones — increments the counter before the auth check runs. +- `zeph-mcp`: `tool_list_locked` entries could outlive their server, leaking orphaned + locks (#6139, follow-up to #6118). `remove_server` and `shutdown_all_shared` + (`crates/zeph-mcp/src/manager/server.rs`) cleared `server_tools`, `server_trust`, + `server_fingerprints`, and `last_refresh` on disconnect but never removed the + corresponding `tool_list_locked` entry; and `handle_connect_result` + (`crates/zeph-mcp/src/manager/connect.rs`), used by the `connect_all`/ + `connect_oauth_deferred` path, released the lock on connection or `list_tools` + failure but not when the pre-connect probe blocked the connection — asymmetric with + the equivalent `add_server` path (`probe_or_cleanup`), which already cleaned up + correctly. All four cleanup sites now release `tool_list_locked` consistently. - `zeph-mcp`: `lock_tool_list` hardening silently exempted OAuth-transport MCP servers from post-attestation tool-injection protection (#6118). `tool_list_locked` was only populated by the two non-OAuth connection paths (`spawn_non_oauth_connections`, diff --git a/crates/zeph-mcp/src/manager/connect.rs b/crates/zeph-mcp/src/manager/connect.rs index e9a20a690..c047f19ec 100644 --- a/crates/zeph-mcp/src/manager/connect.rs +++ b/crates/zeph-mcp/src/manager/connect.rs @@ -626,6 +626,8 @@ impl McpManager { // Phase 1: run pre-connect probe if configured. if let Err(e) = self.run_probe(&server_id, &client).await { client.shutdown().await; + // Probe blocked — remove lock so the server is not left permanently locked. + self.tool_list_locked.remove(&server_id); return fail(format!("{e:#}")); } diff --git a/crates/zeph-mcp/src/manager/server.rs b/crates/zeph-mcp/src/manager/server.rs index a2a3496d8..6a0037c83 100644 --- a/crates/zeph-mcp/src/manager/server.rs +++ b/crates/zeph-mcp/src/manager/server.rs @@ -322,6 +322,7 @@ impl McpManager { self.server_trust.write().await.remove(server_id); self.server_fingerprints.write().await.remove(server_id); self.last_refresh.remove(server_id); + self.tool_list_locked.remove(server_id); // Release the serialization lock before the potentially slow shutdown call. drop(add_remove_guard); client.shutdown().await; @@ -400,6 +401,7 @@ impl McpManager { self.server_tools.write().await.clear(); self.server_fingerprints.write().await.clear(); self.last_refresh.clear(); + self.tool_list_locked.clear(); for (id, client) in drained { tracing::info!(server_id = id, "shutting down MCP client"); if tokio::time::timeout(Duration::from_secs(5), client.shutdown()) diff --git a/crates/zeph-mcp/src/manager/tests.rs b/crates/zeph-mcp/src/manager/tests.rs index 871481ee0..65007f9dd 100644 --- a/crates/zeph-mcp/src/manager/tests.rs +++ b/crates/zeph-mcp/src/manager/tests.rs @@ -452,6 +452,16 @@ impl McpManager { async fn has_server_tools_for_test(&self, server_id: &str) -> bool { self.server_tools.read().await.contains_key(server_id) } + + /// Insert a lock entry directly, bypassing the real connect path. + fn inject_tool_list_locked_for_test(&self, server_id: &str) { + self.tool_list_locked.insert(server_id.to_owned(), ()); + } + + /// Return `true` if `tool_list_locked` still contains an entry for `server_id`. + fn is_tool_list_locked_for_test(&self, server_id: &str) -> bool { + self.tool_list_locked.contains_key(server_id) + } } // --- commit_added_server --- @@ -1651,6 +1661,94 @@ async fn concurrent_remove_server_calls_are_serialized() { ); } +// ── tool_list_locked orphan cleanup (#6139) ──────────────────────────────────────────────── + +/// `remove_server` must clean up the `tool_list_locked` entry it may hold. +/// +/// Before the fix, `remove_server` never removed the server from `tool_list_locked`, +/// so a server connected with `lock_tool_list = true` and later removed at runtime +/// stayed permanently locked if an ID with the same name was ever reconnected. +#[tokio::test] +async fn remove_server_cleans_up_tool_list_locked_when_client_present() { + let mgr = McpManager::new(vec![], vec![], PolicyEnforcer::new(vec![])); + let entry = make_entry("locked-srv"); + + // Simulate the post-connect state: a real client entry plus a lock, as + // `connect_and_list_tools` would have set up for a `lock_tool_list = true` server. + let client = McpClient::new_disconnected_for_test("locked-srv"); + mgr.commit_added_server(&entry, client, vec![], None) + .await + .expect("commit must succeed"); + mgr.inject_tool_list_locked_for_test("locked-srv"); + assert!( + mgr.is_tool_list_locked_for_test("locked-srv"), + "precondition: lock entry must exist before removal" + ); + + mgr.remove_server("locked-srv") + .await + .expect("remove must succeed for a server with a real client entry"); + + assert!( + !mgr.is_tool_list_locked_for_test("locked-srv"), + "tool_list_locked entry must be removed by remove_server" + ); +} + +/// `shutdown_all_shared` must clear the entire `tool_list_locked` map. +/// +/// Before the fix, `tool_list_locked` was never cleared on shutdown, so any locked +/// server ID would appear pre-locked if the manager were rebuilt and reconnected +/// with the same server IDs still tracked by a surviving `Arc>` clone. +#[tokio::test] +async fn shutdown_all_shared_clears_tool_list_locked() { + let mgr = McpManager::new(vec![], vec![], PolicyEnforcer::new(vec![])); + mgr.inject_tool_list_locked_for_test("srv1"); + mgr.inject_tool_list_locked_for_test("srv2"); + assert!(mgr.is_tool_list_locked_for_test("srv1")); + assert!(mgr.is_tool_list_locked_for_test("srv2")); + + mgr.shutdown_all_shared().await; + + assert!( + !mgr.is_tool_list_locked_for_test("srv1"), + "tool_list_locked must be cleared by shutdown_all_shared" + ); + assert!( + !mgr.is_tool_list_locked_for_test("srv2"), + "tool_list_locked must be cleared by shutdown_all_shared" + ); +} + +/// `connect_all` must not leave an orphaned `tool_list_locked` entry when a +/// `lock_tool_list = true` server fails to connect. +/// +/// `spawn_non_oauth_connections` inserts the lock *before* spawning the connection +/// task (MF-2: no window for a refresh event to slip through). `handle_connect_result` +/// is responsible for removing it again on every failure branch — including the +/// pre-connect-probe-blocked branch added by this fix (`connect.rs` `handle_connect_result`, +/// "Probe blocked" arm). This test drives the connection-failure branch (nonexistent +/// binary), which is reachable without a live MCP server and exercises the same +/// insert-before-spawn / remove-on-any-failure invariant the probe-blocked arm relies on. +#[tokio::test] +async fn connect_all_does_not_orphan_tool_list_locked_on_connect_failure() { + let mgr = McpManager::new( + vec![make_entry("locked-fail")], + vec![], + PolicyEnforcer::new(vec![]), + ) + .with_lock_tool_list(true); + + let (tools, outcomes) = mgr.connect_all().await; + + assert!(tools.is_empty()); + assert!(outcomes.iter().all(|o| !o.connected)); + assert!( + !mgr.is_tool_list_locked_for_test("locked-fail"), + "tool_list_locked must not retain a stale entry after connect_all fails to connect" + ); +} + /// `commit_added_server` must return `ServerAlreadyConnected` when called for a /// server ID that already has a client entry. ///