diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 15125354..961d0927 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -526,9 +526,11 @@ name = "codex-exec" version = "0.1.0" dependencies = [ "anyhow", + "chrono", "clap", "codex-core", - "owo-colors", + "owo-colors 4.2.0", + "shlex", "tokio", "tracing", "tracing-subscriber", @@ -561,7 +563,7 @@ dependencies = [ "anyhow", "clap", "codex-core", - "owo-colors", + "owo-colors 4.2.0", "rand", "tokio", "tracing", @@ -599,7 +601,7 @@ dependencies = [ "eyre", "indenter", "once_cell", - "owo-colors", + "owo-colors 3.5.0", "tracing-error", ] @@ -610,7 +612,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd6be1b2a7e382e2b98b43b2adcca6bb0e465af0bdd38123873ae61eb17a72c2" dependencies = [ "once_cell", - "owo-colors", + "owo-colors 3.5.0", "tracing-core", "tracing-error", ] @@ -2225,6 +2227,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "owo-colors" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f" + [[package]] name = "owo-colors" version = "4.2.0" diff --git a/codex-rs/exec/Cargo.toml b/codex-rs/exec/Cargo.toml index 05f9ea40..a6c16977 100644 --- a/codex-rs/exec/Cargo.toml +++ b/codex-rs/exec/Cargo.toml @@ -13,8 +13,11 @@ path = "src/lib.rs" [dependencies] anyhow = "1" +chrono = "0.4.40" clap = { version = "4", features = ["derive"] } codex-core = { path = "../core", features = ["cli"] } +owo-colors = "4.2.0" +shlex = "1.3.0" tokio = { version = "1", features = [ "io-std", "macros", @@ -24,4 +27,3 @@ tokio = { version = "1", features = [ ] } tracing = { version = "0.1.41", features = ["log"] } tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } -owo-colors = "4.2.0" diff --git a/codex-rs/exec/src/cli.rs b/codex-rs/exec/src/cli.rs index 1613845a..f5917a77 100644 --- a/codex-rs/exec/src/cli.rs +++ b/codex-rs/exec/src/cli.rs @@ -1,4 +1,5 @@ use clap::Parser; +use clap::ValueEnum; use codex_core::SandboxModeCliArg; use std::path::PathBuf; @@ -27,6 +28,19 @@ pub struct Cli { #[arg(long = "disable-response-storage", default_value_t = false)] pub disable_response_storage: bool, + /// Specifies color settings for use in the output. + #[arg(long = "color", value_enum, default_value_t = Color::Auto)] + pub color: Color, + /// Initial instructions for the agent. - pub prompt: Option, + pub prompt: String, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)] +#[value(rename_all = "kebab-case")] +pub enum Color { + Always, + Never, + #[default] + Auto, } diff --git a/codex-rs/exec/src/event_processor.rs b/codex-rs/exec/src/event_processor.rs new file mode 100644 index 00000000..9abdc96a --- /dev/null +++ b/codex-rs/exec/src/event_processor.rs @@ -0,0 +1,307 @@ +use chrono::Utc; +use codex_core::protocol::Event; +use codex_core::protocol::EventMsg; +use codex_core::protocol::FileChange; +use owo_colors::OwoColorize; +use owo_colors::Style; +use shlex::try_join; +use std::collections::HashMap; + +/// This should be configurable. When used in CI, users may not want to impose +/// a limit so they can see the full transcript. +const MAX_OUTPUT_LINES_FOR_EXEC_TOOL_CALL: usize = 20; + +pub(crate) struct EventProcessor { + call_id_to_command: HashMap, + call_id_to_patch: HashMap, + + // To ensure that --color=never is respected, ANSI escapes _must_ be added + // using .style() with one of these fields. If you need a new style, add a + // new field here. + bold: Style, + dimmed: Style, + + magenta: Style, + red: Style, + green: Style, +} + +impl EventProcessor { + pub(crate) fn create_with_ansi(with_ansi: bool) -> Self { + let call_id_to_command = HashMap::new(); + let call_id_to_patch = HashMap::new(); + + if with_ansi { + Self { + call_id_to_command, + call_id_to_patch, + bold: Style::new().bold(), + dimmed: Style::new().dimmed(), + magenta: Style::new().magenta(), + red: Style::new().red(), + green: Style::new().green(), + } + } else { + Self { + call_id_to_command, + call_id_to_patch, + bold: Style::new(), + dimmed: Style::new(), + magenta: Style::new(), + red: Style::new(), + green: Style::new(), + } + } + } +} + +struct ExecCommandBegin { + command: Vec, + start_time: chrono::DateTime, +} + +struct PatchApplyBegin { + start_time: chrono::DateTime, + auto_approved: bool, +} + +macro_rules! ts_println { + ($($arg:tt)*) => {{ + let now = Utc::now(); + let formatted = now.format("%Y-%m-%dT%H:%M:%S").to_string(); + print!("[{}] ", formatted); + println!($($arg)*); + }}; +} + +impl EventProcessor { + pub(crate) fn process_event(&mut self, event: Event) { + let Event { id, msg } = event; + match msg { + EventMsg::Error { message } => { + let prefix = "ERROR:".style(self.red); + ts_println!("{prefix} {message}"); + } + EventMsg::BackgroundEvent { message } => { + ts_println!("{}", message.style(self.dimmed)); + } + EventMsg::TaskStarted => { + let msg = format!("Task started: {id}"); + ts_println!("{}", msg.style(self.dimmed)); + } + EventMsg::TaskComplete => { + let msg = format!("Task complete: {id}"); + ts_println!("{}", msg.style(self.bold)); + } + EventMsg::AgentMessage { message } => { + let prefix = "Agent message:".style(self.bold); + ts_println!("{prefix} {message}"); + } + EventMsg::ExecCommandBegin { + call_id, + command, + cwd, + } => { + self.call_id_to_command.insert( + call_id.clone(), + ExecCommandBegin { + command: command.clone(), + start_time: Utc::now(), + }, + ); + ts_println!( + "{} {} in {}", + "exec".style(self.magenta), + escape_command(&command).style(self.bold), + cwd, + ); + } + EventMsg::ExecCommandEnd { + call_id, + stdout, + stderr, + exit_code, + } => { + let exec_command = self.call_id_to_command.remove(&call_id); + let (duration, call) = if let Some(ExecCommandBegin { + command, + start_time, + }) = exec_command + { + ( + format_duration(start_time), + format!("{}", escape_command(&command).style(self.bold)), + ) + } else { + ("".to_string(), format!("exec('{call_id}')")) + }; + + let output = if exit_code == 0 { stdout } else { stderr }; + let truncated_output = output + .lines() + .take(MAX_OUTPUT_LINES_FOR_EXEC_TOOL_CALL) + .collect::>() + .join("\n"); + match exit_code { + 0 => { + let title = format!("{call} succeded{duration}:"); + ts_println!("{}", title.style(self.green)); + } + _ => { + let title = format!("{call} exited {exit_code}{duration}:"); + ts_println!("{}", title.style(self.red)); + } + } + println!("{}", truncated_output.style(self.dimmed)); + } + EventMsg::PatchApplyBegin { + call_id, + auto_approved, + changes, + } => { + // Store metadata so we can calculate duration later when we + // receive the corresponding PatchApplyEnd event. + self.call_id_to_patch.insert( + call_id.clone(), + PatchApplyBegin { + start_time: Utc::now(), + auto_approved, + }, + ); + + ts_println!( + "{} auto_approved={}:", + "apply_patch".style(self.magenta), + auto_approved, + ); + + // Pretty-print the patch summary with colored diff markers so + // it’s easy to scan in the terminal output. + for (path, change) in changes.iter() { + match change { + FileChange::Add { content } => { + let header = format!( + "{} {}", + format_file_change(change), + path.to_string_lossy() + ); + println!("{}", header.style(self.magenta)); + for line in content.lines() { + println!("{}", line.style(self.green)); + } + } + FileChange::Delete => { + let header = format!( + "{} {}", + format_file_change(change), + path.to_string_lossy() + ); + println!("{}", header.style(self.magenta)); + } + FileChange::Update { + unified_diff, + move_path, + } => { + let header = if let Some(dest) = move_path { + format!( + "{} {} -> {}", + format_file_change(change), + path.to_string_lossy(), + dest.to_string_lossy() + ) + } else { + format!("{} {}", format_file_change(change), path.to_string_lossy()) + }; + println!("{}", header.style(self.magenta)); + + // Colorize diff lines. We keep file header lines + // (--- / +++) without extra coloring so they are + // still readable. + for diff_line in unified_diff.lines() { + if diff_line.starts_with('+') && !diff_line.starts_with("+++") { + println!("{}", diff_line.style(self.green)); + } else if diff_line.starts_with('-') + && !diff_line.starts_with("---") + { + println!("{}", diff_line.style(self.red)); + } else { + println!("{diff_line}"); + } + } + } + } + } + } + EventMsg::PatchApplyEnd { + call_id, + stdout, + stderr, + success, + } => { + let patch_begin = self.call_id_to_patch.remove(&call_id); + + // Compute duration and summary label similar to exec commands. + let (duration, label) = if let Some(PatchApplyBegin { + start_time, + auto_approved, + }) = patch_begin + { + ( + format_duration(start_time), + format!("apply_patch(auto_approved={})", auto_approved), + ) + } else { + (String::new(), format!("apply_patch('{call_id}')")) + }; + + let (exit_code, output, title_style) = if success { + (0, stdout, self.green) + } else { + (1, stderr, self.red) + }; + + let title = format!("{label} exited {exit_code}{duration}:"); + ts_println!("{}", title.style(title_style)); + for line in output.lines() { + println!("{}", line.style(self.dimmed)); + } + } + EventMsg::ExecApprovalRequest { .. } => { + // Should we exit? + } + EventMsg::ApplyPatchApprovalRequest { .. } => { + // Should we exit? + } + _ => { + // Ignore event. + } + } + } +} + +fn escape_command(command: &[String]) -> String { + try_join(command.iter().map(|s| s.as_str())).unwrap_or_else(|_| command.join(" ")) +} + +fn format_file_change(change: &FileChange) -> &'static str { + match change { + FileChange::Add { .. } => "A", + FileChange::Delete => "D", + FileChange::Update { + move_path: Some(_), .. + } => "R", + FileChange::Update { + move_path: None, .. + } => "M", + } +} + +fn format_duration(start_time: chrono::DateTime) -> String { + let elapsed = Utc::now().signed_duration_since(start_time); + let millis = elapsed.num_milliseconds(); + if millis < 1000 { + format!(" in {}ms", millis) + } else { + format!(" in {:.2}s", millis as f64 / 1000.0) + } +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 7874f44b..51e17267 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1,4 +1,7 @@ mod cli; +mod event_processor; + +use std::io::IsTerminal; use std::sync::Arc; pub use cli::Cli; @@ -8,76 +11,59 @@ use codex_core::config::ConfigOverrides; use codex_core::protocol::AskForApproval; use codex_core::protocol::Event; use codex_core::protocol::EventMsg; -use codex_core::protocol::FileChange; use codex_core::protocol::InputItem; use codex_core::protocol::Op; use codex_core::util::is_inside_git_repo; +use event_processor::EventProcessor; use owo_colors::OwoColorize; +use owo_colors::Style; use tracing::debug; use tracing::error; use tracing::info; use tracing_subscriber::EnvFilter; -/// Returns `true` if a recognised API key is present in the environment. -/// -/// At present we only support `OPENAI_API_KEY`, mirroring the behaviour of the -/// Node-based `codex-cli`. Additional providers can be added here when the -/// Rust implementation gains first-class support for them. -fn has_api_key() -> bool { - std::env::var("OPENAI_API_KEY") - .map(|s| !s.trim().is_empty()) - .unwrap_or(false) -} - pub async fn run_main(cli: Cli) -> anyhow::Result<()> { - // TODO(mbolin): Take a more thoughtful approach to logging. - let default_level = "error"; - let allow_ansi = true; - let _ = tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env() - .or_else(|_| EnvFilter::try_new(default_level)) - .unwrap(), - ) - .with_ansi(allow_ansi) - .with_writer(std::io::stderr) - .try_init(); - let Cli { images, model, sandbox_policy, skip_git_repo_check, disable_response_storage, + color, prompt, - .. } = cli; - // --------------------------------------------------------------------- - // API key handling - // --------------------------------------------------------------------- + let (stdout_with_ansi, stderr_with_ansi) = match color { + cli::Color::Always => (true, true), + cli::Color::Never => (false, false), + cli::Color::Auto => ( + std::io::stdout().is_terminal(), + std::io::stderr().is_terminal(), + ), + }; - if !has_api_key() { - eprintln!( - "\n{msg}\n\nSet the environment variable {var} and re-run this command.\nYou can create a key here: {url}\n", - msg = "Missing OpenAI API key.".red(), - var = "OPENAI_API_KEY".bold(), - url = "https://platform.openai.com/account/api-keys".bold().underline(), - ); - std::process::exit(1); - } + assert_api_key(stderr_with_ansi); if !skip_git_repo_check && !is_inside_git_repo() { eprintln!("Not inside a Git repo and --skip-git-repo-check was not specified."); std::process::exit(1); - } else if images.is_empty() && prompt.is_none() { - eprintln!("No images or prompt specified."); - std::process::exit(1); } + // TODO(mbolin): Take a more thoughtful approach to logging. + let default_level = "error"; + let _ = tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env() + .or_else(|_| EnvFilter::try_new(default_level)) + .unwrap(), + ) + .with_ansi(stderr_with_ansi) + .with_writer(std::io::stderr) + .try_init(); + // Load configuration and determine approval policy let overrides = ConfigOverrides { - model: model.clone(), + model, // This CLI is intended to be headless and has no affordances for asking // the user for approval. approval_policy: Some(AskForApproval::Never), @@ -115,7 +101,6 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { res = codex.next_event() => match res { Ok(event) => { debug!("Received event: {event:?}"); - process_event(&event); if let Err(e) = tx.send(event) { error!("Error sending event: {e:?}"); break; @@ -131,8 +116,8 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { }); } + // Send images first, if any. if !images.is_empty() { - // Send images first. let items: Vec = images .into_iter() .map(|path| InputItem::LocalImage { path }) @@ -146,101 +131,56 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { } } - if let Some(prompt) = prompt { - // Send the prompt. - let items: Vec = vec![InputItem::Text { text: prompt }]; - let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?; - info!("Sent prompt with event ID: {initial_prompt_task_id}"); - while let Some(event) = rx.recv().await { - if event.id == initial_prompt_task_id && matches!(event.msg, EventMsg::TaskComplete) { - break; - } + // Send the prompt. + let items: Vec = vec![InputItem::Text { text: prompt }]; + let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?; + info!("Sent prompt with event ID: {initial_prompt_task_id}"); + + // Run the loop until the task is complete. + let mut event_processor = EventProcessor::create_with_ansi(stdout_with_ansi); + while let Some(event) = rx.recv().await { + let last_event = + event.id == initial_prompt_task_id && matches!(event.msg, EventMsg::TaskComplete); + event_processor.process_event(event); + if last_event { + break; } } Ok(()) } -fn process_event(event: &Event) { - let Event { id, msg } = event; - match msg { - EventMsg::Error { message } => { - println!("Error: {message}"); - } - EventMsg::BackgroundEvent { .. } => { - // Ignore these for now. - } - EventMsg::TaskStarted => { - println!("Task started: {id}"); - } - EventMsg::TaskComplete => { - println!("Task complete: {id}"); - } - EventMsg::AgentMessage { message } => { - println!("Agent message: {message}"); - } - EventMsg::ExecCommandBegin { - call_id, - command, - cwd, - } => { - println!("exec('{call_id}'): {:?} in {cwd}", command); - } - EventMsg::ExecCommandEnd { - call_id, - stdout, - stderr, - exit_code, - } => { - let output = if *exit_code == 0 { stdout } else { stderr }; - let truncated_output = output.lines().take(5).collect::>().join("\n"); - println!("exec('{call_id}') exited {exit_code}:\n{truncated_output}"); - } - EventMsg::PatchApplyBegin { - call_id, - auto_approved, - changes, - } => { - let changes = changes - .iter() - .map(|(path, change)| { - format!("{} {}", format_file_change(change), path.to_string_lossy()) - }) - .collect::>() - .join("\n"); - println!("apply_patch('{call_id}') auto_approved={auto_approved}:\n{changes}"); - } - EventMsg::PatchApplyEnd { - call_id, - stdout, - stderr, - success, - } => { - let (exit_code, output) = if *success { (0, stdout) } else { (1, stderr) }; - let truncated_output = output.lines().take(5).collect::>().join("\n"); - println!("apply_patch('{call_id}') exited {exit_code}:\n{truncated_output}"); - } - EventMsg::ExecApprovalRequest { .. } => { - // Should we exit? - } - EventMsg::ApplyPatchApprovalRequest { .. } => { - // Should we exit? - } - _ => { - // Ignore event. - } +/// If a valid API key is not present in the environment, print an error to +/// stderr and exits with 1; otherwise, does nothing. +fn assert_api_key(stderr_with_ansi: bool) { + if !has_api_key() { + let (msg_style, var_style, url_style) = if stderr_with_ansi { + ( + Style::new().red(), + Style::new().bold(), + Style::new().bold().underline(), + ) + } else { + (Style::new(), Style::new(), Style::new()) + }; + + eprintln!( + "\n{msg}\n\nSet the environment variable {var} and re-run this command.\nYou can create a key here: {url}\n", + msg = "Missing OpenAI API key.".style(msg_style), + var = "OPENAI_API_KEY".style(var_style), + url = "https://platform.openai.com/account/api-keys".style(url_style), + ); + std::process::exit(1); } } -fn format_file_change(change: &FileChange) -> &'static str { - match change { - FileChange::Add { .. } => "A", - FileChange::Delete => "D", - FileChange::Update { - move_path: Some(_), .. - } => "R", - FileChange::Update { - move_path: None, .. - } => "M", - } +/// Returns `true` if a recognized API key is present in the environment. +/// +/// At present we only support `OPENAI_API_KEY`, mirroring the behavior of the +/// Node-based `codex-cli`. Additional providers can be added here when the +/// Rust implementation gains first-class support for them. +fn has_api_key() -> bool { + std::env::var("OPENAI_API_KEY") + .map(|s| !s.trim().is_empty()) + .unwrap_or(false) }