Preserve suffix history for prefix compaction

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<Session>,
     turn_context: &Arc<TurnContext>,
-    initial_context_injection: InitialContextInjection,
     candidate: PrefixCompactCandidate,
 ) -> CodexResult<bool> {
     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<ResponseItem> = 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;
             }

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<ResponseItem>,
+    captured_context: Vec<ResponseItem>,
     retained_suffix: Vec<ResponseItem>,
-    initial_context: Vec<ResponseItem>,
 ) -> Vec<ResponseItem> {
-    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<I>(items: I) -> Vec<ResponseItem>
-where
-    I: IntoIterator<Item = ResponseItem>,
-{
-    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<ResponseItem>,
     user_messages: &[String],

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", "<permissions instructions>\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: "<permissions instructions>\nworkspace-write".to_string(),
-            },
-        ],
-        end_turn: None,
-        phase: None,
-    };
-    let user_instructions = input_message(
-        "user",
-        r#"# AGENTS.md instructions for /repo
-
-<INSTRUCTIONS>
-Use cargo test.
-</INSTRUCTIONS>"#,
-    );
-    let contextual_user = input_message(
-        "user",
-        r#"<environment_context>
-  <cwd>/repo</cwd>
-  <shell>zsh</shell>
-</environment_context>"#,
-    );
-    let skill_payload = input_message(
-        "user",
-        r#"<skill>
-<name>fix-ci</name>
-<path>/skills/fix-ci/SKILL.md</path>
-Debug the failing check.
-</skill>"#,
-    );
-    let user_shell_command = input_message(
-        "user",
-        r#"<user_shell_command>
-<command>
-echo hi
-</command>
-<result>
-Exit code: 0
-Output:
-hi
-</result>
-</user_shell_command>"#,
-    );
-    let subagent_notification = input_message(
-        "user",
-        r#"<subagent_notification>{"agent_path":"agent-1","status":"completed"}</subagent_notification>"#,
-    );
-    let turn_aborted = input_message("user", "<turn_aborted>interrupted</turn_aborted>");
-    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", "<permissions instructions>\nnew sandbox");
+    let suffix_user_context = input_message(
+        "user",
+        r#"<environment_context>
+  <cwd>/repo</cwd>
+  <shell>zsh</shell>
+</environment_context>"#,
+    );
+    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,
+        ]
     );
 }

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<ResponseItem>,
+    pub(crate) captured_context: Vec<ResponseItem>,
+    pub(crate) captured_reference_context_item: Option<TurnContextItem>,
 }
 
 #[derive(Debug, Clone)]
@@ -309,6 +313,8 @@ pub(crate) struct PrefixCompactCandidate {
     pub(crate) model_slug: String,
     pub(crate) base_history: Vec<ResponseItem>,
     pub(crate) replacement_prefix: Vec<ResponseItem>,
+    pub(crate) captured_context: Vec<ResponseItem>,
+    pub(crate) captured_reference_context_item: Option<TurnContextItem>,
 }
 
 // Sometimes new snapshots don't include credits or plan information.

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(())