tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 31 Packages Projects Releases Wiki Activity

preserve extensions when adding MCP servers (#11339)

Browse Source
This commit is contained in:
Lifei Zhou
2026-08-19 20:36:24 +00:00
committed by GitHub
parent 81578789e2
commit b832501b6f
3 changed files with 249 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+1 -11
View File
@@ -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::<String>("ADDITIONAL_AGENT_SOURCE_ROOTS")
crates/goose/src/acp/server.rs
+209 -55
View File
@@ -200,6 +200,22 @@ pub struct AcpBuiltinSelection {
pub explicit: Vec<String>,
}
impl AcpBuiltinSelection {
pub fn from_requested(builtins: Vec<String>) -> 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<ExtensionConf
}
}
fn add_mcp_servers(
extensions: &mut Vec<ExtensionConfig>,
mcp_servers: Vec<McpServer>,
) -> 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<ExtensionConfig>,
) -> Result<ExtensionData, 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")?;
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<McpServer>,
goose_extensions: Option<Vec<GooseExtension>>,
recipe_extensions: Option<&[ExtensionConfig]>,
) -> Result<Vec<ExtensionConfig>, 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<ExtensionConfig>, 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<McpServer>,
goose_extensions: Option<Vec<GooseExtension>>,
recipe_extensions: Option<&[ExtensionConfig]>,
) -> Result<Vec<ExtensionConfig>, 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<Client>,
@@ -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<Vec<GooseExtension>>,
recipe_extensions: Option<&[ExtensionConfig]>,
) -> Result<ExtensionData, agent_client_protocol::Error> {
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<Agent>) {
@@ -2361,10 +2393,7 @@ pub async fn run(builtins: Vec<String>, 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("");
crates/goose/tests/acp_common_tests/mod.rs
+39
View File
@@ -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<C: Connection>() {
}
pub async fn run_load_session_mcp<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
assert_eq!(output.text, FAKE_CODE, "tool call failed in loaded session");
}
async fn assert_stored_extensions<C: Connection>(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::<Vec<_>>();
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<C: Connection>() {
let expected_session_id = C::expected_session_id();
let openai = OpenAiFixture::new(