diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 3325e9ae8e..4f8e1f469d 100644 --- a/crates/goose/src/acp/server.rs +++ b/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, disable_session_naming: bool, provider_inventory: ProviderInventoryService, + #[allow(dead_code)] goose_platform: GoosePlatform, additional_source_roots: Vec, } @@ -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:' | 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 { } } -fn session_meta(session: &Session) -> serde_json::Map { +pub(super) fn session_meta(session: &Session) -> serde_json::Map { 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 { @@ -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 { }) } -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, session: &Session, - tool_requests: HashMap, ) -> Result<(Arc, Vec), 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, + tool_requests: HashMap, + ) { 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, + session: &Session, + tool_requests: HashMap, + ) -> Result<(Arc, Vec), 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, args: NewSessionRequest, ) -> Result { - 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, mcp_servers: Vec, @@ -2791,7 +2700,7 @@ impl GooseAcpAgent { cx: &ConnectionTo, args: LoadSessionRequest, ) -> Result { - 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, args: ForkSessionRequest, ) -> Result { - 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( diff --git a/crates/goose/src/acp/server/fork_session.rs b/crates/goose/src/acp/server/fork_session.rs new file mode 100644 index 0000000000..b6d3d36a68 --- /dev/null +++ b/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, + args: ForkSessionRequest, + ) -> Result { + 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) + } +} diff --git a/crates/goose/src/acp/server/load_session.rs b/crates/goose/src/acp/server/load_session.rs index 1bf41d2db3..649d40cb3c 100644 --- a/crates/goose/src/acp/server/load_session.rs +++ b/crates/goose/src/acp/server/load_session.rs @@ -191,6 +191,7 @@ fn collect_submitted_elicitation_ids(messages: &[Message]) -> HashSet { } impl GooseAcpAgent { + #[allow(dead_code)] pub(super) async fn build_agent_for_session( &self, cx: &ConnectionTo, @@ -403,6 +404,7 @@ impl GooseAcpAgent { Ok((agent, extension_results)) } + #[allow(dead_code)] pub(super) async fn handle_load_session( &self, cx: &ConnectionTo, @@ -523,4 +525,100 @@ impl GooseAcpAgent { ); Ok(response) } + + #[allow(dead_code)] + pub(super) async fn handle_load_session_refactor( + &self, + cx: &ConnectionTo, + args: LoadSessionRequest, + ) -> Result { + 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) + } } diff --git a/crates/goose/src/acp/server/new_session.rs b/crates/goose/src/acp/server/new_session.rs index 3443430ce8..480204b67b 100644 --- a/crates/goose/src/acp/server/new_session.rs +++ b/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) diff --git a/crates/goose/src/acp/server/session_setup.rs b/crates/goose/src/acp/server/session_setup.rs new file mode 100644 index 0000000000..2167148a02 --- /dev/null +++ b/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, +) -> Result { + 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, + include_messages_on_reload: bool, +) -> Result { + 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) +}