From 4e9d85bde5dcd06b79092ba22616ee49abfce9db Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Thu, 23 Jan 2025 14:09:04 -0800 Subject: [PATCH] fix: collection of CLI updates (#719) --- crates/goose-cli/src/commands/session.rs | 82 ++++++++++++++++++++++-- crates/goose-cli/src/logging.rs | 2 +- crates/goose-cli/src/main.rs | 49 +++++++++++++- crates/goose-cli/src/prompt/rustyline.rs | 5 +- crates/mcp-client/src/transport/stdio.rs | 2 +- 5 files changed, 129 insertions(+), 11 deletions(-) diff --git a/crates/goose-cli/src/commands/session.rs b/crates/goose-cli/src/commands/session.rs index 0365e168e..2148ca2ad 100644 --- a/crates/goose-cli/src/commands/session.rs +++ b/crates/goose-cli/src/commands/session.rs @@ -3,14 +3,21 @@ use std::process; use crate::prompt::rustyline::RustylinePrompt; use crate::session::{ensure_session_dir, get_most_recent_session, Session}; -use goose::agents::extension::ExtensionError; +use console::style; +use goose::agents::extension::{Envs, ExtensionError}; use goose::agents::AgentFactory; -use goose::config::{Config, ExtensionManager}; +use goose::config::{Config, ExtensionConfig, ExtensionManager}; use goose::providers::create; +use std::path::Path; use mcp_client::transport::Error as McpClientError; -pub async fn build_session(name: Option, resume: bool) -> Session<'static> { +pub async fn build_session( + name: Option, + resume: bool, + extension: Option, + builtin: Option, +) -> Session<'static> { // Load config and get provider/model let config = Config::global(); @@ -19,10 +26,10 @@ pub async fn build_session(name: Option, resume: bool) -> Session<'stati .expect("No provider configured. Run 'goose configure' first"); let session_dir = ensure_session_dir().expect("Failed to create session directory"); - let model = config + let model: String = config .get("GOOSE_MODEL") .expect("No model configured. Run 'goose configure' first"); - let model_config = goose::model::ModelConfig::new(model); + let model_config = goose::model::ModelConfig::new(model.clone()); let provider = create(&provider_name, model_config).expect("Failed to create provider"); // Create the agent @@ -53,6 +60,48 @@ pub async fn build_session(name: Option, resume: bool) -> Session<'stati } } + // Add extension if provided + if let Some(extension_str) = extension { + let mut parts: Vec<&str> = extension_str.split_whitespace().collect(); + let mut envs = std::collections::HashMap::new(); + + // Parse environment variables (format: KEY=value) + while let Some(part) = parts.first() { + if !part.contains('=') { + break; + } + let env_part = parts.remove(0); + let (key, value) = env_part.split_once('=').unwrap(); + envs.insert(key.to_string(), value.to_string()); + } + + if parts.is_empty() { + eprintln!("No command provided in extension string"); + process::exit(1); + } + + let cmd = parts.remove(0).to_string(); + let config = ExtensionConfig::Stdio { + cmd, + args: parts.iter().map(|s| s.to_string()).collect(), + envs: Envs::new(envs), + }; + + agent.add_extension(config).await.unwrap_or_else(|e| { + eprintln!("Failed to start extension: {}", e); + process::exit(1); + }); + } + + // Add builtin extension if provided + if let Some(name) = builtin { + let config = ExtensionConfig::Builtin { name }; + agent.add_extension(config).await.unwrap_or_else(|e| { + eprintln!("Failed to start builtin extension: {}", e); + process::exit(1); + }); + } + // If resuming, try to find the session if resume { if let Some(ref session_name) = name { @@ -91,5 +140,28 @@ pub async fn build_session(name: Option, resume: bool) -> Session<'stati } let prompt = Box::new(RustylinePrompt::new()); + + display_session_info(resume, &provider_name, &model, &session_file); Session::new(agent, prompt, session_file) } + +fn display_session_info(resume: bool, provider: &str, model: &str, session_file: &Path) { + let start_session_msg = if resume { + "resuming session |" + } else { + "starting session |" + }; + println!( + "{} {} {} {} {}", + style(start_session_msg).dim(), + style("provider:").dim(), + style(provider).cyan().dim(), + style("model:").dim(), + style(model).cyan().dim(), + ); + println!( + " {} {}", + style("logging to").dim(), + style(session_file.display()).dim().cyan(), + ); +} diff --git a/crates/goose-cli/src/logging.rs b/crates/goose-cli/src/logging.rs index 36bf91c39..f7ebc5347 100644 --- a/crates/goose-cli/src/logging.rs +++ b/crates/goose-cli/src/logging.rs @@ -87,7 +87,7 @@ pub fn setup_logging(name: Option<&str>) -> Result<()> { // Build the subscriber with required layers let subscriber = Registry::default() .with(file_layer.with_filter(env_filter)) // Gets all logs - .with(console_layer.with_filter(LevelFilter::INFO)); // Controls log levels + .with(console_layer.with_filter(LevelFilter::WARN)); // Controls log levels // Initialize with Langfuse if available if let Some(langfuse) = langfuse_layer::create_langfuse_observer() { diff --git a/crates/goose-cli/src/main.rs b/crates/goose-cli/src/main.rs index 897da1ab1..84f0a9cac 100644 --- a/crates/goose-cli/src/main.rs +++ b/crates/goose-cli/src/main.rs @@ -61,6 +61,24 @@ enum Command { long_help = "Continue from a previous chat session. If --session is provided, resumes that specific session. Otherwise resumes the last used session." )] resume: bool, + + /// Add a stdio extension with environment variables and command + #[arg( + long = "with-extension", + value_name = "COMMAND", + help = "Add a stdio extension (e.g., 'GITHUB_TOKEN=xyz npx -y @modelcontextprotocol/server-github')", + long_help = "Add a stdio extension from a full command with environment variables. Format: 'ENV1=val1 ENV2=val2 command args...'" + )] + extension: Option, + + /// Add a builtin extension by name + #[arg( + long = "with-builtin", + value_name = "NAME", + help = "Add a builtin extension by name (e.g., 'developer')", + long_help = "Add a builtin extension that is bundled with goose by specifying its name" + )] + builtin: Option, }, /// Execute commands from an instruction file @@ -106,6 +124,24 @@ enum Command { long_help = "Continue from a previous run, maintaining the execution state and context." )] resume: bool, + + /// Add a stdio extension with environment variables and command + #[arg( + long = "with-extension", + value_name = "COMMAND", + help = "Add a stdio extension with environment variables and command (e.g., 'GITHUB_TOKEN=xyz npx -y @modelcontextprotocol/server-github')", + long_help = "Add a stdio extension with environment variables and command. Format: 'ENV1=val1 ENV2=val2 command args...'" + )] + extension: Option, + + /// Add a builtin extension by name + #[arg( + long = "with-builtin", + value_name = "NAME", + help = "Add a builtin extension by name (e.g., 'developer')", + long_help = "Add a builtin extension that is compiled into goose by specifying its name" + )] + builtin: Option, }, /// List available agent versions @@ -136,8 +172,13 @@ async fn main() -> Result<()> { Some(Command::Mcp { name }) => { let _ = run_server(&name).await; } - Some(Command::Session { name, resume }) => { - let mut session = build_session(name, resume).await; + Some(Command::Session { + name, + resume, + extension, + builtin, + }) => { + let mut session = build_session(name, resume, extension, builtin).await; setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?; let _ = session.start().await; @@ -148,6 +189,8 @@ async fn main() -> Result<()> { input_text, name, resume, + extension, + builtin, }) => { // Validate that we have some input source if instructions.is_none() && input_text.is_none() { @@ -167,7 +210,7 @@ async fn main() -> Result<()> { .expect("Failed to read from stdin"); stdin }; - let mut session = build_session(name, resume).await; + let mut session = build_session(name, resume, extension, builtin).await; let _ = session.headless_start(contents.clone()).await; return Ok(()); } diff --git a/crates/goose-cli/src/prompt/rustyline.rs b/crates/goose-cli/src/prompt/rustyline.rs index e7986c120..c196bb325 100644 --- a/crates/goose-cli/src/prompt/rustyline.rs +++ b/crates/goose-cli/src/prompt/rustyline.rs @@ -94,7 +94,10 @@ impl Prompt for RustylinePrompt { }; message_text = message_text.trim().to_string(); - if message_text.eq_ignore_ascii_case("/exit") || message_text.eq_ignore_ascii_case("/quit") + if message_text.eq_ignore_ascii_case("/exit") + || message_text.eq_ignore_ascii_case("/quit") + || message_text.eq_ignore_ascii_case("exit") + || message_text.eq_ignore_ascii_case("quit") { Ok(Input { input_type: InputType::Exit, diff --git a/crates/mcp-client/src/transport/stdio.rs b/crates/mcp-client/src/transport/stdio.rs index 5789640f6..2ee4234d7 100644 --- a/crates/mcp-client/src/transport/stdio.rs +++ b/crates/mcp-client/src/transport/stdio.rs @@ -60,7 +60,7 @@ impl StdioActor { "Process ended unexpectedly".to_string() }; - tracing::error!("Process stderr: {}", err_msg); + tracing::info!("Process stderr: {}", err_msg); let _ = self .error_sender .send(Error::StdioProcessError(err_msg))