From b6b6e3e485606e8e8c6d36232ac81cc3bfe3be3e Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Tue, 14 Apr 2026 05:55:52 -0700 Subject: [PATCH] Refine prefix compaction context placement --- codex-rs/core/src/codex.rs | 20 ++++----- codex-rs/core/src/compact.rs | 35 ++++++++++++++++ codex-rs/core/src/compact_tests.rs | 65 ++++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 548e22edc1..4839967b54 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6826,25 +6826,25 @@ async fn apply_prefix_compact_candidate( sess.emit_turn_item_started(turn_context, &compaction_item) .await; - let mut new_history = candidate.replacement_prefix; let retained_suffix = compact::strip_injected_context_from_retained_suffix( current_items[candidate.base_history.len()..] .iter() .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) .cloned(), ); - new_history.extend(retained_suffix); - - if matches!( + let initial_context = if matches!( initial_context_injection, InitialContextInjection::BeforeLastUserMessage ) { - let initial_context = sess.build_initial_context(turn_context.as_ref()).await; - new_history = compact::insert_initial_context_before_last_real_user_or_summary( - new_history, - initial_context, - ); - } + sess.build_initial_context(turn_context.as_ref()).await + } else { + Vec::new() + }; + let mut new_history = compact::build_prefix_compacted_history( + candidate.replacement_prefix, + retained_suffix, + initial_context, + ); new_history.extend( current_items diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 52731471f8..ec55257e24 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -481,6 +481,41 @@ 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. +/// +/// 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. +pub(crate) fn build_prefix_compacted_history( + mut processed_prefix: 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(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 diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 15737d976e..c831d830b6 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -38,6 +38,18 @@ fn input_message(role: &str, text: &str) -> ResponseItem { } } +fn output_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } +} + #[test] fn remote_compact_task_supports_openai_provider() { let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); @@ -729,3 +741,56 @@ fn insert_initial_context_before_last_real_user_or_summary_keeps_compaction_last ]; assert_eq!(refreshed, expected); } + +#[test] +fn build_prefix_compacted_history_injects_before_last_suffix_user() { + let prefix_user = input_message("user", "older user in compacted prefix"); + 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()], + ); + + assert_eq!( + refreshed, + vec![ + prefix_user, + suffix_assistant, + initial_context, + suffix_user, + suffix_after_user, + ] + ); +} + +#[test] +fn build_prefix_compacted_history_injects_after_prefix_when_suffix_has_no_user() { + 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 refreshed = build_prefix_compacted_history( + vec![prefix_user.clone()], + vec![suffix_assistant.clone(), hook_prompt.clone()], + vec![initial_context.clone()], + ); + + assert_eq!( + refreshed, + vec![prefix_user, initial_context, suffix_assistant, hook_prompt] + ); +}