fix(routing): treat unknown roots-capability as pending, not rootless

Stamp the session roots capability synchronously before the first await
in on_initialized so a tools/list racing the notification sees the right
flag. In the resolver, treat unknown capability (None — never observed
notifications/initialized for this session yet) like roots-capable:
return PendingRoots so the next request retries via the on-demand probe
instead of falling through to client-grants and getting permanently
denied (the "only meta tools until reconnect" bug).

Also fixes an incidental clippy lint (unnecessary_sort_by) in
prefix_cache.rs surfaced by the newer toolchain.

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

Author: its-mash

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

@@ -211,14 +211,15 @@ impl McpMuxGatewayHandler {
         if self.services.session_roots.get(sid).is_some() {
             return;
         }
-        // Not roots-capable → resolver routes via client grants, no
-        // probe useful.
-        if !self
-            .services
-            .session_roots
-            .is_roots_capable(sid)
-            .unwrap_or(false)
-        {
+        // Skip the probe only when we *know* this session is rootless
+        // (`Some(false)`). When capability is unknown (`None` — we
+        // haven't observed `notifications/initialized` for this session
+        // yet, e.g. tools/list racing in before the notification's
+        // handler completed), still try to probe: the worst case is one
+        // wasted call on a genuinely rootless client where peer.list_roots()
+        // returns method-not-found, vs. a stuck PendingRoots / empty
+        // response on a client that *would* report roots if asked.
+        if self.services.session_roots.is_roots_capable(sid) == Some(false) {
             return;
         }
         // Cool-down after a recent failed probe so we don't hammer a
@@ -392,11 +393,41 @@ impl ServerHandler for McpMuxGatewayHandler {
             }
         };
 
+        let peer = std::sync::Arc::new(context.peer);
+        let session_id_for_register = extract_session_id(&context.extensions);
+
+        // CRITICAL: stamp the roots capability **before any await** so the
+        // resolver / probe paths see the right answer if a request from
+        // this session arrives while we're still partway through this
+        // handler. The race we hit before this reordering:
+        //
+        //     on_initialized starts
+        //     register_session ✓
+        //     await prime_hashes_for_space ← yields ~5ms
+        //                                    tools/list races in here,
+        //                                    is_roots_capable() == None,
+        //                                    → resolver falls to "no roots
+        //                                    + no grants — deny" (Tier 2),
+        //                                    returns 4 meta tools
+        //     await returns, now set_roots_capable(true) — too late
+        //
+        // Stamping synchronously up front guarantees that whatever else
+        // tokio decides to schedule between here and the spawned
+        // list_roots() task at the bottom, the resolver has the right
+        // capability flag.
+        if let Some(sid) = session_id_for_register.as_deref() {
+            let declares_roots = peer
+                .peer_info()
+                .map(|info| info.capabilities.roots.is_some())
+                .unwrap_or(false);
+            self.services
+                .session_roots
+                .set_roots_capable(sid, declares_roots);
+        }
+
         // Register the *session* with MCPNotifier so subsequent fanout can
         // re-resolve per session (a single OAuth client can hold multiple
         // sessions on different folders, each routing independently).
-        let peer = std::sync::Arc::new(context.peer);
-        let session_id_for_register = extract_session_id(&context.extensions);
         if let Some(sid) = session_id_for_register.as_deref() {
             self.notification_bridge.register_session(
                 sid.to_string(),
@@ -422,18 +453,14 @@ impl ServerHandler for McpMuxGatewayHandler {
         // workspace roots into the session registry so the resolver can pick
         // a binding. Then log + (if no binding matched) prompt the UI.
         if let Some(session_id) = extract_session_id(&context.extensions) {
-            let declares_roots = peer
-                .peer_info()
-                .map(|info| info.capabilities.roots.is_some())
-                .unwrap_or(false);
-            // Stash the capability so the resolver can branch between
-            // workspace-binding routing (capable) and the per-client grant
-            // fallback (rootless). Done unconditionally so the registry has
-            // a definitive answer for every session, not just those with
-            // roots declared.
-            self.services
+            // Capability already stamped at the top of this handler — read
+            // it back rather than re-deriving so we stay consistent if the
+            // peer_info() ever flapped during the await above.
+            let declares_roots = self
+                .services
                 .session_roots
-                .set_roots_capable(&session_id, declares_roots);
+                .is_roots_capable(&session_id)
+                .unwrap_or(false);
             // Persist the bit on the client row, *always* — the Clients UI
             // needs to distinguish "never observed" from "explicitly
             // rootless" so its capability badge isn't misleading on

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

@@ -152,7 +152,20 @@ impl FeatureSetResolverService {
         if let Some(sid) = session_id {
             let roots = self.session_roots.get(sid);
             let has_roots = roots.as_ref().is_some_and(|r| !r.is_empty());
-            let roots_capable = self.session_roots.is_roots_capable(sid).unwrap_or(false);
+            // Three states for capability:
+            //   `Some(true)`  — declared roots on initialize
+            //   `Some(false)` — explicitly didn't declare; treat as rootless
+            //   `None`        — never observed `notifications/initialized`
+            //                   for this session yet. Treat as PROBABLY
+            //                   capable (most modern MCP clients are) so
+            //                   the resolver returns PendingRoots instead
+            //                   of falling all the way through to the
+            //                   rootless client_grants tier — otherwise a
+            //                   tools/list that races on_initialized
+            //                   yields "no roots + no grants — deny" and
+            //                   the user sees only meta tools until
+            //                   reconnect.
+            let roots_capable_known = self.session_roots.is_roots_capable(sid);
 
             // Tier 1: session reported roots — try a binding match.
             if has_roots {
@@ -183,14 +196,19 @@ impl FeatureSetResolverService {
                 });
             }
 
-            // Tier 1c: client declared `roots` but they haven't shown up yet.
-            // Don't fall through to client grants — that's the leak the old
-            // Tier-2 fallback caused. Return empty; we'll fire `list_changed`
-            // when roots actually arrive.
-            if roots_capable {
+            // Tier 1c: client declared `roots` but they haven't shown up
+            // yet, OR we haven't observed `initialize` yet so we don't
+            // know either way. Returning PendingRoots (empty) means the
+            // first response is empty if the on-demand probe loses the
+            // race, but the next request retries via the probe + the
+            // on_initialized list_roots task fires `list_changed` once
+            // roots actually land. Beats falling through to grants and
+            // getting permanently denied.
+            if !matches!(roots_capable_known, Some(false)) {
                 debug!(
                     session_id = %sid,
-                    "[FeatureSetResolver] roots-capable, roots pending — empty until they arrive",
+                    capability = ?roots_capable_known,
+                    "[FeatureSetResolver] roots-capable (or unknown), roots pending — empty until they arrive",
                 );
                 return Ok(ResolvedFeatureSet {
                     feature_set_ids: vec![],

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

@@ -138,7 +138,7 @@ impl PrefixCacheService {
 
         // Sort by created_at (earliest first)
         // TODO: Add verified status priority when registry supports it
-        servers.sort_by(|a, b| a.created_at.cmp(&b.created_at));
+        servers.sort_by_key(|a| a.created_at);
 
         // Clear existing cache for this space
         self.clear_space(space_id).await;

tests/rust/tests/integration/feature_set_resolver.rs

@@ -128,12 +128,28 @@ async fn deny_when_no_session_id_and_no_grants() {
 }
 
 #[tokio::test]
-async fn deny_when_session_has_no_roots_and_not_capable() {
-    // Default capability state is "unknown" (None). The resolver treats
-    // missing capability info as rootless, so this falls through to Tier 2
-    // (no client_id supplied → Deny).
+async fn pending_when_session_has_no_roots_and_capability_unknown() {
+    // Default capability state for a session we've never seen
+    // `notifications/initialized` for is `None` (unknown). The resolver
+    // treats unknown like roots-capable: returns `PendingRoots` so the
+    // *next* request retries via the on-demand probe instead of being
+    // permanently denied. This was the bug where a tools/list racing
+    // on_initialized resolved to "no roots + no grants — deny" and the
+    // user saw only meta tools until reconnect.
     let f = Fixture::new().await;
     let r = f.resolver.resolve(Some("orphan"), None).await.unwrap();
+    assert_eq!(r.source, ResolutionSource::PendingRoots);
+    assert!(r.feature_set_ids.is_empty());
+}
+
+#[tokio::test]
+async fn deny_when_session_explicitly_rootless_and_no_grants() {
+    // Explicit Some(false) capability — client told us it doesn't
+    // support roots — and no client grants. This is the only path where
+    // the resolver legitimately lands on Deny without a session id.
+    let f = Fixture::new().await;
+    f.session_roots.set_roots_capable("rootless", false);
+    let r = f.resolver.resolve(Some("rootless"), None).await.unwrap();
     assert_eq!(r.source, ResolutionSource::Deny);
 }