diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index c0ba3c1b5..a144512cf 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -1411,17 +1411,7 @@ async fn handle_serve_command(args: ServeCommandArgs) -> Result<()> { enable_scheduler, } = args; - let builtins = if builtins.is_empty() { - AcpBuiltinSelection { - defaults: vec!["developer".to_string()], - explicit: Vec::new(), - } - } else { - AcpBuiltinSelection { - defaults: Vec::new(), - explicit: builtins, - } - }; + let builtins = AcpBuiltinSelection::from_requested(builtins); let additional_source_roots = Config::global() .get_param::("ADDITIONAL_AGENT_SOURCE_ROOTS") diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 165052123..80f90c50f 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -200,6 +200,22 @@ pub struct AcpBuiltinSelection { pub explicit: Vec, } +impl AcpBuiltinSelection { + pub fn from_requested(builtins: Vec) -> Self { + if builtins.is_empty() { + Self { + defaults: vec!["developer".to_string()], + explicit: Vec::new(), + } + } else { + Self { + defaults: Vec::new(), + explicit: builtins, + } + } + } +} + pub struct GooseAcpAgentOptions { pub provider_factory: AcpProviderFactory, pub builtin_selection: AcpBuiltinSelection, @@ -413,6 +429,29 @@ fn mcp_server_to_extension_config(mcp_server: McpServer) -> Result, + mcp_servers: Vec, +) -> Result<(), agent_client_protocol::Error> { + for mcp_server in mcp_servers { + let extension = mcp_server_to_extension_config(mcp_server) + .map_err(|message| agent_client_protocol::Error::invalid_params().data(message))?; + push_or_replace_extension(extensions, extension); + } + Ok(()) +} + +fn enabled_extensions_data( + session: &Session, + extensions: Vec, +) -> Result { + 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) +} + fn selected_builtin_extensions( config: &Config, builtin_selection: &AcpBuiltinSelection, @@ -432,6 +471,38 @@ fn selected_builtin_extensions( extensions } +fn initial_session_extensions( + config: &Config, + builtin_selection: &AcpBuiltinSelection, + project_root: &Path, + mcp_servers: Vec, + goose_extensions: Option>, + recipe_extensions: Option<&[ExtensionConfig]>, +) -> Result, agent_client_protocol::Error> { + let mut extensions = selected_builtin_extensions(config, builtin_selection); + + if let Some(recipe_extensions) = recipe_extensions { + for extension in recipe_extensions { + push_or_replace_extension(&mut extensions, extension.clone()); + } + } else if let Some(goose_extensions) = goose_extensions { + for extension in extensions::goose_extensions_to_configs(goose_extensions)? { + push_or_replace_extension(&mut extensions, extension); + } + } else { + for extension in get_enabled_extensions_with_config(config) { + push_or_replace_extension(&mut extensions, extension); + } + for extension in crate::plugins::mcp_servers::enabled_plugin_mcp_servers(Some(project_root)) + { + push_or_replace_extension(&mut extensions, extension); + } + add_mcp_servers(&mut extensions, mcp_servers)?; + } + + Ok(extensions) +} + fn push_or_replace_extension(extensions: &mut Vec, extension: ExtensionConfig) { let name = extension.name().to_string(); if let Some(index) = extensions @@ -809,45 +880,6 @@ impl GooseAcpAgent { .map_err(|error| agent_creation_error(error, "Failed to create agent")) } - fn initial_session_extensions( - &self, - config: &Config, - project_root: &Path, - mcp_servers: Vec, - goose_extensions: Option>, - recipe_extensions: Option<&[ExtensionConfig]>, - ) -> Result, agent_client_protocol::Error> { - let mut extensions = selected_builtin_extensions(config, &self.builtin_selection); - - if let Some(recipe_extensions) = recipe_extensions { - for extension in recipe_extensions { - push_or_replace_extension(&mut extensions, extension.clone()); - } - } else if let Some(goose_extensions) = goose_extensions { - for extension in extensions::goose_extensions_to_configs(goose_extensions)? { - push_or_replace_extension(&mut extensions, extension); - } - } else if mcp_servers.is_empty() { - for extension in get_enabled_extensions_with_config(config) { - push_or_replace_extension(&mut extensions, extension); - } - for extension in - crate::plugins::mcp_servers::enabled_plugin_mcp_servers(Some(project_root)) - { - push_or_replace_extension(&mut extensions, extension); - } - } else { - for mcp_server in mcp_servers { - let extension = mcp_server_to_extension_config(mcp_server).map_err(|message| { - agent_client_protocol::Error::invalid_params().data(message) - })?; - push_or_replace_extension(&mut extensions, extension); - } - } - - Ok(extensions) - } - async fn apply_acp_extension_overrides( &self, cx: &ConnectionTo, @@ -952,12 +984,15 @@ impl GooseAcpAgent { session_needs_update = true; } - if !mcp_servers.is_empty() - || EnabledExtensionsState::from_extension_data(&session.extension_data).is_none() - { - let extension_data = - self.build_enabled_extensions_data(config, &session, mcp_servers, None, None)?; - builder = builder.extension_data(extension_data); + if !mcp_servers.is_empty() { + let mut stored_extensions = + EnabledExtensionsState::from_extension_data(&session.extension_data) + .unwrap_or_else(|| EnabledExtensionsState::new(Vec::new())); + add_mcp_servers(&mut stored_extensions.extensions, mcp_servers)?; + builder = builder.extension_data(enabled_extensions_data( + &session, + stored_extensions.extensions, + )?); session_needs_update = true; } @@ -991,18 +1026,15 @@ impl GooseAcpAgent { goose_extensions: Option>, recipe_extensions: Option<&[ExtensionConfig]>, ) -> Result { - let extensions = self.initial_session_extensions( + let extensions = initial_session_extensions( config, + &self.builtin_selection, &session.working_dir, mcp_servers, goose_extensions, recipe_extensions, )?; - 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) + enabled_extensions_data(session, extensions) } async fn register_acp_session(&self, session_id: String, agent: Arc) { @@ -2361,10 +2393,7 @@ pub async fn run(builtins: Vec, enable_scheduler: bool) -> Result<()> { let server = crate::acp::server_factory::AcpServer::new( crate::acp::server_factory::AcpServerFactoryConfig { - builtins: AcpBuiltinSelection { - explicit: builtins, - ..Default::default() - }, + builtins: AcpBuiltinSelection::from_requested(builtins), data_dir: Paths::data_dir(), config_dir: Paths::config_dir(), goose_platform: GoosePlatform::GooseCli, @@ -2442,6 +2471,131 @@ mod tests { } } + #[test] + fn requested_builtins_default_to_developer() { + let selected = AcpBuiltinSelection::from_requested(Vec::new()); + assert_eq!(selected.defaults, vec!["developer"]); + assert!(selected.explicit.is_empty()); + } + + #[test] + fn requested_builtins_replace_default_selection() { + let selected = AcpBuiltinSelection::from_requested(vec!["github".to_string()]); + assert!(selected.defaults.is_empty()); + assert_eq!(selected.explicit, vec!["github"]); + } + + #[test] + fn new_session_mcp_is_additive_to_enabled_config_extensions() { + let (config, _c, _s) = config_with_yaml( + r#" +extensions: + developer: + enabled: true + type: builtin + name: developer +"#, + ); + let project_root = tempfile::tempdir().unwrap(); + let extensions = initial_session_extensions( + &config, + &AcpBuiltinSelection::default(), + project_root.path(), + vec![McpServer::Http(McpServerHttp::new( + "zed-mcp", + "http://localhost/mcp", + ))], + None, + None, + ) + .unwrap(); + + assert!(extensions + .iter() + .any(|extension| extension.name() == "developer")); + assert!(extensions + .iter() + .any(|extension| extension.name() == "zed-mcp")); + } + + #[test] + fn new_session_mcp_does_not_enable_disabled_default_builtin() { + let (config, _c, _s) = config_with_yaml( + r#" +extensions: + developer: + enabled: false + type: builtin + name: developer +"#, + ); + let project_root = tempfile::tempdir().unwrap(); + let extensions = initial_session_extensions( + &config, + &AcpBuiltinSelection::from_requested(Vec::new()), + project_root.path(), + vec![McpServer::Http(McpServerHttp::new( + "zed-mcp", + "http://localhost/mcp", + ))], + None, + None, + ) + .unwrap(); + + assert!(!has_developer(&extensions)); + assert!(extensions + .iter() + .any(|extension| extension.name() == "zed-mcp")); + } + + #[test] + fn acp_mcp_is_additive_to_stored_extensions() { + let mut stored_extensions = vec![builtin_to_extension_config("developer")]; + add_mcp_servers( + &mut stored_extensions, + vec![McpServer::Http(McpServerHttp::new( + "zed-mcp", + "http://localhost/mcp", + ))], + ) + .unwrap(); + + assert!(has_developer(&stored_extensions)); + assert!(stored_extensions + .iter() + .any(|extension| extension.name() == "zed-mcp")); + } + + #[test] + fn acp_mcp_replaces_same_named_extension() { + let mut extensions = + vec![ + mcp_server_to_extension_config(McpServer::Http(McpServerHttp::new( + "zed-mcp", + "http://localhost/old", + ))) + .unwrap(), + ]; + add_mcp_servers( + &mut extensions, + vec![McpServer::Http(McpServerHttp::new( + "zed-mcp", + "http://localhost/new", + ))], + ) + .unwrap(); + + assert_eq!(extensions.len(), 1); + match &extensions[0] { + ExtensionConfig::StreamableHttp { name, uri, .. } => { + assert_eq!(name, "zed-mcp"); + assert_eq!(uri, "http://localhost/new"); + } + extension => panic!("expected streamable HTTP extension, got {extension:?}"), + } + } + #[test] fn default_builtin_developer_loads_when_config_is_empty() { let (config, _c, _s) = config_with_yaml(""); diff --git a/crates/goose/tests/acp_common_tests/mod.rs b/crates/goose/tests/acp_common_tests/mod.rs index aa97f1df3..42c811f1b 100644 --- a/crates/goose/tests/acp_common_tests/mod.rs +++ b/crates/goose/tests/acp_common_tests/mod.rs @@ -16,6 +16,7 @@ use fs_err as fs; use goose::acp::server::AcpProviderFactory; use goose::config::base::CONFIG_YAML_NAME; use goose::config::GooseMode; +use goose::session::{EnabledExtensionsState, SessionManager}; use goose_test_support::{McpFixture, FAKE_CODE, TEST_IMAGE_B64, TEST_MODEL}; use sqlx::sqlite::SqlitePoolOptions; use std::sync::Arc; @@ -560,6 +561,15 @@ pub async fn run_load_model() { } pub async fn run_load_session_mcp() { + let temp_dir = tempfile::tempdir().unwrap(); + fs::write( + temp_dir.path().join(CONFIG_YAML_NAME), + format!( + "GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\nextensions:\n developer:\n enabled: true\n type: platform\n name: developer\n description: Developer\n display_name: Developer\n bundled: true\n available_tools: []\n" + ), + ) + .unwrap(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the get_code tool and output only its result."; let mcp = McpFixture::new().await; @@ -592,12 +602,19 @@ pub async fn run_load_session_mcp() { let mcp_servers = vec![McpServer::Http(McpServerHttp::new("mcp-fixture", &mcp_url))]; let config = TestConnectionConfig { + data_root: temp_dir.path().to_path_buf(), mcp_servers: mcp_servers.clone(), ..Default::default() }; let mut conn = C::new(config, openai).await; let SessionData { mut session, .. } = conn.new_session().await.unwrap(); expected_session_id.set(&session.session_id().0); + assert_stored_extensions( + &conn, + &session.session_id().0, + &["developer", "mcp-fixture"], + ) + .await; // First prompt: tool should work in the new session. let output = session @@ -612,6 +629,7 @@ pub async fn run_load_session_mcp() { session: mut loaded_session, .. } = conn.load_session(&session_id, mcp_servers).await.unwrap(); + assert_stored_extensions(&conn, &session_id, &["developer", "mcp-fixture"]).await; // Second prompt: tool should work in the loaded session. let output = loaded_session @@ -621,6 +639,27 @@ pub async fn run_load_session_mcp() { assert_eq!(output.text, FAKE_CODE, "tool call failed in loaded session"); } +async fn assert_stored_extensions(conn: &C, session_id: &str, expected: &[&str]) { + let session = SessionManager::new(conn.data_root()) + .get_session(session_id, false) + .await + .unwrap(); + let stored_extensions = + EnabledExtensionsState::from_extension_data(&session.extension_data).unwrap(); + let names = stored_extensions + .extensions + .iter() + .map(|extension| extension.name()) + .collect::>(); + + for expected_name in expected { + assert!( + names.iter().any(|name| name == expected_name), + "expected stored extension {expected_name} in {names:?}" + ); + } +} + pub async fn run_load_session_replays_image_attachment() { let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new(