From bbb5ff106f7759dcdf64759c94f7f905e92a2817 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Fri, 13 Mar 2026 13:08:30 -0700 Subject: [PATCH] Trim synthetic context during rollback When a rollback targets an earlier user turn, also trim the contiguous developer and contextual user messages injected immediately before that turn. This keeps stale per-turn context out of rebuilt history and adds unit plus integration coverage for adjacent and non-adjacent cases. Co-authored-by: Codex --- codex-rs/core/src/codex_tests.rs | 48 ++++++++++++++ codex-rs/core/src/context_manager/history.rs | 35 ++++++++++ .../core/src/context_manager/history_tests.rs | 65 +++++++++++++++++++ 3 files changed, 148 insertions(+) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index a1d88ff3dd..bc79b3d252 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -7,6 +7,7 @@ use crate::config_loader::ConfigLayerStackOrdering; use crate::config_loader::NetworkConstraints; use crate::config_loader::RequirementSource; use crate::config_loader::Sourced; +use crate::contextual_user_message::ENVIRONMENT_CONTEXT_FRAGMENT; use crate::exec::ExecToolCallOutput; use crate::function_tool::FunctionCallError; use crate::mcp_connection_manager::ToolInfo; @@ -1079,6 +1080,53 @@ async fn thread_rollback_clears_history_when_num_turns_exceeds_existing_turns() assert_eq!(initial_context, history.raw_items()); } +#[tokio::test] +async fn thread_rollback_trims_adjacent_synthetic_context_before_removed_turn() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + attach_rollout_recorder(&sess).await; + + let initial_context = sess.build_initial_context(tc.as_ref()).await; + let turn_1 = vec![ + user_message("turn 1 user"), + assistant_message("turn 1 assistant"), + ]; + let turn_2_context = vec![ + ResponseItem::from(DeveloperInstructions::new( + "switched to gpt-test".to_string(), + )), + ENVIRONMENT_CONTEXT_FRAGMENT + .into_message(ENVIRONMENT_CONTEXT_FRAGMENT.wrap("/tmp/updated".to_string())), + ]; + let turn_2 = vec![ + user_message("turn 2 user"), + assistant_message("turn 2 assistant"), + ]; + let mut full_history = Vec::new(); + full_history.extend(initial_context.clone()); + full_history.extend(turn_1.clone()); + full_history.extend(turn_2_context); + full_history.extend(turn_2); + sess.replace_history(full_history.clone(), Some(tc.to_turn_context_item())) + .await; + let rollout_items: Vec = full_history + .into_iter() + .map(RolloutItem::ResponseItem) + .collect(); + sess.persist_rollout_items(&rollout_items).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); + + let mut expected = Vec::new(); + expected.extend(initial_context); + expected.extend(turn_1); + + let history = sess.clone_history().await; + assert_eq!(expected, history.raw_items()); +} + #[tokio::test] async fn thread_rollback_fails_without_persisted_rollout_path() { 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 19dbcf7293..6955214fd7 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -212,6 +212,9 @@ 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. + /// - when rolling back to an earlier turn, contiguous developer/contextual messages inserted + /// immediately before the removed user turn are also trimmed so stale per-turn context does + /// not survive above the restored boundary. pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) { if num_turns == 0 { return; @@ -230,6 +233,8 @@ impl ContextManager { } else { user_positions[user_positions.len() - n_from_end] }; + let cut_idx = + trim_adjacent_rollback_context_before_user_boundary(&snapshot, first_user_idx, cut_idx); self.replace(snapshot[..cut_idx].to_vec()); } @@ -637,6 +642,36 @@ pub(crate) fn is_user_turn_boundary(item: &ResponseItem) -> bool { role == "user" && !is_contextual_user_message_content(content) } +fn trim_adjacent_rollback_context_before_user_boundary( + items: &[ResponseItem], + first_user_idx: usize, + cut_idx: usize, +) -> usize { + if cut_idx <= first_user_idx { + return cut_idx; + } + + let mut adjusted_cut_idx = cut_idx; + while adjusted_cut_idx > first_user_idx + && items + .get(adjusted_cut_idx.saturating_sub(1)) + .is_some_and(is_adjacent_rollback_context_item) + { + adjusted_cut_idx -= 1; + } + adjusted_cut_idx +} + +fn is_adjacent_rollback_context_item(item: &ResponseItem) -> bool { + match item { + ResponseItem::Message { role, .. } if role == "developer" => true, + ResponseItem::Message { role, content, .. } if role == "user" => { + is_contextual_user_message_content(content) + } + _ => false, + } +} + fn user_message_positions(items: &[ResponseItem]) -> Vec { let mut positions = Vec::new(); for (idx, item) in items.iter().enumerate() { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 066272748f..1195c9b1e9 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::contextual_user_message::ENVIRONMENT_CONTEXT_FRAGMENT; use crate::truncate; use crate::truncate::TruncationPolicy; use base64::Engine; @@ -6,6 +7,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_git::GhostCommit; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; +use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -70,6 +72,10 @@ fn user_input_text_msg(text: &str) -> ResponseItem { } } +fn developer_msg(text: &str) -> ResponseItem { + DeveloperInstructions::new(text.to_string()).into() +} + fn custom_tool_call_output(call_id: &str, output: &str) -> ResponseItem { ResponseItem::CustomToolCallOutput { call_id: call_id.to_string(), @@ -794,6 +800,65 @@ 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_adjacent_developer_and_contextual_messages() { + let contextual_update = ENVIRONMENT_CONTEXT_FRAGMENT + .into_message(ENVIRONMENT_CONTEXT_FRAGMENT.wrap("/tmp/updated".to_string())); + let items = vec![ + user_input_text_msg("session prefix"), + user_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg("switched to gpt-test"), + contextual_update, + user_msg("turn 2 user"), + assistant_msg("turn 2 assistant"), + ]; + + let modalities = default_input_modalities(); + let mut history = create_history_with_items(items); + history.drop_last_n_user_turns(1); + + assert_eq!( + history.for_prompt(&modalities), + vec![ + user_input_text_msg("session prefix"), + user_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + ] + ); +} + +#[test] +fn drop_last_n_user_turns_keeps_non_adjacent_synthetic_context() { + let contextual_update = ENVIRONMENT_CONTEXT_FRAGMENT + .into_message(ENVIRONMENT_CONTEXT_FRAGMENT.wrap("/tmp/updated".to_string())); + let items = vec![ + user_input_text_msg("session prefix"), + user_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg("switched to gpt-test"), + assistant_msg("separator"), + contextual_update, + user_msg("turn 2 user"), + assistant_msg("turn 2 assistant"), + ]; + + let modalities = default_input_modalities(); + let mut history = create_history_with_items(items); + history.drop_last_n_user_turns(1); + + assert_eq!( + history.for_prompt(&modalities), + vec![ + user_input_text_msg("session prefix"), + user_msg("turn 1 user"), + assistant_msg("turn 1 assistant"), + developer_msg("switched to gpt-test"), + assistant_msg("separator"), + ] + ); +} + #[test] fn remove_first_item_handles_custom_tool_pair() { let items = vec![