Skip to content

fix: make feature discovery capability-aware and bounded #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
its-mash merged 4 commits into from
Jun 27, 2026
Merged

fix: make feature discovery capability-aware and bounded #194

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e0a1b39
fix: avoid hanging feature discovery on unsupported capabilities
muxammadreza Jun 25, 2026
e01ef9e
Merge branch 'main' into pr/fix-feature-discovery-timeouts
its-mash Jun 26, 2026
715c003
Merge branch 'main' into pr/fix-feature-discovery-timeouts
its-mash Jun 27, 2026
dfac298
test(gateway): cover with_list_timeout (ok / err / timeout paths)
its-mash Jun 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 90 additions & 36 deletions crates/mcpmux-gateway/src/pool/features/discovery.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
//! Feature Discovery Service - SRP: Discovery & caching

use anyhow::Result;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, info, warn};

use super::{convert_to_feature, resource_to_feature, CachedFeatures};
Expand All @@ -14,10 +16,29 @@ pub struct FeatureDiscoveryService {
}

impl FeatureDiscoveryService {
const LIST_TIMEOUT: Duration = Duration::from_secs(10);

pub fn new(feature_repo: Arc<dyn ServerFeatureRepository>) -> Self {
Self { feature_repo }
}

async fn with_list_timeout<T, E, F>(label: &str, fut: F) -> Option<Result<T, E>>
where
F: Future<Output = Result<T, E>>,
{
match tokio::time::timeout(Self::LIST_TIMEOUT, fut).await {
Ok(result) => Some(result),
Err(_) => {
warn!(
"[FeatureDiscovery] {} timed out after {:?}",
label,
Self::LIST_TIMEOUT
);
None
}
}
}

/// Discover features from a connected MCP client and cache them
pub async fn discover_and_cache(
&self,
Expand All @@ -31,50 +52,83 @@ impl FeatureDiscoveryService {
);

let mut discovered = CachedFeatures::default();
let capabilities = client.peer_info().map(|info| info.capabilities.clone());
let capabilities_known = capabilities.is_some();
let has_tools = capabilities
.as_ref()
.and_then(|c| c.tools.as_ref())
.is_some();
let has_prompts = capabilities
.as_ref()
.and_then(|c| c.prompts.as_ref())
.is_some();
let has_resources = capabilities
.as_ref()
.and_then(|c| c.resources.as_ref())
.is_some();

// Discover tools
match client.list_all_tools().await {
Ok(tools) => {
discovered.tools = tools
.into_iter()
.map(|t| convert_to_feature(space_id, server_id, t))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} tools",
discovered.tools.len()
);
debug!(
"[FeatureDiscovery] Capability gates for {}/{}: known={}, tools={}, prompts={}, resources={}",
space_id, server_id, capabilities_known, has_tools, has_prompts, has_resources
);

if !capabilities_known || has_tools {
match Self::with_list_timeout("tools/list", client.list_all_tools()).await {
Some(Ok(tools)) => {
discovered.tools = tools
.into_iter()
.map(|t| convert_to_feature(space_id, server_id, t))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} tools",
discovered.tools.len()
);
}
Some(Err(e)) => warn!("[FeatureDiscovery] Failed to list tools: {}", e),
None => {}
}
Err(e) => warn!("[FeatureDiscovery] Failed to list tools: {}", e),
} else {
debug!(
"[FeatureDiscovery] Skipping tools/list: server explicitly did not advertise tools capability"
);
}

// Discover prompts
match client.list_all_prompts().await {
Ok(prompts) => {
discovered.prompts = prompts
.into_iter()
.map(|p| convert_to_feature(space_id, server_id, p))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} prompts",
discovered.prompts.len()
);
if !capabilities_known || has_prompts {
match Self::with_list_timeout("prompts/list", client.list_all_prompts()).await {
Some(Ok(prompts)) => {
discovered.prompts = prompts
.into_iter()
.map(|p| convert_to_feature(space_id, server_id, p))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} prompts",
discovered.prompts.len()
);
}
Some(Err(e)) => warn!("[FeatureDiscovery] Failed to list prompts: {}", e),
None => {}
}
Err(e) => warn!("[FeatureDiscovery] Failed to list prompts: {}", e),
} else {
debug!("[FeatureDiscovery] Skipping prompts/list: server explicitly did not advertise prompts capability");
}

// Discover resources
match client.list_all_resources().await {
Ok(resources) => {
discovered.resources = resources
.into_iter()
.map(|r| resource_to_feature(space_id, server_id, r))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} resources",
discovered.resources.len()
);
if !capabilities_known || has_resources {
match Self::with_list_timeout("resources/list", client.list_all_resources()).await {
Some(Ok(resources)) => {
discovered.resources = resources
.into_iter()
.map(|r| resource_to_feature(space_id, server_id, r))
.collect();
debug!(
"[FeatureDiscovery] Discovered {} resources",
discovered.resources.len()
);
}
Some(Err(e)) => warn!("[FeatureDiscovery] Failed to list resources: {}", e),
None => {}
}
Err(e) => warn!("[FeatureDiscovery] Failed to list resources: {}", e),
} else {
debug!("[FeatureDiscovery] Skipping resources/list: server explicitly did not advertise resources capability");
}

// Cache all features in database
Expand Down
Loading