refactor(mcp): SharedConnection — remove mem::forget leak; tighten classification

Self-review of 0362660 surfaced five issues. This commit closes all of them.

M-R1 (critical) — `mem::forget(RunningService)` leaked the child
process AND the background tokio task. rmcp's `RunningService` holds
a `DropGuard` that cancels the cancellation_token on drop; forgetting
the value bypassed that, so the spawned MCP child process kept
running even after dirge exited. Confirmed by reading
`rmcp-1.7.0/src/service.rs:621-632` (the Drop impl is what reaps
the child).

Fix: new `SharedConnection { peer: RwLock<Peer>, running_service:
Mutex<Option<RunningService>> }`. Every McpTool from the same server
holds the same `Arc<SharedConnection>`. `replace(new_peer, new_rs)`
swaps both fields atomically; the OLD RunningService falls out of
scope at the end of `replace` and its DropGuard correctly cancels
the dead transport. No more mem::forget.

M-R2 — `reconnect_lock` was constructed per `collect_tools` call,
so different agent rebuilds got different locks and the gen counter
was reset on each rebuild. Fix: store `reconnect_locks: HashMap<String,
Arc<Mutex<u64>>>` on `McpClientManager` itself; cloned into every
McpTool, canonical for the process lifetime.

M-R3 — both attempts of try-reconnect-retry used the full 120s
timeout (worst case 240s total). Fix: thread `started: Instant` +
`total_budget: Duration` through; each `call_once` + the reconnect
get whatever budget remains. New `remaining_budget` helper saturates
at zero past the deadline (with a unit test).

M-R4 — manager's handles became stale after tool-side reconnect
(manager pointed at the OLD dead RunningService while the tool had
swapped in a new one). Fix falls out of M-R1: manager.connections
and tool.connection are the SAME Arc<SharedConnection>; whichever
side reconnects, both see the new state.

M-R5 — `is_transport_failure` was too aggressive: included
`UnexpectedResponse` (protocol mismatch, server alive but buggy)
and `Timeout` (slow tool legitimately running). Both would tear
down healthy connections. Fix: narrowed to `TransportSend` and
`TransportClosed` only. Test matrix updated to lock the contract.

Also removed `McpClientHandle` entirely (was a thin wrapper that
became dead weight); manager now uses
`HashMap<String, Arc<SharedConnection>>` directly. `connect()` is a
free function returning the Arc. `raw_connect` returns the
(peer, running_service) pair for the swap path. `list_tools` is a
free function over `&SharedConnection`.

Tests: 1 → 12 mcp tests (1104 with plugin / 899 without).
- `is_transport_failure_classifies_correctly` updated to assert
  the narrower set + locks the M-R5 contract.
- `remaining_budget_decays_and_saturates` covers the M-R3 deadline
  arithmetic.
- Existing tests confirm the refactor preserved behaviour.

No new beads filed — this closes the review's recommended follow-up.

Author: Yogthos

src/extras/mcp/client.rs

@@ -4,159 +4,180 @@ use std::sync::Arc;
 
 use rmcp::service::{Peer, RoleClient, RunningService, serve_client};
 use tokio::process::{ChildStderr, Command};
-use tokio::sync::RwLock;
+use tokio::sync::{Mutex, RwLock};
 
 use super::config::McpServerConfig;
 
