From 0b46df26770ca321ec28289590deaed954c1fabe Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 6 May 2025 21:03:59 -0700 Subject: [PATCH] feat: save rollouts in Rust CLI --- codex-rs/Cargo.lock | 11 +++ codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/codex.rs | 35 ++++++++++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/rollout.rs | 128 +++++++++++++++++++++++++++++++++++ 5 files changed, 177 insertions(+) create mode 100644 codex-rs/core/src/rollout.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 77a9ff74b3..34eeb74612 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -524,12 +524,14 @@ dependencies = [ "serde_json", "tempfile", "thiserror 2.0.12", + "time", "tokio", "tokio-util", "toml", "tracing", "tree-sitter", "tree-sitter-bash", + "uuid", "wiremock", ] @@ -3838,6 +3840,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9" +dependencies = [ + "getrandom 0.3.2", +] + [[package]] name = "valuable" version = "0.1.1" diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 9e0105082d..614c4350a2 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -29,6 +29,7 @@ reqwest = { version = "0.12", features = ["json", "stream"] } serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2.0.12" +time = { version = "0.3", features = ["formatting", "macros"] } tokio = { version = "1", features = [ "io-std", "macros", @@ -41,6 +42,7 @@ toml = "0.8.20" tracing = { version = "0.1.41", features = ["log"] } tree-sitter = "0.25.3" tree-sitter-bash = "0.23.3" +uuid = { version = "1", features = ["v4"] } [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2.172" diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b5c04ddda6..3442464991 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -57,6 +57,7 @@ use crate::protocol::Op; use crate::protocol::ReviewDecision; use crate::protocol::SandboxPolicy; use crate::protocol::Submission; +use crate::rollout::RolloutRecorder; use crate::safety::assess_command_safety; use crate::safety::assess_patch_safety; use crate::safety::SafetyCheck; @@ -213,6 +214,10 @@ pub(crate) struct Session { /// External notifier command (will be passed as args to exec()). When /// `None` this feature is disabled. notify: Option>, + + /// Optional rollout recorder for persisting the conversation transcript so + /// sessions can be replayed or inspected later. + rollout: Mutex>, state: Mutex, } @@ -321,6 +326,17 @@ impl Session { state.approved_commands.insert(cmd); } + /// Append the given items to the session's rollout transcript (if enabled) + /// and persist them to disk. + fn record_rollout_items(&self, items: &[ResponseItem]) { + let mut guard = self.rollout.lock().unwrap(); + if let Some(recorder) = guard.as_mut() { + if let Err(e) = recorder.record_items(items) { + error!("failed to record rollout items: {e:#}"); + } + } + } + async fn notify_exec_command_begin(&self, sub_id: &str, call_id: &str, params: &ExecParams) { let event = Event { id: sub_id.to_string(), @@ -601,6 +617,16 @@ async fn submission_loop( } }; + // Attempt to create a RolloutRecorder *before* moving the + // `instructions` value into the Session struct. + let rollout_recorder = match RolloutRecorder::new(instructions.clone()) { + Ok(r) => Some(r), + Err(e) => { + tracing::warn!("failed to initialise rollout recorder: {e}"); + None + } + }; + sess = Some(Arc::new(Session { client, tx_event: tx_event.clone(), @@ -613,6 +639,7 @@ async fn submission_loop( mcp_connection_manager, notify, state: Mutex::new(state), + rollout: Mutex::new(rollout_recorder), })); // ack @@ -711,6 +738,10 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { net_new_turn_input }; + // Persist the input part of the turn to the rollout (user messages / + // function_call_output from previous step). + sess.record_rollout_items(&turn_input); + let turn_input_messages: Vec = turn_input .iter() .filter_map(|item| match item { @@ -738,6 +769,10 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { // Only attempt to take the lock if there is something to record. if !items.is_empty() { + // First persist model-generated output to the rollout file – this only borrows. + sess.record_rollout_items(&items); + + // For ZDR we also need to keep a transcript clone. if let Some(transcript) = sess.state.lock().unwrap().zdr_transcript.as_mut() { transcript.record_items(items); } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 919d05f154..d274f50e20 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -24,5 +24,6 @@ mod safety; mod user_notification; pub mod util; mod zdr_transcript; +mod rollout; pub use codex::Codex; diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs new file mode 100644 index 0000000000..a47be5729b --- /dev/null +++ b/codex-rs/core/src/rollout.rs @@ -0,0 +1,128 @@ +//! Functionality to persist a Codex conversation *rollout* – a linear list of +//! [`ResponseItem`] objects exchanged during a session – to disk so that +//! sessions can be replayed or inspected later (mirrors the behaviour of the +//! upstream TypeScript implementation). + +use std::fs::File; +use std::fs::{self}; +use std::io::Write; +use time::format_description::FormatItem; +use time::macros::format_description; +use time::OffsetDateTime; + +use serde::Serialize; +use uuid::Uuid; + +use crate::config::codex_dir; +use crate::models::ResponseItem; + +/// Folder inside `~/.codex` that holds saved rollouts. +const SESSIONS_SUBDIR: &str = "sessions"; + +#[derive(Serialize)] +struct SessionMeta { + id: String, + timestamp: String, + #[serde(skip_serializing_if = "Option::is_none")] + instructions: Option, +} + +/// Records all [`ResponseItem`]s for a session and flushes them to disk after +/// every update. +pub(crate) struct RolloutRecorder { + file: File, +} + +impl RolloutRecorder { + /// Attempt to create a new [`RolloutRecorder`]. If the sessions directory + /// cannot be created or the rollout file cannot be opened we return the + /// error so the caller can decide whether to disable persistence. + pub fn new(instructions: Option) -> std::io::Result { + let LogFileInfo { + file, + session_id, + timestamp, + } = create_log_file()?; + + // Build the static session metadata JSON first. + let timestamp_format: &[FormatItem] = format_description!( + "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" + ); + let timestamp = timestamp.format(timestamp_format).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::Other, + format!("failed to format timestamp: {e}"), + ) + })?; + + let meta = SessionMeta { + timestamp, + id: session_id.to_string(), + instructions, + }; + + let mut recorder = Self { file }; + recorder.record_item(&meta)?; + + Ok(recorder) + } + + pub(crate) fn record_items(&mut self, items: &[ResponseItem]) -> std::io::Result<()> { + for item in items { + self.record_item(item)?; + } + Ok(()) + } + + fn record_item(&mut self, item: &impl Serialize) -> std::io::Result<()> { + // Serialize the items to JSON and write them to the file. + let json = serde_json::to_string(item).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::Other, + format!("failed to serialize response items: {e}"), + ) + })?; + writeln!(self.file, "{json}")?; + self.file.flush()?; + + Ok(()) + } +} + +struct LogFileInfo { + /// Opened file handle to the rollout file. + file: File, + + /// Session ID (also embedded in filename). + session_id: Uuid, + + timestamp: OffsetDateTime, +} + +fn create_log_file() -> std::io::Result { + // Resolve ~/.codex/sessions and create it if missing. + let mut dir = codex_dir()?; + dir.push(SESSIONS_SUBDIR); + fs::create_dir_all(&dir)?; + + // Generate a v4 UUID – matches the JS CLI implementation. + let session_id = Uuid::new_v4(); + let timestamp = OffsetDateTime::now_utc(); + // Custom format for YYYY-MM-DD + let format: &[FormatItem] = format_description!("[year]-[month]-[day]"); + let date_str = timestamp.format(format).unwrap(); + + let filename = format!("rollout-{date_str}-{session_id}.jsonl"); + + let path = dir.join(filename); + let file = std::fs::OpenOptions::new() + .append(true) + .create(true) + .open(&path)?; + + Ok(LogFileInfo { + file, + session_id, + timestamp, + }) +}