extract activate agent

crates/goose/src/acp/server.rs

@@ -12,7 +12,9 @@ use crate::agents::extension::{Envs, PLATFORM_EXTENSIONS};
 use crate::agents::extension_manager::TRUSTED_TOOL_UPDATE_META_KEY;
 use crate::agents::mcp_client::{GooseMcpHostInfo, McpClientTrait};
 use crate::agents::platform_extensions::developer::DeveloperClient;
-use crate::agents::{Agent, AgentConfig, ExtensionConfig, GoosePlatform, SessionConfig};
+use crate::agents::{
+    Agent, AgentConfig, ExtensionConfig, ExtensionLoadResult, GoosePlatform, SessionConfig,
+};
 use crate::config::base::CONFIG_YAML_NAME;
 use crate::config::extensions::get_enabled_extensions_with_config;
 use crate::config::paths::Paths;
@@ -32,7 +34,7 @@ use crate::providers::inventory::{
     RefreshPlan, RefreshSkipReason,
 };
 use crate::session::session_manager::{SessionListCursor, SessionType};
-use crate::session::{EnabledExtensionsState, ExtensionState, Session, SessionManager};
+use crate::session::{EnabledExtensionsState, Session, SessionManager};
 use crate::source_roots::SourceRoot;
 use crate::utils::sanitize_unicode_tags;
 use agent_client_protocol::schema::{
@@ -1331,24 +1333,13 @@ impl GooseAcpAgent {
             .await;
     }
 
-    async fn get_or_create_session_agent(
-        &self,
-        cx: &ConnectionTo<Client>,
-        session_id: String,
-    ) -> Result<Arc<Agent>, agent_client_protocol::Error> {
-        Ok(self
-            .get_or_create_session_agent_with_results(cx, session_id)
-            .await?
-            .agent)
-    }
-
     async fn get_or_create_session_agent_with_results(
         &self,
         cx: &ConnectionTo<Client>,
         session_id: String,
     ) -> Result<AgentManagerGetResult, agent_client_protocol::Error> {
         self.agent_manager
-            .get_or_create_agent_with_runtime_context_result(
+            .get_or_create_agent_with_runtime_context(
                 session_id,
                 RuntimeContext {
                     mcp_host_info: self.client_mcp_host_info.get().cloned(),
@@ -1387,29 +1378,6 @@ impl GooseAcpAgent {
         Ok(extensions)
     }
 
-    async fn persist_session_extensions(
-        &self,
-        session: &Session,
-        extensions: Vec<ExtensionConfig>,
-    ) -> Result<Session, agent_client_protocol::Error> {
-        let mut extension_data = session.extension_data.clone();
-        EnabledExtensionsState::new(extensions)
-            .to_extension_data(&mut extension_data)
-            .internal_err_ctx("Failed to initialize session extensions")?;
-
-        self.session_manager
-            .update(&session.id)
-            .extension_data(extension_data)
-            .apply()
-            .await
-            .internal_err_ctx("Failed to update session extensions")?;
-
-        self.session_manager
-            .get_session(&session.id, false)
-            .await
-            .internal_err_ctx("Failed to reload session")
-    }
-
     async fn apply_acp_extension_overrides(
         &self,
         cx: &ConnectionTo<Client>,
@@ -1470,6 +1438,38 @@ impl GooseAcpAgent {
             .await;
     }
 
+    async fn activate_acp_session(
+        &self,
+        cx: &ConnectionTo<Client>,
+        session: &Session,
+        tool_requests: HashMap<String, ToolRequest>,
+    ) -> Result<(Arc<Agent>, Vec<ExtensionLoadResult>), agent_client_protocol::Error> {
+        let agent_result = self
+            .get_or_create_session_agent_with_results(cx, session.id.clone())
+            .await?;
+        let agent = agent_result.agent.clone();
+        self.apply_acp_extension_overrides(cx, &agent, session)
+            .await;
+
+        let acp_session = GooseAcpSession {
+            agent: agent.clone(),
+            tool_requests,
+            chain_membership: HashMap::new(),
+            responded_tool_ids: HashSet::new(),
+            summarized_chains: HashSet::new(),
+            cancel_token: None,
+        };
+        self.sessions
+            .lock()
+            .await
+            .insert(session.id.clone(), acp_session);
+
+        self.maybe_refresh_provider_inventory_with_agent(session, &agent)
+            .await;
+
+        Ok((agent, agent_result.extension_results))
+    }
+
     async fn build_eager_session_config(
         &self,
         mode_state: &SessionModeState,

crates/goose/src/acp/server/new_session.rs

@@ -4,13 +4,14 @@ use crate::acp::server::{
 use crate::config::{Config, GooseMode};
 use crate::model::ModelConfig;
 use crate::session::SessionType;
+use crate::session::{EnabledExtensionsState, ExtensionState};
 
-use super::{GooseAcpAgent, GooseAcpSession};
+use super::GooseAcpAgent;
 use agent_client_protocol::schema::{
     NewSessionRequest, NewSessionResponse, SessionId, SessionNotification, SessionUpdate,
 };
 use agent_client_protocol::{Client, ConnectionTo};
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use tracing::debug;
 
 impl GooseAcpAgent {
@@ -56,9 +57,15 @@ impl GooseAcpAgent {
             .await
             .internal_err_ctx("Failed to create session")?;
         let mut builder = self.session_manager.update(&goose_session.id);
+        let extensions = self.initial_session_extensions(config, args.mcp_servers)?;
+        let mut extension_data = goose_session.extension_data.clone();
+        EnabledExtensionsState::new(extensions)
+            .to_extension_data(&mut extension_data)
+            .internal_err_ctx("Failed to initialize session extensions")?;
         builder = builder
             .provider_name(resolved_provider)
-            .model_config(resolved_model_config);
+            .model_config(resolved_model_config)
+            .extension_data(extension_data);
         if let Some(pid) = project_id {
             builder = builder.project_id(Some(pid));
         }
@@ -72,36 +79,13 @@ impl GooseAcpAgent {
             .get_session(&goose_session.id, false)
             .await
             .internal_err_ctx("Failed to reload session")?;
-        let extensions = self.initial_session_extensions(config, args.mcp_servers)?;
-        goose_session = self
-            .persist_session_extensions(&goose_session, extensions)
-            .await?;
         let session_id_str = goose_session.id.clone();
         let sid = sid_short(&session_id_str);
         debug!(target: "perf", sid = %sid, ms = t0.elapsed().as_millis() as u64, "perf: new_session create_session");
 
-        let agent_result = self
-            .get_or_create_session_agent_with_results(cx, session_id_str.clone())
+        let (_agent, extension_results) = self
+            .activate_acp_session(cx, &goose_session, HashMap::new())
             .await?;
-        let agent = agent_result.agent.clone();
-        self.apply_acp_extension_overrides(cx, &agent, &goose_session)
-            .await;
-
-        let acp_session = GooseAcpSession {
-            agent: agent.clone(),
-            tool_requests: HashMap::new(),
-            chain_membership: HashMap::new(),
-            responded_tool_ids: HashSet::new(),
-            summarized_chains: HashSet::new(),
-            cancel_token: None,
-        };
-        self.sessions
-            .lock()
-            .await
-            .insert(session_id_str.clone(), acp_session);
-
-        self.maybe_refresh_provider_inventory_with_agent(&goose_session, &agent)
-            .await;
 
         let goose_session = self
             .session_manager
@@ -126,7 +110,7 @@ impl GooseAcpAgent {
         if let Some(co) = config_options {
             response = response.config_options(co);
         }
-        if let Ok(extension_results) = serde_json::to_value(&agent_result.extension_results) {
+        if let Ok(extension_results) = serde_json::to_value(&extension_results) {
             let mut meta = serde_json::Map::new();
             meta.insert("extensionResults".to_string(), extension_results);
             response = response.meta(meta);

crates/goose/src/execution/manager.rs

@@ -112,22 +112,13 @@ impl AgentManager {
     }
 
     pub async fn get_or_create_agent(&self, session_id: String) -> Result<Arc<Agent>> {
-        self.get_or_create_agent_with_runtime_context(session_id, RuntimeContext::default())
-            .await
-    }
-
-    pub async fn get_or_create_agent_with_runtime_context(
-        &self,
-        session_id: String,
-        runtime_context: RuntimeContext,
-    ) -> Result<Arc<Agent>> {
         Ok(self
-            .get_or_create_agent_with_runtime_context_result(session_id, runtime_context)
+            .get_or_create_agent_with_runtime_context(session_id, RuntimeContext::default())
             .await?
             .agent)
     }
 
-    pub async fn get_or_create_agent_with_runtime_context_result(
+    pub async fn get_or_create_agent_with_runtime_context(
         &self,
         session_id: String,
         runtime_context: RuntimeContext,