Preserve retained answers across steer rollbacks (#42298)

## Why

A steer creates a new instruction boundary without creating a new turn ID. Retained `request_user_input` answers that share that turn ID must not all be discarded when only the later steer is rolled back.

## What changed

- Associate retained answers with their originating function call and instruction boundary.
- Keep answers whose source calls survive rollback, including through compaction, resume, and rollout migration.
- Remove answers whose source calls belong to the rolled-back instruction.

## Testing

- Cover live rollback and resume with answers before and after a steer.
- Cover rollout migration when delayed answers share a turn ID but originate from different calls.

GitOrigin-RevId: 225e7eb7b6a3ae5a5e7fa6fcb80cb7d259db90d0

codex-rs/core/src/context_manager/history.rs

@@ -1,6 +1,6 @@
 //! Parent model history and bounded host-owned context facts.
 //! Compaction replaces only the model window. Snapshots include retained facts atomically;
-//! checkpoint replay and source-turn rollback share their live lifecycle.
+//! checkpoint replay and source-call rollback share their live lifecycle.
 
 use crate::context::ContextualUserFragment;
 use crate::context::ModelSwitchInstructions;
@@ -514,7 +514,18 @@ impl ContextManager {
             .iter()
             .filter_map(|item| item.turn_id())
             .collect::<Vec<_>>();
-        Arc::make_mut(&mut retained_context).remove_turns(&removed_turns);
+        Arc::make_mut(&mut retained_context).retain_answers(|answer| {
+            // A steer creates an instruction boundary, not a new turn ID. Replay exposes
+            // the original calls from rolled-back checkpoints, so prefer the exact source.
+            if let Some(source_index) = snapshot.iter().rposition(|item| {
+                item.turn_id() == Some(answer.turn_id.as_str())
+                    && matches!(&item.item, ResponseItem::FunctionCall { call_id, .. }
+                        if call_id == &answer.call_id)
+            }) {
+                return source_index < cut_idx;
+            }
+            !removed_turns.contains(&answer.turn_id.as_str())
+        });
         self.replace_annotated(retained_items);
         self.retained_context = retained_context;
         self.review_history = review_history;

codex-rs/core/tests/suite/guardian_retained_context.rs

@@ -175,6 +175,147 @@ async fn compact_and_assert_answers(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn retained_answers_rollback_only_the_steered_instruction() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let server = start_mock_server().await;
+    let test = test_codex()
+        .with_history_mode(ThreadHistoryMode::Legacy)
+        .with_config(|config| {
+            config.experimental_thread_store = ThreadStoreConfig::Local;
+            for feature in [Feature::TokenBudget, Feature::DefaultModeRequestUserInput] {
+                config
+                    .features
+                    .enable(feature)
+                    .expect("enable test feature");
+            }
+        })
+        .build_with_auto_env(&server)
+        .await?;
+    let questions = [
+        ("before-steer", "Publish?", "Only privately."),
+        ("after-steer", "Publish the README?", "Do not publish it."),
+    ];
+    let mut responses = questions
+        .iter()
+        .map(|(call_id, question, _)| {
+            sse(vec![
+                ev_function_call(
+                    call_id,
+                    "request_user_input",
+                    &json!({"questions": [{
+                        "id": "publish", "header": "Publish", "question": question,
+                        "options": [
+                            {"label": "Yes", "description": "Publish privately."},
+                            {"label": "No", "description": "Keep local."}
+                        ]
+                    }]})
+                    .to_string(),
+                ),
+                ev_completed(call_id),
+            ])
+        })
+        .collect::<Vec<_>>();
+    responses.push(sse(vec![ev_completed("done")]));
+    let response_mock = mount_sse_sequence(&server, responses).await;
+    test.codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "Check whether to publish.".to_owned(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let mut answers = Vec::new();
+    for (call_id, question, answer) in questions {
+        let request = wait_for_event_match(&test.codex, |event| match event {
+            EventMsg::RequestUserInput(request) => Some(request.clone()),
+            _ => None,
+        })
+        .await;
+        if answers.is_empty() {
+            test.codex
+                .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+                    text: "Also inspect the README.".to_owned(),
+                    text_elements: Vec::new(),
+                }]))
+                .await?;
+        }
+        test.codex
+            .submit(Op::UserInputAnswer {
+                id: request.turn_id.clone(),
+                response: RequestUserInputResponse {
+                    answers: HashMap::from([(
+                        "publish".to_owned(),
+                        RequestUserInputAnswer {
+                            answers: vec![answer.to_owned()],
+                        },
+                    )]),
+                },
+            })
+            .await?;
+        answers.push(VerifiedAnswer {
+            turn_id: request.turn_id,
+            call_id: call_id.to_owned(),
+            questions: vec![VerifiedQuestionAnswer {
+                question: question.to_owned(),
+                answer: answer.to_owned(),
+            }],
+        });
+    }
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(answers[0].turn_id, answers[1].turn_id);
+    let requests = response_mock.requests();
+    assert_eq!(requests.len(), 3);
+    assert!(
+        requests[0].has_message_with_input_texts("user", |texts| {
+            texts == ["Check whether to publish."]
+        })
+    );
+    for (index, request) in requests.iter().skip(1).enumerate() {
+        assert!(request.has_message_with_input_texts("user", |texts| {
+            texts == ["Also inspect the README."]
+        }));
+        for answer in &answers[..=index] {
+            let (output, _) = request
+                .function_call_output_content_and_success(&answer.call_id)
+                .context("answer tool output")?;
+            let output: serde_json::Value =
+                serde_json::from_str(&output.context("answer tool output content")?)?;
+            assert_eq!(
+                output,
+                json!({"answers": {"publish": {"answers": [answer.questions[0].answer]}}})
+            );
+        }
+    }
+    test.codex
+        .append_rollout_items(
+            &answers
+                .iter()
+                .cloned()
+                .map(|answer| {
+                    RolloutItem::RetainedContext(RetainedContextEvent::VerifiedAnswer(answer))
+                })
+                .collect::<Vec<_>>(),
+        )
+        .await?;
+    let mut thread = resume(&test, &test.codex).await?;
+    compact_and_assert_answers(&test, &thread, &answers).await?;
+    for expected in [&answers[..1], &[]] {
+        thread.submit(Op::ThreadRollback { num_turns: 1 }).await?;
+        wait_for_event(&thread, |event| {
+            matches!(event, EventMsg::ThreadRolledBack(_))
+        })
+        .await;
+        compact_and_assert_answers(&test, &thread, expected).await?;
+        thread = resume(&test, &thread).await?;
+        compact_and_assert_answers(&test, &thread, expected).await?;
+    }
+    thread.shutdown_and_wait().await?;
+    Ok(())
+}
+
 #[derive(Clone, Copy)]
 enum LifecycleBoundary {
     Rollback,

codex-rs/history/src/retained_context.rs

@@ -44,7 +44,7 @@ pub enum RetainedContextEvent {
 }
 
 /// Bounded snapshot of retained families, persisted with the parent compaction checkpoint.
-/// Facts live until their source turn is rolled back; compaction does not expire them.
+/// Facts live until their source instruction is rolled back; compaction does not expire them.
 #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
 pub struct RetainedContext {
     verified_answers: VecDeque<VerifiedAnswer>,
@@ -123,10 +123,9 @@ impl RetainedContext {
         }
     }
 
-    /// Removes facts with explicitly removed source turns, not merely compacted tool calls.
-    pub fn remove_turns(&mut self, turn_ids: &[&str]) {
-        self.verified_answers
-            .retain(|answer| !turn_ids.contains(&answer.turn_id.as_str()));
+    /// Retains facts whose sources survive a lifecycle change. Compaction alone is not removal.
+    pub fn retain_answers(&mut self, keep: impl FnMut(&VerifiedAnswer) -> bool) {
+        self.verified_answers.retain(keep);
     }
 }
 

codex-rs/history/src/retained_context_tests.rs

@@ -35,7 +35,7 @@ fn answer_lifecycle_is_bounded_idempotent_and_checkpointed() {
     }
     assert!(!context.is_complete());
     assert_eq!(context.verified_answers().count(), MAX_ANSWERS);
-    context.remove_turns(&["turn-2"]);
+    context.retain_answers(|answer| answer.turn_id != "turn-2");
     assert_eq!(context.verified_answers().count(), 0);
     assert!(!context.is_complete());
 

codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs

@@ -5,6 +5,7 @@
 //! surviving turn after a newer turn has started. This planner keeps compact per-record ownership
 //! metadata for SQLite visibility, then combines it with `rollback_replay`'s cold-resume answer
 //! before the writer makes its second streaming pass.
+//! Retained answers follow their source calls, not later steers sharing the same turn ID.
 
 use std::collections::HashMap;
 
@@ -86,7 +87,7 @@ pub(super) struct RollbackPlanner {
     pending_user_response: Option<PendingUserResponse>,
     pending_delivery_boundary: Option<usize>,
     turn_boundaries: HashMap<String, usize>,
-    retained_fact_sources: Vec<(usize, String)>,
+    call_boundaries: HashMap<(String, String), Option<usize>>,
     compactions: Vec<CompactionFrame>,
     model_replay: ModelReplayPlanner,
 }
@@ -103,7 +104,7 @@ impl RollbackPlanner {
             pending_user_response: None,
             pending_delivery_boundary: None,
             turn_boundaries: HashMap::new(),
-            retained_fact_sources: Vec::new(),
+            call_boundaries: HashMap::new(),
             compactions: Vec::new(),
             model_replay: ModelReplayPlanner::new(),
         }
@@ -154,6 +155,14 @@ impl RollbackPlanner {
                     // no later turn to attach it to.
                     self.pending_context_records.push(index);
                 }
+                if let ResponseItem::FunctionCall { call_id, .. } = &response.item
+                    && let Some(turn_id) = response.turn_id().or(self.active_turn_id.as_deref())
+                {
+                    self.call_boundaries.insert(
+                        (turn_id.to_owned(), call_id.clone()),
+                        self.record_boundaries[index],
+                    );
+                }
             }
             RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => {
                 self.apply_rollback(rollback.num_turns)?;
@@ -228,9 +237,16 @@ impl RollbackPlanner {
             RolloutItem::RetainedContext(codex_rollout::RetainedContextEvent::VerifiedAnswer(
                 answer,
             )) => {
-                self.assign_targeted_record(index, Some(&answer.turn_id));
-                self.retained_fact_sources
-                    .push((index, answer.turn_id.clone()));
+                let source = (answer.turn_id.clone(), answer.call_id.clone());
+                // A late answer still belongs to its call's instruction boundary, even
+                // if the same running turn has since received another user steer.
+                if let Some(boundary) = self.call_boundaries.get(&source) {
+                    self.record_boundaries[index] = *boundary;
+                } else {
+                    self.assign_targeted_record(index, Some(&answer.turn_id));
+                }
+                self.call_boundaries
+                    .insert(source, self.record_boundaries[index]);
             }
             RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None,
         }
@@ -245,7 +261,7 @@ impl RollbackPlanner {
             compactions,
             model_replay,
             turn_boundaries,
-            retained_fact_sources,
+            call_boundaries,
             ..
         } = self;
         let replay_anchor = model_replay.finish().empty_replacement_history_compaction;
@@ -253,14 +269,6 @@ impl RollbackPlanner {
             .iter()
             .filter(|(_, boundary)| !boundary_alive[**boundary])
             .map(|(turn_id, _)| turn_id.as_str())
-            .chain(
-                retained_fact_sources
-                    .iter()
-                    .filter(|(index, _)| {
-                        record_boundaries[*index].is_some_and(|boundary| !boundary_alive[boundary])
-                    })
-                    .map(|(_, turn_id)| turn_id.as_str()),
-            )
             .collect::<Vec<_>>();
         let compacted_items = compactions
             .into_iter()
@@ -268,7 +276,14 @@ impl RollbackPlanner {
                 // Migration removes rollback markers. Checkpoints must therefore carry
                 // the same surviving facts as their standalone source events.
                 if let Some(context) = &mut frame.item.retained_context {
-                    context.remove_turns(&removed_turns);
+                    context.retain_answers(|answer| {
+                        call_boundaries
+                            .get(&(answer.turn_id.clone(), answer.call_id.clone()))
+                            .map_or_else(
+                                || !removed_turns.contains(&answer.turn_id.as_str()),
+                                |boundary| boundary.is_none_or(|boundary| boundary_alive[boundary]),
+                            )
+                    });
                 }
                 if Some(frame.record_index) == replay_anchor {
                     frame.item.replacement_history = Some(Vec::new());

codex-rs/thread-store/src/local/rollout_migration_tests.rs

@@ -1160,6 +1160,79 @@ async fn migration_rolls_back_pre_compaction_turns_from_sqlite_history() {
     );
 }
 
+#[tokio::test]
+async fn migration_preserves_answers_before_a_rolled_back_steer() {
+    let home = TempDir::new().expect("create Codex home");
+    let thread_id = ThreadId::new();
+    let mut items = vec![started("shared-turn")];
+    let mut answers = Vec::new();
+    let RolloutItem::Compacted(mut checkpoint) = compacted(vec![
+        input_response_message("user", "before-steer"),
+        input_response_message("user", "after-steer"),
+    ]) else {
+        unreachable!()
+    };
+    checkpoint.retained_context = Some(Default::default());
+    for call_id in ["before-steer", "after-steer"] {
+        items.push(user_message(call_id));
+        let mut call: ResponseItem = serde_json::from_value(json!({
+            "type": "function_call", "call_id": call_id,
+            "name": "request_user_input", "arguments": "{}"
+        }))
+        .expect("request_user_input call");
+        call.set_turn_id_if_missing("shared-turn");
+        items.push(rollout_response_item(call));
+        let answer: codex_rollout::RetainedContextEvent = serde_json::from_value(json!({
+            "type": "verified_answer", "turn_id": "shared-turn", "call_id": call_id,
+            "questions": [{"question": "Publish?", "answer": "Only privately."}]
+        }))
+        .expect("verified answer");
+        checkpoint
+            .retained_context
+            .as_mut()
+            .expect("retained context")
+            .record(&answer);
+        answers.push(answer);
+    }
+    // Delayed delivery must use the original call boundary, not the turn's latest steer.
+    items.extend(answers.iter().cloned().map(RolloutItem::RetainedContext));
+    items.push(completed("shared-turn"));
+    items.push(RolloutItem::Compacted(checkpoint));
+    items.push(RolloutItem::EventMsg(EventMsg::ThreadRolledBack(
+        ThreadRolledBackEvent { num_turns: 1 },
+    )));
+    let path = write_rollout(home.path(), thread_id, SessionSource::Cli, items);
+    let store = indexed_store(home.path()).await;
+    store
+        .migrate_rollouts(apply_options())
+        .await
+        .expect("migrate rollback of a steer");
+    let migrated = read_rollout(&path);
+    let events = migrated
+        .iter()
+        .filter_map(|line| match &line.item {
+            RolloutItem::RetainedContext(event) => Some(event.clone()),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(events, answers[..1]);
+    let checkpoint = migrated
+        .iter()
+        .find_map(|line| match &line.item {
+            RolloutItem::Compacted(checkpoint) => checkpoint.retained_context.as_ref(),
+            _ => None,
+        })
+        .expect("retained checkpoint");
+    assert_eq!(
+        checkpoint
+            .verified_answers()
+            .cloned()
+            .map(codex_rollout::RetainedContextEvent::VerifiedAnswer)
+            .collect::<Vec<_>>(),
+        answers[..1]
+    );
+}
+
 #[tokio::test]
 async fn migration_preserves_reverse_replay_anchor_after_pre_compaction_rollback() {
     let home = TempDir::new().expect("create Codex home");