Refactor: move persisting extension to session outside of route (#6685)

crates/goose-acp/src/server.rs

@@ -273,7 +273,11 @@ async fn add_builtins(agent: &Agent, builtins: Vec<String>) {
             }
         };
 
-        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)));
             }

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!(

crates/goose-cli/src/commands/web.rs

@@ -186,7 +186,7 @@ async fn create_agent(provider_name: &str, model: &str) -> Result<Agent> {
 
     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);
         }
     }

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<dyn goose::providers::base::Provider>,
     interactive: bool,
+    session_id: &str,
 ) -> Arc<Agent> {
     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<usize>| {
@@ -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;
 

crates/goose-cli/src/session/mod.rs

@@ -336,16 +336,11 @@ impl CliSession {
     async fn add_and_persist_extensions(&mut self, configs: Vec<ExtensionConfig>) -> 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(())

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<RemoveExtensionRequest>,
 ) -> Result<StatusCode, ErrorResponse> {
     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,
             }
         })?;

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 {

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<std::path::PathBuf>,
+        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(())
     }
 

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<Self>,
         config: ExtensionConfig,
         working_dir: Option<PathBuf>,

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!(

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();
 

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(),

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());

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)

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<dyn std::error::Error>> {
         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 {