From 93bde6adb2f48d79fcc7c52f223788cd8c8311bf Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Mon, 23 Mar 2026 21:37:19 -0700 Subject: [PATCH 1/4] Trim rollback context bundles and clear stale baseline Co-authored-by: Codex --- .../core/src/codex/rollout_reconstruction.rs | 26 ++- .../src/codex/rollout_reconstruction_tests.rs | 165 ++++++++++++++++++ codex-rs/core/src/codex_tests.rs | 115 ++++++++++++ codex-rs/core/src/context_manager/history.rs | 34 +++- .../core/src/context_manager/history_tests.rs | 88 ++++++++++ codex-rs/core/src/event_mapping.rs | 33 ++++ .../core/tests/suite/compact_resume_fork.rs | 40 ++--- ..._followup_turn_trims_context_updates.snap} | 8 +- ...fork__rollback_past_compaction_shapes.snap | 4 +- 9 files changed, 475 insertions(+), 38 deletions(-) rename codex-rs/core/tests/suite/snapshots/{all__suite__compact_resume_fork__rollback_followup_turn_duplicates_context_updates.snap => all__suite__compact_resume_fork__rollback_followup_turn_trims_context_updates.snap} (57%) diff --git a/codex-rs/core/src/codex/rollout_reconstruction.rs b/codex-rs/core/src/codex/rollout_reconstruction.rs index a4c042af0c8..a36ea8a91a6 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction.rs @@ -233,6 +233,10 @@ impl Session { let mut history = ContextManager::new(); let mut saw_legacy_compaction_without_replacement_history = false; + // Trimming a mixed rollback bundle removes non-diff developer content, so the final + // baseline must be cleared unless a later real user turn re-establishes it. + let mut should_clear_reference_context_item = false; + let mut saw_user_turn_since_mixed_context_trim = false; if let Some(base_replacement_history) = base_replacement_history { history.replace(base_replacement_history.to_vec()); } @@ -246,6 +250,7 @@ impl Session { std::iter::once(response_item), turn_context.truncation_policy, ); + saw_user_turn_since_mixed_context_trim |= is_user_turn_boundary(response_item); } RolloutItem::Compacted(compacted) => { if let Some(replacement_history) = &compacted.replacement_history { @@ -271,12 +276,21 @@ impl Session { history.replace(rebuilt); } } + RolloutItem::EventMsg(EventMsg::UserMessage(_)) => { + saw_user_turn_since_mixed_context_trim = true; + } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { - history.drop_last_n_user_turns(rollback.num_turns); + should_clear_reference_context_item |= + history.drop_last_n_user_turns(rollback.num_turns); + saw_user_turn_since_mixed_context_trim = false; + } + RolloutItem::TurnContext(_) => { + if should_clear_reference_context_item && saw_user_turn_since_mixed_context_trim + { + should_clear_reference_context_item = false; + } } - RolloutItem::EventMsg(_) - | RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) => {} + RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} } } @@ -286,7 +300,9 @@ impl Session { Some(*turn_reference_context_item) } }; - let reference_context_item = if saw_legacy_compaction_without_replacement_history { + let reference_context_item = if saw_legacy_compaction_without_replacement_history + || should_clear_reference_context_item + { None } else { reference_context_item diff --git a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs index 33fb5f9af71..40684327a49 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs @@ -35,6 +35,21 @@ fn assistant_message(text: &str) -> ResponseItem { } } +fn developer_message_with_fragments(texts: &[&str]) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: texts + .iter() + .map(|text| ContentItem::InputText { + text: (*text).to_string(), + }) + .collect(), + end_turn: None, + phase: None, + } +} + fn inter_agent_assistant_message(text: &str) -> ResponseItem { let communication = InterAgentCommunication::new( AgentPath::root(), @@ -262,6 +277,156 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com ); } +#[tokio::test] +async fn reconstruct_history_rollback_preserves_later_reference_context_item_after_mixed_trim() { + let (session, turn_context) = make_session_and_context().await; + let first_context_item = turn_context.to_turn_context_item(); + let first_turn_id = first_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let mut rolled_back_context_item = first_context_item.clone(); + rolled_back_context_item.turn_id = Some("rolled-back-turn".to_string()); + rolled_back_context_item.model = "rolled-back-model".to_string(); + let rolled_back_turn_id = rolled_back_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let mut follow_up_context_item = first_context_item.clone(); + follow_up_context_item.turn_id = Some("follow-up-turn".to_string()); + follow_up_context_item.model = "follow-up-model".to_string(); + let follow_up_turn_id = follow_up_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let turn_one_user = user_message("turn 1 user"); + let turn_one_assistant = assistant_message("turn 1 assistant"); + let turn_three_contextual_user = + user_message("follow-up-cwd"); + let turn_three_user = user_message("turn 3 user"); + let turn_three_assistant = assistant_message("turn 3 assistant"); + let turn_three_developer = developer_message_with_fragments(&[ + "follow-up contextual permissions", + "follow-up persistent plugin instructions", + ]); + + let rollout_items = vec![ + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: first_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 1 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(first_context_item.clone()), + RolloutItem::ResponseItem(turn_one_user.clone()), + RolloutItem::ResponseItem(turn_one_assistant.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: first_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: rolled_back_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 2 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(rolled_back_context_item), + RolloutItem::ResponseItem(developer_message_with_fragments(&[ + "contextual permissions", + "persistent plugin instructions", + ])), + RolloutItem::ResponseItem(user_message( + "rolled-back-cwd", + )), + RolloutItem::ResponseItem(user_message("turn 2 user")), + RolloutItem::ResponseItem(assistant_message("turn 2 assistant")), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: rolled_back_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack( + codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, + )), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: follow_up_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 3 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(follow_up_context_item.clone()), + RolloutItem::ResponseItem(turn_three_developer.clone()), + RolloutItem::ResponseItem(turn_three_contextual_user.clone()), + RolloutItem::ResponseItem(turn_three_user.clone()), + RolloutItem::ResponseItem(turn_three_assistant.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: follow_up_turn_id, + last_agent_message: None, + }, + )), + ]; + + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + + assert_eq!( + reconstructed.history, + vec![ + turn_one_user, + turn_one_assistant, + turn_three_developer, + turn_three_contextual_user, + turn_three_user, + turn_three_assistant, + ] + ); + assert_eq!( + reconstructed.previous_turn_settings, + Some(PreviousTurnSettings { + model: "follow-up-model".to_string(), + realtime_active: Some(turn_context.realtime_active), + }) + ); + assert_eq!( + serde_json::to_value(reconstructed.reference_context_item) + .expect("serialize reconstructed reference context item"), + serde_json::to_value(Some(follow_up_context_item)) + .expect("serialize expected reconstructed reference context item") + ); +} + #[tokio::test] async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_incomplete_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 5d6769e0eac..8170a8e33a0 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -157,6 +157,21 @@ fn skill_message(text: &str) -> ResponseItem { } } +fn developer_message_with_fragments(texts: &[&str]) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: texts + .iter() + .map(|text| ContentItem::InputText { + text: (*text).to_string(), + }) + .collect(), + end_turn: None, + phase: None, + } +} + #[tokio::test] async fn regular_turn_emits_turn_started_without_waiting_for_startup_prewarm() { let (sess, tc, rx) = make_session_and_context_with_rx().await; @@ -1517,6 +1532,106 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context ); } +#[tokio::test] +async fn thread_rollback_clears_reference_context_item_for_trimmed_mixed_context_bundle() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + attach_rollout_recorder(&sess).await; + + let first_context_item = tc.to_turn_context_item(); + let first_turn_id = first_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let mut rolled_back_context_item = first_context_item.clone(); + rolled_back_context_item.turn_id = Some("rolled-back-turn".to_string()); + rolled_back_context_item.model = "rolled-back-model".to_string(); + let rolled_back_turn_id = rolled_back_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let turn_one_user = user_message("turn 1 user"); + let turn_one_assistant = assistant_message("turn 1 assistant"); + let turn_two_user = user_message("turn 2 user"); + let turn_two_assistant = assistant_message("turn 2 assistant"); + + sess.persist_rollout_items(&[ + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: first_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 1 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(first_context_item.clone()), + RolloutItem::ResponseItem(turn_one_user.clone()), + RolloutItem::ResponseItem(turn_one_assistant.clone()), + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: first_turn_id, + last_agent_message: None, + })), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: rolled_back_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 2 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(rolled_back_context_item), + RolloutItem::ResponseItem(developer_message_with_fragments(&[ + "contextual permissions", + "persistent plugin instructions", + ])), + RolloutItem::ResponseItem(user_message( + "rolled-back-cwd", + )), + RolloutItem::ResponseItem(turn_two_user), + RolloutItem::ResponseItem(turn_two_assistant), + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: rolled_back_turn_id, + last_agent_message: None, + })), + ]) + .await; + sess.replace_history( + vec![assistant_message("stale history")], + Some(first_context_item.clone()), + ) + .await; + + handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; + let rollback_event = wait_for_thread_rolled_back(&rx).await; + assert_eq!(rollback_event.num_turns, 1); + + assert_eq!( + sess.clone_history().await.raw_items(), + vec![turn_one_user, turn_one_assistant] + ); + assert_eq!( + sess.previous_turn_settings().await, + Some(PreviousTurnSettings { + model: tc.model_info.slug.clone(), + realtime_active: Some(tc.realtime_active), + }) + ); + assert!(sess.reference_context_item().await.is_none()); +} + #[tokio::test] async fn thread_rollback_restores_cleared_reference_context_item_after_compaction() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 57807e5f855..c837929919f 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,5 +1,7 @@ use crate::codex::TurnContext; use crate::context_manager::normalize; +use crate::event_mapping::has_non_contextual_dev_message_content; +use crate::event_mapping::is_contextual_dev_message_content; use crate::event_mapping::is_contextual_user_message_content; use crate::truncate::TruncationPolicy; use crate::truncate::approx_bytes_for_tokens; @@ -215,26 +217,50 @@ impl ContextManager { /// - if there are no user turns, this is a no-op /// - if `num_turns` exceeds the number of user turns, all user turns are dropped while /// preserving any items that occurred before the first user message. - pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) { + /// + /// Returns true when rollback trimmed a developer context bundle that also contained + /// non-context fragments, so callers should clear the final `reference_context_item` + /// baseline and force a full reinjection on the next real user turn. + pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) -> bool { if num_turns == 0 { - return; + return false; } let snapshot = self.items.clone(); let user_positions = user_message_positions(&snapshot); let Some(&first_user_idx) = user_positions.first() else { self.replace(snapshot); - return; + return false; }; let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); - let cut_idx = if n_from_end >= user_positions.len() { + let mut cut_idx = if n_from_end >= user_positions.len() { first_user_idx } else { user_positions[user_positions.len() - n_from_end] }; + let mut should_clear_reference_context_item = false; + while cut_idx > first_user_idx { + match &snapshot[cut_idx - 1] { + ResponseItem::Message { role, content, .. } + if role == "developer" && is_contextual_dev_message_content(content) => + { + should_clear_reference_context_item |= + has_non_contextual_dev_message_content(content); + cut_idx -= 1; + } + ResponseItem::Message { role, content, .. } + if role == "user" && is_contextual_user_message_content(content) => + { + cut_idx -= 1; + } + _ => break, + } + } + self.replace(snapshot[..cut_idx].to_vec()); + should_clear_reference_context_item } pub(crate) fn update_token_info( diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 1a4dc0ed8e3..56d0f056a43 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -90,6 +90,33 @@ fn user_input_text_msg(text: &str) -> ResponseItem { } } +fn developer_msg(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } +} + +fn developer_msg_with_fragments(texts: &[&str]) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: texts + .iter() + .map(|text| ContentItem::InputText { + text: (*text).to_string(), + }) + .collect(), + end_turn: None, + phase: None, + } +} + fn custom_tool_call_output(call_id: &str, output: &str) -> ResponseItem { ResponseItem::CustomToolCallOutput { call_id: call_id.to_string(), @@ -851,6 +878,67 @@ fn drop_last_n_user_turns_ignores_session_prefix_user_messages() { assert_eq!(history.for_prompt(&modalities), expected_prefix_only); } +#[test] +fn drop_last_n_user_turns_trims_context_updates_above_rolled_back_turn() { + let items = vec![ + assistant_msg("session prefix item"), + user_input_text_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg("Generated images are saved to /tmp as /tmp/image-1.png by default."), + developer_msg("ROLLED_BACK_DEV_INSTRUCTIONS"), + user_input_text_msg( + "PRETURN_CONTEXT_DIFF_CWD", + ), + user_input_text_msg("turn 2 user"), + assistant_msg("turn 2 assistant"), + ]; + + let modalities = default_input_modalities(); + let mut history = create_history_with_items(items); + let should_clear_reference_context_item = history.drop_last_n_user_turns(1); + + assert_eq!( + history.for_prompt(&modalities), + vec![ + assistant_msg("session prefix item"), + user_input_text_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg("Generated images are saved to /tmp as /tmp/image-1.png by default."), + ] + ); + assert!(!should_clear_reference_context_item); +} + +#[test] +fn drop_last_n_user_turns_flags_mixed_developer_context_bundles_for_reference_clear() { + let items = vec![ + user_input_text_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg_with_fragments(&[ + "contextual permissions", + "persistent plugin instructions", + ]), + user_input_text_msg( + "PRETURN_CONTEXT_DIFF_CWD", + ), + user_input_text_msg("turn 2 user"), + assistant_msg("turn 2 assistant"), + ]; + + let modalities = default_input_modalities(); + let mut history = create_history_with_items(items); + let should_clear_reference_context_item = history.drop_last_n_user_turns(1); + + assert_eq!( + history.for_prompt(&modalities), + vec![ + user_input_text_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + ] + ); + assert!(should_clear_reference_context_item); +} + #[test] fn remove_first_item_handles_custom_tool_pair() { let items = vec![ diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 8f80547a4c6..26173a07bc6 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -14,6 +14,8 @@ use codex_protocol::models::is_image_close_tag_text; use codex_protocol::models::is_image_open_tag_text; use codex_protocol::models::is_local_image_close_tag_text; use codex_protocol::models::is_local_image_open_tag_text; +use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG; +use codex_protocol::protocol::REALTIME_CONVERSATION_OPEN_TAG; use codex_protocol::user_input::UserInput; use tracing::warn; use uuid::Uuid; @@ -22,10 +24,41 @@ use crate::contextual_user_message::is_contextual_user_fragment; use crate::contextual_user_message::parse_visible_hook_prompt_message; use crate::web_search::web_search_action_detail; +const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ + "", + "", + COLLABORATION_MODE_OPEN_TAG, + REALTIME_CONVERSATION_OPEN_TAG, + "", +]; + pub(crate) fn is_contextual_user_message_content(message: &[ContentItem]) -> bool { message.iter().any(is_contextual_user_fragment) } +pub(crate) fn is_contextual_dev_message_content(message: &[ContentItem]) -> bool { + message.iter().any(is_contextual_dev_fragment) +} + +pub(crate) fn has_non_contextual_dev_message_content(message: &[ContentItem]) -> bool { + message + .iter() + .any(|content_item| !is_contextual_dev_fragment(content_item)) +} + +fn is_contextual_dev_fragment(content_item: &ContentItem) -> bool { + let ContentItem::InputText { text } = content_item else { + return false; + }; + + let trimmed = text.trim_start(); + CONTEXTUAL_DEVELOPER_PREFIXES.iter().any(|prefix| { + trimmed + .get(..prefix.len()) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(prefix)) + }) +} + fn parse_user_message(message: &[ContentItem]) -> Option { if is_contextual_user_message_content(message) { return None; diff --git a/codex-rs/core/tests/suite/compact_resume_fork.rs b/codex-rs/core/tests/suite/compact_resume_fork.rs index 2007aab55bd..a3c3f91370d 100644 --- a/codex-rs/core/tests/suite/compact_resume_fork.rs +++ b/codex-rs/core/tests/suite/compact_resume_fork.rs @@ -513,8 +513,9 @@ async fn snapshot_rollback_past_compaction_replays_append_only_history() -> Resu #[tokio::test(flavor = "multi_thread", worker_threads = 2)] /// Scenario: rolling back a turn that introduced persistent pre-turn context -/// diffs currently duplicates those context updates on the next request. -async fn snapshot_rollback_followup_turn_duplicates_context_updates() -> Result<()> { +/// diffs should trim those context updates so the next request includes them +/// only once. +async fn snapshot_rollback_followup_turn_trims_context_updates() -> Result<()> { if network_disabled() { println!("Skipping test because network is disabled in this sandbox"); return Ok(()); @@ -593,14 +594,12 @@ async fn snapshot_rollback_followup_turn_duplicates_context_updates() -> Result< let requests = request_log.requests(); assert_eq!(requests.len(), 3); - assert_eq!( - requests[1] - .message_input_texts("developer") - .iter() - .filter(|text| text.contains(ROLLED_BACK_DEV_INSTRUCTIONS)) - .count(), - 1 - ); + let before_rollback_developer_count = requests[1] + .message_input_texts("developer") + .iter() + .filter(|text| text.contains(ROLLED_BACK_DEV_INSTRUCTIONS)) + .count(); + assert_eq!(before_rollback_developer_count, 1); assert_eq!( requests[1] .message_input_texts("user") @@ -609,14 +608,13 @@ async fn snapshot_rollback_followup_turn_duplicates_context_updates() -> Result< .count(), 1 ); - assert_eq!( - requests[2] - .message_input_texts("developer") - .iter() - .filter(|text| text.contains(ROLLED_BACK_DEV_INSTRUCTIONS)) - .count(), - 2 - ); + + let after_rollback_developer_count = requests[2] + .message_input_texts("developer") + .iter() + .filter(|text| text.contains(ROLLED_BACK_DEV_INSTRUCTIONS)) + .count(); + assert_eq!(after_rollback_developer_count, 1); let after_rollback_user_texts = requests[2].message_input_texts("user"); assert_eq!( @@ -624,7 +622,7 @@ async fn snapshot_rollback_followup_turn_duplicates_context_updates() -> Result< .iter() .filter(|text| text.contains(PRETURN_CONTEXT_DIFF_CWD)) .count(), - 2 + 1 ); assert_eq!( after_rollback_user_texts.last().map(String::as_str), @@ -632,9 +630,9 @@ async fn snapshot_rollback_followup_turn_duplicates_context_updates() -> Result< ); insta::assert_snapshot!( - "rollback_followup_turn_duplicates_context_updates", + "rollback_followup_turn_trims_context_updates", context_snapshot::format_labeled_requests_snapshot( - "rollback currently duplicates pre-turn override context updates on the follow-up request", + "rollback trims pre-turn override context updates before the follow-up request", &[ ("rolled-back turn request", &requests[1]), ("follow-up request after rollback", &requests[2]), diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_duplicates_context_updates.snap b/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_trims_context_updates.snap similarity index 57% rename from codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_duplicates_context_updates.snap rename to codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_trims_context_updates.snap index 676421c3730..925a810a43c 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_duplicates_context_updates.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_followup_turn_trims_context_updates.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/compact_resume_fork.rs -expression: "context_snapshot::format_labeled_requests_snapshot(\"rollback currently duplicates pre-turn override context updates on the follow-up request\",\n&[(\"rolled-back turn request\", &requests[1]),\n(\"follow-up request after rollback\", &requests[2]),],\n&ContextSnapshotOptions::default().strip_capability_instructions().render_mode(ContextSnapshotRenderMode::KindWithTextPrefix\n{ max_chars: 96 }),)" +expression: "context_snapshot::format_labeled_requests_snapshot(\"rollback trims pre-turn override context updates before the follow-up request\",\n&[(\"rolled-back turn request\", &requests[1]),\n(\"follow-up request after rollback\", &requests[2]),],\n&ContextSnapshotOptions::default().strip_capability_instructions().render_mode(ContextSnapshotRenderMode::KindWithTextPrefix\n{ max_chars: 96 }),)" --- -Scenario: rollback currently duplicates pre-turn override context updates on the follow-up request +Scenario: rollback trims pre-turn override context updates before the follow-up request ## rolled-back turn request 00:message/developer: @@ -20,6 +20,4 @@ Scenario: rollback currently duplicates pre-turn override context updates on the 03:message/assistant:turn 1 assistant 04:message/developer:ROLLED_BACK_DEV_INSTRUCTIONS 05:message/user: -06:message/developer:ROLLED_BACK_DEV_INSTRUCTIONS -07:message/user: -08:message/user:follow-up user +06:message/user:follow-up user diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_past_compaction_shapes.snap b/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_past_compaction_shapes.snap index 04e45c3a682..bb694af6934 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_past_compaction_shapes.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__compact_resume_fork__rollback_past_compaction_shapes.snap @@ -23,6 +23,4 @@ Scenario: rollback past compaction replay after rollback 01:message/user:\nSUMMARY_ONLY_CONTEXT 02:message/developer: 03:message/user:> -04:message/developer: -05:message/user:> -06:message/user:AFTER_ROLLBACK +04:message/user:AFTER_ROLLBACK From bc1cfc3178c40d9243e5e9f25f1446fd434ac096 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Mon, 23 Mar 2026 23:17:21 -0700 Subject: [PATCH 2/4] remove rollout reconstruction diff --- .../core/src/codex/rollout_reconstruction.rs | 26 +-- .../src/codex/rollout_reconstruction_tests.rs | 165 ------------------ codex-rs/core/src/codex_tests.rs | 115 ------------ 3 files changed, 5 insertions(+), 301 deletions(-) diff --git a/codex-rs/core/src/codex/rollout_reconstruction.rs b/codex-rs/core/src/codex/rollout_reconstruction.rs index a36ea8a91a6..a4c042af0c8 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction.rs @@ -233,10 +233,6 @@ impl Session { let mut history = ContextManager::new(); let mut saw_legacy_compaction_without_replacement_history = false; - // Trimming a mixed rollback bundle removes non-diff developer content, so the final - // baseline must be cleared unless a later real user turn re-establishes it. - let mut should_clear_reference_context_item = false; - let mut saw_user_turn_since_mixed_context_trim = false; if let Some(base_replacement_history) = base_replacement_history { history.replace(base_replacement_history.to_vec()); } @@ -250,7 +246,6 @@ impl Session { std::iter::once(response_item), turn_context.truncation_policy, ); - saw_user_turn_since_mixed_context_trim |= is_user_turn_boundary(response_item); } RolloutItem::Compacted(compacted) => { if let Some(replacement_history) = &compacted.replacement_history { @@ -276,21 +271,12 @@ impl Session { history.replace(rebuilt); } } - RolloutItem::EventMsg(EventMsg::UserMessage(_)) => { - saw_user_turn_since_mixed_context_trim = true; - } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { - should_clear_reference_context_item |= - history.drop_last_n_user_turns(rollback.num_turns); - saw_user_turn_since_mixed_context_trim = false; - } - RolloutItem::TurnContext(_) => { - if should_clear_reference_context_item && saw_user_turn_since_mixed_context_trim - { - should_clear_reference_context_item = false; - } + history.drop_last_n_user_turns(rollback.num_turns); } - RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::EventMsg(_) + | RolloutItem::TurnContext(_) + | RolloutItem::SessionMeta(_) => {} } } @@ -300,9 +286,7 @@ impl Session { Some(*turn_reference_context_item) } }; - let reference_context_item = if saw_legacy_compaction_without_replacement_history - || should_clear_reference_context_item - { + let reference_context_item = if saw_legacy_compaction_without_replacement_history { None } else { reference_context_item diff --git a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs index 40684327a49..33fb5f9af71 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs @@ -35,21 +35,6 @@ fn assistant_message(text: &str) -> ResponseItem { } } -fn developer_message_with_fragments(texts: &[&str]) -> ResponseItem { - ResponseItem::Message { - id: None, - role: "developer".to_string(), - content: texts - .iter() - .map(|text| ContentItem::InputText { - text: (*text).to_string(), - }) - .collect(), - end_turn: None, - phase: None, - } -} - fn inter_agent_assistant_message(text: &str) -> ResponseItem { let communication = InterAgentCommunication::new( AgentPath::root(), @@ -277,156 +262,6 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com ); } -#[tokio::test] -async fn reconstruct_history_rollback_preserves_later_reference_context_item_after_mixed_trim() { - let (session, turn_context) = make_session_and_context().await; - let first_context_item = turn_context.to_turn_context_item(); - let first_turn_id = first_context_item - .turn_id - .clone() - .expect("turn context should have turn_id"); - let mut rolled_back_context_item = first_context_item.clone(); - rolled_back_context_item.turn_id = Some("rolled-back-turn".to_string()); - rolled_back_context_item.model = "rolled-back-model".to_string(); - let rolled_back_turn_id = rolled_back_context_item - .turn_id - .clone() - .expect("turn context should have turn_id"); - let mut follow_up_context_item = first_context_item.clone(); - follow_up_context_item.turn_id = Some("follow-up-turn".to_string()); - follow_up_context_item.model = "follow-up-model".to_string(); - let follow_up_turn_id = follow_up_context_item - .turn_id - .clone() - .expect("turn context should have turn_id"); - let turn_one_user = user_message("turn 1 user"); - let turn_one_assistant = assistant_message("turn 1 assistant"); - let turn_three_contextual_user = - user_message("follow-up-cwd"); - let turn_three_user = user_message("turn 3 user"); - let turn_three_assistant = assistant_message("turn 3 assistant"); - let turn_three_developer = developer_message_with_fragments(&[ - "follow-up contextual permissions", - "follow-up persistent plugin instructions", - ]); - - let rollout_items = vec![ - RolloutItem::EventMsg(EventMsg::TurnStarted( - codex_protocol::protocol::TurnStartedEvent { - turn_id: first_turn_id.clone(), - model_context_window: Some(128_000), - collaboration_mode_kind: ModeKind::Default, - }, - )), - RolloutItem::EventMsg(EventMsg::UserMessage( - codex_protocol::protocol::UserMessageEvent { - message: "turn 1 user".to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - }, - )), - RolloutItem::TurnContext(first_context_item.clone()), - RolloutItem::ResponseItem(turn_one_user.clone()), - RolloutItem::ResponseItem(turn_one_assistant.clone()), - RolloutItem::EventMsg(EventMsg::TurnComplete( - codex_protocol::protocol::TurnCompleteEvent { - turn_id: first_turn_id, - last_agent_message: None, - }, - )), - RolloutItem::EventMsg(EventMsg::TurnStarted( - codex_protocol::protocol::TurnStartedEvent { - turn_id: rolled_back_turn_id.clone(), - model_context_window: Some(128_000), - collaboration_mode_kind: ModeKind::Default, - }, - )), - RolloutItem::EventMsg(EventMsg::UserMessage( - codex_protocol::protocol::UserMessageEvent { - message: "turn 2 user".to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - }, - )), - RolloutItem::TurnContext(rolled_back_context_item), - RolloutItem::ResponseItem(developer_message_with_fragments(&[ - "contextual permissions", - "persistent plugin instructions", - ])), - RolloutItem::ResponseItem(user_message( - "rolled-back-cwd", - )), - RolloutItem::ResponseItem(user_message("turn 2 user")), - RolloutItem::ResponseItem(assistant_message("turn 2 assistant")), - RolloutItem::EventMsg(EventMsg::TurnComplete( - codex_protocol::protocol::TurnCompleteEvent { - turn_id: rolled_back_turn_id, - last_agent_message: None, - }, - )), - RolloutItem::EventMsg(EventMsg::ThreadRolledBack( - codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, - )), - RolloutItem::EventMsg(EventMsg::TurnStarted( - codex_protocol::protocol::TurnStartedEvent { - turn_id: follow_up_turn_id.clone(), - model_context_window: Some(128_000), - collaboration_mode_kind: ModeKind::Default, - }, - )), - RolloutItem::EventMsg(EventMsg::UserMessage( - codex_protocol::protocol::UserMessageEvent { - message: "turn 3 user".to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - }, - )), - RolloutItem::TurnContext(follow_up_context_item.clone()), - RolloutItem::ResponseItem(turn_three_developer.clone()), - RolloutItem::ResponseItem(turn_three_contextual_user.clone()), - RolloutItem::ResponseItem(turn_three_user.clone()), - RolloutItem::ResponseItem(turn_three_assistant.clone()), - RolloutItem::EventMsg(EventMsg::TurnComplete( - codex_protocol::protocol::TurnCompleteEvent { - turn_id: follow_up_turn_id, - last_agent_message: None, - }, - )), - ]; - - let reconstructed = session - .reconstruct_history_from_rollout(&turn_context, &rollout_items) - .await; - - assert_eq!( - reconstructed.history, - vec![ - turn_one_user, - turn_one_assistant, - turn_three_developer, - turn_three_contextual_user, - turn_three_user, - turn_three_assistant, - ] - ); - assert_eq!( - reconstructed.previous_turn_settings, - Some(PreviousTurnSettings { - model: "follow-up-model".to_string(), - realtime_active: Some(turn_context.realtime_active), - }) - ); - assert_eq!( - serde_json::to_value(reconstructed.reference_context_item) - .expect("serialize reconstructed reference context item"), - serde_json::to_value(Some(follow_up_context_item)) - .expect("serialize expected reconstructed reference context item") - ); -} - #[tokio::test] async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_incomplete_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 8170a8e33a0..5d6769e0eac 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -157,21 +157,6 @@ fn skill_message(text: &str) -> ResponseItem { } } -fn developer_message_with_fragments(texts: &[&str]) -> ResponseItem { - ResponseItem::Message { - id: None, - role: "developer".to_string(), - content: texts - .iter() - .map(|text| ContentItem::InputText { - text: (*text).to_string(), - }) - .collect(), - end_turn: None, - phase: None, - } -} - #[tokio::test] async fn regular_turn_emits_turn_started_without_waiting_for_startup_prewarm() { let (sess, tc, rx) = make_session_and_context_with_rx().await; @@ -1532,106 +1517,6 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context ); } -#[tokio::test] -async fn thread_rollback_clears_reference_context_item_for_trimmed_mixed_context_bundle() { - let (sess, tc, rx) = make_session_and_context_with_rx().await; - attach_rollout_recorder(&sess).await; - - let first_context_item = tc.to_turn_context_item(); - let first_turn_id = first_context_item - .turn_id - .clone() - .expect("turn context should have turn_id"); - let mut rolled_back_context_item = first_context_item.clone(); - rolled_back_context_item.turn_id = Some("rolled-back-turn".to_string()); - rolled_back_context_item.model = "rolled-back-model".to_string(); - let rolled_back_turn_id = rolled_back_context_item - .turn_id - .clone() - .expect("turn context should have turn_id"); - let turn_one_user = user_message("turn 1 user"); - let turn_one_assistant = assistant_message("turn 1 assistant"); - let turn_two_user = user_message("turn 2 user"); - let turn_two_assistant = assistant_message("turn 2 assistant"); - - sess.persist_rollout_items(&[ - RolloutItem::EventMsg(EventMsg::TurnStarted( - codex_protocol::protocol::TurnStartedEvent { - turn_id: first_turn_id.clone(), - model_context_window: Some(128_000), - collaboration_mode_kind: ModeKind::Default, - }, - )), - RolloutItem::EventMsg(EventMsg::UserMessage( - codex_protocol::protocol::UserMessageEvent { - message: "turn 1 user".to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - }, - )), - RolloutItem::TurnContext(first_context_item.clone()), - RolloutItem::ResponseItem(turn_one_user.clone()), - RolloutItem::ResponseItem(turn_one_assistant.clone()), - RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: first_turn_id, - last_agent_message: None, - })), - RolloutItem::EventMsg(EventMsg::TurnStarted( - codex_protocol::protocol::TurnStartedEvent { - turn_id: rolled_back_turn_id.clone(), - model_context_window: Some(128_000), - collaboration_mode_kind: ModeKind::Default, - }, - )), - RolloutItem::EventMsg(EventMsg::UserMessage( - codex_protocol::protocol::UserMessageEvent { - message: "turn 2 user".to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - }, - )), - RolloutItem::TurnContext(rolled_back_context_item), - RolloutItem::ResponseItem(developer_message_with_fragments(&[ - "contextual permissions", - "persistent plugin instructions", - ])), - RolloutItem::ResponseItem(user_message( - "rolled-back-cwd", - )), - RolloutItem::ResponseItem(turn_two_user), - RolloutItem::ResponseItem(turn_two_assistant), - RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: rolled_back_turn_id, - last_agent_message: None, - })), - ]) - .await; - sess.replace_history( - vec![assistant_message("stale history")], - Some(first_context_item.clone()), - ) - .await; - - handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; - let rollback_event = wait_for_thread_rolled_back(&rx).await; - assert_eq!(rollback_event.num_turns, 1); - - assert_eq!( - sess.clone_history().await.raw_items(), - vec![turn_one_user, turn_one_assistant] - ); - assert_eq!( - sess.previous_turn_settings().await, - Some(PreviousTurnSettings { - model: tc.model_info.slug.clone(), - realtime_active: Some(tc.realtime_active), - }) - ); - assert!(sess.reference_context_item().await.is_none()); -} - #[tokio::test] async fn thread_rollback_restores_cleared_reference_context_item_after_compaction() { let (sess, tc, rx) = make_session_and_context_with_rx().await; From a520b7b276e69833d55255b6e9b19561c0fb7f4f Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Mon, 23 Mar 2026 23:48:06 -0700 Subject: [PATCH 3/4] improve docs --- codex-rs/core/src/context_manager/history.rs | 27 ++++++---- .../core/src/context_manager/history_tests.rs | 50 ++++++++++++++++--- codex-rs/core/src/event_mapping.rs | 7 +++ 3 files changed, 67 insertions(+), 17 deletions(-) diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index c837929919f..f4cf39fbf0c 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -42,7 +42,9 @@ pub(crate) struct ContextManager { /// match the current turn after context updates are persisted. /// /// When this is `None`, settings diffing treats the next turn as having no - /// baseline and emits a full reinjection of context state. + /// baseline and emits a full reinjection of context state. Rollback may + /// also clear this when it trims a mixed initial-context developer bundle + /// whose non-diff fragments no longer exist in the surviving history. reference_context_item: Option, } @@ -218,19 +220,21 @@ impl ContextManager { /// - if `num_turns` exceeds the number of user turns, all user turns are dropped while /// preserving any items that occurred before the first user message. /// - /// Returns true when rollback trimmed a developer context bundle that also contained - /// non-context fragments, so callers should clear the final `reference_context_item` - /// baseline and force a full reinjection on the next real user turn. - pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) -> bool { + /// If rollback trims a pre-turn developer message that mixes contextual fragments with + /// persistent developer text from `build_initial_context`, this also clears + /// `reference_context_item`. The surviving history no longer contains the full bundle that + /// established the prior baseline, so future turns must fall back to full reinjection instead + /// of diffing against stale state. + pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) { if num_turns == 0 { - return false; + return; } let snapshot = self.items.clone(); let user_positions = user_message_positions(&snapshot); let Some(&first_user_idx) = user_positions.first() else { self.replace(snapshot); - return false; + return; }; let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); @@ -240,13 +244,13 @@ impl ContextManager { user_positions[user_positions.len() - n_from_end] }; - let mut should_clear_reference_context_item = false; + let mut trimmed_mixed_developer_context_bundle = false; while cut_idx > first_user_idx { match &snapshot[cut_idx - 1] { ResponseItem::Message { role, content, .. } if role == "developer" && is_contextual_dev_message_content(content) => { - should_clear_reference_context_item |= + trimmed_mixed_developer_context_bundle |= has_non_contextual_dev_message_content(content); cut_idx -= 1; } @@ -260,7 +264,10 @@ impl ContextManager { } self.replace(snapshot[..cut_idx].to_vec()); - should_clear_reference_context_item + if trimmed_mixed_developer_context_bundle { + // Must rebuild full context since we trimmed full context reinjection after backtracking + self.reference_context_item = None; + } } pub(crate) fn update_token_info( diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 56d0f056a43..91e7c8546c4 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -5,6 +5,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_git::GhostCommit; use codex_protocol::AgentPath; +use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputBody; @@ -18,12 +19,16 @@ use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::default_input_modalities; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::TurnContextItem; use image::ImageBuffer; use image::ImageFormat; use image::Rgba; use pretty_assertions::assert_eq; use regex_lite::Regex; +use std::path::PathBuf; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; @@ -117,6 +122,29 @@ fn developer_msg_with_fragments(texts: &[&str]) -> ResponseItem { } } +fn reference_context_item() -> TurnContextItem { + TurnContextItem { + turn_id: Some("reference-turn".to_string()), + trace_id: None, + cwd: PathBuf::from("/tmp/reference-cwd"), + current_date: Some("2026-03-23".to_string()), + timezone: Some("America/Los_Angeles".to_string()), + approval_policy: AskForApproval::OnRequest, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + network: None, + model: "gpt-test".to_string(), + personality: None, + collaboration_mode: None, + realtime_active: Some(false), + effort: None, + summary: ReasoningSummary::Auto, + user_instructions: None, + developer_instructions: None, + final_output_json_schema: None, + truncation_policy: Some(codex_protocol::protocol::TruncationPolicy::Tokens(10_000)), + } +} + fn custom_tool_call_output(call_id: &str, output: &str) -> ResponseItem { ResponseItem::CustomToolCallOutput { call_id: call_id.to_string(), @@ -895,10 +923,12 @@ fn drop_last_n_user_turns_trims_context_updates_above_rolled_back_turn() { let modalities = default_input_modalities(); let mut history = create_history_with_items(items); - let should_clear_reference_context_item = history.drop_last_n_user_turns(1); + 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.for_prompt(&modalities), + history.clone().for_prompt(&modalities), vec![ assistant_msg("session prefix item"), user_input_text_msg("turn 1 user"), @@ -906,11 +936,16 @@ fn drop_last_n_user_turns_trims_context_updates_above_rolled_back_turn() { developer_msg("Generated images are saved to /tmp as /tmp/image-1.png by default."), ] ); - assert!(!should_clear_reference_context_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") + ); } #[test] -fn drop_last_n_user_turns_flags_mixed_developer_context_bundles_for_reference_clear() { +fn drop_last_n_user_turns_clears_reference_context_for_mixed_developer_context_bundles() { let items = vec![ user_input_text_msg("turn 1 user"), assistant_msg("turn 1 assistant"), @@ -927,16 +962,17 @@ fn drop_last_n_user_turns_flags_mixed_developer_context_bundles_for_reference_cl let modalities = default_input_modalities(); let mut history = create_history_with_items(items); - let should_clear_reference_context_item = history.drop_last_n_user_turns(1); + history.set_reference_context_item(Some(reference_context_item())); + history.drop_last_n_user_turns(1); assert_eq!( - history.for_prompt(&modalities), + history.clone().for_prompt(&modalities), vec![ user_input_text_msg("turn 1 user"), assistant_msg("turn 1 assistant"), ] ); - assert!(should_clear_reference_context_item); + assert!(history.reference_context_item().is_none()); } #[test] diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 26173a07bc6..5e174944f12 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -36,10 +36,17 @@ pub(crate) fn is_contextual_user_message_content(message: &[ContentItem]) -> boo message.iter().any(is_contextual_user_fragment) } +/// Returns true when a developer message contains any rollback-trimmable contextual fragment. +/// +/// `build_initial_context` can bundle these fragments together with persistent developer text in a +/// single developer message, so callers that care about invalidating a stored reference baseline +/// should pair this with `has_non_contextual_dev_message_content`. pub(crate) fn is_contextual_dev_message_content(message: &[ContentItem]) -> bool { message.iter().any(is_contextual_dev_fragment) } +/// Returns true when a developer message contains any fragment that is not part of the +/// rollback-trimmable contextual prefix set. pub(crate) fn has_non_contextual_dev_message_content(message: &[ContentItem]) -> bool { message .iter() From dcb0a1b579524dd8bace7fcd2a4f4b4658989fbb Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Tue, 24 Mar 2026 12:20:59 -0700 Subject: [PATCH 4/4] Clarify rollback context trimming Co-authored-by: Codex --- codex-rs/core/src/context_manager/history.rs | 75 +++++++++++++------- 1 file changed, 51 insertions(+), 24 deletions(-) diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index f4cf39fbf0c..b518d8ce4ff 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -232,42 +232,22 @@ impl ContextManager { let snapshot = self.items.clone(); let user_positions = user_message_positions(&snapshot); - let Some(&first_user_idx) = user_positions.first() else { + let Some(&first_instruction_turn_idx) = user_positions.first() else { self.replace(snapshot); return; }; let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); let mut cut_idx = if n_from_end >= user_positions.len() { - first_user_idx + first_instruction_turn_idx } else { user_positions[user_positions.len() - n_from_end] }; - let mut trimmed_mixed_developer_context_bundle = false; - while cut_idx > first_user_idx { - match &snapshot[cut_idx - 1] { - ResponseItem::Message { role, content, .. } - if role == "developer" && is_contextual_dev_message_content(content) => - { - trimmed_mixed_developer_context_bundle |= - has_non_contextual_dev_message_content(content); - cut_idx -= 1; - } - ResponseItem::Message { role, content, .. } - if role == "user" && is_contextual_user_message_content(content) => - { - cut_idx -= 1; - } - _ => break, - } - } + cut_idx = + self.trim_pre_turn_context_updates(&snapshot, first_instruction_turn_idx, cut_idx); self.replace(snapshot[..cut_idx].to_vec()); - if trimmed_mixed_developer_context_bundle { - // Must rebuild full context since we trimmed full context reinjection after backtracking - self.reference_context_item = None; - } } pub(crate) fn update_token_info( @@ -415,6 +395,53 @@ impl ContextManager { | ResponseItem::Other => item.clone(), } } + + /// Walk backward from a rollback cut and trim contiguous pre-turn context-update items. + /// + /// 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. + /// + /// 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 + /// stored `reference_context_item` baseline so the next real turn falls back to full + /// reinjection. + 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 { + 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; + } + cut_idx -= 1; + } + ResponseItem::Message { role, content, .. } + if role == "user" && is_contextual_user_message_content(content) => + { + cut_idx -= 1; + } + _ => break, + } + } + cut_idx + } } fn truncate_function_output_payload(