extract common behaviour for setup sessions

crates/goose/src/acp/server.rs

@@ -86,12 +86,14 @@ mod custom_dispatch;
 mod dictation;
 mod dispatch;
 mod extensions;
+mod fork_session;
 mod load_session;
 mod new_session;
 mod new_session_agent_manager;
 mod onboarding;
 mod providers;
 mod resources;
+mod session_setup;
 mod sessions;
 mod sources;
 mod tools;
@@ -202,6 +204,7 @@ struct ToolChain {
     message_id: String,
 }
 
+#[allow(dead_code)]
 struct SessionInitState {
     mode_state: SessionModeState,
     resolved_provider: Result<(String, crate::model::ModelConfig), String>,
@@ -235,6 +238,7 @@ pub struct GooseAcpAgent {
     permission_manager: Arc<PermissionManager>,
     disable_session_naming: bool,
     provider_inventory: ProviderInventoryService,
+    #[allow(dead_code)]
     goose_platform: GoosePlatform,
     additional_source_roots: Vec<SourceRoot>,
 }
@@ -242,7 +246,7 @@ pub struct GooseAcpAgent {
 /// Shorten a session/thread id for perf log correlation.
 /// All `perf:` logs use `sid=<8-char-prefix>` so a single session's activity
 /// can be extracted with `grep 'perf:' <log> | grep 'sid=abc12345'`.
-fn sid_short(id: &str) -> String {
+pub(super) fn sid_short(id: &str) -> String {
     id.chars().take(8).collect()
 }
 
@@ -347,7 +351,7 @@ fn display_title(s: &Session) -> Option<String> {
     }
 }
 
-fn session_meta(session: &Session) -> serde_json::Map<String, serde_json::Value> {
+pub(super) fn session_meta(session: &Session) -> serde_json::Map<String, serde_json::Value> {
     let mut meta = serde_json::Map::new();
     meta.insert(
         "messageCount".to_string(),
@@ -946,6 +950,7 @@ fn builtin_to_extension_config(name: &str) -> ExtensionConfig {
     }
 }
 
+#[allow(dead_code)]
 async fn provider_default_model_config(
     provider_name: &str,
 ) -> Result<crate::model::ModelConfig, String> {
@@ -958,6 +963,7 @@ async fn provider_default_model_config(
         .map(|model_config| model_config.with_canonical_limits(provider_name))
 }
 
+#[allow(dead_code)]
 fn global_model_config(
     config: &Config,
     provider_name: &str,
@@ -968,6 +974,7 @@ fn global_model_config(
         .map(|model_config| model_config.with_canonical_limits(provider_name))
 }
 
+#[allow(dead_code)]
 async fn resolve_provider_and_model_config(
     config: &Config,
     provider_selection: Option<&str>,
@@ -994,6 +1001,7 @@ async fn resolve_provider_and_model_config(
 
 /// Resolve the provider name and model config for a session from an
 /// already-loaded `Config`.
+#[allow(dead_code)]
 async fn resolve_provider_and_model_from_config(
     config: &Config,
     goose_session: &Session,
@@ -1043,7 +1051,8 @@ fn with_preserved_session_request_params(
 
 /// Convenience wrapper: reads config from disk, then resolves provider + model.
 /// Cheap enough to call from `on_new_session` (file + registry reads, no network).
-async fn resolve_provider_and_model(
+#[allow(dead_code)]
+pub(super) async fn resolve_provider_and_model(
     config_dir: &std::path::Path,
     goose_session: &Session,
 ) -> Result<(String, crate::model::ModelConfig), String> {
@@ -1090,7 +1099,7 @@ fn build_usage_updates(session: &Session) -> Option<UsageUpdates> {
     })
 }
 
-fn validate_absolute_cwd(cwd: &Path) -> Result<(), agent_client_protocol::Error> {
+pub(super) fn validate_absolute_cwd(cwd: &Path) -> Result<(), agent_client_protocol::Error> {
     if !cwd.is_absolute() {
         return Err(
             agent_client_protocol::Error::invalid_params().data("cwd must be an absolute path")
@@ -1202,6 +1211,7 @@ impl GooseAcpAgent {
         .await
     }
 
+    #[allow(dead_code)]
     async fn prepare_session_init_config(
         &self,
         resolved: &Result<(String, crate::model::ModelConfig), String>,
@@ -1252,6 +1262,7 @@ impl GooseAcpAgent {
         (Some(model_state), Some(config_options), prebuilt_provider)
     }
 
+    #[allow(dead_code)]
     async fn maybe_refresh_provider_inventory(&self, goose_session: &Session) {
         let Some(provider_name) = goose_session.provider_name.as_deref() else {
             return;
@@ -1438,11 +1449,10 @@ impl GooseAcpAgent {
             .await;
     }
 
-    async fn activate_acp_session(
+    async fn prepare_acp_session_agent(
         &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())
@@ -1450,24 +1460,40 @@ impl GooseAcpAgent {
         let agent = agent_result.agent.clone();
         self.apply_acp_extension_overrides(cx, &agent, session)
             .await;
+        self.maybe_refresh_provider_inventory_with_agent(session, &agent)
+            .await;
 
+        Ok((agent, agent_result.extension_results))
+    }
+
+    async fn register_acp_session(
+        &self,
+        session_id: String,
+        agent: Arc<Agent>,
+        tool_requests: HashMap<String, ToolRequest>,
+    ) {
         let acp_session = GooseAcpSession {
-            agent: agent.clone(),
+            agent,
             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.sessions.lock().await.insert(session_id, acp_session);
+    }
 
-        self.maybe_refresh_provider_inventory_with_agent(session, &agent)
+    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, extension_results) = self.prepare_acp_session_agent(cx, session).await?;
+        self.register_acp_session(session.id.clone(), agent.clone(), tool_requests)
             .await;
 
-        Ok((agent, agent_result.extension_results))
+        Ok((agent, extension_results))
     }
 
     async fn build_eager_session_config(
@@ -1500,6 +1526,7 @@ impl GooseAcpAgent {
         (Some(model_state), Some(config_options))
     }
 
+    #[allow(dead_code)]
     async fn build_session_provider(
         &self,
         provider_name: &str,
@@ -1647,6 +1674,7 @@ impl GooseAcpAgent {
         }
     }
 
+    #[allow(dead_code)]
     async fn prepare_session_init_state(
         &self,
         goose_session: &Session,
@@ -2609,127 +2637,7 @@ impl GooseAcpAgent {
         cx: &ConnectionTo<Client>,
         args: NewSessionRequest,
     ) -> Result<NewSessionResponse, agent_client_protocol::Error> {
-        debug!(?args, "new session request");
-        let t_start = std::time::Instant::now();
-        validate_absolute_cwd(&args.cwd)?;
-
-        let requested_provider = meta_string(args.meta.as_ref(), "provider");
-        let project_id = meta_string(args.meta.as_ref(), "projectId");
-        let config = self.config()?;
-        let (resolved_provider, resolved_model_config) =
-            resolve_provider_and_model_config(&config, requested_provider.as_deref(), None)
-                .await
-                .map_err(|error| {
-                    agent_client_protocol::Error::internal_error()
-                        .data(format!("Failed to resolve provider: {}", error))
-                })?;
-
-        // When _meta.client is set, the session is created by a known client
-        // (e.g. "goose" for the desktop app) and treated as a User session.
-        // Without it, sessions default to Acp for programmatic ACP clients.
-        let session_type = match meta_string(args.meta.as_ref(), "client") {
-            Some(_) => SessionType::User,
-            None => SessionType::Acp,
-        };
-
-        let current_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
-
-        let t0 = std::time::Instant::now();
-        let goose_session = self
-            .session_manager
-            .create_session(
-                args.cwd.clone(),
-                "New Chat".to_string(),
-                session_type,
-                current_mode,
-            )
-            .await
-            .internal_err_ctx("Failed to create session")?;
-
-        let mut builder = self.session_manager.update(&goose_session.id);
-        builder = builder
-            .provider_name(resolved_provider)
-            .model_config(resolved_model_config);
-        if let Some(pid) = project_id {
-            builder = builder.project_id(Some(pid));
-        }
-        builder
-            .apply()
-            .await
-            .internal_err_ctx("Failed to update session")?;
-
-        let goose_session = self
-            .session_manager
-            .get_session(&goose_session.id, false)
-            .await
-            .internal_err_ctx("Failed to reload session")?;
-
-        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 acp_session_id = SessionId::new(session_id_str.clone());
-        let init_state = self.prepare_session_init_state(&goose_session).await?;
-
-        let working_dir = goose_session.working_dir.clone();
-
-        let (agent, _extension_results) = self
-            .build_agent_for_session(
-                cx,
-                &goose_session,
-                init_state.resolved_provider.as_ref().ok().cloned(),
-                init_state.prebuilt_provider,
-            )
-            .await?;
-
-        if let Err(error) =
-            Self::add_mcp_extensions(&agent, args.mcp_servers, &goose_session.id).await
-        {
-            error!(
-                error = %error,
-                "new_session MCP server setup failed; continuing with ready session"
-            );
-        }
-
-        let acp_session = GooseAcpSession {
-            agent,
-            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);
-
-        let mut response =
-            NewSessionResponse::new(acp_session_id.clone()).modes(init_state.mode_state);
-        if let Some(ms) = init_state.model_state {
-            response = response.models(ms);
-        }
-        if let Some(co) = init_state.config_options {
-            response = response.config_options(co);
-        }
-        if let Some(updates) = init_state.usage_updates {
-            cx.send_notification(updates.custom)?;
-            // Legacy ACP notification — emitted alongside the custom one for
-            // backwards compatibility. Remove once all known clients have
-            // migrated to `_goose/unstable/session/update`.
-            cx.send_notification(SessionNotification::new(
-                acp_session_id.clone(),
-                SessionUpdate::UsageUpdate(updates.legacy),
-            ))?;
-        }
-        Self::send_available_commands_update(cx, &acp_session_id, &working_dir)?;
-        debug!(
-            target: "perf",
-            sid = %sid,
-            ms = t_start.elapsed().as_millis() as u64,
-            "perf: new_session done"
-        );
-        Ok(response)
+        self.handle_new_session(cx, args).await
     }
 
     /// Look up the session's agent.  Optionally sets a cancellation token on
@@ -2750,6 +2658,7 @@ impl GooseAcpAgent {
         Ok(session.agent.clone())
     }
 
+    #[allow(dead_code)]
     async fn add_mcp_extensions(
         agent: &Arc<Agent>,
         mcp_servers: Vec<McpServer>,
@@ -2791,7 +2700,7 @@ impl GooseAcpAgent {
         cx: &ConnectionTo<Client>,
         args: LoadSessionRequest,
     ) -> Result<LoadSessionResponse, agent_client_protocol::Error> {
-        self.handle_load_session(cx, args).await
+        self.handle_load_session_refactor(cx, args).await
     }
 
     async fn on_prompt(
@@ -3321,93 +3230,7 @@ impl GooseAcpAgent {
         cx: &ConnectionTo<Client>,
         args: ForkSessionRequest,
     ) -> Result<ForkSessionResponse, agent_client_protocol::Error> {
-        validate_absolute_cwd(&args.cwd)?;
-        let source_session_id = &*args.session_id.0;
-
-        let source = self
-            .session_manager
-            .get_session(source_session_id, false)
-            .await
-            .internal_err()?;
-        let fork_name = if source.name.trim().is_empty() {
-            "(copy)".to_string()
-        } else {
-            format!("{} (copy)", source.name)
-        };
-
-        let new_session = self
-            .session_manager
-            .copy_session(source_session_id, fork_name)
-            .await
-            .internal_err()?;
-        let new_session_id = new_session.id.clone();
-
-        // Update working dir for the fork.
-        self.session_manager
-            .update(&new_session_id)
-            .working_dir(args.cwd.clone())
-            .apply()
-            .await
-            .internal_err()?;
-
-        let goose_session = self
-            .session_manager
-            .get_session(&new_session_id, false)
-            .await
-            .internal_err()?;
-
-        let mode_state = build_mode_state(goose_session.goose_mode)?;
-        let resolved = resolve_provider_and_model(&self.config_dir, &goose_session).await;
-        let (model_state, config_options, prebuilt_provider) = self
-            .prepare_session_init_config(&resolved, &mode_state, &goose_session)
-            .await;
-
-        let (agent, _extension_results) = self
-            .build_agent_for_session(
-                cx,
-                &goose_session,
-                resolved.as_ref().ok().cloned(),
-                prebuilt_provider,
-            )
-            .await?;
-
-        if let Err(error) =
-            Self::add_mcp_extensions(&agent, args.mcp_servers, &goose_session.id).await
-        {
-            error!(
-                error = %error,
-                "fork_session MCP server setup failed; continuing with ready session"
-            );
-        }
-
-        let acp_session_id = SessionId::new(new_session_id.clone());
-        let acp_session = GooseAcpSession {
-            agent,
-            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(new_session_id.clone(), acp_session);
-
-        let meta = session_meta(&new_session);
-
-        let mut response = ForkSessionResponse::new(acp_session_id.clone())
-            .modes(mode_state)
-            .meta(meta);
-
-        if let Some(ms) = model_state {
-            response = response.models(ms);
-        }
-        if let Some(co) = config_options {
-            response = response.config_options(co);
-        }
-        Self::send_available_commands_update(cx, &acp_session_id, &args.cwd)?;
-        Ok(response)
+        self.handle_fork_session(cx, args).await
     }
 
     async fn on_close_session(

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

@@ -0,0 +1,71 @@
+use super::*;
+
+impl GooseAcpAgent {
+    #[allow(dead_code)]
+    pub(super) async fn handle_fork_session(
+        &self,
+        cx: &ConnectionTo<Client>,
+        args: ForkSessionRequest,
+    ) -> Result<ForkSessionResponse, agent_client_protocol::Error> {
+        validate_absolute_cwd(&args.cwd)?;
+        let source_session_id = &*args.session_id.0;
+
+        let source = self
+            .session_manager
+            .get_session(source_session_id, false)
+            .await
+            .internal_err()?;
+        let fork_name = if source.name.trim().is_empty() {
+            "(copy)".to_string()
+        } else {
+            format!("{} (copy)", source.name)
+        };
+
+        let new_session = self
+            .session_manager
+            .copy_session(source_session_id, fork_name)
+            .await
+            .internal_err()?;
+        let new_session_id = new_session.id.clone();
+
+        let goose_session = self
+            .session_manager
+            .get_session(&new_session_id, false)
+            .await
+            .internal_err()?;
+
+        let goose_session = super::session_setup::prepare_session_for_activation(
+            self,
+            goose_session,
+            args.cwd.clone(),
+            args.mcp_servers,
+            false,
+        )
+        .await?;
+
+        let (_agent, _extension_results) = self
+            .activate_acp_session(cx, &goose_session, HashMap::new())
+            .await?;
+
+        let acp_session_id = SessionId::new(new_session_id.clone());
+        let meta = session_meta(&new_session);
+
+        let mode_state = build_mode_state(goose_session.goose_mode)?;
+        let (model_state, config_options) = self
+            .build_eager_session_config(&mode_state, &goose_session)
+            .await;
+
+        let mut response = ForkSessionResponse::new(acp_session_id.clone())
+            .modes(mode_state)
+            .meta(meta);
+
+        if let Some(ms) = model_state {
+            response = response.models(ms);
+        }
+        if let Some(co) = config_options {
+            response = response.config_options(co);
+        }
+        Self::send_available_commands_update(cx, &acp_session_id, &args.cwd)?;
+        Ok(response)
+    }
+}

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

@@ -191,6 +191,7 @@ fn collect_submitted_elicitation_ids(messages: &[Message]) -> HashSet<String> {
 }
 
 impl GooseAcpAgent {
+    #[allow(dead_code)]
     pub(super) async fn build_agent_for_session(
         &self,
         cx: &ConnectionTo<Client>,
@@ -403,6 +404,7 @@ impl GooseAcpAgent {
         Ok((agent, extension_results))
     }
 
+    #[allow(dead_code)]
     pub(super) async fn handle_load_session(
         &self,
         cx: &ConnectionTo<Client>,
@@ -523,4 +525,100 @@ impl GooseAcpAgent {
         );
         Ok(response)
     }
+
+    #[allow(dead_code)]
+    pub(super) async fn handle_load_session_refactor(
+        &self,
+        cx: &ConnectionTo<Client>,
+        args: LoadSessionRequest,
+    ) -> Result<LoadSessionResponse, agent_client_protocol::Error> {
+        debug!(?args, "load session request");
+        validate_absolute_cwd(&args.cwd)?;
+
+        let session_id_str = args.session_id.0.to_string();
+        let sid = sid_short(&session_id_str);
+        let t_start = std::time::Instant::now();
+
+        let mut session = self
+            .session_manager
+            .get_session(&session_id_str, true)
+            .await
+            .map_err(|_| {
+                agent_client_protocol::Error::resource_not_found(Some(session_id_str.clone()))
+                    .data(format!("Session not found: {}", session_id_str))
+            })?;
+
+        session = super::session_setup::prepare_session_for_activation(
+            self,
+            session,
+            args.cwd.clone(),
+            args.mcp_servers,
+            true,
+        )
+        .await?;
+
+        let (agent, extension_results) = self.prepare_acp_session_agent(cx, &session).await?;
+        let replay_tool_requests = replay_conversation_to_client(cx, &session)?;
+        self.register_acp_session(session_id_str.clone(), agent, replay_tool_requests)
+            .await;
+
+        session = self
+            .session_manager
+            .get_session(&session_id_str, true)
+            .await
+            .internal_err_ctx("Failed to reload session")?;
+
+        let mode_state = build_mode_state(session.goose_mode)?;
+        let usage_updates = build_usage_updates(&session);
+        let (model_state, config_options) =
+            self.build_eager_session_config(&mode_state, &session).await;
+
+        if let Some(updates) = usage_updates {
+            cx.send_notification(updates.custom)?;
+            cx.send_notification(SessionNotification::new(
+                args.session_id.clone(),
+                SessionUpdate::UsageUpdate(updates.legacy),
+            ))?;
+        }
+
+        Self::send_available_commands_update(cx, &args.session_id, &session.working_dir)?;
+
+        let mut response = LoadSessionResponse::new().modes(mode_state);
+        if let Some(ms) = model_state {
+            response = response.models(ms);
+        }
+        if let Some(co) = config_options {
+            response = response.config_options(co);
+        }
+
+        let mut meta = serde_json::Map::new();
+        if let Some(recipe) = &session.recipe {
+            if let Ok(v) = serde_json::to_value(recipe) {
+                meta.insert("recipe".to_string(), v);
+            }
+        }
+        if let Some(values) = &session.user_recipe_values {
+            if let Ok(v) = serde_json::to_value(values) {
+                meta.insert("userRecipeValues".to_string(), v);
+            }
+        }
+        if let Ok(v) = serde_json::to_value(&extension_results) {
+            meta.insert("extensionResults".to_string(), v);
+        }
+        meta.insert(
+            "workingDir".to_string(),
+            serde_json::Value::String(session.working_dir.to_string_lossy().to_string()),
+        );
+        if !meta.is_empty() {
+            response = response.meta(meta);
+        }
+
+        debug!(
+            target: "perf",
+            sid = %sid,
+            ms = t_start.elapsed().as_millis() as u64,
+            "perf: load_session_refactor done"
+        );
+        Ok(response)
+    }
 }

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

@@ -2,9 +2,7 @@ use crate::acp::server::{
     build_mode_state, build_usage_updates, meta_string, sid_short, validate_absolute_cwd, ResultExt,
 };
 use crate::config::{Config, GooseMode};
-use crate::model::ModelConfig;
 use crate::session::SessionType;
-use crate::session::{EnabledExtensionsState, ExtensionState};
 
 use super::GooseAcpAgent;
 use agent_client_protocol::schema::{
@@ -30,20 +28,8 @@ impl GooseAcpAgent {
             None => SessionType::Acp,
         };
         let config = Config::global();
-        let resolved_provider = config.get_goose_provider().map_err(|error| {
-            agent_client_protocol::Error::internal_error()
-                .data(format!("Failed to resolve provider: {}", error))
-        })?;
-        let resolved_model = config.get_goose_model().map_err(|error| {
-            agent_client_protocol::Error::internal_error()
-                .data(format!("Failed to resolve model: {}", error))
-        })?;
-        let resolved_model_config = ModelConfig::new(&resolved_model)
-            .map(|model_config| model_config.with_canonical_limits(&resolved_provider))
-            .map_err(|error| {
-                agent_client_protocol::Error::internal_error()
-                    .data(format!("Failed to resolve model: {}", error))
-            })?;
+        let (resolved_provider, resolved_model_config) =
+            super::session_setup::resolve_default_provider_model_config(config)?;
         let current_mode: GooseMode = config.get_goose_mode().unwrap_or_default();
         let t0 = std::time::Instant::now();
         let mut goose_session = self
@@ -57,11 +43,12 @@ 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")?;
+        let extension_data = super::session_setup::build_enabled_extensions_data(
+            self,
+            config,
+            &goose_session,
+            args.mcp_servers,
+        )?;
         builder = builder
             .provider_name(resolved_provider)
             .model_config(resolved_model_config)

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

@@ -0,0 +1,85 @@
+use super::*;
+use crate::session::{ExtensionData, ExtensionState};
+
+pub(super) fn resolve_default_provider_model_config(
+    config: &Config,
+) -> Result<(String, crate::model::ModelConfig), agent_client_protocol::Error> {
+    let resolved_provider = config.get_goose_provider().map_err(|error| {
+        agent_client_protocol::Error::internal_error()
+            .data(format!("Failed to resolve provider: {}", error))
+    })?;
+    let resolved_model = config.get_goose_model().map_err(|error| {
+        agent_client_protocol::Error::internal_error()
+            .data(format!("Failed to resolve model: {}", error))
+    })?;
+    let resolved_model_config = crate::model::ModelConfig::new(&resolved_model)
+        .map(|model_config| model_config.with_canonical_limits(&resolved_provider))
+        .map_err(|error| {
+            agent_client_protocol::Error::internal_error()
+                .data(format!("Failed to resolve model: {}", error))
+        })?;
+    Ok((resolved_provider, resolved_model_config))
+}
+
+pub(super) fn build_enabled_extensions_data(
+    agent: &GooseAcpAgent,
+    config: &Config,
+    session: &Session,
+    mcp_servers: Vec<McpServer>,
+) -> Result<ExtensionData, agent_client_protocol::Error> {
+    let extensions = agent.initial_session_extensions(config, mcp_servers)?;
+    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")?;
+    Ok(extension_data)
+}
+
+pub(super) async fn prepare_session_for_activation(
+    agent: &GooseAcpAgent,
+    mut session: Session,
+    cwd: std::path::PathBuf,
+    mcp_servers: Vec<McpServer>,
+    include_messages_on_reload: bool,
+) -> Result<Session, agent_client_protocol::Error> {
+    let config = Config::global();
+    let mut builder = agent.session_manager.update(&session.id);
+    let mut session_needs_update = false;
+
+    if cwd != session.working_dir {
+        builder = builder.working_dir(cwd);
+        session_needs_update = true;
+    }
+
+    if session.provider_name.is_none() || session.model_config.is_none() {
+        let (resolved_provider, resolved_model_config) =
+            resolve_default_provider_model_config(config)?;
+        builder = builder
+            .provider_name(resolved_provider)
+            .model_config(resolved_model_config);
+        session_needs_update = true;
+    }
+
+    if !mcp_servers.is_empty()
+        || EnabledExtensionsState::from_extension_data(&session.extension_data).is_none()
+    {
+        let extension_data = build_enabled_extensions_data(agent, config, &session, mcp_servers)?;
+        builder = builder.extension_data(extension_data);
+        session_needs_update = true;
+    }
+
+    if session_needs_update {
+        let session_id = session.id.clone();
+        builder
+            .apply()
+            .await
+            .internal_err_ctx("Failed to update session")?;
+        session = agent
+            .session_manager
+            .get_session(&session_id, include_messages_on_reload)
+            .await
+            .internal_err_ctx("Failed to reload session")?;
+    }
+
+    Ok(session)
+}