diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index f9ad24af27..8779f99699 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index f6b459261c..a5b9e05919 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/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::>(); - 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; diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index 2ef4f593b3..8f45117e48 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/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::>(); + 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::>(), + ) + .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, diff --git a/codex-rs/history/src/retained_context.rs b/codex-rs/history/src/retained_context.rs index 3214212a0c..97ed607b5d 100644 --- a/codex-rs/history/src/retained_context.rs +++ b/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, @@ -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); } } diff --git a/codex-rs/history/src/retained_context_tests.rs b/codex-rs/history/src/retained_context_tests.rs index baa34b748f..5129daae0a 100644 --- a/codex-rs/history/src/retained_context_tests.rs +++ b/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()); diff --git a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs index 58d1944407..44bd5c6937 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs +++ b/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, pending_delivery_boundary: Option, turn_boundaries: HashMap, - retained_fact_sources: Vec<(usize, String)>, + call_boundaries: HashMap<(String, String), Option>, compactions: Vec, 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::>(); 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()); diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 3ffc7aa5ef..7ed8810601 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/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::>(); + 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::>(), + answers[..1] + ); +} + #[tokio::test] async fn migration_preserves_reverse_replay_anchor_after_pre_compaction_rollback() { let home = TempDir::new().expect("create Codex home");