fix(handler): single-flight on-demand probe per session

The boolean `claim_probe` rate-limiter let the first list request
enter the probe path, but **followers in the same burst skipped the
probe entirely** and resolved to PendingRoots immediately — returning
empty *before* the first probe's `peer.list_roots()` came back.

Trace from a Claude Code claude-vscode session opening a new
workspace:

  02:17:19.950  resolver: roots-capable, roots pending  (req 1)
  02:17:19.950  resolver: roots-capable, roots pending  (req 2 — skipped probe)
  02:17:19.973  on-demand probe populated roots         (req 1's probe finally landed)

Two of the three list responses (tools/list + prompts/list +
resources/list, fired within 1ms by Claude Code at init) returned
empty. The probe success at +23ms triggered a notifications/
tools/list_changed which the CLI variant honors but the VS Code
extension doesn't refetch on, so the empty initial list is what the
panel kept showing.

Replace the boolean with a per-session `tokio::sync::Mutex` for
single-flight semantics. First request acquires the lock, fires the
probe, populates `session_roots`. Followers await the same lock; on
acquire, they re-check `session_roots.get(sid)` and exit early since
the predecessor already populated. Net result: one upstream
`peer.list_roots()` call, three list responses with the correct
routing decision.

Kept the cool-down semantic too: `should_throttle_probe` +
`mark_probe_completed` rate-limit *sequential* probe attempts when
the previous one errored, so a peer that's failing list_roots
doesn't get hammered. Distinct from the lock — the lock is
concurrency, the throttle is failure recovery.

Doesn't help the upstream client bug (Claude Code claude-vscode
doesn't refetch on tools/list_changed) but means the *first*
response is correct, which is the only one some clients ever read.

cargo check + clippy (-D warnings) clean.

Signed-off-by: Mohammod Al Amin Ashik <maa.ashik00@gmail.com>

Author: its-mash

crates/mcpmux-gateway/src/mcp/handler.rs

@@ -206,7 +206,8 @@ impl McpMuxGatewayHandler {
         client_id: &str,
     ) {
         let Some(sid) = session_id else { return };
-        // Already have a definitive answer (Some(roots) — possibly empty).
+        // Fast path: already have a definitive answer (Some(roots),
+        // possibly empty). No probe needed.
         if self.services.session_roots.get(sid).is_some() {
             return;
         }
@@ -220,18 +221,45 @@ impl McpMuxGatewayHandler {
         {
             return;
         }
-        // Throttle: at most one probe per session per second so a
-        // request burst doesn't multiply upstream calls.
-        if !self
+        // Cool-down after a recent failed probe so we don't hammer a
+        // peer whose previous list_roots() errored. Doesn't apply
+        // when a probe is currently *running* — that's the
+        // probe_lock's job below.
+        if self
             .services
             .session_roots
-            .claim_probe(sid, std::time::Duration::from_secs(1))
+            .should_throttle_probe(sid, std::time::Duration::from_secs(1))
         {
             return;
         }
 
+        // Single-flight: serialize concurrent probes per session so a
+        // burst of three list calls (tools/list + prompts/list +
+        // resources/list within milliseconds) doesn't fan out three
+        // upstream `peer.list_roots()` calls. The first request enters
+        // the critical section, fires the probe, populates
+        // session_roots; the second and third await the same lock,
+        // then re-check session_roots and exit early.
+        //
+        // Without this, the followers used to skip the probe entirely
+        // (boolean `claim_probe` flag) and resolve to PendingRoots —
+        // exactly the empty-tools-list bug Claude Code's VS Code
+        // extension was hitting.
+        let lock = self.services.session_roots.probe_lock(sid);
+        let _guard = lock.lock().await;
+
+        // Recheck after acquiring the lock — the predecessor probe may
+        // have already populated the registry.
+        if self.services.session_roots.get(sid).is_some() {
+            return;
+        }
+
         const PROBE_BUDGET: std::time::Duration = std::time::Duration::from_millis(300);
-        match tokio::time::timeout(PROBE_BUDGET, peer.list_roots()).await {
+        let outcome = tokio::time::timeout(PROBE_BUDGET, peer.list_roots()).await;
+        // Stamp completion regardless of success/failure so the
+        // sequential cool-down kicks in for the next caller.
+        self.services.session_roots.mark_probe_completed(sid);
+        match outcome {
             Ok(Ok(result)) => {
                 let uris: Vec<String> = result.roots.iter().map(|r| r.uri.to_string()).collect();
                 self.services
@@ -276,15 +304,15 @@ impl McpMuxGatewayHandler {
                     %client_id,
                     session_id = %sid,
                     error = %e,
-                    "[FeatureSetResolver] on-demand probe failed (will retry on next request)",
+                    "[FeatureSetResolver] on-demand probe failed (will retry on next request after throttle)",
                 );
             }
             Err(_elapsed) => {
                 debug!(
                     %client_id,
                     session_id = %sid,
                     budget_ms = PROBE_BUDGET.as_millis(),
-                    "[FeatureSetResolver] on-demand probe timed out (will retry on next request)",
+                    "[FeatureSetResolver] on-demand probe timed out (will retry on next request after throttle)",
                 );
             }
         }

crates/mcpmux-gateway/src/services/session_roots.rs

@@ -38,13 +38,27 @@ pub struct SessionRootsRegistry {
     roots_capable: DashMap<String, bool>,
     /// `session_id -> Instant of the last on-demand `list_roots()` probe`.
     ///
-    /// Used by the request-time re-probe path in the MCP handler to avoid
-    /// firing N parallel `list_roots()` calls when a roots-capable session
-    /// hits a burst of `tools/list` / `prompts/list` / `resources/list` in
-    /// quick succession. The handler calls `claim_probe(sid, throttle)`
-    /// before firing; if it returns false, another probe was attempted
-    /// recently and this one is skipped.
+    /// Used by [`Self::should_throttle_probe`] to avoid hammering a
+    /// failing client when its previous probe already errored out
+    /// recently. Only stamped after a probe attempt completes (success
+    /// or failure), not on entry — so concurrent in-flight probes
+    /// coordinate via [`Self::probe_lock`] instead of this throttle.
     last_probe: DashMap<String, Instant>,
+    /// Per-session mutex guarding `peer.list_roots()` probe attempts.
+    ///
+    /// Single-flight semantics: when a burst of three list requests
+    /// (`tools/list` + `prompts/list` + `resources/list`) hits a
+    /// roots-pending session within milliseconds, only one upstream
+    /// `list_roots()` call should be in flight. The other two block on
+    /// the same lock; once the first attempt populates `map`, the
+    /// followers re-check `map.get(sid)` and skip the upstream call
+    /// entirely.
+    ///
+    /// Without this, a boolean "already tried" flag let the followers
+    /// see `roots_pending` and return empty *before* the first probe's
+    /// result landed — exactly the bug that left Claude Code's
+    /// VS Code extension showing only the meta tools.
+    probe_lock: DashMap<String, Arc<tokio::sync::Mutex<()>>>,
 }
 
 impl SessionRootsRegistry {
@@ -54,33 +68,40 @@ impl SessionRootsRegistry {
             last_resolution: DashMap::new(),
             roots_capable: DashMap::new(),
             last_probe: DashMap::new(),
+            probe_lock: DashMap::new(),
         })
     }
 
-    /// Try to claim a probe slot for `session_id`. Returns `true` if it's
-    /// been at least `throttle` since the last attempt for this session
-    /// (or if there's never been one) — and stamps the new attempt
-    /// atomically. Returns `false` if a probe was already attempted
-    /// within the throttle window.
+    /// Get (or create) the per-session probe lock. The returned Arc is
+    /// what the handler awaits to serialize concurrent probes — see
+    /// [`Self::probe_lock`] for the rationale.
+    pub fn probe_lock(&self, session_id: &str) -> Arc<tokio::sync::Mutex<()>> {
+        self.probe_lock
+            .entry(session_id.to_string())
+            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
+            .clone()
+    }
+
+    /// Should we skip an on-demand probe because the previous attempt
+    /// completed (success or failure) within the last `throttle`?
     ///
-    /// The handler calls this before firing `peer.list_roots()` so a
-    /// burst of three `tools/list` / `prompts/list` / `resources/list`
-    /// calls in 50 ms results in at most one upstream probe.
-    pub fn claim_probe(&self, session_id: &str, throttle: Duration) -> bool {
-        let now = Instant::now();
-        match self.last_probe.entry(session_id.to_string()) {
-            dashmap::mapref::entry::Entry::Occupied(mut e) => {
-                if now.duration_since(*e.get()) < throttle {
-                    return false;
-                }
-                e.insert(now);
-                true
-            }
-            dashmap::mapref::entry::Entry::Vacant(e) => {
-                e.insert(now);
-                true
-            }
-        }
+    /// Distinct from `probe_lock`: the lock serializes *concurrent*
+    /// probes; this rate-limit prevents *sequential* probes from
+    /// hammering a peer whose previous attempt errored.
+    pub fn should_throttle_probe(&self, session_id: &str, throttle: Duration) -> bool {
+        let Some(last) = self.last_probe.get(session_id) else {
+            return false;
+        };
+        Instant::now().duration_since(*last) < throttle
+    }
+
+    /// Stamp the completion of an on-demand probe so the next caller
+    /// observes the throttle. Called after the probe returns (regardless
+    /// of success or failure) so successive probes back off only when
+    /// the previous one actually finished.
+    pub fn mark_probe_completed(&self, session_id: &str) {
+        self.last_probe
+            .insert(session_id.to_string(), Instant::now());
     }
 
     /// Record whether a session declared the MCP `roots` capability on
@@ -122,6 +143,7 @@ impl SessionRootsRegistry {
         self.last_resolution.remove(session_id);
         self.roots_capable.remove(session_id);
         self.last_probe.remove(session_id);
+        self.probe_lock.remove(session_id);
     }
 
     /// Compare-and-set the session's resolved feature-set id. Returns `true`