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

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 {