openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-13 11:47:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Optimize remote compaction history handling (#34431)

Browse Source 
## 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
This commit is contained in:
Charlie Marsh
2026-07-20 23:25:53 +00:00
committed by copyberry
parent 56c11cf658
commit fd3c1dc13d
9 changed files with 72 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
codex-rs/core/src/compact_remote.rs
View File

@@ -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::<Vec<_>>();
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)
}

6
codex-rs/core/src/compact_remote_request.rs
View File

@@ -19,7 +19,7 @@ use tracing::info;
pub(super) struct RemoteCompactAttempt {
pub(super) new_history: Vec<ResponseItem>,
pub(super) trace_input_history: Vec<ResponseItem>,
pub(super) trace_input_history: Option<Vec<ResponseItem>>,
}
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(),

10
codex-rs/core/src/compact_remote_v2.rs
View File

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

11
codex-rs/core/src/compact_remote_v2_attempt.rs
View File

@@ -21,7 +21,7 @@ use tokio_util::sync::CancellationToken;
use tracing::info;
pub(super) struct RemoteCompactV2Attempt {
pub(super) trace_input_history: Vec<ResponseItem>,
pub(super) trace_input_history: Option<Vec<ResponseItem>>,
pub(super) prompt_input: Vec<ResponseItem>,
pub(super) compaction_output: ResponseItem,
pub(super) token_usage: Option<TokenUsage>,
@@ -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,

6
codex-rs/core/src/context_manager/history.rs
View File

@@ -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)
}

1
codex-rs/core/src/context_manager/mod.rs
View File

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

5
codex-rs/rollout-trace/src/compaction.rs
View File

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

2
codex-rs/rollout-trace/src/thread_tests.rs
View File

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

2
codex-rs/state/src/migrations_tests.rs
View File

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