Fix fork snapshot truncation boundary

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

codex-rs/app-server-protocol/src/protocol/thread_history.rs

@@ -74,6 +74,8 @@ pub struct ThreadHistoryBuilder {
     turns: Vec<Turn>,
     current_turn: Option<PendingTurn>,
     next_item_index: i64,
+    current_rollout_index: usize,
+    next_rollout_index: usize,
 }
 
 impl Default for ThreadHistoryBuilder {
@@ -88,6 +90,8 @@ impl ThreadHistoryBuilder {
             turns: Vec::new(),
             current_turn: None,
             next_item_index: 1,
+            current_rollout_index: 0,
+            next_rollout_index: 0,
         }
     }
 
@@ -118,6 +122,12 @@ impl ThreadHistoryBuilder {
             .map(|turn| turn.id.clone())
     }
 
+    pub fn active_turn_start_index(&self) -> Option<usize> {
+        self.current_turn
+            .as_ref()
+            .map(|turn| turn.rollout_start_index)
+    }
+
     /// Shared reducer for persisted rollout replay and in-memory current-turn
     /// tracking used by running thread resume/rejoin.
     ///
@@ -189,6 +199,8 @@ impl ThreadHistoryBuilder {
     }
 
     pub fn handle_rollout_item(&mut self, item: &RolloutItem) {
+        self.current_rollout_index = self.next_rollout_index;
+        self.next_rollout_index += 1;
         match item {
             RolloutItem::EventMsg(event) => self.handle_event(event),
             RolloutItem::Compacted(payload) => self.handle_compacted(payload),
@@ -981,6 +993,7 @@ impl ThreadHistoryBuilder {
             status: TurnStatus::Completed,
             opened_explicitly: false,
             saw_compaction: false,
+            rollout_start_index: self.current_rollout_index,
         }
     }
 
@@ -1144,6 +1157,8 @@ struct PendingTurn {
     /// True when this turn includes a persisted `RolloutItem::Compacted`, which
     /// should keep the turn from being dropped even without normal items.
     saw_compaction: bool,
+    /// Index of the rollout item that opened this turn during replay.
+    rollout_start_index: usize,
 }
 
 impl PendingTurn {

codex-rs/core/src/thread_manager.rs

@@ -146,10 +146,10 @@ pub enum ForkSnapshot {
     ///
     /// When `n` is within range, this cuts before that 0-based user-message
     /// boundary. When `n` is out of range and the source thread is currently
-    /// mid-turn, this instead cuts before the last user message so the fork
-    /// drops the unfinished turn suffix. When `n` is out of range and the
-    /// source thread is already at a turn boundary, this returns the full
-    /// committed history unchanged.
+    /// mid-turn, this instead cuts before the active turn's opening boundary
+    /// so the fork drops the unfinished turn suffix. When `n` is out of range
+    /// and the source thread is already at a turn boundary, this returns the
+    /// full committed history unchanged.
     TruncateBeforeNthUserMessage(usize),
 
     /// Fork the current persisted history as if the source thread had been
@@ -162,6 +162,14 @@ pub enum ForkSnapshot {
     Interrupted,
 }
 
+/// Preserve legacy `fork_thread(usize, ...)` callsites by mapping them to the
+/// existing truncate-before-nth-user-message snapshot mode.
+impl From<usize> for ForkSnapshot {
+    fn from(value: usize) -> Self {
+        Self::TruncateBeforeNthUserMessage(value)
+    }
+}
+
 #[derive(Debug, Default, PartialEq, Eq)]
 pub struct ThreadShutdownReport {
     pub completed: Vec<ThreadId>,
@@ -581,23 +589,23 @@ impl ThreadManager {
     /// `snapshot` and starting a new thread with identical configuration
     /// (unless overridden by the caller's `config`). The new thread will have
     /// a fresh id.
-    pub async fn fork_thread(
+    pub async fn fork_thread<S>(
         &self,
-        snapshot: ForkSnapshot,
+        snapshot: S,
         config: Config,
         path: PathBuf,
         persist_extended_history: bool,
         parent_trace: Option<W3cTraceContext>,
-    ) -> CodexResult<NewThread> {
+    ) -> CodexResult<NewThread>
+    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.ends_mid_turn,
-                )
+                truncate_before_nth_user_message(history, nth_user_message, &snapshot_state)
             }
             ForkSnapshot::Interrupted => {
                 let history = match history {
@@ -898,17 +906,21 @@ impl ThreadManagerState {
 /// Return a fork snapshot cut strictly before the nth user message (0-based).
 ///
 /// Out-of-range values keep the full committed history at a turn boundary, but
-/// fall back to cutting before the last user message when the source thread is
-/// currently mid-turn so the fork omits the unfinished suffix.
+/// when the source thread is currently mid-turn they fall back to cutting
+/// before the active turn's opening boundary so the fork omits the unfinished
+/// suffix entirely.
 fn truncate_before_nth_user_message(
     history: InitialHistory,
     n: usize,
-    snapshot_mid_turn: bool,
+    snapshot_state: &SnapshotTurnState,
 ) -> InitialHistory {
     let items: Vec<RolloutItem> = history.get_rollout_items();
     let user_positions = truncation::user_message_positions_in_rollout(&items);
-    let rolled = if snapshot_mid_turn && n >= user_positions.len() {
-        if let Some(cut_idx) = user_positions.last().copied() {
+    let rolled = if snapshot_state.ends_mid_turn && n >= user_positions.len() {
+        if let Some(cut_idx) = snapshot_state
+            .active_turn_start_index
+            .or_else(|| user_positions.last().copied())
+        {
             items[..cut_idx].to_vec()
         } else {
             items
@@ -928,6 +940,7 @@ fn truncate_before_nth_user_message(
 struct SnapshotTurnState {
     ends_mid_turn: bool,
     active_turn_id: Option<String>,
+    active_turn_start_index: Option<usize>,
 }
 
 fn snapshot_turn_state(history: &InitialHistory) -> SnapshotTurnState {
@@ -945,12 +958,14 @@ fn snapshot_turn_state(history: &InitialHistory) -> SnapshotTurnState {
             return SnapshotTurnState {
                 ends_mid_turn: false,
                 active_turn_id: None,
+                active_turn_start_index: None,
             };
         }
 
         return SnapshotTurnState {
             ends_mid_turn: true,
             active_turn_id: builder.active_turn_id_if_explicit(),
+            active_turn_start_index: builder.active_turn_start_index(),
         };
     }
 
@@ -961,6 +976,7 @@ fn snapshot_turn_state(history: &InitialHistory) -> SnapshotTurnState {
         return SnapshotTurnState {
             ends_mid_turn: false,
             active_turn_id: None,
+            active_turn_start_index: None,
         };
     };
 
@@ -975,6 +991,7 @@ fn snapshot_turn_state(history: &InitialHistory) -> SnapshotTurnState {
             )
         }),
         active_turn_id: None,
+        active_turn_start_index: None,
     }
 }
 

codex-rs/core/src/thread_manager_tests.rs

@@ -75,7 +75,11 @@ fn truncates_before_requested_user_message() {
     let truncated = truncate_before_nth_user_message(
         InitialHistory::Forked(initial),
         1,
-        /*snapshot_mid_turn*/ false,
+        &SnapshotTurnState {
+            ends_mid_turn: false,
+            active_turn_id: None,
+            active_turn_start_index: None,
+        },
     );
     let got_items = truncated.get_rollout_items();
     let expected_items = vec![
@@ -96,7 +100,11 @@ fn truncates_before_requested_user_message() {
     let truncated2 = truncate_before_nth_user_message(
         InitialHistory::Forked(initial2.clone()),
         2,
-        /*snapshot_mid_turn*/ false,
+        &SnapshotTurnState {
+            ends_mid_turn: false,
+            active_turn_id: None,
+            active_turn_start_index: None,
+        },
     );
     assert_eq!(
         serde_json::to_value(truncated2.get_rollout_items()).unwrap(),
@@ -116,7 +124,66 @@ fn out_of_range_truncation_drops_only_unfinished_suffix_mid_turn() {
     let truncated = truncate_before_nth_user_message(
         InitialHistory::Forked(items.clone()),
         usize::MAX,
-        /*snapshot_mid_turn*/ true,
+        &SnapshotTurnState {
+            ends_mid_turn: true,
+            active_turn_id: None,
+            active_turn_start_index: None,
+        },
+    );
+
+    assert_eq!(
+        serde_json::to_value(truncated.get_rollout_items()).unwrap(),
+        serde_json::to_value(items[..2].to_vec()).unwrap()
+    );
+}
+
+#[test]
+fn fork_thread_accepts_legacy_usize_snapshot_argument() {
+    fn assert_legacy_snapshot_callsite(
+        manager: &ThreadManager,
+        config: Config,
+        path: std::path::PathBuf,
+    ) {
+        let _future = manager.fork_thread(
+            usize::MAX,
+            config,
+            path,
+            /*persist_extended_history*/ false,
+            /*parent_trace*/ None,
+        );
+    }
+
+    let _: fn(&ThreadManager, Config, std::path::PathBuf) = assert_legacy_snapshot_callsite;
+}
+
+#[test]
+fn out_of_range_truncation_drops_pre_user_active_turn_prefix() {
+    let items = vec![
+        RolloutItem::ResponseItem(user_msg("u1")),
+        RolloutItem::ResponseItem(assistant_msg("a1")),
+        RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
+            turn_id: "turn-2".to_string(),
+            model_context_window: None,
+            collaboration_mode_kind: Default::default(),
+        })),
+        RolloutItem::ResponseItem(user_msg("u2")),
+        RolloutItem::ResponseItem(assistant_msg("partial")),
+    ];
+
+    let snapshot_state = snapshot_turn_state(&InitialHistory::Forked(items.clone()));
+    assert_eq!(
+        snapshot_state,
+        SnapshotTurnState {
+            ends_mid_turn: true,
+            active_turn_id: Some("turn-2".to_string()),
+            active_turn_start_index: Some(2),
+        },
+    );
+
+    let truncated = truncate_before_nth_user_message(
+        InitialHistory::Forked(items.clone()),
+        usize::MAX,
+        &snapshot_state,
     );
 
     assert_eq!(
@@ -143,7 +210,11 @@ async fn ignores_session_prefix_messages_when_truncating() {
     let truncated = truncate_before_nth_user_message(
         InitialHistory::Forked(rollout_items),
         1,
-        /*snapshot_mid_turn*/ false,
+        &SnapshotTurnState {
+            ends_mid_turn: false,
+            active_turn_id: None,
+            active_turn_start_index: None,
+        },
     );
     let got_items = truncated.get_rollout_items();
 
@@ -275,6 +346,7 @@ fn interrupted_snapshot_is_not_mid_turn() {
         SnapshotTurnState {
             ends_mid_turn: false,
             active_turn_id: None,
+            active_turn_start_index: None,
         },
     );
 }
@@ -300,6 +372,7 @@ fn completed_legacy_event_history_is_not_mid_turn() {
         SnapshotTurnState {
             ends_mid_turn: false,
             active_turn_id: None,
+            active_turn_start_index: None,
         },
     );
 }
@@ -449,6 +522,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() {
         SnapshotTurnState {
             ends_mid_turn: true,
             active_turn_id: Some("turn-explicit".to_string()),
+            active_turn_start_index: Some(0),
         },
     );