From 9997539fa1108bb59b18b42cdcd6b03270e30b0c Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Wed, 15 Apr 2026 22:01:13 -0700 Subject: [PATCH] Preserve suffix history for prefix compaction --- codex-rs/core/src/codex.rs | 44 +++-- codex-rs/core/src/compact.rs | 72 +------- codex-rs/core/src/compact_tests.rs | 173 ++++---------------- codex-rs/core/src/state/session.rs | 6 + codex-rs/core/tests/suite/compact_remote.rs | 2 +- 5 files changed, 64 insertions(+), 233 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4839967b54..a49c13cbff 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6760,9 +6760,7 @@ async fn run_auto_compact( .take_ready_prefix_compact(&turn_context.model_info.slug) .await; if let Some(candidate) = prefix_candidate { - if apply_prefix_compact_candidate(sess, turn_context, initial_context_injection, candidate) - .await? - { + if apply_prefix_compact_candidate(sess, turn_context, candidate).await? { return Ok(()); } debug!( @@ -6803,7 +6801,6 @@ async fn run_auto_compact( async fn apply_prefix_compact_candidate( sess: &Arc, turn_context: &Arc, - initial_context_injection: InitialContextInjection, candidate: PrefixCompactCandidate, ) -> CodexResult { let current_history = sess.clone_history().await; @@ -6826,24 +6823,19 @@ async fn apply_prefix_compact_candidate( sess.emit_turn_item_started(turn_context, &compaction_item) .await; - let retained_suffix = compact::strip_injected_context_from_retained_suffix( - current_items[candidate.base_history.len()..] - .iter() - .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) - .cloned(), - ); - let initial_context = if matches!( - initial_context_injection, - InitialContextInjection::BeforeLastUserMessage - ) { - sess.build_initial_context(turn_context.as_ref()).await - } else { - Vec::new() - }; + let retained_suffix: Vec = current_items[candidate.base_history.len()..] + .iter() + .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) + .cloned() + .collect(); + let reference_context_item = sess + .reference_context_item() + .await + .or(candidate.captured_reference_context_item); let mut new_history = compact::build_prefix_compacted_history( candidate.replacement_prefix, + candidate.captured_context, retained_suffix, - initial_context, ); new_history.extend( @@ -6853,10 +6845,6 @@ async fn apply_prefix_compact_candidate( .cloned(), ); - let reference_context_item = match initial_context_injection { - InitialContextInjection::DoNotInject => None, - InitialContextInjection::BeforeLastUserMessage => Some(turn_context.to_turn_context_item()), - }; let compacted_item = CompactedItem { message: String::new(), replacement_history: Some(new_history.clone()), @@ -6917,7 +6905,7 @@ async fn maybe_start_prefix_compact( return; } - let Some(start) = sess + let Some(mut start) = sess .begin_prefix_compact(turn_context.model_info.slug.clone()) .await else { @@ -6930,6 +6918,8 @@ async fn maybe_start_prefix_compact( ); return; }; + start.captured_context = sess.build_initial_context(turn_context.as_ref()).await; + start.captured_reference_context_item = Some(turn_context.to_turn_context_item()); debug!( turn_id = %turn_context.sub_id, @@ -6939,6 +6929,7 @@ async fn maybe_start_prefix_compact( prefix_compact_limit, auto_compact_limit, base_history_len = start.base_history.len(), + captured_context_len = start.captured_context.len(), "starting background prefix compaction" ); spawn_prefix_compact_task(Arc::clone(sess), Arc::clone(turn_context), start); @@ -6954,6 +6945,8 @@ fn spawn_prefix_compact_task( generation, model_slug, base_history, + captured_context, + captured_reference_context_item, } = start; let compacted = run_remote_prefix_compact_task( Arc::clone(&sess), @@ -6969,6 +6962,7 @@ fn spawn_prefix_compact_task( model_slug = %model_slug, base_history_len = base_history.len(), replacement_prefix_len = replacement_prefix.len(), + captured_context_len = captured_context.len(), "background prefix compaction ready" ); sess.finish_prefix_compact(PrefixCompactCandidate { @@ -6976,6 +6970,8 @@ fn spawn_prefix_compact_task( model_slug, base_history, replacement_prefix, + captured_context, + captured_reference_context_item, }) .await; } diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index ec55257e24..a104a3075c 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -481,82 +481,22 @@ pub(crate) fn insert_initial_context_before_last_real_user_or_summary( compacted_history } -/// Builds prefix-compaction replacement history from a processed compacted prefix and the live -/// retained suffix, then inserts canonical context at the active suffix boundary. +/// Builds prefix-compaction replacement history from a processed compacted prefix, the context +/// captured when that prefix snapshot was taken, and the live retained suffix. /// /// The prefix has already gone through the same retention filtering as classic remote compaction. -/// The suffix has not been summarized by `/compact`, so callers should only remove stale injected -/// context from it and preserve live transcript/tool state. -/// -/// Placement rules: -/// - If the suffix has a real user message, inject immediately before the last such suffix message. -/// - If the suffix has no real user message, inject immediately after the compacted prefix. +/// The captured context is inserted immediately after the prefix so the untouched suffix remains +/// under the same chronological context state that existed when the suffix began. pub(crate) fn build_prefix_compacted_history( mut processed_prefix: Vec, + captured_context: Vec, retained_suffix: Vec, - initial_context: Vec, ) -> Vec { - let prefix_len = processed_prefix.len(); - let insertion_index = retained_suffix - .iter() - .enumerate() - .rev() - .find_map(|(i, item)| is_real_user_message(item).then_some(prefix_len + i)) - .unwrap_or(prefix_len); - + processed_prefix.extend(captured_context); processed_prefix.extend(retained_suffix); - processed_prefix.splice(insertion_index..insertion_index, initial_context); processed_prefix } -fn is_real_user_message(item: &ResponseItem) -> bool { - let Some(TurnItem::UserMessage(user)) = crate::event_mapping::parse_turn_item(item) else { - return false; - }; - !is_summary_message(&user.message()) -} - -/// Removes client-injected context scaffolding from retained, uncompacted suffix history. -/// -/// Prefix compaction replaces only an earlier prefix and appends the still-live suffix. The suffix -/// can contain context items that were injected before a later turn and are now stale after we -/// re-inject canonical current context. Unlike compacted prefix cleanup, this intentionally keeps -/// live model/tool state such as reasoning, function calls, and tool outputs. -pub(crate) fn strip_injected_context_from_retained_suffix(items: I) -> Vec -where - I: IntoIterator, -{ - items - .into_iter() - .filter(|item| !is_injected_context_item(item)) - .collect() -} - -fn is_injected_context_item(item: &ResponseItem) -> bool { - match item { - ResponseItem::Message { role, content, .. } if role == "developer" => { - crate::event_mapping::is_contextual_dev_message_content(content) - } - ResponseItem::Message { role, content, .. } if role == "user" => { - is_regenerable_contextual_user_message(content) - } - _ => false, - } -} - -fn is_regenerable_contextual_user_message(content: &[ContentItem]) -> bool { - !content.is_empty() && content.iter().all(is_regenerable_contextual_user_fragment) -} - -fn is_regenerable_contextual_user_fragment(content_item: &ContentItem) -> bool { - let ContentItem::InputText { text } = content_item else { - return false; - }; - - codex_instructions::AGENTS_MD_FRAGMENT.matches_text(text) - || crate::contextual_user_message::ENVIRONMENT_CONTEXT_FRAGMENT.matches_text(text) -} - pub(crate) fn build_compacted_history( initial_context: Vec, user_messages: &[String], diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index c831d830b6..b93fe7d17f 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -2,9 +2,6 @@ use super::*; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; -use codex_protocol::items::HookPromptFragment; -use codex_protocol::items::build_hook_prompt_message; -use codex_protocol::models::FunctionCallOutputPayload; use pretty_assertions::assert_eq; async fn process_compacted_history_with_test_session( @@ -174,123 +171,6 @@ do things assert_eq!(vec!["real user message".to_string()], collected); } -#[test] -fn strip_injected_context_from_retained_suffix_preserves_live_history() { - let contextual_developer = input_message("developer", "\nread only"); - let mixed_developer = ResponseItem::Message { - id: None, - role: "developer".to_string(), - content: vec![ - ContentItem::InputText { - text: "Persistent developer instruction.".to_string(), - }, - ContentItem::InputText { - text: "\nworkspace-write".to_string(), - }, - ], - end_turn: None, - phase: None, - }; - let user_instructions = input_message( - "user", - r#"# AGENTS.md instructions for /repo - - -Use cargo test. -"#, - ); - let contextual_user = input_message( - "user", - r#" - /repo - zsh -"#, - ); - let skill_payload = input_message( - "user", - r#" -fix-ci -/skills/fix-ci/SKILL.md -Debug the failing check. -"#, - ); - let user_shell_command = input_message( - "user", - r#" - -echo hi - - -Exit code: 0 -Output: -hi - -"#, - ); - let subagent_notification = input_message( - "user", - r#"{"agent_path":"agent-1","status":"completed"}"#, - ); - let turn_aborted = input_message("user", "interrupted"); - let real_user = input_message("user", "please continue"); - let assistant = ResponseItem::Message { - id: None, - role: "assistant".to_string(), - content: vec![ContentItem::OutputText { - text: "continuing".to_string(), - }], - end_turn: None, - phase: None, - }; - let function_call = ResponseItem::FunctionCall { - id: None, - name: "read_file".to_string(), - namespace: None, - arguments: "{}".to_string(), - call_id: "call-1".to_string(), - }; - let function_output = ResponseItem::FunctionCallOutput { - call_id: "call-1".to_string(), - output: FunctionCallOutputPayload::from_text("ok".to_string()), - }; - let hook_prompt = build_hook_prompt_message(&[HookPromptFragment::from_single_hook( - "Retry with tests.", - "hook-run-1", - )]) - .expect("hook prompt message"); - - let cleaned = strip_injected_context_from_retained_suffix(vec![ - contextual_developer, - mixed_developer, - user_instructions, - contextual_user, - skill_payload.clone(), - user_shell_command.clone(), - subagent_notification.clone(), - turn_aborted.clone(), - real_user.clone(), - assistant.clone(), - function_call.clone(), - function_output.clone(), - hook_prompt.clone(), - ]); - - assert_eq!( - cleaned, - vec![ - skill_payload, - user_shell_command, - subagent_notification, - turn_aborted, - real_user, - assistant, - function_call, - function_output, - hook_prompt, - ] - ); -} - #[test] fn build_token_limited_compacted_history_truncates_overlong_user_messages() { // Use a small truncation limit so the test remains fast while still validating @@ -743,54 +623,63 @@ fn insert_initial_context_before_last_real_user_or_summary_keeps_compaction_last } #[test] -fn build_prefix_compacted_history_injects_before_last_suffix_user() { +fn build_prefix_compacted_history_inserts_captured_context_after_prefix() { let prefix_user = input_message("user", "older user in compacted prefix"); + let captured_developer = input_message("developer", "captured permissions"); + let captured_user = input_message("user", "captured environment"); let suffix_assistant = output_message("working through suffix"); let suffix_user = input_message("user", "latest user in retained suffix"); - let suffix_after_user = output_message("assistant after suffix user"); - let initial_context = input_message("developer", "fresh current context"); let refreshed = build_prefix_compacted_history( vec![prefix_user.clone()], - vec![ - suffix_assistant.clone(), - suffix_user.clone(), - suffix_after_user.clone(), - ], - vec![initial_context.clone()], + vec![captured_developer.clone(), captured_user.clone()], + vec![suffix_assistant.clone(), suffix_user.clone()], ); assert_eq!( refreshed, vec![ prefix_user, + captured_developer, + captured_user, suffix_assistant, - initial_context, suffix_user, - suffix_after_user, ] ); } #[test] -fn build_prefix_compacted_history_injects_after_prefix_when_suffix_has_no_user() { +fn build_prefix_compacted_history_preserves_suffix_context_items() { let prefix_user = input_message("user", "older user in compacted prefix"); - let suffix_assistant = output_message("continuing without user"); - let hook_prompt = build_hook_prompt_message(&[HookPromptFragment::from_single_hook( - "Retry with tests.", - "hook-run-1", - )]) - .expect("hook prompt message"); - let initial_context = input_message("developer", "fresh current context"); + let captured_context = input_message("developer", "captured permissions"); + let suffix_developer = input_message("developer", "\nnew sandbox"); + let suffix_user_context = input_message( + "user", + r#" + /repo + zsh +"#, + ); + let suffix_assistant = output_message("continuing after context change"); let refreshed = build_prefix_compacted_history( vec![prefix_user.clone()], - vec![suffix_assistant.clone(), hook_prompt.clone()], - vec![initial_context.clone()], + vec![captured_context.clone()], + vec![ + suffix_developer.clone(), + suffix_user_context.clone(), + suffix_assistant.clone(), + ], ); assert_eq!( refreshed, - vec![prefix_user, initial_context, suffix_assistant, hook_prompt] + vec![ + prefix_user, + captured_context, + suffix_developer, + suffix_user_context, + suffix_assistant, + ] ); } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index d2b672058f..d8dfd60105 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -112,6 +112,8 @@ impl SessionState { generation, model_slug, base_history, + captured_context: Vec::new(), + captured_reference_context_item: None, }) } @@ -301,6 +303,8 @@ pub(crate) struct PrefixCompactStart { pub(crate) generation: u64, pub(crate) model_slug: String, pub(crate) base_history: Vec, + pub(crate) captured_context: Vec, + pub(crate) captured_reference_context_item: Option, } #[derive(Debug, Clone)] @@ -309,6 +313,8 @@ pub(crate) struct PrefixCompactCandidate { pub(crate) model_slug: String, pub(crate) base_history: Vec, pub(crate) replacement_prefix: Vec, + pub(crate) captured_context: Vec, + pub(crate) captured_reference_context_item: Option, } // Sometimes new snapshots don't include credits or plan information. diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 4ab5cb6bb0..3da18458c0 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -1379,7 +1379,7 @@ async fn ready_prefix_compact_is_applied_by_pre_turn_auto_compact() -> Result<() .count(); assert_eq!( stale_context_count, 1, - "prefix compaction should strip stale retained suffix context before current context is reinjected" + "prefix compaction should preserve context updates in the retained suffix" ); Ok(())