Update prefix compaction mode and suffix cleanup

This commit is contained in:
Rohit Arunachalam
2026-04-13 17:09:57 -07:00
parent 2dc3358a5d
commit a6c594e296
8 changed files with 169 additions and 18 deletions

View File

@@ -32,8 +32,14 @@ pub struct CompactionInput<'a> {
pub reasoning: Option<Reasoning>,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<TextControls>,
#[serde(skip_serializing_if = "std::ops::Not::not")]
pub prefix_mode: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub mode: Option<CompactionMode>,
}
#[derive(Debug, Clone, Copy, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactionMode {
Prefix,
}
/// Canonical input payload for the memory summarize endpoint.

View File

@@ -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;

View File

@@ -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<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
session_telemetry: &SessionTelemetry,
prefix_mode: bool,
mode: Option<ApiCompactionMode>,
) -> Result<Vec<ResponseItem>> {
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();

View File

@@ -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,

View File

@@ -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<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" => {
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<ResponseItem>,
user_messages: &[String],

View File

@@ -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;

View File

@@ -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", "<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 contextual_user = input_message(
"user",
r#"<environment_context>
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment_context>"#,
);
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

View File

@@ -134,10 +134,7 @@ impl Respond for PrefixRaceCompactResponder {
let body = request
.body_json::<serde_json::Value>()
.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);