From fd3c1dc13d0a0941af406e1bc1f697c9d14110ea Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 20 Jul 2026 23:25:53 +0000 Subject: [PATCH] Optimize remote compaction history handling (#34431) ## Why Remote compaction can process large histories. Repeatedly estimating and replacing the full history, and cloning it when tracing is disabled, adds avoidable CPU and memory overhead. ## What changed - Estimate each history item's token count once, update the total as trailing tool outputs are rewritten, and replace history only after all rewrites are selected. - Preserve unclamped token totals while calculating removed tokens so saturated estimates do not hide overflow. - Snapshot compaction input history only when rollout tracing is enabled. - Reuse the v2 request input instead of cloning it before adding the compaction trigger. ## Testing - Cover enabled and disabled compaction trace contexts through `is_enabled()` assertions. GitOrigin-RevId: 8de9c9704ba29532fba05430eb204bbae2e8bf83 --- codex-rs/core/src/compact_remote.rs | 69 +++++++++++-------- codex-rs/core/src/compact_remote_request.rs | 6 +- codex-rs/core/src/compact_remote_v2.rs | 10 +-- .../core/src/compact_remote_v2_attempt.rs | 11 +-- codex-rs/core/src/context_manager/history.rs | 6 +- codex-rs/core/src/context_manager/mod.rs | 1 + codex-rs/rollout-trace/src/compaction.rs | 5 ++ codex-rs/rollout-trace/src/thread_tests.rs | 2 + codex-rs/state/src/migrations_tests.rs | 2 +- 9 files changed, 72 insertions(+), 40 deletions(-) diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 5c056749f0..16718f3719 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -11,6 +11,7 @@ use crate::compact_model_fallback::record_model_fallback; use crate::compact_model_fallback::should_retry_with_current_model; use crate::context::world_state::WorldState; use crate::context_manager::ContextManager; +use crate::context_manager::estimate_item_token_count; use crate::hook_runtime::PostCompactHookOutcome; use crate::hook_runtime::PreCompactHookOutcome; use crate::hook_runtime::run_post_compact_hooks; @@ -35,6 +36,7 @@ use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::TurnStartedEvent; use codex_rollout_trace::CompactionCheckpointTracePayload; +use codex_utils_output_truncation::approx_token_count; #[path = "compact_remote_request.rs"] mod request; @@ -283,10 +285,12 @@ async fn run_remote_compact_task_inner_impl( // Install is the semantic boundary where the compact endpoint's output becomes live // thread history. Keep it distinct from the later inference request so the reducer can // still represent repeated developer/context prefix items exactly as the model saw them. - compaction_trace.record_installed(&CompactionCheckpointTracePayload { - input_history: &trace_input_history, - replacement_history: &new_history, - }); + if let Some(trace_input_history) = trace_input_history.as_deref() { + compaction_trace.record_installed(&CompactionCheckpointTracePayload { + input_history: trace_input_history, + replacement_history: &new_history, + }); + } sess.replace_compacted_history( compaction_turn_context.as_ref(), new_history, @@ -374,37 +378,46 @@ pub(crate) fn trim_function_call_history_to_fit_context_window( let Some(context_window) = turn_context.model_context_window() else { return (0, 0); }; - let mut rewritten_outputs = 0usize; - let mut estimated_deleted_tokens = 0i64; - let item_count = history.raw_items().len(); + // Keep the unclamped total so replacing an item cannot lose an overflow hidden by i64 + // saturation in the normal history estimator. + 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) + .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(); - for index in (0..item_count).rev() { - let Some(estimated_tokens_before) = - history.estimate_token_count_with_base_instructions(base_instructions) - else { - break; - }; - if estimated_tokens_before <= context_window { + for (item, item_tokens) in original_items.iter().zip(item_token_estimates).rev() { + if i64::try_from(estimated_tokens).unwrap_or(i64::MAX) <= context_window { break; } - let Some(rewritten_item) = history - .raw_items() - .get(index) - .and_then(rewritten_output_for_context_window) - else { + let Some(rewritten_item) = rewritten_output_for_context_window(item) else { break; }; - let mut items = history.raw_items().to_vec(); - items[index] = rewritten_item; - history.replace(items); - let estimated_tokens_after = history - .estimate_token_count_with_base_instructions(base_instructions) - .unwrap_or_default(); - rewritten_outputs += 1; - estimated_deleted_tokens = estimated_deleted_tokens - .saturating_add(estimated_tokens_before.saturating_sub(estimated_tokens_after)); + estimated_tokens = estimated_tokens + .saturating_sub(i128::from(item_tokens)) + .saturating_add(i128::from(estimate_item_token_count(&rewritten_item))); + rewritten_items.push(rewritten_item); } + let rewritten_outputs = rewritten_items.len(); + if rewritten_outputs > 0 { + let retained_len = original_items.len() - rewritten_outputs; + let mut items = original_items[..retained_len].to_vec(); + items.extend(rewritten_items.into_iter().rev()); + history.replace(items); + } + + let final_estimated_tokens = i64::try_from(estimated_tokens).unwrap_or(i64::MAX); + let estimated_deleted_tokens = initial_estimated_tokens.saturating_sub(final_estimated_tokens); (rewritten_outputs, estimated_deleted_tokens) } diff --git a/codex-rs/core/src/compact_remote_request.rs b/codex-rs/core/src/compact_remote_request.rs index bbf2f04446..f3c5347fe6 100644 --- a/codex-rs/core/src/compact_remote_request.rs +++ b/codex-rs/core/src/compact_remote_request.rs @@ -19,7 +19,7 @@ use tracing::info; pub(super) struct RemoteCompactAttempt { pub(super) new_history: Vec, - pub(super) trace_input_history: Vec, + pub(super) trace_input_history: Option>, } pub(super) async fn run_remote_compact_attempt( @@ -57,7 +57,9 @@ pub(super) async fn run_remote_compact_attempt( .saturating_sub(estimated_deleted_tokens.min(max_local_deleted_tokens)) }); } - let trace_input_history = history.raw_items().to_vec(); + let trace_input_history = compaction_trace + .is_enabled() + .then(|| history.raw_items().to_vec()); let prompt_input = history.for_prompt(&turn_context.model_info.input_modalities); let tool_router = built_tools( sess.as_ref(), diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index a0bb9e11ae..9a39f481df 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -305,10 +305,12 @@ async fn run_remote_compact_task_inner_impl( previous_window_id: new_window_ids.previous_window_id.map(|id| id.to_string()), window_id: Some(new_window_ids.window_id.to_string()), }; - compaction_trace.record_installed(&CompactionCheckpointTracePayload { - input_history: &trace_input_history, - replacement_history: &new_history, - }); + if let Some(trace_input_history) = trace_input_history.as_deref() { + compaction_trace.record_installed(&CompactionCheckpointTracePayload { + input_history: trace_input_history, + replacement_history: &new_history, + }); + } sess.replace_compacted_history( compaction_turn_context.as_ref(), new_history, diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 5ebe128e3b..224d96b1d8 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -21,7 +21,7 @@ use tokio_util::sync::CancellationToken; use tracing::info; pub(super) struct RemoteCompactV2Attempt { - pub(super) trace_input_history: Vec, + pub(super) trace_input_history: Option>, pub(super) prompt_input: Vec, pub(super) compaction_output: ResponseItem, pub(super) token_usage: Option, @@ -65,15 +65,16 @@ pub(super) async fn run_remote_compact_v2_attempt( }); } - let trace_input_history = history.raw_items().to_vec(); - let prompt_input = history.for_prompt(&turn_context.model_info.input_modalities); + let trace_input_history = compaction_trace + .is_enabled() + .then(|| history.raw_items().to_vec()); + let mut input = history.for_prompt(&turn_context.model_info.input_modalities); let tool_router = built_tools( sess.as_ref(), step_context.as_ref(), &CancellationToken::new(), ) .await?; - let mut input = prompt_input.clone(); input.push(ResponseItem::CompactionTrigger {}); let prompt = Prompt { input, @@ -129,6 +130,8 @@ pub(super) async fn run_remote_compact_v2_attempt( }), ) .await; + let mut prompt_input = prompt.input; + prompt_input.pop(); Ok(RemoteCompactV2Attempt { trace_input_history, prompt_input, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 1cd6587249..d4d2b67185 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -490,7 +490,11 @@ fn estimate_encrypted_function_output_length(encoded_len: usize) -> usize { encoded_len.saturating_mul(9).div_ceil(16) } -fn estimate_item_token_count(item: &ResponseItem) -> i64 { +/// Returns the same coarse, model-visible token estimate used for full history estimates. +/// +/// Ordinary items are JSON-serialized, so callers estimating many items should reuse these +/// results instead of repeatedly estimating the full history. +pub(crate) fn estimate_item_token_count(item: &ResponseItem) -> i64 { let model_visible_bytes = estimate_response_item_model_visible_bytes(item); approx_tokens_from_byte_count_i64(model_visible_bytes) } diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index f811c8b373..f2bdd89c41 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -3,5 +3,6 @@ mod normalize; pub(crate) mod updates; pub(crate) use history::ContextManager; +pub(crate) use history::estimate_item_token_count; pub(crate) use history::is_user_turn_boundary; pub(crate) use history::truncate_function_output_payload; diff --git a/codex-rs/rollout-trace/src/compaction.rs b/codex-rs/rollout-trace/src/compaction.rs index 5dffad2462..c8f18795ab 100644 --- a/codex-rs/rollout-trace/src/compaction.rs +++ b/codex-rs/rollout-trace/src/compaction.rs @@ -115,6 +115,11 @@ impl CompactionTraceContext { } } + /// Returns whether this context records compaction traces. + pub fn is_enabled(&self) -> bool { + matches!(self.state, CompactionTraceContextState::Enabled(_)) + } + /// Starts a new upstream attempt and records the exact compact endpoint request. pub fn start_attempt(&self, request: &impl Serialize) -> CompactionTraceAttempt { let CompactionTraceContextState::Enabled(context) = &self.state else { diff --git a/codex-rs/rollout-trace/src/thread_tests.rs b/codex-rs/rollout-trace/src/thread_tests.rs index 6a41e0beeb..e0a5547c33 100644 --- a/codex-rs/rollout-trace/src/thread_tests.rs +++ b/codex-rs/rollout-trace/src/thread_tests.rs @@ -142,6 +142,7 @@ fn disabled_thread_context_accepts_trace_calls_without_writing() -> anyhow::Resu "gpt-test", "test-provider", ); + assert!(!compaction_trace.is_enabled()); let compaction_attempt = compaction_trace.start_attempt(&serde_json::json!({ "kind": "compaction" })); compaction_attempt.record_completed(&[]); @@ -175,6 +176,7 @@ fn compaction_contexts_share_identity_across_models() -> anyhow::Result<()> { for model in ["gpt-previous", "gpt-selected"] { let compaction_trace = thread_trace.compaction_trace_context("turn-1", "compaction-1", model, "test-provider"); + assert!(compaction_trace.is_enabled()); compaction_trace .start_attempt(&serde_json::json!({ "model": model })) .record_failed("test failure"); diff --git a/codex-rs/state/src/migrations_tests.rs b/codex-rs/state/src/migrations_tests.rs index 05c30e2d1a..13da97f134 100644 --- a/codex-rs/state/src/migrations_tests.rs +++ b/codex-rs/state/src/migrations_tests.rs @@ -36,7 +36,7 @@ async fn agent_job_tables_are_dropped_when_upgrading() { let _cleanup = scopeguard::guard(sqlite_home.clone(), |sqlite_home| { let _ = std::fs::remove_dir_all(sqlite_home); }); - let sqlite = crate::SqliteConfig::new_for_testing(sqlite_home.clone()); + let sqlite = crate::SqliteConfig::new_for_testing(sqlite_home.as_path().abs()); let pool = sqlite .open_read_write_pool(&state_db_path(&sqlite_home)) .await