-/// Shared, swappable peer reference for an MCP server. Cloned by
-/// every `McpTool` from the same server so that an auto-reconnect
-/// triggered by ANY tool call updates the peer for ALL tools.
-/// Previously each `McpTool` stored a direct `Peer<RoleClient>`
-/// clone, leaving them orphaned the moment the underlying transport
-/// died (audit dirge-dvi).
-pub type SharedPeer = Arc<RwLock<Peer<RoleClient>>>;
-
-pub struct McpClientHandle {
+/// Co-owned (peer, running_service) pair for one MCP server.
+///
+/// Every `McpTool` from the same server holds the same
+/// `Arc<SharedConnection>`. On reconnect — manager-side OR
+/// tool-side — `replace` atomically swaps in a fresh peer +
+/// running_service. The OLD `RunningService` drops at the end of
+/// the swap, which cancels its cancellation_token, closes the
+/// transport, and (for child-process transports) kills the dead
+/// child. This was the M-R1 review finding: the prior code did
+/// `mem::forget(RunningService)` which leaked the spawned process.
+///
+/// Lock order to avoid deadlock: always take `running_service`
+/// before `peer`. Readers take a single lock at a time.
+pub struct SharedConnection {
+    /// Kept for debugging / tracing — every error path logs the
+    /// server name, so the structured field stays for log
+    /// correlation even when no code reads it directly.
+    #[allow(dead_code)]
     pub server_name: String,
-    pub running_service: RunningService<RoleClient, ()>,
-    /// The shared peer ref. Updated in place by `replace_peer`
-    /// when the manager reconnects, so already-handed-out McpTool
-    /// instances pick up the new peer transparently.
-    peer_ref: SharedPeer,
+    peer: RwLock<Peer<RoleClient>>,
+    /// `Option` so the consuming-`Drop` of `RunningService::cancel`
+    /// can take it out cleanly during shutdown. `None` after the
+    /// connection has been explicitly shut down.
+    running_service: Mutex<Option<RunningService<RoleClient, ()>>>,
 }
 
-/// Upper bound on how long we'll wait for an MCP server to complete
-/// initialization. Command-based servers that hang on `initialize`
-/// (e.g. waiting for stdin that never comes) would otherwise pin
-/// startup indefinitely. 10s is generous for legitimate inits — npm
-/// install-on-first-run servers take a few seconds; locally-running
-/// binaries respond in <100ms. Past the cap we abort and log.
-const MCP_INIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
-
-impl McpClientHandle {
-    pub async fn connect(server_name: String, config: &McpServerConfig) -> anyhow::Result<Self> {
-        // Wrap the entire connect in a timeout so a wedged server
-        // doesn't block startup forever. Returns a clean
-        // "init timeout" error past the cap.
-        let inner = Self::connect_inner(server_name.clone(), config);
-        match tokio::time::timeout(MCP_INIT_TIMEOUT, inner).await {
-            Ok(result) => result,
-            Err(_) => Err(anyhow::anyhow!(
-                "MCP server {server_name:?} did not initialize within {}s — skipping",
-                MCP_INIT_TIMEOUT.as_secs(),
-            )),
+impl SharedConnection {
+    /// Wrap a freshly-built (peer, running_service) pair. Pub-crate so
+    /// the manager's reconnect path can create a SharedConnection for
+    /// servers that failed initial connect but succeed later.
+    pub(crate) fn new(
+        server_name: String,
+        peer: Peer<RoleClient>,
+        rs: RunningService<RoleClient, ()>,
+    ) -> Self {
+        Self {
+            server_name,
+            peer: RwLock::new(peer),
+            running_service: Mutex::new(Some(rs)),
         }
     }
 
-    async fn connect_inner(server_name: String, config: &McpServerConfig) -> anyhow::Result<Self> {
-        match config {
-            McpServerConfig::Command { command, args, env } => {
-                let mut cmd = Command::new(command);
-                cmd.args(args);
-                for (k, v) in env {
-                    cmd.env(k, v);
-                }
-                // CRITICAL: capture stderr instead of inheriting it.
-                // rmcp's default `TokioChildProcess::new` uses
-                // `Stdio::inherit()` for stderr, which gives the MCP
-                // server (and its descendants) direct access to
-                // dirge's controlling terminal. If the server (or
-                // any library it uses) emits terminal queries — OSC
-                // 11 for bg-color detection, `\x1b[c` for DA1,
-                // `\x1b[6n` for CPR — those queries reach the
-                // terminal, which replies via the TTY's INPUT side
-                // (dirge's stdin). Crossterm's event parser doesn't
-                // recognize those reply shapes, so the bytes sit in
-                // the OS stdin buffer until exit, when the shell
-                // inherits them and renders the literal escape
-                // payload as visible garbage at the prompt.
-                //
-                // Pipe stderr instead. The child's logs are still
-                // surfaced — we line-read them and forward to
-                // dirge's stderr via tracing — but the child no
-                // longer has a route to send escape queries that
-                // can elicit a reply on dirge's stdin.
-                let (transport, stderr) =
-                    rmcp::transport::child_process::TokioChildProcess::builder(cmd)
-                        .stderr(Stdio::piped())
-                        .spawn()?;
-                if let Some(child_stderr) = stderr {
-                    spawn_stderr_forwarder(server_name.clone(), child_stderr);
-                }
-                let running_service = serve_client((), transport).await.map_err(|e| {
-                    anyhow::anyhow!("MCP connection failed for '{server_name}': {e}")
-                })?;
-                let peer = running_service.peer().clone();
-                Ok(Self {
-                    server_name,
-                    running_service,
-                    peer_ref: Arc::new(RwLock::new(peer)),
-                })
-            }
-            McpServerConfig::Url { url, headers } => {
-                let custom_headers = parse_headers(headers)?;
-                let cfg = rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig::with_uri(url.as_str())
-                    .custom_headers(custom_headers);
-                type HttpClient = rmcp::transport::StreamableHttpClientTransport<reqwest::Client>;
-                let transport = HttpClient::from_config(cfg);
-                let running_service = serve_client((), transport).await.map_err(|e| {
-                    anyhow::anyhow!("MCP HTTP connection failed for '{server_name}': {e}")
-                })?;
-                let peer = running_service.peer().clone();
-                Ok(Self {
-                    server_name,
-                    running_service,
-                    peer_ref: Arc::new(RwLock::new(peer)),
-                })
-            }
-        }
+    /// Snapshot the current peer. Cheap — `Peer` is an `mpsc::Sender`
+    /// wrapper, cloning bumps a refcount.
+    pub async fn current_peer(&self) -> Peer<RoleClient> {
+        self.peer.read().await.clone()
     }
 
-    /// Direct peer clone (legacy path). Prefer [`shared_peer`] in
-    /// new code so auto-reconnect updates flow through.
-    #[allow(dead_code)]
-    pub fn peer(&self) -> rmcp::service::Peer<RoleClient> {
-        self.running_service.peer().clone()
+    /// Atomically swap in a fresh (peer, running_service). The OLD
+    /// `RunningService` drops as `_old` falls out of scope, cancelling
+    /// its background task + closing its transport.
+    pub async fn replace(
+        &self,
+        new_peer: Peer<RoleClient>,
+        new_rs: RunningService<RoleClient, ()>,
+    ) {
+        // Order: take running_service first (Option swap), then peer
+        // (RwLock write). Both consumed before the OLD running_service
+        // is dropped so the new one is fully wired before the cleanup
+        // signal fires on the old transport.
+        let _old = {
+            let mut rs_guard = self.running_service.lock().await;
+            std::mem::replace(&mut *rs_guard, Some(new_rs))
+        };
+        *self.peer.write().await = new_peer;
+        // `_old` drops here. If it was `Some`, that `RunningService`'s
+        // `DropGuard` cancels its cancellation_token; the background
+        // task observes the cancel + closes the transport; the
+        // TokioChildProcess transport's drop reaps the child.
     }
 
-    /// Shared, swappable peer ref. Cloning this `Arc` is cheap; the
-    /// inner peer is mutated in place when the server is
-    /// reconnected, so every cloned holder sees the fresh peer
-    /// on its next read.
-    pub fn shared_peer(&self) -> SharedPeer {
-        Arc::clone(&self.peer_ref)
+    /// Explicit shutdown — drops the running service synchronously
+    /// (via `Drop`'s async-cancellation guard) and renders the
+    /// connection dead. Called by `McpClientManager::shutdown`.
+    pub async fn shutdown(&self) {
+        let mut rs = self.running_service.lock().await;
+        rs.take(); // dropping the Some(...) here triggers cleanup
     }
+}
 
-    /// Replace the inner peer with a fresh one. Called when the
-    /// manager (or a self-reconnecting McpTool) builds a new
-    /// connection — write-locks briefly, swaps, drops the lock.
-    /// Existing McpTool clones holding the same Arc see the new
-    /// peer on their next read.
-    pub async fn replace_peer(&self, new_peer: Peer<RoleClient>) {
-        let mut guard = self.peer_ref.write().await;
-        *guard = new_peer;
-    }
+/// Upper bound on how long we'll wait for an MCP server to complete
+/// initialization. Command-based servers that hang on `initialize`
+/// (e.g. waiting for stdin that never comes) would otherwise pin
+/// startup indefinitely. 10s is generous for legitimate inits — npm
+/// install-on-first-run servers take a few seconds; locally-running
+/// binaries respond in <100ms. Past the cap we abort and log.
+const MCP_INIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
 
-    pub async fn list_tools(&self) -> Result<Vec<rmcp::model::Tool>, rmcp::ServiceError> {
-        self.running_service.peer().list_all_tools().await
+/// Connect to one MCP server and wrap the connection in a
+/// shared, swappable container. Returns the `Arc<SharedConnection>`
+/// the manager + every McpTool clone holds.
+pub async fn connect(
+    server_name: String,
+    config: &McpServerConfig,
+) -> anyhow::Result<Arc<SharedConnection>> {
+    let inner = connect_inner(server_name.clone(), config);
+    match tokio::time::timeout(MCP_INIT_TIMEOUT, inner).await {
+        Ok(result) => result,
+        Err(_) => Err(anyhow::anyhow!(
+            "MCP server {server_name:?} did not initialize within {}s — skipping",
+            MCP_INIT_TIMEOUT.as_secs(),
+        )),
     }
 }
 
-/// Build a fresh `RunningService` for an existing server using its
-/// config. Used by the McpTool auto-reconnect path on transport
-/// failure: returns just the new peer; the caller swaps it into a
-/// `SharedPeer` and drops the old RunningService.
+async fn connect_inner(
+    server_name: String,
+    config: &McpServerConfig,
+) -> anyhow::Result<Arc<SharedConnection>> {
+    let (peer, rs) = raw_connect(&server_name, config).await?;
+    Ok(Arc::new(SharedConnection::new(server_name, peer, rs)))
+}
+
+/// Build a new `RunningService` + extract its peer, without wrapping
+/// in `SharedConnection`. Used by `SharedConnection::replace` callers
+/// (manager + tool-side auto-reconnect) which already own the
+/// container they want to swap into.
 ///
-/// This is `pub` so the tool side can call it without owning a full
-/// `McpClientHandle`. Wraps `connect_inner` directly to skip the
-/// init-timeout layer (caller already times out the whole reconnect).
-pub async fn fresh_peer_for(
+/// Does NOT wrap in `MCP_INIT_TIMEOUT` — the caller times out the
+/// whole reconnect operation.
+pub async fn raw_connect(
     server_name: &str,
     config: &McpServerConfig,
 ) -> anyhow::Result<(Peer<RoleClient>, RunningService<RoleClient, ()>)> {
-    let handle = McpClientHandle::connect(server_name.to_string(), config).await?;
-    let peer = handle.running_service.peer().clone();
-    Ok((peer, handle.running_service))
+    match config {
+        McpServerConfig::Command { command, args, env } => {
+            let mut cmd = Command::new(command);
+            cmd.args(args);
+            for (k, v) in env {
+                cmd.env(k, v);
+            }
+            // CRITICAL: capture stderr instead of inheriting it.
+            // See lengthy explanation below at the original
+            // call site — terminal-query bytes from the child must
+            // not reach dirge's stdin via the controlling TTY.
+            let (transport, stderr) =
+                rmcp::transport::child_process::TokioChildProcess::builder(cmd)
+                    .stderr(Stdio::piped())
+                    .spawn()?;
+            if let Some(child_stderr) = stderr {
+                spawn_stderr_forwarder(server_name.to_string(), child_stderr);
+            }
+            let rs = serve_client((), transport)
+                .await
+                .map_err(|e| anyhow::anyhow!("MCP connection failed for '{server_name}': {e}"))?;
+            let peer = rs.peer().clone();
+            Ok((peer, rs))
+        }
+        McpServerConfig::Url { url, headers } => {
+            let custom_headers = parse_headers(headers)?;
+            let cfg = rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig::with_uri(url.as_str())
+                .custom_headers(custom_headers);
+            type HttpClient = rmcp::transport::StreamableHttpClientTransport<reqwest::Client>;
+            let transport = HttpClient::from_config(cfg);
+            let rs = serve_client((), transport).await.map_err(|e| {
+                anyhow::anyhow!("MCP HTTP connection failed for '{server_name}': {e}")
+            })?;
+            let peer = rs.peer().clone();
+            Ok((peer, rs))
+        }
+    }
+}
+
+/// List the tools the server advertises. Called once at startup
+/// (or after manual reconnect) to build the agent's tool registry.
+pub async fn list_tools(
+    conn: &SharedConnection,
+) -> Result<Vec<rmcp::model::Tool>, rmcp::ServiceError> {
+    let peer = conn.current_peer().await;
+    peer.list_all_tools().await
 }
 
 /// Forward an MCP child's stderr line-by-line to dirge's tracing