Skip to content

Commit e9d7737

Browse files
codefromthecryptcodefromthecrypt
committed
Resolve ENV stuff prior to passing to CLI providers
Signed-off-by: Adrian Cole <adrian@tetrate.io>
1 parent f2f29d6 commit e9d7737

File tree

9 files changed

+450
-255
lines changed

9 files changed

+450
-255
lines changed

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/goose-cli/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ indicatif = "0.18.1"
5858
tokio-util = { version = "0.7.15", features = ["compat", "rt"] }
5959
anstream = "0.6.18"
6060
open = "5.3.2"
61+
url = { workspace = true }
6162
urlencoding = "2.1"
6263
clap_complete = "4.5.62"
6364

@@ -70,5 +71,6 @@ disable-update = []
7071

7172
[dev-dependencies]
7273
tempfile = "3"
74+
test-case = { workspace = true }
7375
tokio = { workspace = true }
7476

‎crates/goose-cli/src/session/builder.rs‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::cli::StreamableHttpOptions;
33
use super::output;
44
use super::CliSession;
55
use console::style;
6-
use goose::agents::{Agent, Container};
6+
use goose::agents::{Agent, Container, ExtensionError};
77
use goose::config::get_enabled_extensions;
88
use goose::config::resolve_extensions_for_new_session;
99
use goose::config::{get_all_extensions, Config, ExtensionConfig};
@@ -490,14 +490,12 @@ async fn handle_resumed_session_workdir(agent: &Agent, session_id: &str, interac
490490
}
491491
}
492492

