From a6c594e296f87298e57089be0be43f577ba08054 Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Mon, 13 Apr 2026 17:09:57 -0700 Subject: [PATCH] Update prefix compaction mode and suffix cleanup --- codex-rs/codex-api/src/common.rs | 10 ++- codex-rs/codex-api/src/lib.rs | 1 + codex-rs/core/src/client.rs | 5 +- codex-rs/core/src/codex.rs | 3 +- codex-rs/core/src/compact.rs | 32 ++++++++ codex-rs/core/src/compact_remote.rs | 4 +- codex-rs/core/src/compact_tests.rs | 89 +++++++++++++++++++++ codex-rs/core/tests/suite/compact_remote.rs | 43 +++++++--- 8 files changed, 169 insertions(+), 18 deletions(-) diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 40bdec0ddd..7e77721ada 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -32,8 +32,14 @@ pub struct CompactionInput<'a> { pub reasoning: Option, #[serde(skip_serializing_if = "Option::is_none")] pub text: Option, - #[serde(skip_serializing_if = "std::ops::Not::not")] - pub prefix_mode: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub mode: Option, +} + +#[derive(Debug, Clone, Copy, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum CompactionMode { + Prefix, } /// Canonical input payload for the memory summarize endpoint. diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index ac26d3cdba..350473295b 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -19,6 +19,7 @@ pub use crate::api_bridge::CoreAuthProvider; pub use crate::api_bridge::map_api_error; pub use crate::auth::AuthProvider; pub use crate::common::CompactionInput; +pub use crate::common::CompactionMode; pub use crate::common::MemorySummarizeInput; pub use crate::common::MemorySummarizeOutput; pub use crate::common::OpenAiVerbosity; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 3fc2fbc0fb..94a6aae547 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -34,6 +34,7 @@ use std::sync::atomic::Ordering; use codex_api::ApiError; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; +use codex_api::CompactionMode as ApiCompactionMode; use codex_api::Compression; use codex_api::MemoriesClient as ApiMemoriesClient; use codex_api::MemorySummarizeInput as ApiMemorySummarizeInput; @@ -417,7 +418,7 @@ impl ModelClient { effort: Option, summary: ReasoningSummaryConfig, session_telemetry: &SessionTelemetry, - prefix_mode: bool, + mode: Option, ) -> Result> { if prompt.input.is_empty() { return Ok(Vec::new()); @@ -462,7 +463,7 @@ impl ModelClient { parallel_tool_calls: prompt.parallel_tool_calls, reasoning, text, - prefix_mode, + mode, }; let mut extra_headers = ApiHeaderMap::new(); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 90dc029b58..548e22edc1 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6827,12 +6827,13 @@ async fn apply_prefix_compact_candidate( .await; let mut new_history = candidate.replacement_prefix; - new_history.extend( + 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!( initial_context_injection, diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 580713ab4d..4afcaa235a 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -481,6 +481,38 @@ pub(crate) fn insert_initial_context_before_last_real_user_or_summary( compacted_history } +/// 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" => { + crate::event_mapping::is_contextual_user_message_content(content) + && !matches!( + crate::event_mapping::parse_turn_item(item), + Some(TurnItem::HookPrompt(_)) + ) + } + _ => false, + } +} + pub(crate) fn build_compacted_history( initial_context: Vec, user_messages: &[String], diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index d63b3e949f..e57fdc9750 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -122,7 +122,7 @@ pub(crate) async fn run_remote_prefix_compact_task( turn_context.reasoning_effort, turn_context.reasoning_summary, &turn_context.session_telemetry, - /*prefix_mode*/ true, + Some(codex_api::CompactionMode::Prefix), ) .or_else(|err| async { let total_usage_breakdown = sess.get_total_token_usage_breakdown().await; @@ -241,7 +241,7 @@ async fn run_remote_compact_task_inner_impl( turn_context.reasoning_effort, turn_context.reasoning_summary, &turn_context.session_telemetry, - /*prefix_mode*/ false, + None, ) .or_else(|err| async { let total_usage_breakdown = sess.get_total_token_usage_breakdown().await; diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 74154f6d3d..541f67eb7b 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -2,6 +2,9 @@ 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( @@ -23,6 +26,18 @@ async fn process_compacted_history_with_test_session( (refreshed, initial_context) } +fn input_message(role: &str, text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: role.to_string(), + content: vec![ContentItem::InputText { + 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); @@ -147,6 +162,80 @@ 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 contextual_user = input_message( + "user", + r#" + /repo + zsh +"#, + ); + 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, + contextual_user, + real_user.clone(), + assistant.clone(), + function_call.clone(), + function_output.clone(), + hook_prompt.clone(), + ]); + + assert_eq!( + cleaned, + vec![ + 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 diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index dca37d427f..4ab5cb6bb0 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -134,10 +134,7 @@ impl Respond for PrefixRaceCompactResponder { let body = request .body_json::() .expect("compact request body should be JSON"); - let prefix_mode = body - .get("prefix_mode") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); + let prefix_mode = body.get("mode").and_then(serde_json::Value::as_str) == Some("prefix"); self.requests .lock() .expect("request lock poisoned") @@ -1262,7 +1259,7 @@ async fn prefix_compact_uses_configured_threshold_percent() -> Result<()> { .expect("request lock poisoned") .clone(); assert_eq!(compact_requests.len(), 1); - assert_eq!(compact_requests[0]["prefix_mode"], json!(true)); + assert_eq!(compact_requests[0]["mode"], json!("prefix")); Ok(()) } @@ -1320,10 +1317,25 @@ async fn ready_prefix_compact_is_applied_by_pre_turn_auto_compact() -> Result<() wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; wait_for_response_requests(&compact_mock, /*count*/ 1).await; assert_eq!( - compact_mock.single_request().body_json()["prefix_mode"], - json!(true) + compact_mock.single_request().body_json()["mode"], + json!("prefix") ); + codex + .submit(Op::OverrideTurnContext { + cwd: Some(PathBuf::from(PRETURN_CONTEXT_DIFF_CWD)), + approval_policy: None, + approvals_reviewer: None, + sandbox_policy: None, + windows_sandbox_level: None, + model: None, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; codex .submit(Op::UserInput { items: vec![UserInput::Text { @@ -1360,6 +1372,15 @@ async fn ready_prefix_compact_is_applied_by_pre_turn_auto_compact() -> Result<() assert!(!third_request_body.contains("FIRST_REMOTE_USER")); assert!(third_request_body.contains("SECOND_REMOTE_USER")); assert!(third_request_body.contains("THIRD_REMOTE_USER")); + let stale_context_count = requests[2] + .message_input_texts("user") + .iter() + .filter(|text| text.contains(PRETURN_CONTEXT_DIFF_CWD)) + .count(); + assert_eq!( + stale_context_count, 1, + "prefix compaction should strip stale retained suffix context before current context is reinjected" + ); Ok(()) } @@ -1452,8 +1473,8 @@ async fn running_prefix_compact_is_abandoned_when_auto_compact_fires() -> Result .expect("request lock poisoned") .clone(); assert_eq!(compact_requests.len(), 2); - assert_eq!(compact_requests[0]["prefix_mode"], json!(true)); - assert!(compact_requests[1].get("prefix_mode").is_none()); + assert_eq!(compact_requests[0]["mode"], json!("prefix")); + assert!(compact_requests[1].get("mode").is_none()); let requests = responses_mock.requests(); assert_eq!(requests.len(), 3); @@ -1528,8 +1549,8 @@ async fn running_prefix_compact_started_on_follow_up_boundary_is_abandoned_when_ .expect("request lock poisoned") .clone(); assert_eq!(compact_requests.len(), 2); - assert_eq!(compact_requests[0]["prefix_mode"], json!(true)); - assert!(compact_requests[1].get("prefix_mode").is_none()); + assert_eq!(compact_requests[0]["mode"], json!("prefix")); + assert!(compact_requests[1].get("mode").is_none()); let requests = responses_mock.requests(); assert_eq!(requests.len(), 3);