Skip to content

Commit

Permalink
fix: collection of CLI updates (#719)
Browse files Browse the repository at this point in the history
  • Loading branch information
baxen authored Jan 23, 2025
1 parent 357ae3c commit 4e9d85b
Show file tree
Hide file tree
Showing 5 changed files with 129 additions and 11 deletions.
82 changes: 77 additions & 5 deletions crates/goose-cli/src/commands/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>, resume: bool) -> Session<'static> {
pub async fn build_session(
name: Option<String>,
resume: bool,
extension: Option<String>,
builtin: Option<String>,
) -> Session<'static> {
// Load config and get provider/model
let config = Config::global();

Expand All @@ -19,10 +26,10 @@ pub async fn build_session(name: Option<String>, 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
Expand Down Expand Up @@ -53,6 +60,48 @@ pub async fn build_session(name: Option<String>, 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 {
Expand Down Expand Up @@ -91,5 +140,28 @@ pub async fn build_session(name: Option<String>, 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(),
);
}
2 changes: 1 addition & 1 deletion crates/goose-cli/src/logging.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
49 changes: 46 additions & 3 deletions crates/goose-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,

/// 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<String>,
},

/// Execute commands from an instruction file
Expand Down Expand Up @@ -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<String>,

/// 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<String>,
},

/// List available agent versions
Expand Down Expand Up @@ -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;
Expand All @@ -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() {
Expand All @@ -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(());
}
Expand Down
5 changes: 4 additions & 1 deletion crates/goose-cli/src/prompt/rustyline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion crates/mcp-client/src/transport/stdio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit 4e9d85b

Please sign in to comment.