diff --git a/codex-rs/core/src/guardian/review_session/execution.rs b/codex-rs/core/src/guardian/review_session/execution.rs index 252f375195..76536002ef 100644 --- a/codex-rs/core/src/guardian/review_session/execution.rs +++ b/codex-rs/core/src/guardian/review_session/execution.rs @@ -16,7 +16,6 @@ use codex_protocol::protocol::Op; use crate::codex::Codex; use crate::protocol::SandboxPolicy; -use crate::rollout::recorder::RolloutRecorder; use super::GUARDIAN_FOLLOWUP_REVIEW_REMINDER; use super::GUARDIAN_INTERRUPT_DRAIN_TIMEOUT; @@ -25,21 +24,6 @@ use super::GuardianReviewSession; use super::GuardianReviewSessionOutcome; use super::GuardianReviewSessionParams; -/// Captures the trunk rollout items that a later parallel fork should inherit. -/// -/// The manager stores only the latest committed snapshot; loading it from rollout storage lives -/// here so the session/orchestration layer does not need to know about recorder details. -pub(super) async fn load_rollout_items_for_fork( - session: &crate::codex::Session, -) -> anyhow::Result>> { - session.flush_rollout().await; - let Some(rollout_path) = session.current_rollout_path().await else { - return Ok(None); - }; - let history = RolloutRecorder::get_rollout_history(rollout_path.as_path()).await?; - Ok(Some(history.get_rollout_items())) -} - pub(super) async fn run_review_on_session( review_session: &GuardianReviewSession, params: &GuardianReviewSessionParams, @@ -70,6 +54,7 @@ pub(super) async fn run_review_on_session( items: params.prompt_items.clone(), cwd: params.parent_turn.cwd.clone(), approval_policy: AskForApproval::Never, + approvals_reviewer: None, sandbox_policy: SandboxPolicy::new_read_only_policy(), model: params.model.clone(), effort: params.reasoning_effort, diff --git a/codex-rs/core/src/guardian/review_session/mod.rs b/codex-rs/core/src/guardian/review_session/mod.rs index d49d831c46..95f23eab7f 100644 --- a/codex-rs/core/src/guardian/review_session/mod.rs +++ b/codex-rs/core/src/guardian/review_session/mod.rs @@ -15,7 +15,6 @@ use std::time::Duration; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::InitialHistory; -use codex_protocol::protocol::RolloutItem; use codex_protocol::user_input::UserInput; use serde_json::Value; use tokio::sync::Mutex; @@ -26,8 +25,9 @@ use crate::codex::Codex; use crate::codex::Session; use crate::codex::TurnContext; use crate::config::Config; +use crate::thread_manager::ForkSnapshot; +use crate::thread_manager::snapshot_rollout_history; -use self::execution::load_rollout_items_for_fork; use self::execution::run_before_review_deadline; use self::execution::run_review_on_session; use self::spawn::GuardianReviewSessionReuseKey; @@ -134,7 +134,8 @@ struct GuardianReviewSessionState { /// Runtime state for one guardian sub-session. /// /// The trunk persists across approvals, while forked sessions are short-lived and always shut down -/// after the review that spawned them. +/// after the review that spawned them. Parallel forks snapshot the trunk through the generic +/// `ForkSnapshot` path instead of maintaining a guardian-specific committed-history cache. struct GuardianReviewSession { /// Child Codex session running the guardian prompt. codex: Codex, @@ -146,8 +147,6 @@ struct GuardianReviewSession { has_prior_review: AtomicBool, /// Prevents overlapping reviews on the same guardian session. review_lock: Mutex<()>, - /// Snapshot used when forking a parallel review from the cached trunk. - last_committed_rollout_items: Mutex>>, } impl GuardianReviewSession { @@ -163,7 +162,6 @@ impl GuardianReviewSession { reuse_key, has_prior_review: AtomicBool::new(has_prior_review), review_lock: Mutex::new(()), - last_committed_rollout_items: Mutex::new(None), } } @@ -179,23 +177,27 @@ impl GuardianReviewSession { })); } + /// Snapshot the trunk as a forkable committed prefix. + /// + /// `usize::MAX` with `TruncateBeforeNthUserMessage` means "keep everything committed so far, + /// but if the session is currently mid-turn, drop that unfinished turn suffix." That matches + /// the guardian policy of ignoring the in-flight review and forking from the last stable trunk + /// state without synthesizing an interrupt marker. async fn fork_initial_history(&self) -> Option { - self.last_committed_rollout_items - .lock() - .await - .clone() - .filter(|items| !items.is_empty()) - .map(InitialHistory::Forked) - } - - async fn refresh_last_committed_rollout_items(&self) { - match load_rollout_items_for_fork(&self.codex.session).await { - Ok(Some(items)) => { - *self.last_committed_rollout_items.lock().await = Some(items); - } - Ok(None) => {} + self.codex.session.ensure_rollout_materialized().await; + self.codex.session.flush_rollout().await; + let rollout_path = self.codex.session.current_rollout_path().await?; + match snapshot_rollout_history( + rollout_path.as_path(), + ForkSnapshot::TruncateBeforeNthUserMessage(usize::MAX), + ) + .await + { + Ok(InitialHistory::New) => None, + Ok(initial_history) => Some(initial_history), Err(err) => { - warn!("failed to refresh guardian trunk rollout snapshot: {err}"); + warn!("failed to snapshot guardian trunk for fork: {err}"); + None } } } @@ -355,14 +357,6 @@ impl GuardianReviewSessionManager { }; let execution_result = run_review_on_session(trunk.as_ref(), ¶ms, deadline).await; - if execution_result.session_healthy - && matches!( - execution_result.outcome, - GuardianReviewSessionOutcome::Completed(_) - ) - { - trunk.refresh_last_committed_rollout_items().await; - } drop(trunk_guard); if execution_result.session_healthy { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 6ce55c6229..77ff54f489 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -44,6 +44,7 @@ use codex_protocol::protocol::W3cTraceContext; use futures::StreamExt; use futures::stream::FuturesUnordered; use std::collections::HashMap; +use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -598,26 +599,7 @@ impl ThreadManager { where S: Into, { - let snapshot = snapshot.into(); - let history = RolloutRecorder::get_rollout_history(&path).await?; - let snapshot_state = snapshot_turn_state(&history); - let history = match snapshot { - ForkSnapshot::TruncateBeforeNthUserMessage(nth_user_message) => { - truncate_before_nth_user_message(history, nth_user_message, &snapshot_state) - } - ForkSnapshot::Interrupted => { - let history = match history { - InitialHistory::New => InitialHistory::New, - InitialHistory::Forked(history) => InitialHistory::Forked(history), - InitialHistory::Resumed(resumed) => InitialHistory::Forked(resumed.history), - }; - if snapshot_state.ends_mid_turn { - append_interrupted_boundary(history, snapshot_state.active_turn_id) - } else { - history - } - } - }; + let history = snapshot_rollout_history(path.as_path(), snapshot).await?; Box::pin(self.state.spawn_thread( config, history, @@ -646,6 +628,40 @@ impl ThreadManager { } } +/// Load persisted rollout history and convert it into a reusable fork snapshot. +/// +/// This is the common snapshotting primitive used by generic thread forking. Callers that already +/// know how they want to spawn the child thread can use it directly without reimplementing rollout +/// parsing and snapshot semantics. +pub(crate) async fn snapshot_rollout_history( + path: &Path, + snapshot: S, +) -> CodexResult +where + S: Into, +{ + let snapshot = snapshot.into(); + let history = RolloutRecorder::get_rollout_history(path).await?; + let snapshot_state = snapshot_turn_state(&history); + Ok(match snapshot { + ForkSnapshot::TruncateBeforeNthUserMessage(nth_user_message) => { + truncate_before_nth_user_message(history, nth_user_message, &snapshot_state) + } + ForkSnapshot::Interrupted => { + let history = match history { + InitialHistory::New => InitialHistory::New, + InitialHistory::Forked(history) => InitialHistory::Forked(history), + InitialHistory::Resumed(resumed) => InitialHistory::Forked(resumed.history), + }; + if snapshot_state.ends_mid_turn { + append_interrupted_boundary(history, snapshot_state.active_turn_id) + } else { + history + } + } + }) +} + impl ThreadManagerState { pub(crate) async fn list_thread_ids(&self) -> Vec { self.threads.read().await.keys().copied().collect()