diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index bc071ecf..15623e2f 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -12,16 +12,15 @@ use rmcp::transport::{ ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport, TokioChildProcess, }; use std::collections::HashMap; -use std::option::Option; use std::path::PathBuf; use std::process::Stdio; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; use tempfile::{tempdir, TempDir}; use tokio::io::AsyncReadExt; use tokio::process::Command; use tokio::sync::Mutex; -use tokio::task; use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; use tracing::{error, warn}; @@ -96,6 +95,8 @@ pub struct ExtensionManager { extensions: Mutex>, context: Mutex, provider: SharedProvider, + tools_cache: Mutex>>>, + tools_cache_version: AtomicU64, } /// A flattened representation of a resource used by the agent to prepare inference @@ -453,6 +454,8 @@ impl ExtensionManager { extension_manager: None, }), provider, + tools_cache: Mutex::new(None), + tools_cache_version: AtomicU64::new(0), } } @@ -648,6 +651,8 @@ impl ExtensionManager { final_name, Extension::new(config, Arc::new(Mutex::new(client)), server_info, temp_dir), ); + drop(extensions); + self.invalidate_tools_cache_and_bump_version().await; Ok(()) } @@ -664,6 +669,7 @@ impl ExtensionManager { .lock() .await .insert(name, Extension::new(config, client, info, temp_dir)); + self.invalidate_tools_cache_and_bump_version().await; } /// Get extensions info for building the system prompt @@ -686,6 +692,7 @@ impl ExtensionManager { pub async fn remove_extension(&self, name: &str) -> ExtensionResult<()> { let sanitized_name = normalize(name); self.extensions.lock().await.remove(&sanitized_name); + self.invalidate_tools_cache_and_bump_version().await; Ok(()) } @@ -723,49 +730,99 @@ impl ExtensionManager { &self, extension_name: Option, ) -> ExtensionResult> { - self.get_prefixed_tools_impl(extension_name, None).await + let all_tools = self.get_all_tools_cached().await?; + Ok(self.filter_tools(&all_tools, extension_name.as_deref(), None)) } - async fn get_prefixed_tools_impl( + pub async fn get_prefixed_tools_excluding(&self, exclude: &str) -> ExtensionResult> { + let all_tools = self.get_all_tools_cached().await?; + Ok(self.filter_tools(&all_tools, None, Some(exclude))) + } + + fn filter_tools( &self, - extension_name: Option, + tools: &[Tool], + extension_name: Option<&str>, exclude: Option<&str>, - ) -> ExtensionResult> { - // Filter clients based on the provided extension_name or include all if None - let filtered_clients: Vec<_> = self - .extensions - .lock() - .await + ) -> Vec { + tools .iter() - .filter(|(name, _ext)| { + .filter(|tool| { + let tool_prefix = tool.name.as_ref().split("__").next().unwrap_or(""); + if let Some(excluded) = exclude { - if name.as_str() == excluded { + if tool_prefix == excluded { return false; } } - if let Some(ref name_filter) = extension_name { - *name == name_filter + if let Some(name_filter) = extension_name { + tool_prefix == name_filter } else { true } }) + .cloned() + .collect() + } + + async fn get_all_tools_cached(&self) -> ExtensionResult>> { + { + let cache = self.tools_cache.lock().await; + if let Some(ref tools) = *cache { + return Ok(Arc::clone(tools)); + } + } + + let version_before = self.tools_cache_version.load(Ordering::SeqCst); + let tools = Arc::new(self.fetch_all_tools().await?); + + { + let mut cache = self.tools_cache.lock().await; + let version_after = self.tools_cache_version.load(Ordering::SeqCst); + if version_after == version_before && cache.is_none() { + *cache = Some(Arc::clone(&tools)); + } + } + + Ok(tools) + } + + async fn invalidate_tools_cache_and_bump_version(&self) { + self.tools_cache_version.fetch_add(1, Ordering::SeqCst); + *self.tools_cache.lock().await = None; + } + + async fn fetch_all_tools(&self) -> ExtensionResult> { + let clients: Vec<_> = self + .extensions + .lock() + .await + .iter() .map(|(name, ext)| (name.clone(), ext.config.clone(), ext.get_client())) .collect(); let cancel_token = CancellationToken::default(); - let client_futures = filtered_clients.into_iter().map(|(name, config, client)| { + let client_futures = clients.into_iter().map(|(name, config, client)| { let cancel_token = cancel_token.clone(); - task::spawn(async move { + let ext_name = name.clone(); + async move { let mut tools = Vec::new(); let client_guard = client.lock().await; - let mut client_tools = client_guard.list_tools(None, cancel_token).await?; + let mut client_tools = match client_guard + .list_tools(None, cancel_token.clone()) + .await + { + Ok(t) => t, + Err(e) => { + warn!(extension = %ext_name, error = %e, "Failed to list tools"); + return (name, vec![]); + } + }; loop { for tool in client_tools.tools { - let is_available = config.is_tool_available(&tool.name); - - if is_available { + if config.is_tool_available(&tool.name) { tools.push(Tool { name: format!("{}__{}", name, tool.name).into(), description: tool.description, @@ -783,35 +840,32 @@ impl ExtensionManager { break; } - client_tools = client_guard - .list_tools(client_tools.next_cursor, CancellationToken::default()) - .await?; + client_tools = match client_guard + .list_tools(client_tools.next_cursor, cancel_token.clone()) + .await + { + Ok(t) => t, + Err(e) => { + warn!(extension = %ext_name, error = %e, "Failed to list tools (pagination)"); + break; + } + }; } - Ok::, ExtensionError>(tools) - }) + (name, tools) + } }); - // Collect all results concurrently let results = future::join_all(client_futures).await; - // Aggregate tools and handle errors let mut tools = Vec::new(); - for result in results { - match result { - Ok(Ok(client_tools)) => tools.extend(client_tools), - Ok(Err(err)) => return Err(err), - Err(join_err) => return Err(ExtensionError::from(join_err)), - } + for (_, client_tools) in results { + tools.extend(client_tools); } Ok(tools) } - pub async fn get_prefixed_tools_excluding(&self, exclude: &str) -> ExtensionResult> { - self.get_prefixed_tools_impl(None, Some(exclude)).await - } - /// Get the extension prompt including client instructions pub async fn get_planning_prompt(&self, tools_info: Vec) -> String { let mut context: HashMap<&str, Value> = HashMap::new(); @@ -1380,6 +1434,7 @@ mod tests { .lock() .await .insert(sanitized_name, extension); + self.invalidate_tools_cache_and_bump_version().await; } } @@ -1828,4 +1883,123 @@ mod tests { ); } } + + #[tokio::test] + async fn test_tools_cache_invalidated_on_add_extension() { + let extension_manager = ExtensionManager::new_without_provider(); + + extension_manager + .add_mock_extension( + "ext_a".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + + let tools_after_first = extension_manager.get_prefixed_tools(None).await.unwrap(); + let tool_names: Vec = tools_after_first + .iter() + .map(|t| t.name.to_string()) + .collect(); + assert!(tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__"))); + + extension_manager + .add_mock_extension( + "ext_b".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + + let tools_after_second = extension_manager.get_prefixed_tools(None).await.unwrap(); + let tool_names: Vec = tools_after_second + .iter() + .map(|t| t.name.to_string()) + .collect(); + assert!(tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(tool_names.iter().any(|n| n.starts_with("ext_b__"))); + } + + #[tokio::test] + async fn test_tools_cache_invalidated_on_remove_extension() { + let extension_manager = ExtensionManager::new_without_provider(); + + extension_manager + .add_mock_extension( + "ext_a".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + extension_manager + .add_mock_extension( + "ext_b".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + + let tools_before = extension_manager.get_prefixed_tools(None).await.unwrap(); + let tool_names: Vec = tools_before.iter().map(|t| t.name.to_string()).collect(); + assert!(tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(tool_names.iter().any(|n| n.starts_with("ext_b__"))); + + extension_manager.remove_extension("ext_b").await.unwrap(); + + let tools_after = extension_manager.get_prefixed_tools(None).await.unwrap(); + let tool_names: Vec = tools_after.iter().map(|t| t.name.to_string()).collect(); + assert!(tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__"))); + } + + #[tokio::test] + async fn test_get_prefixed_tools_excluding() { + let extension_manager = ExtensionManager::new_without_provider(); + + extension_manager + .add_mock_extension( + "ext_a".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + extension_manager + .add_mock_extension( + "ext_b".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + + let tools = extension_manager + .get_prefixed_tools_excluding("ext_a") + .await + .unwrap(); + let tool_names: Vec = tools.iter().map(|t| t.name.to_string()).collect(); + + assert!(!tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(tool_names.iter().any(|n| n.starts_with("ext_b__"))); + } + + #[tokio::test] + async fn test_get_prefixed_tools_by_extension_name() { + let extension_manager = ExtensionManager::new_without_provider(); + + extension_manager + .add_mock_extension( + "ext_a".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + extension_manager + .add_mock_extension( + "ext_b".to_string(), + Arc::new(Mutex::new(Box::new(MockClient {}))), + ) + .await; + + let tools = extension_manager + .get_prefixed_tools(Some("ext_a".to_string())) + .await + .unwrap(); + let tool_names: Vec = tools.iter().map(|t| t.name.to_string()).collect(); + + assert!(tool_names.iter().any(|n| n.starts_with("ext_a__"))); + assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__"))); + } }