From 4bd5b9fd0933c4399351bd59364a96a8fb99cd33 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Tue, 4 Aug 2026 19:10:49 +0000 Subject: [PATCH] Keep image resize notices attached during remote compaction (#36956) ## Why Image resize notices describe the history item immediately before them, so remote compaction must not preserve or discard them independently of that item. ## What changed - Group recognized `image_resize_notice` developer messages with their source history items when filtering compacted history. - Account for each group together when trimming or truncating history to a token budget, preserving the notice only when its source is retained. ## Testing - Extend the remote compaction v2 follow-up test to verify that a retained user image keeps its resize notice, while notices for discarded items and unrecognized developer notices are omitted. GitOrigin-RevId: 82bb70a46eb68b07b4601e7fb345921d0789f442 --- codex-rs/core/src/compact_remote.rs | 33 ++++++----- codex-rs/core/src/compact_remote_history.rs | 59 ++++++++++++++++++ codex-rs/core/src/compact_remote_v2.rs | 39 +++++++++--- codex-rs/core/src/lib.rs | 1 + codex-rs/core/tests/suite/compact_remote.rs | 66 ++++++++++++++++++++- 5 files changed, 174 insertions(+), 24 deletions(-) create mode 100644 codex-rs/core/src/compact_remote_history.rs diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 26e783026f..5b515105d5 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -10,6 +10,8 @@ use crate::compact::compaction_status_from_result; use crate::compact::insert_initial_context_before_last_real_user_or_summary; use crate::compact_model_fallback::record_model_fallback; use crate::compact_model_fallback::should_retry_with_current_model; +use crate::compact_remote_history::HistoryItemGroup; +use crate::compact_remote_history::history_item_groups; use crate::context::world_state::WorldState; use crate::context_manager::ContextManager; use crate::context_manager::estimate_item_token_count; @@ -301,7 +303,7 @@ async fn run_remote_compact_task_inner_impl( pub(crate) async fn process_compacted_history( sess: &Session, - mut compacted_history: Vec, + compacted_history: Vec, initial_context_injection: &InitialContextInjection, ) -> (Vec, Option>) { // Mid-turn compaction is the only path that must inject initial context above the last user @@ -310,7 +312,10 @@ pub(crate) async fn process_compacted_history( let (initial_context, world_state_baseline) = build_compaction_initial_context(sess, initial_context_injection).await; - compacted_history.retain(should_keep_compacted_history_item); + let compacted_history = history_item_groups(compacted_history) + .filter(|group| should_keep_compacted_history_item(&group.source)) + .flat_map(HistoryItemGroup::into_items) + .collect(); ( insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context), world_state_baseline, @@ -375,34 +380,34 @@ pub(crate) fn trim_function_call_history_to_fit_context_window( let base_tokens = i128::try_from(approx_token_count(&base_instructions.text)).unwrap_or(i128::MAX); let original_items = history.raw_items(); - let item_token_estimates = original_items - .iter() - .map(estimate_item_token_count) - .collect::>(); - let mut estimated_tokens = item_token_estimates - .iter() - .copied() - .map(i128::from) + let mut estimated_tokens = history_item_groups(original_items) + .map(|group| group.estimated_token_count()) .fold(base_tokens, i128::saturating_add); let initial_estimated_tokens = i64::try_from(estimated_tokens).unwrap_or(i64::MAX); let mut rewritten_items = Vec::new(); + let mut consumed_items = 0; - for (item, item_tokens) in original_items.iter().zip(item_token_estimates).rev() { + for group in history_item_groups(original_items) + .collect::>() + .into_iter() + .rev() + { if i64::try_from(estimated_tokens).unwrap_or(i64::MAX) <= context_window { break; } - let Some(rewritten_item) = rewritten_output_for_context_window(item) else { + let Some(rewritten_item) = rewritten_output_for_context_window(group.source) else { break; }; estimated_tokens = estimated_tokens - .saturating_sub(i128::from(item_tokens)) + .saturating_sub(group.estimated_token_count()) .saturating_add(i128::from(estimate_item_token_count(&rewritten_item))); + consumed_items += 1 + usize::from(group.attached_notice.is_some()); rewritten_items.push(rewritten_item); } let rewritten_outputs = rewritten_items.len(); if rewritten_outputs > 0 { - let retained_len = original_items.len() - rewritten_outputs; + let retained_len = original_items.len() - consumed_items; let mut items = original_items[..retained_len].to_vec(); items.extend(rewritten_items.into_iter().rev()); history.replace(items); diff --git a/codex-rs/core/src/compact_remote_history.rs b/codex-rs/core/src/compact_remote_history.rs new file mode 100644 index 0000000000..15706cd69d --- /dev/null +++ b/codex-rs/core/src/compact_remote_history.rs @@ -0,0 +1,59 @@ +use std::borrow::Borrow; + +use crate::context_manager::estimate_item_token_count; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; + +const ATTACHED_NOTICE_MARKERS: &[(&str, &str)] = + &[("", "")]; + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct HistoryItemGroup { + pub(crate) source: T, + pub(crate) attached_notice: Option, +} + +impl> HistoryItemGroup { + pub(crate) fn into_items(self) -> impl Iterator { + std::iter::once(self.source).chain(self.attached_notice) + } + + pub(crate) fn estimated_token_count(&self) -> i128 { + let source_tokens = i128::from(estimate_item_token_count(self.source.borrow())); + let notice_tokens = self.attached_notice.as_ref().map_or(0, |notice| { + i128::from(estimate_item_token_count(notice.borrow())) + }); + source_tokens.saturating_add(notice_tokens) + } +} + +pub(crate) fn history_item_groups(items: I) -> impl Iterator> +where + I: IntoIterator, + I::Item: Borrow, +{ + let mut items = items.into_iter().peekable(); + std::iter::from_fn(move || { + let source = items.next()?; + let attached_notice = items.next_if(|notice| is_attached_notice(notice.borrow())); + Some(HistoryItemGroup { + source, + attached_notice, + }) + }) +} + +fn is_attached_notice(notice: &ResponseItem) -> bool { + matches!( + notice, + ResponseItem::Message { role, content, .. } + if role == "developer" + && matches!( + content.as_slice(), + [ContentItem::InputText { text }] + if ATTACHED_NOTICE_MARKERS.iter().any(|(opening, closing)| { + text.starts_with(opening) && text.ends_with(closing) + }) + ) + ) +} diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 34f48513ac..995cd2b700 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -13,6 +13,8 @@ use crate::compact_model_fallback::record_model_fallback; use crate::compact_model_fallback::should_retry_with_current_model; use crate::compact_remote::process_compacted_history; use crate::compact_remote::should_keep_compacted_history_item; +use crate::compact_remote_history::HistoryItemGroup; +use crate::compact_remote_history::history_item_groups; use crate::context_manager::estimate_item_token_count; use crate::hook_runtime::PostCompactHookOutcome; use crate::hook_runtime::PreCompactHookOutcome; @@ -443,10 +445,10 @@ fn build_v2_compacted_history( prompt_input: &[ResponseItem], compaction_output: ResponseItem, ) -> (Vec, usize) { - let retained = prompt_input - .iter() - .filter(|item| is_retained_for_remote_compaction_v2(item)) - .filter(|item| should_keep_compacted_history_item(item)) + let retained = history_item_groups(prompt_input) + .filter(|group| is_retained_for_remote_compaction_v2(group.source)) + .filter(|group| should_keep_compacted_history_item(group.source)) + .flat_map(HistoryItemGroup::into_items) .cloned() .collect::>(); let mut retained = @@ -494,18 +496,37 @@ fn truncate_retained_messages_for_remote_compaction( ) -> Vec { let mut remaining = max_tokens; let mut truncated_reversed = Vec::with_capacity(items.len()); - for item in items.into_iter().rev() { + for group in history_item_groups(items) + .collect::>() + .into_iter() + .rev() + { if remaining == 0 { continue; } - let token_count = message_text_token_count(&item).max(1); + let notice_tokens = group + .attached_notice + .as_ref() + .map_or(0, |notice| message_text_token_count(notice).max(1)); + let token_count = message_text_token_count(&group.source) + .max(1) + .saturating_add(notice_tokens); if token_count <= remaining { - truncated_reversed.push(item); + if let Some(notice) = group.attached_notice { + truncated_reversed.push(notice); + } + truncated_reversed.push(group.source); remaining = remaining.saturating_sub(token_count); - } else if let Some(truncated_item) = - truncate_message_text_to_token_budget(item, /*max_tokens*/ remaining) + } else if remaining > notice_tokens + && let Some(truncated_item) = truncate_message_text_to_token_budget( + group.source, + /*max_tokens*/ remaining - notice_tokens, + ) { + if let Some(notice) = group.attached_notice { + truncated_reversed.push(notice); + } truncated_reversed.push(truncated_item); remaining = 0; } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f1ba9c705e..4cecf46aed 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -22,6 +22,7 @@ pub use turn_metadata::detached_memory_responses_metadata; mod codex_thread; mod compact_model_fallback; mod compact_remote; +mod compact_remote_history; mod compact_remote_v2; mod compact_token_budget; mod config_lock; diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 00efbb24c9..8a1c5ad16c 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -28,6 +28,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::InputModality; use codex_protocol::protocol::ConversationStartParams; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ItemStartedEvent; @@ -928,7 +929,53 @@ async fn remote_compact_v2_reuses_compaction_trigger_for_followups() -> Result<( }), ) .await?; - let codex = harness.test().codex.clone(); + let image_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR4nGNgYAAAAAMAASsJTYQAAAAASUVORK5CYII="; + let user_notice = "retained user image"; + let tool_notice = "discarded tool image"; + let unlisted_notice = "discarded developer notice"; + let developer_message = |text| { + json!({ + "type": "message", + "role": "developer", + "content": [{ "type": "input_text", "text": text }] + }) + }; + let initial_history = [ + json!({ + "type": "message", + "role": "user", + "content": [ + { "type": "input_text", "text": "retained image source" }, + { "type": "input_image", "image_url": image_url } + ] + }), + developer_message(user_notice), + developer_message(unlisted_notice), + json!({ + "type": "function_call", + "name": "view_image", + "arguments": "{}", + "call_id": "image-call" + }), + json!({ + "type": "function_call_output", + "call_id": "image-call", + "output": [{ "type": "input_image", "image_url": image_url }] + }), + developer_message(tool_notice), + ] + .into_iter() + .map(|item| serde_json::from_value(item).map(RolloutItem::ResponseItem)) + .collect::>>()?; + let codex = harness + .test() + .thread_manager + .start_thread(StartThreadOptions { + initial_history: InitialHistory::Forked(initial_history), + ..StartThreadOptions::new(harness.test().config.clone()) + }) + .await? + .thread; let responses_mock = responses::mount_sse_sequence( harness.server(), @@ -1109,6 +1156,23 @@ async fn remote_compact_v2_reuses_compaction_trigger_for_followups() -> Result<( follow_up_body.contains("hello remote compact"), "expected v2 follow-up request to preserve retained original user messages" ); + assert!( + follow_up_request.input().windows(2).any(|items| { + items[0]["role"] == "user" + && items[0]["content"][0]["text"] == "retained image source" + && items[1]["role"] == "developer" + && items[1]["content"][0]["text"] == user_notice + }), + "expected v2 compaction to retain the user image and its adjacent resize notice" + ); + assert!( + !follow_up_body.contains(unlisted_notice), + "expected v2 compaction to drop unlisted developer notices" + ); + assert!( + !follow_up_body.contains(tool_notice), + "expected v2 compaction to drop the resize notice with its tool output" + ); Ok(()) }