From bff03ecce582d55ff62d8ee37e576b63c24bc0a9 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 13 Aug 2026 22:02:50 +0000 Subject: [PATCH] Retain client developer messages across context compaction (#38445) ## Why Client-authored developer instructions should remain available after a context window is compacted when `retain_client_developer_messages` is enabled. ## What changed - Preserve annotated client-authored developer messages during remote compaction v2 and local token-budget context resets. - Apply the existing retained-message token budget when carrying those messages into the new context window. ## Testing - Cover enabled and disabled retention for remote compaction v2, manual token-budget compaction, and mid-turn automatic compaction. GitOrigin-RevId: 4cb3782cb86ef46458413a17a3255d0016fd675f --- codex-rs/core/src/compact_remote_v2.rs | 175 +++++++++++++++----- codex-rs/core/src/session/mod.rs | 18 ++ codex-rs/core/tests/suite/compact_remote.rs | 97 +++++++++++ codex-rs/core/tests/suite/token_budget.rs | 53 +++++- 4 files changed, 304 insertions(+), 39 deletions(-) diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 070986af63..89d5c5a1b1 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -8,10 +8,11 @@ use crate::compact::CompactedHistoryMetadata; use crate::compact::CompactionAnalyticsAttempt; use crate::compact::CompactionAnalyticsDetails; use crate::compact::InitialContextInjection; +use crate::compact::build_compaction_initial_context; 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::process_annotated_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; @@ -32,6 +33,7 @@ use codex_analytics::CompactionImplementation; use codex_analytics::CompactionPhase; use codex_analytics::CompactionReason; use codex_analytics::CompactionTrigger; +use codex_features::Feature; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; use codex_protocol::error::CodexErr; @@ -60,7 +62,7 @@ use attempt::run_remote_compact_v2_attempt; // 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) const RETAINED_MESSAGE_TOKEN_BUDGET: usize = 64_000; const MAX_RETAINED_AGENT_MESSAGE_TOKENS: i64 = 10_000; // Compact attempts can run much longer than normal turns, so keep the per-transport // retry budget smaller than the general Responses stream retry budget. @@ -292,16 +294,18 @@ async fn run_remote_compact_task_inner_impl( analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens); analytics_details.cache_write_input_tokens = Some(token_usage.cache_write_input_tokens); } - let (compacted_history, retained_images) = - build_v2_compacted_history(prompt_input, prompt_input_metadata, compaction_output); + let (compacted_history, retained_images) = build_v2_compacted_history( + prompt_input, + prompt_input_metadata, + compaction_output, + sess.enabled(Feature::RetainClientDeveloperMessages), + ); analytics_details.retained_image_count = Some(retained_images); let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; - let (new_history, world_state_baseline) = process_annotated_compacted_history( - sess.as_ref(), - compacted_history, - &initial_context_injection, - ) - .await; + let (initial_context, world_state_baseline) = + build_compaction_initial_context(sess.as_ref(), &initial_context_injection).await; + let new_history = + insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context); let reference_context_item = match initial_context_injection { InitialContextInjection::DoNotInject => None, @@ -456,6 +460,7 @@ fn build_v2_compacted_history( prompt_input: Vec, prompt_input_metadata: Vec>, compaction_output: ResponseItem, + retain_client_developer_messages: bool, ) -> (Vec, usize) { debug_assert_eq!(prompt_input.len(), prompt_input_metadata.len()); let prompt_input = prompt_input @@ -463,9 +468,13 @@ fn build_v2_compacted_history( .zip(prompt_input_metadata) .map(|(item, metadata)| ResponseItemEnvelope { item, metadata }) .collect::>(); - let retained = history_item_groups(prompt_input) + let retained = v2_history_item_groups(prompt_input) .filter(|group| is_retained_for_remote_compaction_v2(&group.source.item)) - .filter(|group| should_keep_compacted_history_item(&group.source.item)) + .filter(|group| { + should_keep_compacted_history_item(&group.source.item) + || (retain_client_developer_messages + && is_client_authored_developer_message(&group.source)) + }) .flat_map(HistoryItemGroup::into_items) .collect::>(); let mut retained = @@ -478,6 +487,28 @@ fn build_v2_compacted_history( (retained, retained_image_count) } +pub(crate) fn is_client_authored_developer_message(item: &ResponseItemEnvelope) -> bool { + item.metadata + .as_ref() + .is_some_and(|metadata| metadata.client_authored) + && matches!(&item.item, ResponseItem::Message { role, .. } if role == "developer") +} + +fn v2_history_item_groups( + items: Vec, +) -> impl Iterator> { + history_item_groups(items).flat_map(|mut group| { + let client_message = group + .attached_notice + .take_if(|item| is_client_authored_developer_message(item)) + .map(|source| HistoryItemGroup { + source, + attached_notice: None, + }); + std::iter::once(group).chain(client_message) + }) +} + fn is_retained_for_remote_compaction_v2(item: &ResponseItem) -> bool { if let ResponseItem::AgentMessage { content, .. } = item { let is_completion = matches!( @@ -507,13 +538,13 @@ fn retained_input_image_count(item: &ResponseItem) -> usize { .count() } -fn truncate_retained_messages_for_remote_compaction( +pub(crate) fn truncate_retained_messages_for_remote_compaction( items: Vec, max_tokens: usize, ) -> Vec { let mut remaining = max_tokens; let mut truncated_reversed = Vec::with_capacity(items.len()); - for group in history_item_groups(items) + for group in v2_history_item_groups(items) .collect::>() .into_iter() .rev() @@ -522,25 +553,58 @@ fn truncate_retained_messages_for_remote_compaction( continue; } + let client_developer = is_client_authored_developer_message(&group.source); let notice_tokens = group .attached_notice .as_ref() .map_or(0, |notice| message_text_token_count(¬ice.item).max(1)); - let token_count = message_text_token_count(&group.source.item) - .max(1) - .saturating_add(notice_tokens); + let source_tokens = if client_developer { + usize::try_from(estimate_item_token_count(&group.source.item)).unwrap_or(usize::MAX) + } else { + message_text_token_count(&group.source.item).max(1) + }; + let token_count = source_tokens.saturating_add(notice_tokens); if token_count <= remaining { if let Some(notice) = group.attached_notice { truncated_reversed.push(notice); } truncated_reversed.push(group.source); remaining = remaining.saturating_sub(token_count); - } else if remaining > notice_tokens - && let Some(truncated_item) = truncate_message_text_to_token_budget( - group.source, - /*max_tokens*/ remaining - notice_tokens, - ) - { + } else if remaining > notice_tokens { + let available_tokens = remaining - notice_tokens; + let text_budget = if client_developer { + available_tokens.saturating_sub( + source_tokens.saturating_sub(message_text_token_count(&group.source.item)), + ) + } else { + available_tokens + }; + let Some(mut truncated_item) = + truncate_message_text_to_token_budget(group.source, text_budget) + else { + continue; + }; + if client_developer { + let item_tokens = usize::try_from(estimate_item_token_count(&truncated_item.item)) + .unwrap_or(usize::MAX); + if item_tokens > available_tokens { + let adjusted_budget = text_budget + .saturating_sub(item_tokens - available_tokens) + .saturating_sub(1); + let Some(adjusted) = + truncate_message_text_to_token_budget(truncated_item, adjusted_budget) + else { + continue; + }; + if usize::try_from(estimate_item_token_count(&adjusted.item)) + .unwrap_or(usize::MAX) + > available_tokens + { + continue; + } + truncated_item = adjusted; + } + } if let Some(notice) = group.attached_notice { truncated_reversed.push(notice); } @@ -655,7 +719,9 @@ mod tests { output: ResponseItem, ) -> (Vec, usize) { let metadata = vec![None; input.len()]; - build_v2_compacted_history(input, metadata, output) + build_v2_compacted_history( + input, metadata, output, /*retain_client_developer_messages*/ false, + ) } fn annotated(items: Vec) -> Vec { @@ -730,28 +796,63 @@ mod tests { #[test] fn build_v2_compacted_history_preserves_retained_metadata_sidecar() { let retained = message("user", "keep", /*phase*/ None); - let filtered = message("developer", "drop", /*phase*/ None); + let generated_notice = message( + "developer", + "generated", + /*phase*/ None, + ); + let harness = message("developer", "drop", /*phase*/ None); + let client = message( + "developer", + "client", + /*phase*/ None, + ); let output = ResponseItem::Compaction { id: None, encrypted_content: "new".to_string(), internal_chat_message_metadata_passthrough: None, }; - let (history, _) = build_v2_compacted_history( - vec![filtered, retained.clone()], - vec![None, Some(CodexHarnessMetadata::default())], - output.clone(), - ); - assert_eq!( - history, - vec![ + for enabled in [false, true] { + let (history, _) = build_v2_compacted_history( + vec![ + harness.clone(), + client.clone(), + retained.clone(), + generated_notice.clone(), + ], + vec![ + None, + Some(CodexHarnessMetadata { + client_authored: true, + }), + Some(CodexHarnessMetadata::default()), + None, + ], + output.clone(), + enabled, + ); + let mut expected = vec![ ResponseItemEnvelope { - item: retained, + item: retained.clone(), metadata: Some(CodexHarnessMetadata::default()), }, - ResponseItemEnvelope::new(output), - ] - ); + ResponseItemEnvelope::new(generated_notice.clone()), + ResponseItemEnvelope::new(output.clone()), + ]; + if enabled { + expected.insert( + 0, + ResponseItemEnvelope { + item: client.clone(), + metadata: Some(CodexHarnessMetadata { + client_authored: true, + }), + }, + ); + } + assert_eq!(history, expected); + } } #[test] diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 10f3411cd4..2dc42fef6d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3695,6 +3695,23 @@ impl Session { world_state: Arc, ) -> u64 { let turn_context = step_context.turn.as_ref(); + let retained_client_developer_messages = + if self.enabled(Feature::RetainClientDeveloperMessages) { + let history = self.clone_history().await; + crate::compact_remote_v2::truncate_retained_messages_for_remote_compaction( + history + .annotated_items() + .iter() + .filter(|item| { + crate::compact_remote_v2::is_client_authored_developer_message(item) + }) + .cloned() + .collect(), + crate::compact_remote_v2::RETAINED_MESSAGE_TOKEN_BUDGET, + ) + } else { + Vec::new() + }; let window = { let mut state = self.state.lock().await; state.start_new_context_window() @@ -3705,6 +3722,7 @@ impl Session { .await .into_iter() .map(ResponseItemEnvelope::new) + .chain(retained_client_developer_messages) .collect(); let turn_context_item = turn_context.to_turn_context_item(); self.replace_compacted_history( diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 629dff8768..451ec20959 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -69,6 +69,7 @@ use core_test_support::wait_for_event_with_timeout; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use test_case::test_case; use tokio::time::Duration; use wiremock::ResponseTemplate; @@ -571,6 +572,102 @@ async fn remote_compact_v2_retains_metadata_from_resumed_history() -> Result<()> Ok(()) } +#[test_case(false; "feature_disabled")] +#[test_case(true; "feature_enabled")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_compact_v2_retains_only_client_developer_messages_when_enabled( + enabled: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_auto_env_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config + .features + .enable(Feature::RemoteCompactionV2) + .expect("remote compaction v2 should be configurable"); + if enabled { + config + .features + .enable(Feature::RetainClientDeveloperMessages) + .expect("client developer retention should be configurable"); + } + }), + ) + .await?; + let developer = |text: &str| ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + let codex = &harness.test().codex; + let rollout_path = codex.rollout_path().context("rollout path")?; + codex + .inject_response_items(vec![developer("INJECTED_CLIENT_DEVELOPER")]) + .await?; + let response_mock = responses::mount_sse_sequence( + harness.server(), + vec![ + sse(vec![responses::ev_completed("before-compact")]), + sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "CLIENT_RETENTION_SUMMARY", + }, + }), + responses::ev_completed("compact"), + ]), + sse(vec![responses::ev_completed("after-compact")]), + ], + ) + .await; + + harness.test().submit_turn("before compact").await?; + codex.submit(Op::Compact).await?; + wait_for_turn_complete(codex).await; + harness.test().submit_turn("after compact").await?; + + let requests = response_mock.requests(); + let follow_up = requests.last().context("follow-up request")?; + assert_eq!( + follow_up.body_contains_text("INJECTED_CLIENT_DEVELOPER"), + enabled + ); + requests + .iter() + .for_each(assert_compact_request_omits_harness_metadata); + + codex.shutdown_and_wait().await?; + let replacement_history = fs::read_to_string(&rollout_path)? + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .filter_map(|line| match line.item { + RolloutItem::Compacted(compacted) => compacted.replacement_history, + _ => None, + }) + .next_back() + .context("remote compaction should persist a checkpoint")?; + let retained_client_developers = replacement_history + .iter() + .filter(|item| { + item.metadata + .as_ref() + .is_some_and(|metadata| metadata.client_authored) + }) + .count(); + assert_eq!(retained_client_developers, usize::from(enabled)); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn amazon_bedrock_automatic_compaction_uses_v1_endpoint_when_v2_is_enabled() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 93193e7e0d..ff7572c2eb 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -56,6 +56,7 @@ use std::collections::HashMap; use std::path::Path; use std::sync::Arc; use std::time::Duration; +use test_case::test_case; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; const AUTO_COMPACT_FALLBACK_PROMPT: &str = "Save the important state before rollover."; @@ -951,8 +952,12 @@ async fn get_context_remaining_returns_unknown_when_threshold_is_unbounded() -> Ok(()) } +#[test_case(false; "token_budget_only")] +#[test_case(true; "with_client_developer_retention")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { +async fn token_budget_context_uses_new_window_after_compaction( + retain_client_developer_messages: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -982,10 +987,25 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { .features .enable(Feature::TokenBudget) .expect("test config should allow token budget"); + if retain_client_developer_messages { + config + .features + .enable(Feature::RetainClientDeveloperMessages) + .expect("test config should allow client developer retention"); + } }) .build(&server) .await?; + if retain_client_developer_messages { + test.codex + .inject_response_items(vec![serde_json::from_value(json!({ + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "CLIENT_DEVELOPER_INSTRUCTIONS"}] + }))?]) + .await?; + } test.submit_turn("before compact").await?; test.codex.submit(Op::Compact).await?; assert_context_compaction_item_lifecycle(&test.codex).await; @@ -1025,6 +1045,11 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { !requests[1].body_contains_text("assistant before compact"), "token budget compaction should drop prior assistant messages" ); + assert_eq!( + requests[1].body_contains_text("CLIENT_DEVELOPER_INSTRUCTIONS"), + retain_client_developer_messages, + "token budget compaction should retain client-authored developer messages when enabled" + ); assert!( requests[1].body_contains_text("after compact"), "follow-up should still include the new turn input" @@ -1082,8 +1107,12 @@ async fn token_budget_compaction_runs_compact_hooks() -> Result<()> { Ok(()) } +#[test_case(false; "token_budget_only")] +#[test_case(true; "with_client_developer_retention")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() -> Result<()> { +async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up( + retain_client_developer_messages: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -1129,10 +1158,25 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() .features .enable(Feature::TokenBudget) .expect("test config should allow token budget"); + if retain_client_developer_messages { + config + .features + .enable(Feature::RetainClientDeveloperMessages) + .expect("test config should allow client developer retention"); + } }) .build(&server) .await?; + if retain_client_developer_messages { + test.codex + .inject_response_items(vec![serde_json::from_value(json!({ + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "MID_TURN_CLIENT_INSTRUCTIONS"}] + }))?]) + .await?; + } test.submit_turn("trigger mid-turn auto compaction").await?; let requests = responses.requests(); @@ -1179,6 +1223,11 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() !requests[1].body_contains_text("trigger mid-turn auto compaction"), "fresh token-budget windows should drop prior user messages" ); + assert_eq!( + requests[1].body_contains_text("MID_TURN_CLIENT_INSTRUCTIONS"), + retain_client_developer_messages, + "mid-turn token-budget compaction should retain client-authored developer messages when enabled" + ); assert_ne!( follow_up_window_id, initial_window_id, "mid-turn token-budget auto-compaction should reset the context window"