From 3b4093de654ece715128b0ce9bb6b6a7c1f66062 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 28 Apr 2025 16:35:36 -0700 Subject: [PATCH] feat: improve output of exec subcommand --- codex-rs/exec/src/cli.rs | 16 ++- codex-rs/exec/src/console_writer.rs | 76 +++++++++++++ codex-rs/exec/src/event_processor.rs | 107 ++++++++++++++++++ codex-rs/exec/src/lib.rs | 159 ++++++++------------------- 4 files changed, 243 insertions(+), 115 deletions(-) create mode 100644 codex-rs/exec/src/console_writer.rs create mode 100644 codex-rs/exec/src/event_processor.rs diff --git a/codex-rs/exec/src/cli.rs b/codex-rs/exec/src/cli.rs index 1613845a89..f5917a7794 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/console_writer.rs b/codex-rs/exec/src/console_writer.rs new file mode 100644 index 0000000000..5b6ef9bf44 --- /dev/null +++ b/codex-rs/exec/src/console_writer.rs @@ -0,0 +1,76 @@ +/// Trait for writing console messages. +pub trait ConsoleWriter { + fn exec_command_succeed(&mut self, call_id: &str, truncated_output: &str); + fn exec_command_fail(&mut self, call_id: &str, exit_code: i32, truncated_output: &str); +} + +/// Macro to generate both ANSI and Plain ConsoleWriters +macro_rules! console_writer_impl { + ( + $StyledWriter:ident, $PlainWriter:ident, $out_field:ident, + { + $( + fn $method:ident(&mut self, $($arg_name:ident: $arg_ty:ty),*) { + styled: $styled_fmt:expr, + plain: $plain_fmt:expr + } + )* + } + ) => { + pub struct $StyledWriter { + $out_field: W, + } + + pub struct $PlainWriter { + $out_field: W, + } + + impl $StyledWriter { + pub fn new($out_field: W) -> Self { + Self { $out_field } + } + } + + impl $PlainWriter { + pub fn new($out_field: W) -> Self { + Self { $out_field } + } + } + + impl ConsoleWriter for $StyledWriter { + $( + fn $method(&mut self, $($arg_name: $arg_ty),*) { + let _ = writeln!(self.$out_field, $styled_fmt, $($arg_name),*); + } + )* + } + + impl ConsoleWriter for $PlainWriter { + $( + fn $method(&mut self, $($arg_name: $arg_ty),*) { + let _ = writeln!(self.$out_field, $plain_fmt, $($arg_name),*); + } + )* + } + }; +} + +const BOLD_RED: &str = "\x1b[1;31m"; +const BOLD_GREEN: &str = "\x1b[1;32m"; +const DIM: &str = "\x1b[2m"; +const RESET: &str = "\x1b[0m"; + +// TODO(mbolin): Escape ANSI codes in plain text output. + +console_writer_impl!( + AnsiConsoleWriter, PlainConsoleWriter, out, { + fn exec_command_succeed(&mut self, call_id: &str, truncated_output: &str) { + styled: "{BOLD_GREEN}exec({}) succeeded:{RESET}\n{DIM}{}{RESET}", + plain: "exec({}) succeeded:\n{}" + } + fn exec_command_fail(&mut self, call_id: &str, exit_code: i32, truncated_output: &str) { + styled: "{BOLD_RED}exec({}) failed ({}):{RESET}\n{DIM}{}{RESET}", + plain: "exec({}) exited {}:\n{}" + } + } +); diff --git a/codex-rs/exec/src/event_processor.rs b/codex-rs/exec/src/event_processor.rs new file mode 100644 index 0000000000..bce4faeb36 --- /dev/null +++ b/codex-rs/exec/src/event_processor.rs @@ -0,0 +1,107 @@ +use codex_core::protocol::Event; +use codex_core::protocol::EventMsg; +use codex_core::protocol::FileChange; + +use crate::console_writer::ConsoleWriter; + +pub(crate) struct EventProcessor { + writer: Box, +} + +impl EventProcessor { + pub(crate) fn new(writer: Box) -> Self { + EventProcessor { writer } + } + + pub(crate) fn process_event(&mut self, 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"); + match exit_code { + 0 => { + self.writer.exec_command_succeed(call_id, &truncated_output); + } + _ => { + self.writer + .exec_command_fail(call_id, *exit_code, &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. + } + } + } +} + +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", + } +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index daa07e4629..201c3de3c3 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1,4 +1,8 @@ mod cli; +mod console_writer; +mod event_processor; + +use std::io::IsTerminal; use std::sync::Arc; pub use cli::Cli; @@ -8,19 +12,42 @@ 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 console_writer::AnsiConsoleWriter; +use console_writer::ConsoleWriter; +use console_writer::PlainConsoleWriter; use tracing::debug; use tracing::error; use tracing::info; use tracing_subscriber::EnvFilter; pub async fn run_main(cli: Cli) -> anyhow::Result<()> { + let Cli { + images, + model, + sandbox_policy, + skip_git_repo_check, + disable_response_storage, + color, + prompt, + } = cli; + + 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); + } + + let stdout = std::io::stdout(); + let allow_ansi = match color { + cli::Color::Always => true, + cli::Color::Never => false, + cli::Color::Auto => stdout.is_terminal(), + }; + // 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() @@ -31,27 +58,15 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { .with_writer(std::io::stderr) .try_init(); - let Cli { - images, - model, - sandbox_policy, - skip_git_repo_check, - disable_response_storage, - prompt, - .. - } = cli; - - 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); - } + let writer: Box = if allow_ansi { + Box::new(AnsiConsoleWriter::new(stdout)) + } else { + Box::new(PlainConsoleWriter::new(stdout)) + }; // 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), @@ -85,7 +100,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; @@ -116,101 +130,18 @@ 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}"); + + let mut event_processor = event_processor::EventProcessor::new(writer); + while let Some(event) = rx.recv().await { + event_processor.process_event(&event); + if event.id == initial_prompt_task_id && matches!(event.msg, EventMsg::TaskComplete) { + 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. - } - } -} - -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", - } -}