From 5e70ae9de08c87f5120fe6f8919a5e115e5f0cfd Mon Sep 17 00:00:00 2001
From: Charles Cunningham <ccunningham@openai.com>
Date: Wed, 25 Mar 2026 11:40:08 -0700
Subject: [PATCH] Preserve contextual session prefix on rollback

Co-authored-by: Codex <noreply@openai.com>
---
 codex-rs/core/src/context_manager/history.rs | 31 +++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs
index e56942471a..c99b9db803 100644
--- a/codex-rs/core/src/context_manager/history.rs
+++ b/codex-rs/core/src/context_manager/history.rs
@@ -244,7 +244,8 @@ impl ContextManager {
             user_positions[user_positions.len() - n_from_end]
         };
 
-        cut_idx = self.trim_pre_turn_context_updates(&snapshot, cut_idx);
+        cut_idx =
+            self.trim_pre_turn_context_updates(&snapshot, first_instruction_turn_idx, cut_idx);
 
         self.replace(snapshot[..cut_idx].to_vec());
     }
@@ -400,10 +401,14 @@ 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`. When rolling back the first real turn, a fully contextual prefix is treated
+    /// as session bootstrap and is therefore preserved.
+    ///
     /// `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. 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.
+    /// that boundary. The trim walk 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
@@ -412,18 +417,18 @@ impl ContextManager {
     fn trim_pre_turn_context_updates(
         &mut self,
         snapshot: &[ResponseItem],
+        first_instruction_turn_idx: usize,
         mut cut_idx: usize,
     ) -> usize {
+        let original_cut_idx = cut_idx;
+        let mut trimmed_mixed_dev_bundle = false;
         while cut_idx > 0 {
             match &snapshot[cut_idx - 1] {
                 ResponseItem::Message { role, content, .. }
                     if role == "developer" && is_contextual_dev_message_content(content) =>
                 {
                     if has_non_contextual_dev_message_content(content) {
-                        // Mixed `build_initial_context` bundles are not reconstructible from
-                        // steady-state diffs once trimmed, so the next real turn must fully
-                        // reinject context instead of diffing against a stale baseline.
-                        self.reference_context_item = None;
+                        trimmed_mixed_dev_bundle = true;
                     }
                     cut_idx -= 1;
                 }
@@ -435,6 +440,18 @@ impl ContextManager {
                 _ => break,
             }
         }
+
+        if original_cut_idx == first_instruction_turn_idx && cut_idx == 0 {
+            return original_cut_idx;
+        }
+
+        if trimmed_mixed_dev_bundle {
+            // Mixed `build_initial_context` bundles are not reconstructible from steady-state
+            // diffs once trimmed, so the next real turn must fully reinject context instead of
+            // diffing against a stale baseline.
+            self.reference_context_item = None;
+        }
+
         cut_idx
     }
 }