feat: improve output of exec subcommand (#719)

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"

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"

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<String>,
+    pub prompt: String,
+}
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)]
+#[value(rename_all = "kebab-case")]
+pub enum Color {
+    Always,
+    Never,
+    #[default]
+    Auto,
 }

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<String, ExecCommandBegin>,
+    call_id_to_patch: HashMap<String, PatchApplyBegin>,
+
+    // 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<String>,
+    start_time: chrono::DateTime<Utc>,
+}
+
+struct PatchApplyBegin {
+    start_time: chrono::DateTime<Utc>,
+    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::<Vec<_>>()
+                    .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<Utc>) -> 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)
+    }
+}

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;
 
-    // ---------------------------------------------------------------------
+    let (stdout_with_ansi, stderr_with_ansi) = match color {
-    // API key handling
+        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() {
+    assert_api_key(stderr_with_ansi);
-        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);
-    }
 
     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<InputItem> = 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.
-        // Send the prompt.
+    let items: Vec<InputItem> = vec![InputItem::Text { text: prompt }];
-        let items: Vec<InputItem> = vec![InputItem::Text { text: prompt }];
+    let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?;
-        let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?;
+    info!("Sent prompt with event ID: {initial_prompt_task_id}");
-        info!("Sent prompt with event ID: {initial_prompt_task_id}");
+
-        while let Some(event) = rx.recv().await {
+    // Run the loop until the task is complete.
-            if event.id == initial_prompt_task_id && matches!(event.msg, EventMsg::TaskComplete) {
+    let mut event_processor = EventProcessor::create_with_ansi(stdout_with_ansi);
-                break;
+    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) {
+/// If a valid API key is not present in the environment, print an error to
-    let Event { id, msg } = event;
+/// stderr and exits with 1; otherwise, does nothing.
-    match msg {
+fn assert_api_key(stderr_with_ansi: bool) {
-        EventMsg::Error { message } => {
+    if !has_api_key() {
-            println!("Error: {message}");
+        let (msg_style, var_style, url_style) = if stderr_with_ansi {
-        }
+            (
-        EventMsg::BackgroundEvent { .. } => {
+                Style::new().red(),
-            // Ignore these for now.
+                Style::new().bold(),
-        }
+                Style::new().bold().underline(),
-        EventMsg::TaskStarted => {
+            )
-            println!("Task started: {id}");
+        } else {
-        }
+            (Style::new(), Style::new(), Style::new())
-        EventMsg::TaskComplete => {
+        };
-            println!("Task complete: {id}");
+
-        }
+        eprintln!(
-        EventMsg::AgentMessage { message } => {
+            "\n{msg}\n\nSet the environment variable {var} and re-run this command.\nYou can create a key here: {url}\n",
-            println!("Agent message: {message}");
+            msg = "Missing OpenAI API key.".style(msg_style),
-        }
+            var = "OPENAI_API_KEY".style(var_style),
-        EventMsg::ExecCommandBegin {
+            url = "https://platform.openai.com/account/api-keys".style(url_style),
-            call_id,
+        );
-            command,
+        std::process::exit(1);
-            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::<Vec<_>>().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::<Vec<_>>()
-                .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::<Vec<_>>().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 {
+/// Returns `true` if a recognized API key is present in the environment.
-    match change {
+///
-        FileChange::Add { .. } => "A",
+/// At present we only support `OPENAI_API_KEY`, mirroring the behavior of the
-        FileChange::Delete => "D",
+/// Node-based `codex-cli`. Additional providers can be added here when the
-        FileChange::Update {
+/// Rust implementation gains first-class support for them.
-            move_path: Some(_), ..
+fn has_api_key() -> bool {
-        } => "R",
+    std::env::var("OPENAI_API_KEY")
-        FileChange::Update {
+        .map(|s| !s.trim().is_empty())
-            move_path: None, ..
+        .unwrap_or(false)
-        } => "M",
-    }
 }