feat(gateway): add SessionOverrideRegistry and list-path composition

Phase 1 of dynamic MCP toggling: per-session enable/disable sets wired
through FeatureService list paths, session GC, and composition tests.
Also fixes a DashMap deadlock in SessionRootsRegistry::record_resolution.

Signed-off-by: crimsonsunset <jsangio1@gmail.com>

Author: crimsonsunset

crates/mcpmux-gateway/src/consumers/mcp_notifier.rs

@@ -26,7 +26,7 @@ use tracing::{debug, info, trace, warn};
 use uuid::Uuid;
 
 use crate::pool::FeatureService;
-use crate::services::FeatureSetResolverService;
+use crate::services::{FeatureSetResolverService, SessionOverrideRegistry};
 
 /// MCP Notifier — sends `list_changed` notifications to connected sessions.
 ///
@@ -58,6 +58,8 @@ pub struct MCPNotifier {
     feature_set_resolver: Arc<FeatureSetResolverService>,
     /// Feature service for calculating content hashes
     feature_service: Arc<FeatureService>,
+    /// Session override registry — reaped alongside session roots.
+    session_overrides: Arc<SessionOverrideRegistry>,
     /// Throttle tracker: (space_id, notification_type) -> last_sent_timestamp
     /// Prevents sending duplicate notifications within THROTTLE_WINDOW
     throttle_tracker: Arc<RwLock<HashMap<(Uuid, NotificationType), Instant>>>,
@@ -110,11 +112,13 @@ impl MCPNotifier {
     pub fn new(
         feature_set_resolver: Arc<FeatureSetResolverService>,
         feature_service: Arc<FeatureService>,
+        session_overrides: Arc<SessionOverrideRegistry>,
     ) -> Self {
         Self {
             sessions: Arc::new(RwLock::new(HashMap::new())),
             feature_set_resolver,
             feature_service,
+            session_overrides,
             throttle_tracker: Arc::new(RwLock::new(HashMap::new())),
             state_hashes: Arc::new(RwLock::new(HashMap::new())),
         }
@@ -354,6 +358,7 @@ impl MCPNotifier {
         // sessions that no longer exist.
         for sid in &dead {
             self.feature_set_resolver.session_roots().remove(sid);
+            self.session_overrides.remove(sid);
         }
         info!(
             reaped = dead.len(),

crates/mcpmux-gateway/src/lib.rs

@@ -76,7 +76,7 @@ pub use pool::{
 };
 
 // Services module
-pub use services::{EventEmitter, GrantService, PrefixCacheService};
+pub use services::{EventEmitter, GrantService, PrefixCacheService, SessionOverrideRegistry};
 
 // MCP module (rmcp-based implementation)
 pub use mcp::McpMuxGatewayHandler;

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

@@ -676,7 +676,7 @@ impl ServerHandler for McpMuxGatewayHandler {
             .services
             .pool_services
             .feature_service
-            .get_tools_for_grants(&space_id.to_string(), &feature_set_ids)
+            .get_tools_for_grants(&space_id.to_string(), &feature_set_ids, session_id_owned.as_deref())
             .await
             .map_err(|e| McpError::internal_error(format!("Failed to get tools: {}", e), None))?;
 
@@ -861,7 +861,11 @@ impl ServerHandler for McpMuxGatewayHandler {
             .services
             .pool_services
             .feature_service
-            .get_prompts_for_grants(&space_id.to_string(), &feature_set_ids)
+            .get_prompts_for_grants(
+                &space_id.to_string(),
+                &feature_set_ids,
+                session_id_owned.as_deref(),
+            )
             .await
             .map_err(|e| McpError::internal_error(format!("Failed to get prompts: {}", e), None))?;
 
@@ -897,11 +901,9 @@ impl ServerHandler for McpMuxGatewayHandler {
         let oauth_ctx = self
             .get_oauth_context(&context.extensions)
             .map_err(|e| McpError::invalid_params(e.to_string(), None))?;
+        let session_id_owned = extract_session_id(&context.extensions);
         let (space_id, feature_set_ids) = self
-            .resolve_routing(
-                extract_session_id(&context.extensions).as_deref(),
-                &oauth_ctx.client_id,
-            )
+            .resolve_routing(session_id_owned.as_deref(), &oauth_ctx.client_id)
             .await?;
 
         let (server_id, prompt_name) = self
@@ -916,7 +918,11 @@ impl ServerHandler for McpMuxGatewayHandler {
             .services
             .pool_services
             .feature_service
-            .get_prompts_for_grants(&space_id.to_string(), &feature_set_ids)
+            .get_prompts_for_grants(
+                &space_id.to_string(),
+                &feature_set_ids,
+                session_id_owned.as_deref(),
+            )
             .await
             .map_err(|e| {
                 McpError::internal_error(format!("Failed to verify authorization: {}", e), None)
@@ -972,7 +978,11 @@ impl ServerHandler for McpMuxGatewayHandler {
             .services
             .pool_services
             .feature_service
-            .get_resources_for_grants(&space_id.to_string(), &feature_set_ids)
+            .get_resources_for_grants(
+                &space_id.to_string(),
+                &feature_set_ids,
+                session_id_owned.as_deref(),
+            )
             .await
             .map_err(|e| {
                 McpError::internal_error(format!("Failed to get resources: {}", e), None)
@@ -1006,11 +1016,9 @@ impl ServerHandler for McpMuxGatewayHandler {
         let oauth_ctx = self
             .get_oauth_context(&context.extensions)
             .map_err(|e| McpError::invalid_params(e.to_string(), None))?;
+        let session_id_owned = extract_session_id(&context.extensions);
         let (space_id, feature_set_ids) = self
-            .resolve_routing(
-                extract_session_id(&context.extensions).as_deref(),
-                &oauth_ctx.client_id,
-            )
+            .resolve_routing(session_id_owned.as_deref(), &oauth_ctx.client_id)
             .await?;
 
         let server_id = self
@@ -1030,7 +1038,11 @@ impl ServerHandler for McpMuxGatewayHandler {
             .services
             .pool_services
             .feature_service
-            .get_resources_for_grants(&space_id.to_string(), &feature_set_ids)
+            .get_resources_for_grants(
+                &space_id.to_string(),
+                &feature_set_ids,
+                session_id_owned.as_deref(),
+            )
             .await
             .map_err(|e| {
                 McpError::internal_error(format!("Failed to verify authorization: {}", e), None)

crates/mcpmux-gateway/src/pool/features/facade.rs

@@ -1,10 +1,11 @@
 //! Feature Service Facade - Unified API delegating to specialized services
 
 use anyhow::Result;
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use crate::pool::instance::McpClient;
-use crate::services::PrefixCacheService;
+use crate::services::{PrefixCacheService, SessionOverrideRegistry};
 use mcpmux_core::{FeatureSetRepository, FeatureType, ServerFeature, ServerFeatureRepository};
 
 use super::{
@@ -16,13 +17,15 @@ pub struct FeatureService {
     discovery: Arc<FeatureDiscoveryService>,
     resolution: Arc<FeatureResolutionService>,
     routing: Arc<FeatureRoutingService>,
+    session_overrides: Arc<SessionOverrideRegistry>,
 }
 
 impl FeatureService {
     pub fn new(
         feature_repo: Arc<dyn ServerFeatureRepository>,
         feature_set_repo: Arc<dyn FeatureSetRepository>,
         prefix_cache: Arc<PrefixCacheService>,
+        session_overrides: Arc<SessionOverrideRegistry>,
     ) -> Self {
         let discovery = Arc::new(FeatureDiscoveryService::new(feature_repo.clone()));
 
@@ -41,6 +44,7 @@ impl FeatureService {
             discovery,
             resolution,
             routing,
+            session_overrides,
         }
     }
 
@@ -86,35 +90,98 @@ impl FeatureService {
             .await
     }
 
-    // Type-specific helpers
+    /// Resolve granted feature sets to tools, applying session server overrides.
     pub async fn get_tools_for_grants(
         &self,
         space_id: &str,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<ServerFeature>> {
-        self.resolution
-            .resolve_feature_sets(space_id, feature_set_ids, Some(FeatureType::Tool))
-            .await
+        self.get_features_for_grants(
+            space_id,
+            feature_set_ids,
+            session_id,
+            Some(FeatureType::Tool),
+        )
+        .await
     }
 
+    /// Resolve granted feature sets to prompts, applying session server overrides.
     pub async fn get_prompts_for_grants(
         &self,
         space_id: &str,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<ServerFeature>> {
-        self.resolution
-            .resolve_feature_sets(space_id, feature_set_ids, Some(FeatureType::Prompt))
-            .await
+        self.get_features_for_grants(
+            space_id,
+            feature_set_ids,
+            session_id,
+            Some(FeatureType::Prompt),
+        )
+        .await
     }
 
+    /// Resolve granted feature sets to resources, applying session server overrides.
     pub async fn get_resources_for_grants(
         &self,
         space_id: &str,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<ServerFeature>> {
-        self.resolution
-            .resolve_feature_sets(space_id, feature_set_ids, Some(FeatureType::Resource))
-            .await
+        self.get_features_for_grants(
+            space_id,
+            feature_set_ids,
+            session_id,
+            Some(FeatureType::Resource),
+        )
+        .await
+    }
+
+    /// Shared list materialization: binding FS resolution + session overrides.
+    async fn get_features_for_grants(
+        &self,
+        space_id: &str,
+        feature_set_ids: &[String],
+        session_id: Option<&str>,
+        filter_type: Option<FeatureType>,
+    ) -> Result<Vec<ServerFeature>> {
+        let binding_features = self
+            .resolution
+            .resolve_feature_sets(space_id, feature_set_ids, filter_type.clone())
+            .await?;
+
+        let Some(session_id) = session_id else {
+            return Ok(binding_features);
+        };
+
+        let enabled = self.session_overrides.enabled_set(session_id);
+        let disabled = self.session_overrides.disabled_set(session_id);
+
+        if enabled.is_empty() && disabled.is_empty() {
+            return Ok(binding_features);
+        }
+
+        let mut binding_servers: HashSet<String> = binding_features
+            .iter()
+            .map(|f| f.server_id.clone())
+            .collect();
+        binding_servers.extend(enabled.iter().cloned());
+        binding_servers.retain(|server_id| !disabled.contains(server_id));
+
+        if binding_servers.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        let all_features = self
+            .resolution
+            .get_all_features_for_space(space_id, filter_type)
+            .await?;
+
+        Ok(all_features
+            .into_iter()
+            .filter(|f| f.is_available && binding_servers.contains(&f.server_id))
+            .collect())
     }
 
     // Delegate to FeatureRoutingService (with type-specific helpers)

crates/mcpmux-gateway/src/pool/routing.rs

@@ -84,13 +84,14 @@ impl RoutingService {
         &self,
         space_id: Uuid,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<RoutedTool>> {
         let space_id_str = space_id.to_string();
 
         // Resolve feature sets to allowed features
         let allowed_features = self
             .feature_service
-            .get_tools_for_grants(&space_id_str, feature_set_ids)
+            .get_tools_for_grants(&space_id_str, feature_set_ids, session_id)
             .await?;
 
         // Filter to just tools
@@ -119,12 +120,13 @@ impl RoutingService {
         &self,
         space_id: Uuid,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<RoutedPrompt>> {
         let space_id_str = space_id.to_string();
 
         let allowed_features = self
             .feature_service
-            .get_prompts_for_grants(&space_id_str, feature_set_ids)
+            .get_prompts_for_grants(&space_id_str, feature_set_ids, session_id)
             .await?;
 
         let prompts: Vec<RoutedPrompt> = allowed_features
@@ -151,12 +153,13 @@ impl RoutingService {
         &self,
         space_id: Uuid,
         feature_set_ids: &[String],
+        session_id: Option<&str>,
     ) -> Result<Vec<RoutedResource>> {
         let space_id_str = space_id.to_string();
 
         let allowed_features = self
             .feature_service
-            .get_resources_for_grants(&space_id_str, feature_set_ids)
+            .get_resources_for_grants(&space_id_str, feature_set_ids, session_id)
             .await?;
 
         let resources: Vec<RoutedResource> = allowed_features

crates/mcpmux-gateway/src/pool/service_factory.rs

@@ -44,6 +44,7 @@ impl ServiceFactory {
         deps: &GatewayDependencies,
         event_tx: tokio::sync::broadcast::Sender<DomainEvent>,
         prefix_cache: Arc<crate::services::PrefixCacheService>,
+        session_overrides: Arc<crate::services::SessionOverrideRegistry>,
     ) -> PoolServices {
         // TokenService - single source of truth for token management
         let token_service = Arc::new(TokenService::new(
@@ -80,7 +81,8 @@ impl ServiceFactory {
         let feature_service = Arc::new(FeatureService::new(
             deps.feature_repo.clone(),
             deps.feature_set_repo.clone(),
-            prefix_cache.clone(), // Clone here since we use it again below
+            prefix_cache.clone(),
+            session_overrides,
         ));
 
         // ServerManager - event-driven orchestrator for server state

crates/mcpmux-gateway/src/server/handlers.rs

@@ -1087,23 +1087,23 @@ pub async fn oauth_get_client_features(
         .services
         .pool_services
         .feature_service
-        .get_tools_for_grants(&space_id_str, &feature_set_ids)
+        .get_tools_for_grants(&space_id_str, &feature_set_ids, None)
         .await
         .unwrap_or_default();
 
     let prompts = state
         .services
         .pool_services
         .feature_service
-        .get_prompts_for_grants(&space_id_str, &feature_set_ids)
+        .get_prompts_for_grants(&space_id_str, &feature_set_ids, None)
         .await
         .unwrap_or_default();
 
     let resources = state
         .services
         .pool_services
         .feature_service
-        .get_resources_for_grants(&space_id_str, &feature_set_ids)
+        .get_resources_for_grants(&space_id_str, &feature_set_ids, None)
         .await
         .unwrap_or_default();
 

crates/mcpmux-gateway/src/server/mod.rs

@@ -231,6 +231,7 @@ impl GatewayServer {
         let notification_bridge = Arc::new(MCPNotifier::new(
             self.services.feature_set_resolver.clone(),
             self.services.pool_services.feature_service.clone(),
+            self.services.session_overrides.clone(),
         ));
 
         // Start listening to DomainEvents