493-
// Resolve separately from loading so Agentic providers (claude-code, codex, etc.) can
494-
// parse MCP extensions for session initialization.
495-
async fn resolve_extension_configs(
493+
async fn collect_extension_configs(
496494
agent: &Agent,
497495
session_config: &SessionBuilderConfig,
498496
recipe: Option<&Recipe>,
499497
session_id: &str,
500-
) -> Vec<ExtensionConfig> {
498+
) -> Result<Vec<ExtensionConfig>, ExtensionError> {
501499
let configured_extensions: Vec<ExtensionConfig> = if session_config.resume {
502500
agent
503501
.config
@@ -522,21 +520,21 @@ async fn resolve_extension_configs(
522520

523521
let mut all: Vec<ExtensionConfig> = configured_extensions;
524522
all.extend(cli_flag_extensions.into_iter().map(|(_, cfg)| cfg));
525-
all
523+
524+
Ok(all)
526525
}
527526

528527
async fn resolve_and_load_extensions(
529528
agent: Agent,
530-
session_config: &SessionBuilderConfig,
531-
recipe: Option<&Recipe>,
532-
session_id: &str,
529+
extensions: Vec<ExtensionConfig>,
533530
provider_for_debug: Arc<dyn goose::providers::base::Provider>,
531+
interactive: bool,
532+
session_id: &str,
534533
) -> Arc<Agent> {
535534
for warning in goose::config::get_warnings() {
536535
eprintln!("{}", style(format!("Warning: {}", warning)).yellow());
537536
}
538537

539-
let extensions = resolve_extension_configs(&agent, session_config, recipe, session_id).await;
540538
let extensions_to_load: Vec<(String, ExtensionConfig)> = extensions
541539
.into_iter()
542540
.map(|cfg| (cfg.name(), cfg))
@@ -546,7 +544,7 @@ async fn resolve_and_load_extensions(
546544
agent,
547545
extensions_to_load,
548546
provider_for_debug,
549-
session_config.interactive,
547+
interactive,
550548
session_id,
551549
)
552550
.await
@@ -623,12 +621,18 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
623621
}
624622

625623
let extensions_for_provider =
626-
resolve_extension_configs(&agent, &session_config, recipe, &session_id).await;
624+
match collect_extension_configs(&agent, &session_config, recipe, &session_id).await {
625+
Ok(exts) => exts,
626+
Err(e) => {
627+
output::render_error(&format!("Failed to collect extensions: {}", e));
628+
process::exit(1);
629+
}
630+
};
627631

628632
let new_provider = match create(
629633
&resolved.provider_name,
630634
resolved.model_config,
631-
extensions_for_provider,
635+
extensions_for_provider.clone(),
632636
)
633637
.await
634638
{
@@ -679,10 +683,10 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
679683
// Extensions are loaded after session creation because we may change directory when resuming
680684
let agent_ptr = resolve_and_load_extensions(
681685
agent,
682-
&session_config,
683-
recipe,
684-
&session_id,
686+
extensions_for_provider,
685687
Arc::clone(&provider_for_display),
688+
session_config.interactive,
689+
&session_id,
686690
)
687691
.await;
688692

‎crates/goose-cli/src/session/mod.rs‎

Lines changed: 122 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,9 +287,14 @@ impl CliSession {
287287
}
288288

289289
let cmd = parts.remove(0).to_string();
290+
let name = std::path::Path::new(&cmd)
291+
.file_name()
292+
.and_then(|f| f.to_str())
293+
.unwrap_or("unnamed")
294+
.to_string();
290295

291296
Ok(ExtensionConfig::Stdio {
292-
name: String::new(),
297+
name,
293298
cmd,
294299
args: parts.iter().map(|s| s.to_string()).collect(),
295300
envs: Envs::new(envs),
@@ -302,8 +307,29 @@ impl CliSession {
302307
}
303308

304309
pub fn parse_streamable_http_extension(extension_url: &str, timeout: u64) -> ExtensionConfig {
310+
let name = url::Url::parse(extension_url)
311+
.ok()
312+
.map(|u| {
313+
let mut s = String::new();
314+
if let Some(host) = u.host_str() {
315+
s.push_str(host);
316+
}
317+
if let Some(port) = u.port() {
318+
s.push('_');
319+
s.push_str(&port.to_string());
320+
}
321+
let path = u.path().trim_matches('/');
322+
if !path.is_empty() {
323+
s.push('_');
324+
s.push_str(path);
325+
}
326+
s
327+
})
328+
.filter(|s| !s.is_empty())
329+
.unwrap_or_else(|| "unnamed".to_string());
330+
305331
ExtensionConfig::StreamableHttp {
306-
name: String::new(),
332+
name,
307333
uri: extension_url.to_string(),
308334
envs: Envs::new(HashMap::new()),
309335
env_keys: Vec::new(),
@@ -1864,7 +1890,10 @@ fn format_elapsed_time(duration: std::time::Duration) -> String {
18641890
#[cfg(test)]
18651891
mod tests {
18661892
use super::*;
1893+
use goose::agents::extension::Envs;
1894+
use goose::config::ExtensionConfig;
18671895
use std::time::Duration;
1896+
use test_case::test_case;
18681897

18691898
#[test]
18701899
fn test_format_elapsed_time_under_60_seconds() {
@@ -1931,4 +1960,95 @@ mod tests {
19311960
let duration = Duration::from_millis(60500);
19321961
assert_eq!(format_elapsed_time(duration), "1m 00s");
19331962
}
1963+
1964+
#[test_case(
1965+
"/usr/bin/my-server",
1966+
ExtensionConfig::Stdio {
1967+
name: "my-server".into(),
1968+
cmd: "/usr/bin/my-server".into(),
1969+
args: vec![],
1970+
envs: Envs::default(),
1971+
env_keys: vec![],
1972+
description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
1973+
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
1974+
bundled: None,
1975+
available_tools: vec![],
1976+
}
1977+
; "name_from_cmd_basename"
1978+
)]
1979+
#[test_case(
1980+
"MY_SECRET=s3cret npx -y @modelcontextprotocol/server-everything",
1981+
ExtensionConfig::Stdio {
1982+
name: "npx".into(),
1983+
cmd: "npx".into(),
1984+
args: vec!["-y".into(), "@modelcontextprotocol/server-everything".into()],
1985+
envs: Envs::new([("MY_SECRET".into(), "s3cret".into())].into()),
1986+
env_keys: vec![],
1987+
description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
1988+
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
1989+
bundled: None,
1990+
available_tools: vec![],
1991+
}
1992+
; "env_prefix_name_from_cmd"
1993+
)]
1994+
fn test_parse_stdio_extension(input: &str, expected: ExtensionConfig) {
1995+
assert_eq!(CliSession::parse_stdio_extension(input).unwrap(), expected);
1996+
}
1997+
1998+
#[test]
1999+
fn test_parse_stdio_extension_no_command() {
2000+
assert!(CliSession::parse_stdio_extension("").is_err());
2001+
}
2002+
2003+
#[test_case(
2004+
"https://mcp.kiwi.com", 300,
2005+
ExtensionConfig::StreamableHttp {
2006+
name: "mcp.kiwi.com".into(),
2007+
uri: "https://mcp.kiwi.com".into(),
2008+
envs: Envs::default(),
2009+
env_keys: vec![],
2010+
headers: HashMap::new(),
2011+
description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
2012+
timeout: Some(300),
2013+
bundled: None,
2014+
available_tools: vec![],
2015+
}
2016+
; "name_from_host"
2017+
)]
2018+
#[test_case(
2019+
"http://localhost:8080/api", 300,
2020+
ExtensionConfig::StreamableHttp {
2021+
name: "localhost_8080_api".into(),
2022+
uri: "http://localhost:8080/api".into(),
2023+
envs: Envs::default(),
2024+
env_keys: vec![],
2025+
headers: HashMap::new(),
2026+
description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
2027+
timeout: Some(300),
2028+
bundled: None,
2029+
available_tools: vec![],
2030+
}
2031+
; "port_and_path"
2032+
)]
2033+
#[test_case(
2034+
"http://localhost:9090/other", 300,
2035+
ExtensionConfig::StreamableHttp {
2036+
name: "localhost_9090_other".into(),
2037+
uri: "http://localhost:9090/other".into(),
2038+
envs: Envs::default(),
2039+
env_keys: vec![],
2040+
headers: HashMap::new(),
2041+
description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
2042+
timeout: Some(300),
2043+
bundled: None,
2044+
available_tools: vec![],
2045+
}
2046+
; "different_port_and_path"
2047+
)]
2048+
fn test_parse_streamable_http_extension(url: &str, timeout: u64, expected: ExtensionConfig) {
2049+
assert_eq!(
2050+
CliSession::parse_streamable_http_extension(url, timeout),
2051+
expected
2052+
);
2053+
}
19342054
}

0 commit comments

Comments
 (0)