From 120251eb5a4bcd25636805bc74903902383e9beb Mon Sep 17 00:00:00 2001
From: Charles Cunningham <ccunningham@openai.com>
Date: Wed, 25 Mar 2026 10:47:32 -0700
Subject: [PATCH] Trim first-turn rollback context scaffolding

Co-authored-by: Codex <noreply@openai.com>
---
 codex-rs/core/src/context_manager/history.rs  | 13 +++------
 .../core/src/context_manager/history_tests.rs | 29 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs
index ec2df30d3b..e56942471a 100644
--- a/codex-rs/core/src/context_manager/history.rs
+++ b/codex-rs/core/src/context_manager/history.rs
@@ -244,8 +244,7 @@ impl ContextManager {
             user_positions[user_positions.len() - n_from_end]
         };
 
-        cut_idx =
-            self.trim_pre_turn_context_updates(&snapshot, first_instruction_turn_idx, cut_idx);
+        cut_idx = self.trim_pre_turn_context_updates(&snapshot, cut_idx);
 
         self.replace(snapshot[..cut_idx].to_vec());
     }
@@ -401,13 +400,10 @@ impl ContextManager {
     /// Returns the adjusted cut index after removing contextual developer/user items immediately
     /// above the rolled-back turn boundary.
     ///
-    /// `first_instruction_turn_idx` is the earliest rollback-eligible instruction-turn boundary
-    /// in `snapshot`; the trim walk never crosses it so any session-prefix items that predate the
-    /// first real turn survive rollback.
-    ///
     /// `cut_idx` is the tentative slice boundary after dropping the requested number of
     /// instruction turns, before stripping contextual pre-turn items that sit immediately above
-    /// that boundary.
+    /// that boundary. The trim walk may continue all the way to index `0`, but it stops as soon
+    /// as it encounters a non-contextual item, so only actual per-turn scaffolding is removed.
     ///
     /// If any trimmed developer message was a mixed `build_initial_context` bundle containing both
     /// rollback-trimmable contextual fragments and persistent developer text, this also clears the
@@ -416,10 +412,9 @@ impl ContextManager {
     fn trim_pre_turn_context_updates(
         &mut self,
         snapshot: &[ResponseItem],
-        first_instruction_turn_idx: usize,
         mut cut_idx: usize,
     ) -> usize {
-        while cut_idx > first_instruction_turn_idx {
+        while cut_idx > 0 {
             match &snapshot[cut_idx - 1] {
                 ResponseItem::Message { role, content, .. }
                     if role == "developer" && is_contextual_dev_message_content(content) =>
diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs
index 3c508e05ff..633452c204 100644
--- a/codex-rs/core/src/context_manager/history_tests.rs
+++ b/codex-rs/core/src/context_manager/history_tests.rs
@@ -954,7 +954,34 @@ fn drop_last_n_user_turns_trims_context_updates_above_rolled_back_turn() {
     );
 }
 
-#[test]
+fn drop_last_n_user_turns_trims_context_updates_above_first_rolled_back_turn() {
+    let items = vec![
+        assistant_msg("session prefix item"),
+        developer_msg("<collaboration_mode>ROLLED_BACK_DEV_INSTRUCTIONS</collaboration_mode>"),
+        user_input_text_msg(
+            "<environment_context><cwd>PRETURN_CONTEXT_DIFF_CWD</cwd></environment_context>",
+        ),
+        user_input_text_msg("turn 1 user"),
+        assistant_msg("turn 1 assistant"),
+    ];
+
+    let modalities = default_input_modalities();
+    let mut history = create_history_with_items(items);
+    let reference_context_item = reference_context_item();
+    history.set_reference_context_item(Some(reference_context_item.clone()));
+    history.drop_last_n_user_turns(1);
+
+    assert_eq!(
+        history.clone().for_prompt(&modalities),
+        vec![assistant_msg("session prefix item")]
+    );
+    assert_eq!(
+        serde_json::to_value(history.reference_context_item())
+            .expect("serialize retained reference context item"),
+        serde_json::to_value(Some(reference_context_item))
+            .expect("serialize expected reference context item")
+    );
+}
 fn drop_last_n_user_turns_clears_reference_context_for_mixed_developer_context_bundles() {
     let items = vec![
         user_input_text_msg("turn 1 user"),