feat(goose-acp): enable parallel sessions with isolated agent state (#6392)

Signed-off-by: Adrian Cole <adrian@tetrate.io>

crates/goose/src/agents/code_execution_extension.rs

@@ -1,6 +1,6 @@
 use crate::agents::extension::PlatformExtensionContext;
 use crate::agents::extension_manager::get_parameter_names;
-use crate::agents::mcp_client::{Error, McpClientTrait};
+use crate::agents::mcp_client::{Error, McpClientTrait, McpMeta};
 use anyhow::Result;
 use async_trait::async_trait;
 use boa_engine::builtins::promise::PromiseState;
@@ -9,10 +9,9 @@ use boa_engine::{js_string, Context, JsNativeError, JsString, JsValue, NativeFun
 use indoc::indoc;
 use regex::Regex;
 use rmcp::model::{
-    CallToolRequestParam, CallToolResult, Content, GetPromptResult, Implementation,
-    InitializeResult, JsonObject, ListPromptsResult, ListResourcesResult, ListToolsResult,
-    ProtocolVersion, RawContent, ReadResourceResult, ServerCapabilities, ServerNotification,
-    Tool as McpTool, ToolAnnotations, ToolsCapability,
+    CallToolRequestParam, CallToolResult, Content, Implementation, InitializeResult, JsonObject,
+    ListToolsResult, ProtocolVersion, RawContent, ServerCapabilities, Tool as McpTool,
+    ToolAnnotations, ToolsCapability,
 };
 use schemars::{schema_for, JsonSchema};
 use serde::{Deserialize, Serialize};
@@ -478,6 +477,7 @@ impl CodeExecutionClient {
 
     async fn handle_execute_code(
         &self,
+        session_id: &str,
         arguments: Option<JsonObject>,
     ) -> Result<Vec<Content>, String> {
         let code = arguments
@@ -490,6 +490,7 @@ impl CodeExecutionClient {
         let tools = self.get_tool_infos().await;
         let (call_tx, call_rx) = mpsc::unbounded_channel();
         let tool_handler = tokio::spawn(Self::run_tool_handler(
+            session_id.to_string(),
             call_rx,
             self.context.extension_manager.clone(),
         ));
@@ -658,6 +659,7 @@ impl CodeExecutionClient {
     }
 
     async fn run_tool_handler(
+        session_id: String,
         mut call_rx: mpsc::UnboundedReceiver<ToolCallRequest>,
         extension_manager: Option<std::sync::Weak<crate::agents::ExtensionManager>>,
     ) {
@@ -669,7 +671,7 @@ impl CodeExecutionClient {
                         arguments: serde_json::from_str(&arguments).ok(),
                     };
                     match manager
-                        .dispatch_tool_call(tool_call, CancellationToken::new())
+                        .dispatch_tool_call(&session_id, tool_call, CancellationToken::new())
                         .await
                     {
                         Ok(dispatch_result) => match dispatch_result.result.await {
@@ -700,22 +702,6 @@ impl CodeExecutionClient {
 
 #[async_trait]
 impl McpClientTrait for CodeExecutionClient {
-    async fn list_resources(
-        &self,
-        _next_cursor: Option<String>,
-        _cancellation_token: CancellationToken,
-    ) -> Result<ListResourcesResult, Error> {
-        Err(Error::TransportClosed)
-    }
-
-    async fn read_resource(
-        &self,
-        _uri: &str,
-        _cancellation_token: CancellationToken,
-    ) -> Result<ReadResourceResult, Error> {
-        Err(Error::TransportClosed)
-    }
-
     #[allow(clippy::too_many_lines)]
     async fn list_tools(
         &self,
@@ -845,10 +831,11 @@ impl McpClientTrait for CodeExecutionClient {
         &self,
         name: &str,
         arguments: Option<JsonObject>,
+        meta: McpMeta,
         _cancellation_token: CancellationToken,
     ) -> Result<CallToolResult, Error> {
         let content = match name {
-            "execute_code" => self.handle_execute_code(arguments).await,
+            "execute_code" => self.handle_execute_code(&meta.session_id, arguments).await,
             "read_module" => self.handle_read_module(arguments).await,
             "search_modules" => self.handle_search_modules(arguments).await,
             _ => Err(format!("Unknown tool: {name}")),
@@ -862,32 +849,11 @@ impl McpClientTrait for CodeExecutionClient {
         }
     }
 
-    async fn list_prompts(
-        &self,
-        _next_cursor: Option<String>,
-        _cancellation_token: CancellationToken,
-    ) -> Result<ListPromptsResult, Error> {
-        Err(Error::TransportClosed)
-    }
-
-    async fn get_prompt(
-        &self,
-        _name: &str,
-        _arguments: Value,
-        _cancellation_token: CancellationToken,
-    ) -> Result<GetPromptResult, Error> {
-        Err(Error::TransportClosed)
-    }
-
-    async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
-        mpsc::channel(1).1
-    }
-
     fn get_info(&self) -> Option<&InitializeResult> {
         Some(&self.info)
     }
 
-    async fn get_moim(&self) -> Option<String> {
+    async fn get_moim(&self, _session_id: &str) -> Option<String> {
         let tools = self.get_tool_infos().await;
         if tools.is_empty() {
             return None;
@@ -923,9 +889,13 @@ mod tests {
 
     #[tokio::test]
     async fn test_execute_code_simple() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let session_manager = Arc::new(crate::session::SessionManager::new(
+            temp_dir.path().to_path_buf(),
+        ));
         let context = PlatformExtensionContext {
-            session_id: None,
             extension_manager: None,
+            session_manager,
         };
         let client = CodeExecutionClient::new(context).unwrap();
 
@@ -936,7 +906,12 @@ mod tests {
         );
 
         let result = client
-            .call_tool("execute_code", Some(args), CancellationToken::new())
+            .call_tool(
+                "execute_code",
+                Some(args),
+                McpMeta::new("test-session-id"),
+                CancellationToken::new(),
+            )
             .await
             .unwrap();
 
@@ -982,9 +957,13 @@ mod tests {
 
     #[tokio::test]
     async fn test_read_module_not_found() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let session_manager = Arc::new(crate::session::SessionManager::new(
+            temp_dir.path().to_path_buf(),
+        ));
         let context = PlatformExtensionContext {
-            session_id: None,
             extension_manager: None,
+            session_manager,
         };
         let client = CodeExecutionClient::new(context).unwrap();