core: reuse generic fork snapshots for guardian

Co-authored-by: Codex <noreply@openai.com>

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<Option<Vec<codex_protocol::protocol::RolloutItem>>> {
-    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,

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<Option<Vec<RolloutItem>>>,
 }
 
 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<InitialHistory> {
-        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(), &params, 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 {

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<ForkSnapshot>,
     {
-        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<S>(
+    path: &Path,
+    snapshot: S,
+) -> CodexResult<InitialHistory>
+where
+    S: Into<ForkSnapshot>,
+{
+    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<ThreadId> {
         self.threads.read().await.keys().copied().collect()