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

feat: MCP support for agentic CLI providers (#6972)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-02-12 07:37:52 +08:00
committed by GitHub
parent c6e27b8f84
commit 8d59c2dced
57 changed files with 1196 additions and 254 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+4 -1
View File
@@ -1584,7 +1584,10 @@ impl Agent {
}
};
let provider = crate::providers::create(&provider_name, model_config)
let extensions =
EnabledExtensionsState::extensions_or_default(Some(&session.extension_data), config);
let provider = crate::providers::create(&provider_name, model_config, extensions)
.await
.map_err(|e| anyhow!("Could not create provider: {}", e))?;
crates/goose/src/agents/extension.rs
+264 -3
View File
@@ -11,6 +11,7 @@ use crate::agents::mcp_client::McpClientTrait;
use crate::config;
use crate::config::extensions::name_to_key;
use crate::config::permission::PermissionLevel;
use crate::config::Config;
use once_cell::sync::Lazy;
use rmcp::model::Tool;
use rmcp::service::ClientInitializeError;
@@ -532,11 +533,9 @@ impl ExtensionConfig {
}
pub fn key(&self) -> String {
let name = self.name();
name_to_key(&name)
name_to_key(&self.name())
}
/// Get the extension name regardless of variant
pub fn name(&self) -> String {
match self {
Self::Sse { name, .. } => name,
@@ -578,6 +577,69 @@ impl ExtensionConfig {
// If tools are specified, only those tools are available
available_tools.is_empty() || available_tools.contains(&tool_name.to_string())
}
pub async fn resolve(self, config: &Config) -> ExtensionResult<Self> {
use crate::agents::extension_manager::{merge_environments, substitute_env_vars};
match self {
Self::Stdio {
name,
description,
cmd,
args,
envs,
env_keys,
timeout,
bundled,
available_tools,
} => {
let merged = merge_environments(&envs, &env_keys, &name, config).await?;
Ok(Self::Stdio {
name,
description,
cmd,
args,
envs: Envs::new(merged),
env_keys: vec![],
timeout,
bundled,
available_tools,
})
}
Self::StreamableHttp {
name,
description,
uri,
envs,
env_keys,
headers,
timeout,
bundled,
available_tools,
} => {
let merged = merge_environments(&envs, &env_keys, &name, config).await?;
let headers = headers
.into_iter()
.map(|(k, v)| {
let v = substitute_env_vars(&v, &merged);
(k, v)
})
.collect();
Ok(Self::StreamableHttp {
name,
description,
uri,
envs: Envs::new(merged),
env_keys: vec![],
headers,
timeout,
bundled,
available_tools,
})
}
other => Ok(other),
}
}
}
impl std::fmt::Display for ExtensionConfig {
@@ -661,6 +723,8 @@ impl ToolInfo {
#[cfg(test)]
mod tests {
use crate::agents::*;
use crate::config;
use test_case::test_case;
#[test]
fn test_deserialize_missing_description() {
@@ -722,4 +786,201 @@ available_tools: []
panic!("unexpected result of deserialization: {}", config)
}
}
#[test_case(
ExtensionConfig::Builtin {
name: "developer".into(),
description: "dev".into(),
display_name: None,
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::Builtin {
name: "developer".into(),
description: "dev".into(),
display_name: None,
timeout: None,
bundled: None,
available_tools: vec![],
}
; "builtin_unchanged"
)]
#[test_case(
ExtensionConfig::StreamableHttp {
name: "test".into(),
description: String::new(),
uri: "https://example.com".into(),
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("AUTH_TOKEN".to_string(), "secret".to_string());
m
}),
env_keys: vec![],
headers: [(
"Authorization".to_string(),
"Bearer $AUTH_TOKEN".to_string(),
)]
.into_iter()
.collect(),
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::StreamableHttp {
name: "test".into(),
description: String::new(),
uri: "https://example.com".into(),
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("AUTH_TOKEN".to_string(), "secret".to_string());
m
}),
env_keys: vec![],
headers: [(
"Authorization".to_string(),
"Bearer secret".to_string(),
)]
.into_iter()
.collect(),
timeout: None,
bundled: None,
available_tools: vec![],
}
; "header_substitution"
)]
#[test_case(
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::default(),
env_keys: vec![],
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::default(),
env_keys: vec![],
timeout: None,
bundled: None,
available_tools: vec![],
}
; "env_keys_cleared"
)]
#[test_case(
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::default(),
env_keys: vec!["MY_SECRET".into()],
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("MY_SECRET".to_string(), "secret_value".to_string());
m
}),
env_keys: vec![],
timeout: None,
bundled: None,
available_tools: vec![],
}
; "env_key_resolved"
)]
#[test_case(
ExtensionConfig::StreamableHttp {
name: "test".into(),
description: String::new(),
uri: "https://example.com".into(),
envs: extension::Envs::default(),
env_keys: vec!["MY_SECRET".into()],
headers: [(
"Authorization".to_string(),
"Bearer $MY_SECRET".to_string(),
)]
.into_iter()
.collect(),
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::StreamableHttp {
name: "test".into(),
description: String::new(),
uri: "https://example.com".into(),
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("MY_SECRET".to_string(), "secret_value".to_string());
m
}),
env_keys: vec![],
headers: [("Authorization".to_string(), "Bearer secret_value".to_string())]
.into_iter()
.collect(),
timeout: None,
bundled: None,
available_tools: vec![],
}
; "http_env_key_and_header_substitution"
)]
#[test_case(
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("MY_SECRET".to_string(), "original".to_string());
m
}),
env_keys: vec!["MY_SECRET".into()],
timeout: None,
bundled: None,
available_tools: vec![],
},
ExtensionConfig::Stdio {
name: "test".into(),
description: String::new(),
cmd: "echo".into(),
args: vec![],
envs: extension::Envs::new({
let mut m = std::collections::HashMap::new();
m.insert("MY_SECRET".to_string(), "original".to_string());
m
}),
env_keys: vec![],
timeout: None,
bundled: None,
available_tools: vec![],
}
; "env_key_skipped_when_already_in_envs"
)]
#[tokio::test]
async fn test_resolve(config: ExtensionConfig, expected: ExtensionConfig) {
let dir = tempfile::tempdir().unwrap();
let cfg = config::Config::new_with_file_secrets(
dir.path().join("config.yaml"),
dir.path().join("secrets.yaml"),
)
.unwrap();
cfg.set("MY_SECRET", &"secret_value", true).unwrap();
assert_eq!(config.resolve(&cfg).await.unwrap(), expected);
}
}
crates/goose/src/agents/extension_manager.rs
+17 -75
View File
@@ -3,7 +3,6 @@ use axum::http::{HeaderMap, HeaderName};
use chrono::{DateTime, Utc};
use futures::stream::{FuturesUnordered, StreamExt};
use futures::{future, FutureExt};
use rand::{distributions::Alphanumeric, Rng};
use rmcp::service::{ClientInitializeError, ServiceError};
use rmcp::transport::streamable_http_client::{
AuthRequiredError, StreamableHttpClientTransportConfig, StreamableHttpError,
@@ -135,31 +134,6 @@ impl ResourceItem {
}
}
/// Generates extension name from server info; adds random suffix on collision.
fn generate_extension_name(
server_info: Option<&ServerInfo>,
name_exists: impl Fn(&str) -> bool,
) -> String {
let base = server_info
.and_then(|info| {
let name = info.server_info.name.as_str();
(!name.is_empty()).then(|| name_to_key(name))
})
.unwrap_or_else(|| "unnamed".to_string());
if !name_exists(&base) {
return base;
}
let suffix: String = rand::thread_rng()
.sample_iter(Alphanumeric)
.take(6)
.map(char::from)
.collect();
format!("{base}_{suffix}")
}
fn resolve_command(cmd: &str) -> PathBuf {
SearchPaths::builder()
.with_npm()
@@ -315,20 +289,20 @@ fn extract_auth_error(
}
/// Merge environment variables from direct envs and keychain-stored env_keys
async fn merge_environments(
pub(crate) async fn merge_environments(
envs: &Envs,
env_keys: &[String],
ext_name: &str,
config: &Config,
) -> Result<HashMap<String, String>, ExtensionError> {
let mut all_envs = envs.get_env();
let config_instance = Config::global();
for key in env_keys {
if all_envs.contains_key(key) {
continue;
}
match config_instance.get(key, true) {
match config.get(key, true) {
Ok(value) => {
if value.is_null() {
warn!(
@@ -369,7 +343,7 @@ async fn merge_environments(
}
/// Substitute environment variables in a string. Supports both ${VAR} and $VAR syntax.
fn substitute_env_vars(value: &str, env_map: &HashMap<String, String>) -> String {
pub(crate) fn substitute_env_vars(value: &str, env_map: &HashMap<String, String>) -> String {
let mut result = value.to_string();
let re_braces =
@@ -404,7 +378,6 @@ async fn create_streamable_http_client(
timeout: Option<u64>,
headers: &HashMap<String, String>,
name: &str,
all_envs: &HashMap<String, String>,
provider: SharedProvider,
) -> ExtensionResult<Box<dyn McpClientTrait>> {
let mut default_headers = HeaderMap::new();
@@ -412,11 +385,10 @@ async fn create_streamable_http_client(
default_headers.insert(reqwest::header::USER_AGENT, GOOSE_USER_AGENT);
for (key, value) in headers {
let substituted_value = substitute_env_vars(value, all_envs);
default_headers.insert(
HeaderName::try_from(key)
.map_err(|_| ExtensionError::ConfigError(format!("invalid header: {}", key)))?,
substituted_value.parse().map_err(|_| {
value.parse().map_err(|_| {
ExtensionError::ConfigError(format!("invalid header value: {}", key))
})?,
);
@@ -517,8 +489,7 @@ impl ExtensionManager {
container: Option<&Container>,
session_id: Option<&str>,
) -> ExtensionResult<()> {
let config_name = config.key().to_string();
let sanitized_name = name_to_key(&config_name);
let sanitized_name = config.key();
if self.extensions.lock().await.contains_key(&sanitized_name) {
return Ok(());
@@ -545,13 +516,17 @@ impl ExtensionManager {
env_keys,
..
} => {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let config = Config::global();
let all_envs = merge_environments(envs, env_keys, &sanitized_name, config).await?;
let resolved_headers = headers
.iter()
.map(|(k, v)| (k.clone(), substitute_env_vars(v, &all_envs)))
.collect();
create_streamable_http_client(
uri,
*timeout,
headers,
&resolved_headers,
name,
&all_envs,
self.provider.clone(),
)
.await?
@@ -564,7 +539,9 @@ impl ExtensionManager {
timeout,
..
} => {
let mut all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let config = Config::global();
let mut all_envs =
merge_environments(envs, env_keys, &sanitized_name, config).await?;
if let Some(sid) = session_id {
all_envs.insert("AGENT_SESSION_ID".to_string(), sid.to_string());
@@ -706,15 +683,9 @@ impl ExtensionManager {
let server_info = client.get_info().cloned();
// Only generate name from server info when config has no name (e.g., CLI --with-*-extension args)
let mut extensions = self.extensions.lock().await;
let final_name = if sanitized_name.is_empty() {
generate_extension_name(server_info.as_ref(), |n| extensions.contains_key(n))
} else {
sanitized_name
};
extensions.insert(
final_name,
sanitized_name,
Extension::new(config, Arc::new(Mutex::new(client)), server_info, temp_dir),
);
drop(extensions);
@@ -2020,35 +1991,6 @@ mod tests {
assert_eq!(result, "Authorization: Bearer secret123 and API key456");
}
mod generate_extension_name_tests {
use super::*;
use rmcp::model::Implementation;
use test_case::test_case;
fn make_info(name: &str) -> ServerInfo {
ServerInfo {
server_info: Implementation {
name: name.into(),
..Default::default()
},
..Default::default()
}
}
#[test_case(Some("kiwi-mcp-server"), None, "^kiwi-mcp-server$" ; "already normalized server name")]
#[test_case(Some("Context7"), None, "^context7$" ; "mixed case normalized")]
#[test_case(Some("@huggingface/mcp-services"), None, "^_huggingface_mcp-services$" ; "special chars normalized")]
#[test_case(None, None, "^unnamed$" ; "no server info falls back")]
#[test_case(Some(""), None, "^unnamed$" ; "empty server name falls back")]
#[test_case(Some("github-mcp-server"), Some("github-mcp-server"), r"^github-mcp-server_[A-Za-z0-9]{6}$" ; "duplicate adds suffix")]
fn test_generate_name(server_name: Option<&str>, collision: Option<&str>, expected: &str) {
let info = server_name.map(make_info);
let result = generate_extension_name(info.as_ref(), |n| collision == Some(n));
let re = regex::Regex::new(expected).unwrap();
assert!(re.is_match(&result));
}
}
#[tokio::test]
async fn test_collect_moim_uses_minute_granularity() {
let temp_dir = tempfile::tempdir().unwrap();
crates/goose/src/agents/mod.rs
+1 -1
View File
@@ -30,7 +30,7 @@ pub mod types;
pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult};
pub use container::Container;
pub use execute_commands::COMPACT_TRIGGERS;
pub use extension::ExtensionConfig;
pub use extension::{ExtensionConfig, ExtensionError};
pub use extension_manager::ExtensionManager;
pub use prompt_manager::PromptManager;
pub use subagent_handler::SUBAGENT_TOOL_REQUEST_TYPE;
crates/goose/src/agents/summon_extension.rs
+7 -14
View File
@@ -13,11 +13,12 @@ use crate::agents::subagent_handler::{
use crate::agents::subagent_task_config::{TaskConfig, DEFAULT_SUBAGENT_MAX_TURNS};
use crate::agents::AgentConfig;
use crate::config::paths::Paths;
use crate::config::Config;
use crate::providers;
use crate::recipe::build_recipe::build_recipe_from_template;
use crate::recipe::local_recipes::load_local_recipe_file;
use crate::recipe::{Recipe, Settings, RECIPE_FILE_EXTENSIONS};
use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
use crate::session::extension_data::EnabledExtensionsState;
use crate::session::SessionType;
use anyhow::Result;
use async_trait::async_trait;
@@ -1290,7 +1291,10 @@ impl SummonClient {
) -> Result<TaskConfig, anyhow::Error> {
let provider = self.resolve_provider(params, recipe, session).await?;
let mut extensions = self.resolve_extensions(session)?;
let mut extensions = EnabledExtensionsState::extensions_or_default(
Some(&session.extension_data),
Config::global(),
);
if let Some(filter) = &params.extensions {
if filter.is_empty() {
@@ -1349,18 +1353,7 @@ impl SummonClient {
model_config = model_config.with_temperature(Some(temp));
}
providers::create(&provider_name, model_config).await
}
fn resolve_extensions(
&self,
session: &crate::session::Session,
) -> Result<Vec<crate::agents::ExtensionConfig>, anyhow::Error> {
let extensions = EnabledExtensionsState::from_extension_data(&session.extension_data)
.map(|s| s.extensions)
.unwrap_or_else(crate::config::get_enabled_extensions);
Ok(extensions)
providers::create(&provider_name, model_config, Vec::new()).await
}
fn resolve_max_turns(&self, session: &crate::session::Session) -> usize {