mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: retain remote compaction truncation parity in v2 (#23728)
## Why Remote compaction now has two implementations: the existing server-rebuilt v1 path and the newer client-rebuilt v2 path behind `remote_compaction_v2`. The v1 path bounds retained user/developer/system history before installing the compaction item, while v2 was previously carrying the full retained history forward. That made the two paths diverge for large pre-compaction transcripts even though they are meant to preserve the same compaction contract. This aligns v2 with the retained-history budget expected from v1 so switching the feature flag does not materially change which pre-compaction messages survive into the rebuilt history. ## What changed - Apply a retained-message character budget while rebuilding v2 compacted history in `core/src/compact_remote_v2.rs`. - Keep newest retained messages first, truncate the boundary message with the shared `truncate_text(...)` helper, and drop older retained messages once the budget is exhausted. - Preserve non-text retained message content such as images while truncating text content. - Use the current `64_000` token retained-message default translated to the existing `4x` character budget. ## Testing - `cargo test -p codex-core compact_remote_v2::tests::` - Added focused coverage for newest-first retention and truncating multipart retained messages without dropping images.
This commit is contained in:
committed by
GitHub
Unverified
parent
a6bedc8a7c
commit
94442b7f95
@@ -286,7 +286,7 @@ pub(crate) async fn process_compacted_history(
|
||||
/// - `user`-role warnings that parse as `TurnItem::UserMessage` and compaction-generated summary
|
||||
/// messages. Legacy warning fragments are filtered by `parse_turn_item` before they reach this
|
||||
/// check.
|
||||
fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {
|
||||
pub(crate) fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {
|
||||
match item {
|
||||
ResponseItem::Message { role, .. } if role == "developer" => false,
|
||||
ResponseItem::Message { role, .. } if role == "user" => {
|
||||
|
||||
@@ -10,6 +10,7 @@ use crate::compact::compaction_status_from_result;
|
||||
use crate::compact_remote::build_compact_request_log_data;
|
||||
use crate::compact_remote::log_remote_compact_failure;
|
||||
use crate::compact_remote::process_compacted_history;
|
||||
use crate::compact_remote::should_keep_compacted_history_item;
|
||||
use crate::compact_remote::trim_function_call_history_to_fit_context_window;
|
||||
use crate::hook_runtime::PostCompactHookOutcome;
|
||||
use crate::hook_runtime::PreCompactHookOutcome;
|
||||
@@ -27,17 +28,25 @@ use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::error::Result as CodexResult;
|
||||
use codex_protocol::items::ContextCompactionItem;
|
||||
use codex_protocol::items::TurnItem;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::CompactedItem;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::TruncationPolicy;
|
||||
use codex_protocol::protocol::TurnStartedEvent;
|
||||
use codex_rollout_trace::CompactionCheckpointTracePayload;
|
||||
use codex_rollout_trace::InferenceTraceContext;
|
||||
use codex_utils_output_truncation::approx_token_count;
|
||||
use codex_utils_output_truncation::truncate_text;
|
||||
use futures::StreamExt;
|
||||
use futures::TryFutureExt;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::info;
|
||||
|
||||
// Mirror the current /responses/compact retained-message default while the
|
||||
// server-side path remains the reference implementation.
|
||||
const RETAINED_MESSAGE_TOKEN_BUDGET: usize = 64_000;
|
||||
|
||||
pub(crate) async fn run_inline_remote_auto_compact_task(
|
||||
sess: Arc<Session>,
|
||||
turn_context: Arc<TurnContext>,
|
||||
@@ -343,11 +352,14 @@ fn build_v2_compacted_history(
|
||||
prompt_input: &[ResponseItem],
|
||||
compaction_output: ResponseItem,
|
||||
) -> Vec<ResponseItem> {
|
||||
let mut retained = prompt_input
|
||||
let retained = prompt_input
|
||||
.iter()
|
||||
.filter(|item| is_retained_for_remote_compaction_v2(item))
|
||||
.filter(|item| should_keep_compacted_history_item(item))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let mut retained =
|
||||
truncate_retained_messages_for_remote_compaction(retained, RETAINED_MESSAGE_TOKEN_BUDGET);
|
||||
retained.push(compaction_output);
|
||||
retained
|
||||
}
|
||||
@@ -360,6 +372,98 @@ fn is_retained_for_remote_compaction_v2(item: &ResponseItem) -> bool {
|
||||
matches!(role.as_str(), "user" | "developer" | "system")
|
||||
}
|
||||
|
||||
fn truncate_retained_messages_for_remote_compaction(
|
||||
items: Vec<ResponseItem>,
|
||||
max_tokens: usize,
|
||||
) -> Vec<ResponseItem> {
|
||||
let mut remaining = max_tokens;
|
||||
let mut truncated_reversed = Vec::with_capacity(items.len());
|
||||
for item in items.into_iter().rev() {
|
||||
if remaining == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let token_count = message_text_token_count(&item).max(1);
|
||||
if token_count <= remaining {
|
||||
truncated_reversed.push(item);
|
||||
remaining = remaining.saturating_sub(token_count);
|
||||
} else if let Some(truncated_item) =
|
||||
truncate_message_text_to_token_budget(item, /*max_tokens*/ remaining)
|
||||
{
|
||||
truncated_reversed.push(truncated_item);
|
||||
remaining = 0;
|
||||
}
|
||||
}
|
||||
truncated_reversed.reverse();
|
||||
truncated_reversed
|
||||
}
|
||||
|
||||
fn message_text_token_count(item: &ResponseItem) -> usize {
|
||||
let ResponseItem::Message { content, .. } = item else {
|
||||
return 0;
|
||||
};
|
||||
|
||||
content
|
||||
.iter()
|
||||
.map(|item| match item {
|
||||
ContentItem::InputText { text } | ContentItem::OutputText { text } => {
|
||||
approx_token_count(text)
|
||||
}
|
||||
ContentItem::InputImage { .. } => 0,
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
fn truncate_message_text_to_token_budget(
|
||||
item: ResponseItem,
|
||||
max_tokens: usize,
|
||||
) -> Option<ResponseItem> {
|
||||
let ResponseItem::Message {
|
||||
id,
|
||||
role,
|
||||
content,
|
||||
phase,
|
||||
} = item
|
||||
else {
|
||||
return Some(item);
|
||||
};
|
||||
|
||||
let mut remaining = max_tokens;
|
||||
let mut truncated_content = Vec::with_capacity(content.len());
|
||||
for mut content_item in content {
|
||||
match &mut content_item {
|
||||
ContentItem::InputText { text } | ContentItem::OutputText { text } => {
|
||||
if remaining == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let token_count = approx_token_count(text);
|
||||
if token_count <= remaining {
|
||||
remaining = remaining.saturating_sub(token_count);
|
||||
} else {
|
||||
*text = truncate_text(text, TruncationPolicy::Tokens(remaining));
|
||||
remaining = 0;
|
||||
}
|
||||
if !text.is_empty() {
|
||||
truncated_content.push(content_item);
|
||||
}
|
||||
}
|
||||
ContentItem::InputImage { .. } => truncated_content.push(content_item),
|
||||
}
|
||||
}
|
||||
|
||||
if truncated_content.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(ResponseItem::Message {
|
||||
id,
|
||||
role,
|
||||
content: truncated_content,
|
||||
phase,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -395,7 +499,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_v2_compacted_history_matches_prod_retention_shape() {
|
||||
fn build_v2_compacted_history_filters_to_installed_retention_shape() {
|
||||
let input = vec![
|
||||
message("developer", "dev", /*phase*/ None),
|
||||
message("system", "sys", /*phase*/ None),
|
||||
@@ -421,15 +525,145 @@ mod tests {
|
||||
|
||||
assert_eq!(
|
||||
history,
|
||||
vec![message("user", "user", /*phase*/ None), output]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_v2_compacted_history_discards_messages_before_truncating() {
|
||||
let old = message("user", "old", /*phase*/ None);
|
||||
let new = message("user", "new", /*phase*/ None);
|
||||
let huge_developer_message = "d".repeat((RETAINED_MESSAGE_TOKEN_BUDGET + 1) * 4);
|
||||
let huge_contextual_message = format!(
|
||||
"<environment_context>\n{}\n</environment_context>",
|
||||
"c".repeat((RETAINED_MESSAGE_TOKEN_BUDGET + 1) * 4)
|
||||
);
|
||||
let input = vec![
|
||||
old.clone(),
|
||||
message("developer", &huge_developer_message, /*phase*/ None),
|
||||
message("user", &huge_contextual_message, /*phase*/ None),
|
||||
new.clone(),
|
||||
];
|
||||
let output = ResponseItem::Compaction {
|
||||
encrypted_content: "new".to_string(),
|
||||
};
|
||||
|
||||
let history = build_v2_compacted_history(&input, output.clone());
|
||||
|
||||
assert_eq!(history, vec![old, new, output]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_history_truncation_keeps_newest_messages_first() {
|
||||
let middle = message("user", "middle1234", /*phase*/ None);
|
||||
let new = message("user", "new", /*phase*/ None);
|
||||
let retained = vec![
|
||||
message("user", "old-old", /*phase*/ None),
|
||||
middle,
|
||||
new.clone(),
|
||||
];
|
||||
|
||||
let truncated =
|
||||
truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 3);
|
||||
|
||||
assert_eq!(
|
||||
truncated,
|
||||
vec![
|
||||
message("developer", "dev", /*phase*/ None),
|
||||
message("system", "sys", /*phase*/ None),
|
||||
message("user", "user", /*phase*/ None),
|
||||
output,
|
||||
message("user", "midd…1 tokens truncated…1234", /*phase*/ None),
|
||||
new,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_history_truncation_preserves_images_and_truncates_later_text_parts() {
|
||||
let item = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![
|
||||
ContentItem::InputText {
|
||||
text: "abcdef".to_string(),
|
||||
},
|
||||
ContentItem::InputImage {
|
||||
image_url: "data:image/png;base64,abc".to_string(),
|
||||
detail: None,
|
||||
},
|
||||
ContentItem::OutputText {
|
||||
text: "uvwxyz".to_string(),
|
||||
},
|
||||
],
|
||||
phase: None,
|
||||
};
|
||||
|
||||
let truncated =
|
||||
truncate_retained_messages_for_remote_compaction(vec![item], /*max_tokens*/ 3);
|
||||
|
||||
assert_eq!(
|
||||
truncated,
|
||||
vec![ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![
|
||||
ContentItem::InputText {
|
||||
text: "abcdef".to_string(),
|
||||
},
|
||||
ContentItem::InputImage {
|
||||
image_url: "data:image/png;base64,abc".to_string(),
|
||||
detail: None,
|
||||
},
|
||||
ContentItem::OutputText {
|
||||
text: "uv…1 tokens truncated…yz".to_string(),
|
||||
},
|
||||
],
|
||||
phase: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_history_truncation_charges_image_only_messages() {
|
||||
let image_only_message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputImage {
|
||||
image_url: "data:image/png;base64,abc".to_string(),
|
||||
detail: None,
|
||||
}],
|
||||
phase: None,
|
||||
};
|
||||
let newest = message("user", "new", /*phase*/ None);
|
||||
let retained = vec![
|
||||
message("user", "old", /*phase*/ None),
|
||||
image_only_message.clone(),
|
||||
newest.clone(),
|
||||
];
|
||||
|
||||
let truncated =
|
||||
truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 2);
|
||||
|
||||
assert_eq!(truncated, vec![image_only_message, newest]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_history_truncation_drops_image_only_messages_after_budget_is_spent() {
|
||||
let image_only_message = ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputImage {
|
||||
image_url: "data:image/png;base64,abc".to_string(),
|
||||
detail: None,
|
||||
}],
|
||||
phase: None,
|
||||
};
|
||||
let newest = message("user", "new", /*phase*/ None);
|
||||
let retained = vec![image_only_message, newest.clone()];
|
||||
|
||||
let truncated =
|
||||
truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 1);
|
||||
|
||||
assert_eq!(truncated, vec![newest]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn collect_compaction_output_accepts_additional_output_items() {
|
||||
let compaction = ResponseItem::Compaction {
|
||||
|
||||
Reference in New Issue
Block a user