Handle MCP tool list change notifications

Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
Co-authored-by: Alexis Rohou <arohou@gmail.com>
Co-authored-by: Alexis Rohou <a.rohou@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

crates/goose/src/agents/extension_manager.rs

@@ -17,7 +17,7 @@ use std::path::PathBuf;
 use std::pin::Pin;
 use std::process::Stdio;
 use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::task::{Context, Poll};
 use std::time::Duration;
 use tempfile::{tempdir, TempDir};
@@ -404,6 +404,7 @@ struct ResolvedTool {
     resource_uri: Option<String>,
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn child_process_client(
     mut command: Command,
     timeout: &Option<u64>,
@@ -412,6 +413,7 @@ async fn child_process_client(
     docker_container: Option<String>,
     client_name: String,
     capabilities: GooseMcpClientCapabilities,
+    extension_manager: Weak<ExtensionManager>,
 ) -> ExtensionResult<McpClient> {
     configure_subprocess(&mut command);
 
@@ -450,6 +452,7 @@ async fn child_process_client(
         client_name,
         capabilities,
         working_dir.clone(),
+        extension_manager,
     )
     .await;
 
@@ -620,6 +623,7 @@ async fn connect_with_auth(
     client_name: String,
     capabilities: GooseMcpClientCapabilities,
     roots_dir: &std::path::Path,
+    extension_manager: Weak<ExtensionManager>,
 ) -> ExtensionResult<Box<dyn McpClientTrait>> {
     let mut auth_headers = HeaderMap::new();
     auth_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT);
@@ -654,6 +658,7 @@ async fn connect_with_auth(
             client_name,
             capabilities,
             roots_dir.to_path_buf(),
+            extension_manager,
         )
         .await?,
     ))
@@ -671,6 +676,7 @@ async fn create_streamable_http_client(
     client_name: String,
     capabilities: GooseMcpClientCapabilities,
     roots_dir: &std::path::Path,
+    extension_manager: Weak<ExtensionManager>,
 ) -> ExtensionResult<Box<dyn McpClientTrait>> {
     #[cfg(unix)]
     if let Some(socket_path) = socket {
@@ -684,6 +690,7 @@ async fn create_streamable_http_client(
             client_name,
             capabilities,
             roots_dir,
+            extension_manager,
         )
         .await;
     }
@@ -739,6 +746,7 @@ async fn create_streamable_http_client(
                     client_name.clone(),
                     capabilities.clone(),
                     roots_dir,
+                    extension_manager.clone(),
                 )
                 .await;
 
@@ -777,6 +785,7 @@ async fn create_streamable_http_client(
         client_name.clone(),
         capabilities.clone(),
         roots_dir.to_path_buf(),
+        extension_manager.clone(),
     )
     .await;
 
@@ -792,6 +801,7 @@ async fn create_streamable_http_client(
                     client_name,
                     capabilities,
                     roots_dir,
+                    extension_manager,
                 )
                 .await
             }
@@ -814,6 +824,7 @@ async fn create_unix_socket_http_client(
     client_name: String,
     capabilities: GooseMcpClientCapabilities,
     roots_dir: &std::path::Path,
+    extension_manager: Weak<ExtensionManager>,
 ) -> ExtensionResult<Box<dyn McpClientTrait>> {
     use rmcp::transport::UnixSocketHttpClient;
 
@@ -851,6 +862,7 @@ async fn create_unix_socket_http_client(
         client_name.clone(),
         capabilities.clone(),
         roots_dir.to_path_buf(),
+        extension_manager,
     )
     .await;
 
@@ -995,6 +1007,7 @@ impl ExtensionManager {
                     self.client_name.clone(),
                     self.mcp_client_capabilities(),
                     &effective_working_dir,
+                    Arc::downgrade(self),
                 )
                 .await?
             }
@@ -1052,6 +1065,7 @@ impl ExtensionManager {
                             Some(container_id.to_string()),
                             self.client_name.clone(),
                             self.mcp_client_capabilities(),
+                            Arc::downgrade(self),
                         )
                         .await?;
                         Box::new(client)
@@ -1068,6 +1082,7 @@ impl ExtensionManager {
                                 self.client_name.clone(),
                                 self.mcp_client_capabilities(),
                                 effective_working_dir.clone(),
+                                Arc::downgrade(self),
                             )
                             .await?,
                         )
@@ -1129,6 +1144,7 @@ impl ExtensionManager {
                     container.map(|c| c.id().to_string()),
                     self.client_name.clone(),
                     self.mcp_client_capabilities(),
+                    Arc::downgrade(self),
                 )
                 .await?;
                 Box::new(client)
@@ -1161,6 +1177,7 @@ impl ExtensionManager {
                     container.map(|c| c.id().to_string()),
                     self.client_name.clone(),
                     self.mcp_client_capabilities(),
+                    Arc::downgrade(self),
                 )
                 .await?;
 
@@ -1385,7 +1402,7 @@ impl ExtensionManager {
         Some(attachment)
     }
 
