diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index ae63e2bcbd..9a5dbc85be 100644 --- a/crates/goose-acp/src/server.rs +++ b/crates/goose-acp/src/server.rs @@ -273,7 +273,11 @@ async fn add_builtins(agent: &Agent, builtins: Vec) { } }; - match agent.add_extension(config).await { + match agent + .extension_manager + .add_extension(config, None, None) + .await + { Ok(_) => info!(extension = %builtin, "extension loaded"), Err(e) => warn!(extension = %builtin, error = %e, "extension load failed"), } @@ -728,7 +732,7 @@ impl GooseAcpAgent { } }; let name = config.name().to_string(); - if let Err(e) = self.agent.add_extension(config).await { + if let Err(e) = self.agent.add_extension(config, &goose_session.id).await { return Err(sacp::Error::internal_error() .data(format!("Failed to add MCP server '{}': {}", name, e))); } diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index d058f25413..5f5e581b8b 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -1448,7 +1448,7 @@ pub async fn configure_tool_permissions_dialog() -> anyhow::Result<()> { agent.update_provider(new_provider, &session.id).await?; if let Some(config) = get_extension_by_name(&selected_extension_name) { agent - .add_extension(config.clone()) + .add_extension(config.clone(), &session.id) .await .unwrap_or_else(|_| { println!( diff --git a/crates/goose-cli/src/commands/web.rs b/crates/goose-cli/src/commands/web.rs index f916bb8618..143c753ff3 100644 --- a/crates/goose-cli/src/commands/web.rs +++ b/crates/goose-cli/src/commands/web.rs @@ -186,7 +186,7 @@ async fn create_agent(provider_name: &str, model: &str) -> Result { let enabled_configs = goose::config::get_enabled_extensions(); for config in enabled_configs { - if let Err(e) = agent.add_extension(config.clone()).await { + if let Err(e) = agent.add_extension(config.clone(), &init_session.id).await { eprintln!("Warning: Failed to load extension {}: {}", config.name(), e); } } diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index e31c0b0209..7599ce37e8 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/crates/goose-cli/src/session/builder.rs @@ -211,7 +211,10 @@ async fn offer_extension_debugging_help( let extensions = get_all_extensions(); for ext_wrapper in extensions { if ext_wrapper.enabled && ext_wrapper.config.name() == "developer" { - if let Err(e) = debug_agent.add_extension(ext_wrapper.config).await { + if let Err(e) = debug_agent + .add_extension(ext_wrapper.config, &session.id) + .await + { // If we can't add developer extension, continue without it eprintln!( "Note: Could not load developer extension for debugging: {}", @@ -258,6 +261,7 @@ async fn load_extensions( extensions_to_load: Vec<(String, ExtensionConfig)>, provider_for_debug: Arc, interactive: bool, + session_id: &str, ) -> Arc { let mut set = JoinSet::new(); let agent_ptr = Arc::new(agent); @@ -266,7 +270,8 @@ async fn load_extensions( for (id, (_label, extension)) in extensions_to_load.iter().enumerate() { let agent_ptr = agent_ptr.clone(); let cfg = extension.clone(); - set.spawn(async move { (id, agent_ptr.add_extension(cfg).await) }); + let sid = session_id.to_string(); + set.spawn(async move { (id, agent_ptr.add_extension(cfg, &sid).await) }); } let get_message = |waiting_ids: &BTreeSet| { @@ -554,6 +559,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { extensions_to_load, Arc::clone(&provider_for_display), session_config.interactive, + &session_id, ) .await; diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 3697bd6499..26447c5d78 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -336,16 +336,11 @@ impl CliSession { async fn add_and_persist_extensions(&mut self, configs: Vec) -> Result<()> { for config in configs { self.agent - .add_extension(config) + .add_extension(config, &self.session_id) .await .map_err(|e| anyhow::anyhow!("Failed to start extension: {}", e))?; } - self.agent - .persist_extension_state(&self.session_id) - .await - .map_err(|e| anyhow::anyhow!("Failed to save extension state: {}", e))?; - self.invalidate_completion_cache().await; Ok(()) diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 588a2dfb5c..17061cf684 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -586,23 +586,15 @@ async fn agent_add_extension( let extension_name = request.config.name(); let agent = state.get_agent(request.session_id.clone()).await?; - agent.add_extension(request.config).await.map_err(|e| { - goose::posthog::emit_error( - "extension_add_failed", - &format!("{}: {}", extension_name, e), - ); - ErrorResponse::internal(format!("Failed to add extension: {}", e)) - })?; - - // Persist here rather than in add_extension to ensure we only save state - // after the extension successfully loads. This prevents failed extensions - // from being persisted as enabled in the session. agent - .persist_extension_state(&request.session_id) + .add_extension(request.config, &request.session_id) .await .map_err(|e| { - error!("Failed to persist extension state: {}", e); - ErrorResponse::internal(format!("Failed to persist extension state: {}", e)) + goose::posthog::emit_error( + "extension_add_failed", + &format!("{}: {}", extension_name, e), + ); + ErrorResponse::internal(format!("Failed to add extension: {}", e)) })?; Ok(StatusCode::OK) @@ -624,15 +616,14 @@ async fn agent_remove_extension( Json(request): Json, ) -> Result { let agent = state.get_agent(request.session_id.clone()).await?; - agent.remove_extension(&request.name).await?; agent - .persist_extension_state(&request.session_id) + .remove_extension(&request.name, &request.session_id) .await .map_err(|e| { - error!("Failed to persist extension state: {}", e); + error!("Failed to remove extension: {}", e); ErrorResponse { - message: format!("Failed to persist extension state: {}", e), + message: format!("Failed to remove extension: {}", e), status: StatusCode::INTERNAL_SERVER_ERROR, } })?; diff --git a/crates/goose/examples/agent.rs b/crates/goose/examples/agent.rs index 6274fd1da0..41c61408c6 100644 --- a/crates/goose/examples/agent.rs +++ b/crates/goose/examples/agent.rs @@ -35,7 +35,7 @@ async fn main() -> anyhow::Result<()> { DEFAULT_EXTENSION_TIMEOUT, ) .with_args(vec!["mcp", "developer"]); - agent.add_extension(config).await?; + agent.add_extension(config, &session.id).await?; println!("Extensions:"); for extension in agent.list_extensions().await { diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 95e3eabbfb..798f54e25f 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -678,15 +678,14 @@ impl Agent { } }; - // Capture the session's working_dir to pass to extensions - let working_dir = session.working_dir.clone(); + let session_id = session.id.clone(); let extension_futures = enabled_configs .into_iter() .map(|config| { let config_clone = config.clone(); let agent_ref = self.clone(); - let working_dir_clone = working_dir.clone(); + let session_id_clone = session_id.clone(); async move { let name = config_clone.name().to_string(); @@ -705,7 +704,7 @@ impl Agent { } match agent_ref - .add_extension_with_working_dir(config_clone, Some(working_dir_clone)) + .add_extension(config_clone, &session_id_clone) .await { Ok(_) => ExtensionLoadResult { @@ -730,15 +729,24 @@ impl Agent { futures::future::join_all(extension_futures).await } - pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> { - self.add_extension_with_working_dir(extension, None).await - } - - pub async fn add_extension_with_working_dir( + pub async fn add_extension( &self, extension: ExtensionConfig, - working_dir: Option, + session_id: &str, ) -> ExtensionResult<()> { + let session = self + .config + .session_manager + .get_session(session_id, false) + .await + .map_err(|e| { + crate::agents::extension::ExtensionError::SetupError(format!( + "Failed to get session '{}': {}", + session_id, e + )) + })?; + let working_dir = Some(session.working_dir); + match &extension { ExtensionConfig::Frontend { tools, @@ -768,15 +776,22 @@ impl Agent { _ => { let container = self.container.lock().await; self.extension_manager - .add_extension_with_working_dir( - extension.clone(), - working_dir, - container.as_ref(), - ) + .add_extension(extension.clone(), working_dir, container.as_ref()) .await?; } } + // Persist extension state after successful add + self.persist_extension_state(session_id) + .await + .map_err(|e| { + error!("Failed to persist extension state: {}", e); + crate::agents::extension::ExtensionError::SetupError(format!( + "Failed to persist extension state: {}", + e + )) + })?; + Ok(()) } @@ -833,8 +848,17 @@ impl Agent { prefixed_tools } - pub async fn remove_extension(&self, name: &str) -> Result<()> { + pub async fn remove_extension(&self, name: &str, session_id: &str) -> Result<()> { self.extension_manager.remove_extension(name).await?; + + // Persist extension state after successful removal + self.persist_extension_state(session_id) + .await + .map_err(|e| { + error!("Failed to persist extension state: {}", e); + anyhow!("Failed to persist extension state: {}", e) + })?; + Ok(()) } diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index c67b8fac08..2f7695f3e8 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -484,7 +484,7 @@ impl ExtensionManager { /// Add an extension with an optional working directory. /// If working_dir is None, falls back to current_dir. #[allow(clippy::too_many_lines)] - pub async fn add_extension_with_working_dir( + pub async fn add_extension( self: &Arc, config: ExtensionConfig, working_dir: Option, diff --git a/crates/goose/src/agents/extension_manager_extension.rs b/crates/goose/src/agents/extension_manager_extension.rs index fbb4f064aa..cf163e4d8c 100644 --- a/crates/goose/src/agents/extension_manager_extension.rs +++ b/crates/goose/src/agents/extension_manager_extension.rs @@ -211,7 +211,7 @@ impl ExtensionManagerClient { }; extension_manager - .add_extension_with_working_dir(config, None, None) + .add_extension(config, None, None) .await .map(|_| { vec![Content::text(format!( diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 658d35472f..854d87eca1 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -491,14 +491,17 @@ mod tests { ]; agent - .add_extension(crate::agents::extension::ExtensionConfig::Frontend { - name: "frontend".to_string(), - description: "desc".to_string(), - tools: frontend_tools, - instructions: None, - bundled: None, - available_tools: vec![], - }) + .add_extension( + crate::agents::extension::ExtensionConfig::Frontend { + name: "frontend".to_string(), + description: "desc".to_string(), + tools: frontend_tools, + instructions: None, + bundled: None, + available_tools: vec![], + }, + &session.id, + ) .await .unwrap(); diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index 54e0d2dab8..d4ee6fba7a 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -133,7 +133,7 @@ fn get_agent_messages( .map_err(|e| anyhow!("Failed to set provider on sub agent: {}", e))?; for extension in task_config.extensions { - if let Err(e) = agent.add_extension(extension.clone()).await { + if let Err(e) = agent.add_extension(extension.clone(), &session_id).await { debug!( "Failed to add extension '{}' to subagent: {}", extension.name(), diff --git a/crates/goose/src/scheduler.rs b/crates/goose/src/scheduler.rs index c3d2a46a0a..31c41ad88a 100644 --- a/crates/goose/src/scheduler.rs +++ b/crates/goose/src/scheduler.rs @@ -743,11 +743,6 @@ async fn execute_job( let agent_provider = create(&provider_name, model_config).await?; - let extensions = resolve_extensions_for_new_session(recipe.extensions.as_deref(), None); - for ext in extensions { - agent.add_extension(ext.clone()).await?; - } - let session = agent .config .session_manager @@ -760,6 +755,11 @@ async fn execute_job( agent.update_provider(agent_provider, &session.id).await?; + let extensions = resolve_extensions_for_new_session(recipe.extensions.as_deref(), None); + for ext in extensions { + agent.add_extension(ext.clone(), &session.id).await?; + } + let mut jobs_guard = jobs.lock().await; if let Some((_, job_def)) = jobs_guard.get_mut(job_id.as_str()) { job_def.current_session_id = Some(session.id.clone()); diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index ee6493d788..4f353e6c2f 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -503,6 +503,8 @@ mod tests { use goose::session::SessionManager; async fn setup_agent_with_extension_manager() -> (Agent, String) { + use goose::session::session_manager::SessionType; + // Add the TODO extension to the config so it can be discovered by search_available_extensions // Set it as disabled initially so tests can enable it let todo_extension_entry = ExtensionEntry { @@ -522,9 +524,8 @@ mod tests { // Create agent with session_id from the start let temp_dir = tempfile::tempdir().unwrap(); let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf())); - let session_id = "test-session-id".to_string(); let config = AgentConfig::new( - session_manager, + session_manager.clone(), PermissionManager::instance(), None, GooseMode::Auto, @@ -532,6 +533,16 @@ mod tests { let agent = Agent::with_config(config); + let session = session_manager + .create_session( + std::path::PathBuf::from("."), + "Test Session".to_string(), + SessionType::Hidden, + ) + .await + .expect("Failed to create session"); + let session_id = session.id; + // Now add the extension manager platform extension let ext_config = ExtensionConfig::Platform { name: "extensionmanager".to_string(), @@ -542,7 +553,7 @@ mod tests { }; agent - .add_extension(ext_config) + .add_extension(ext_config, &session_id) .await .expect("Failed to add extension manager"); (agent, session_id) diff --git a/crates/goose/tests/mcp_integration_test.rs b/crates/goose/tests/mcp_integration_test.rs index ecc11cde05..635c8d6602 100644 --- a/crates/goose/tests/mcp_integration_test.rs +++ b/crates/goose/tests/mcp_integration_test.rs @@ -248,7 +248,7 @@ async fn test_replayed_session( #[allow(clippy::redundant_closure_call)] let result = (async || -> Result<(), Box> { extension_manager - .add_extension_with_working_dir(extension_config, None, None) + .add_extension(extension_config, None, None) .await?; let mut results = Vec::new(); for tool_call in tool_calls {