-    async fn invalidate_tools_cache_and_bump_version(&self) {
+    pub(crate) async fn invalidate_tools_cache_and_bump_version(&self) {
         self.tools_cache_version.fetch_add(1, Ordering::SeqCst);
         *self.tools_cache.lock().await = None;
     }
@@ -3329,6 +3346,7 @@ mod tests {
             "goose-test".to_string(),
             capabilities,
             temp_dir.path(),
+            Weak::new(),
         )
         .await;
 
@@ -3364,6 +3382,7 @@ mod tests {
             "goose-test".to_string(),
             capabilities,
             temp_dir.path(),
+            Weak::new(),
         )
         .await;
 
@@ -3410,6 +3429,7 @@ mod tests {
             "goose-test".to_string(),
             capabilities,
             temp_dir.path(),
+            Weak::new(),
         )
         .await;
 
@@ -3491,6 +3511,7 @@ mod tests {
             "goose-test".to_string(),
             capabilities,
             temp_dir.path(),
+            Weak::new(),
         )
         .await;
 

crates/goose/src/agents/mcp_client.rs

@@ -1,4 +1,5 @@
 use crate::action_required_manager::{ActionRequiredManager, ElicitationOutcome};
+use crate::agents::extension_manager::ExtensionManager;
 use crate::agents::tool_execution::ToolCallContext;
 use crate::agents::types::SharedProvider;
 use crate::session_context::{SESSION_ID_HEADER, TOOL_CALL_REQUEST_ID_HEADER, WORKING_DIR_HEADER};
@@ -31,7 +32,10 @@ use rmcp::{
 };
 use serde_json::Value;
 use std::{
-    collections::HashMap, path::PathBuf, sync::Arc, sync::Mutex as StdMutex, time::Duration,
+    collections::HashMap,
+    path::PathBuf,
+    sync::{Arc, Mutex as StdMutex, Weak},
+    time::Duration,
 };
 use tokio::sync::{
     mpsc::{self, Sender},
@@ -185,6 +189,7 @@ pub struct GooseClient {
     client_name: String,
     capabilities: GooseMcpClientCapabilities,
     working_dir: Arc<tokio::sync::RwLock<PathBuf>>,
+    extension_manager: Weak<ExtensionManager>,
 }
 
 impl GooseClient {
@@ -194,6 +199,7 @@ impl GooseClient {
         client_name: String,
         capabilities: GooseMcpClientCapabilities,
         working_dir: PathBuf,
+        extension_manager: Weak<ExtensionManager>,
     ) -> Self {
         GooseClient {
             notification_handlers: handlers,
@@ -203,6 +209,7 @@ impl GooseClient {
             client_name,
             capabilities,
             working_dir: Arc::new(tokio::sync::RwLock::new(working_dir)),
+            extension_manager,
         }
     }
 
@@ -219,6 +226,14 @@ impl GooseClient {
         *slot = Some(session_id.to_string());
     }
 
+    async fn handle_tool_list_changed(&self) {
+        if let Some(extension_manager) = self.extension_manager.upgrade() {
+            extension_manager
+                .invalidate_tools_cache_and_bump_version()
+                .await;
+        }
+    }
+
     async fn current_session_id(&self) -> Option<String> {
         self.session_id.lock().await.clone()
     }
@@ -362,6 +377,10 @@ impl ClientHandler for GooseClient {
             });
     }
 
+    async fn on_tool_list_changed(&self, _context: rmcp::service::NotificationContext<RoleClient>) {
+        self.handle_tool_list_changed().await;
+    }
+
     #[expect(deprecated)]
     async fn on_logging_message(
         &self,
@@ -570,6 +589,7 @@ impl McpClient {
         client_name: String,
         capabilities: GooseMcpClientCapabilities,
         working_dir: PathBuf,
+        extension_manager: Weak<ExtensionManager>,
     ) -> Result<Self, ClientInitializeError>
     where
         T: IntoTransport<RoleClient, E, A>,
@@ -583,10 +603,12 @@ impl McpClient {
             client_name,
             capabilities,
             working_dir,
+            extension_manager,
         )
         .await
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub async fn connect_with_container<T, E, A>(
         transport: T,
         timeout: std::time::Duration,
@@ -595,6 +617,7 @@ impl McpClient {
         client_name: String,
         capabilities: GooseMcpClientCapabilities,
         working_dir: PathBuf,
+        extension_manager: Weak<ExtensionManager>,
     ) -> Result<Self, ClientInitializeError>
     where
         T: IntoTransport<RoleClient, E, A>,
@@ -609,6 +632,7 @@ impl McpClient {
             client_name.clone(),
             capabilities.clone(),
             working_dir,
+            extension_manager,
         );
         let client: rmcp::service::RunningService<rmcp::RoleClient, GooseClient> =
             client.serve(transport).await?;
@@ -1004,9 +1028,62 @@ fn inject_session_context_into_request(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::agents::extension::ExtensionConfig;
     use crate::agents::GoosePlatform;
+    use rmcp::model::Tool;
     use serde_json::json;
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use test_case::test_case;
+    use tokio::sync::Semaphore;
+
+    struct BlockingToolsClient {
+        calls: AtomicUsize,
+        first_fetch_started: Semaphore,
+        release_first_fetch: Semaphore,
+    }
+
+    #[async_trait::async_trait]
+    impl McpClientTrait for BlockingToolsClient {
+        async fn list_tools(
+            &self,
+            _session_id: &str,
+            _next_cursor: Option<String>,
+            _cancel_token: CancellationToken,
+        ) -> Result<ListToolsResult, Error> {
+            let call = self.calls.fetch_add(1, Ordering::SeqCst);
+            let name = if call == 0 { "old" } else { "new" };
+
+            if call == 0 {
+                self.first_fetch_started.add_permits(1);
+                let _permit = self.release_first_fetch.acquire().await.unwrap();
+            }
+
+            Ok(ListToolsResult {
+                tools: vec![Tool::new(
+                    name,
+                    format!("{name} tool list"),
+                    Arc::new(JsonObject::new()),
+                )],
+                next_cursor: None,
+                meta: None,
+                ..Default::default()
+            })
+        }
+
+        async fn call_tool(
+            &self,
+            _ctx: &ToolCallContext,
+            _name: &str,
+            _arguments: Option<JsonObject>,
+            _cancel_token: CancellationToken,
+        ) -> Result<CallToolResult, Error> {
+            Ok(CallToolResult::success(vec![]))
+        }
+
+        fn get_info(&self) -> Option<&InitializeResult> {
+            None
+        }
+    }
 
     fn new_client(platform: GoosePlatform) -> GooseClient {
         let capabilities = match platform {
@@ -1026,9 +1103,74 @@ mod tests {
             platform.to_string(),
             capabilities,
             std::env::current_dir().unwrap_or_default(),
+            Weak::new(),
         )
     }
 
+    #[tokio::test]
+    async fn tool_list_changed_during_fetch_prevents_stale_cache() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let extension_manager = Arc::new(ExtensionManager::new_without_provider(
+            temp_dir.path().to_path_buf(),
+        ));
+        let tools_client = Arc::new(BlockingToolsClient {
+            calls: AtomicUsize::new(0),
+            first_fetch_started: Semaphore::new(0),
+            release_first_fetch: Semaphore::new(0),
+        });
+        let config = ExtensionConfig::Builtin {
+            name: "dynamic".to_string(),
+            display_name: Some("dynamic".to_string()),
+            description: "dynamic tools".to_string(),
+            timeout: None,
+            bundled: None,
+            available_tools: vec![],
+        };
+        extension_manager
+            .add_client(
+                "dynamic".to_string(),
+                config,
+                tools_client.clone(),
+                None,
+                None,
+            )
+            .await;
+
+        let goose_client = GooseClient::new(
+            Arc::new(Mutex::new(Vec::new())),
+            Arc::new(Mutex::new(None)),
+            "goose-test".to_string(),
+            GooseMcpClientCapabilities {
+                mcpui: false,
+                host_info: None,
+            },
+            temp_dir.path().to_path_buf(),
+            Arc::downgrade(&extension_manager),
+        );
+
+        let manager = extension_manager.clone();
+        let first_fetch = tokio::spawn(async move {
+            manager
+                .get_prefixed_tools("test-session", None)
+                .await
+                .unwrap()
+        });
+
+        let _started = tools_client.first_fetch_started.acquire().await.unwrap();
+        goose_client.handle_tool_list_changed().await;
+        tools_client.release_first_fetch.add_permits(1);
+
+        let stale_result = first_fetch.await.unwrap();
+        assert!(stale_result.iter().any(|tool| tool.name == "dynamic__old"));
+
+        let refreshed = extension_manager
+            .get_prefixed_tools("test-session", None)
+            .await
+            .unwrap();
+        assert!(refreshed.iter().any(|tool| tool.name == "dynamic__new"));
+        assert_eq!(tools_client.calls.load(Ordering::SeqCst), 2);
+    }
+
     fn request_extensions(request: &ClientRequest) -> Option<&Extensions> {
         match request {
             ClientRequest::ListResourcesRequest(req) => Some(&req.extensions),
@@ -1378,6 +1520,7 @@ mod tests {
                 }),
             },
             std::env::current_dir().unwrap_or_default(),
+            Weak::new(),
         );
 
         let info = ClientHandler::get_info(&client);
@@ -1409,6 +1552,7 @@ mod tests {
                 }),
             },
             std::env::current_dir().unwrap_or_default(),
+            Weak::new(),
         );
 
         let info = ClientHandler::get_info(&client);
@@ -1437,6 +1581,7 @@ mod tests {
                 }),
             },
             std::env::current_dir().unwrap_or_default(),
+            Weak::new(),
         );
 
         let info = ClientHandler::get_info(&client);