From 9167a0b2d4a74d44a4d388057d304682e47313c6 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Tue, 13 Oct 2026 09:09:47 +0000 Subject: [PATCH] Reuse background results for remote V2 compaction (#53620) ## What changed - Await a compatible pending background request when foreground compaction is needed, and share the result installation path with foreground requests. - Validate the captured history generation, prefix length, model, and authentication owner before installing a prepared result. Fall back to foreground compaction when required if the result is rejected. - Replace only the captured history prefix, preserving later messages, retained tool outputs, goal edits, and the current world-state baseline. Record installation analytics and completion events only after a result is installed. GitOrigin-RevId: 798963f2fc603c3febba2c21eafa1bcb094b3ebf --- codex-rs/core/src/compact_remote_v2.rs | 213 ++++--------- .../core/src/compact_remote_v2_attempt.rs | 4 +- .../src/compact_remote_v2_installation.rs | 283 ++++++++++++++++++ codex-rs/core/src/session/async_compaction.rs | 40 ++- codex-rs/core/src/session/mod.rs | 83 +++-- codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/session/turn.rs | 59 +++- 7 files changed, 477 insertions(+), 207 deletions(-) create mode 100644 codex-rs/core/src/compact_remote_v2_installation.rs diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 81474d6833..470224bb97 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -11,7 +11,6 @@ use crate::client_common::ResponseEvent; use crate::compact::CompactedHistoryMetadata; use crate::compact::CompactionAnalyticsAttempt; use crate::compact::CompactionAnalyticsDetails; -use crate::compact::build_compaction_replacement_history; use crate::compact::compaction_status_from_result; use crate::compact_model_fallback::record_model_fallback; use crate::compact_model_fallback::should_retry_with_current_model; @@ -57,7 +56,6 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TruncationPolicy; -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; @@ -73,6 +71,10 @@ pub(crate) use attempt::capture_background_input; pub(crate) use attempt::run_background_compaction; use attempt::run_remote_compact_v2_attempt; +#[path = "compact_remote_v2_installation.rs"] +mod installation; +use installation::run_compaction; + #[path = "compact_remote_v2_images.rs"] mod images; @@ -99,6 +101,14 @@ enum CompactionRequestTransport<'a> { Background(Option, Option), } +/// The caller decides whether a declined prepared result needs a foreground request. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum CompactionFallback { + Foreground, + Skip, +} + +#[allow(clippy::too_many_arguments)] pub(crate) async fn run_inline_remote_auto_compact_task( sess: Arc, step_context: Arc, @@ -107,7 +117,9 @@ pub(crate) async fn run_inline_remote_auto_compact_task( world_state: Arc, reason: CompactionReason, phase: CompactionPhase, -) -> CodexResult<()> { + candidate: Option, + fallback: CompactionFallback, +) -> CodexResult { let compaction_metadata = CompactionTurnMetadata::new( CompactionTrigger::Auto, reason, @@ -121,6 +133,8 @@ pub(crate) async fn run_inline_remote_auto_compact_task( Some(client_session), world_state, compaction_metadata, + candidate, + fallback, ) .await } @@ -143,10 +157,14 @@ pub(crate) async fn run_remote_compact_task( /*client_session*/ None, world_state, compaction_metadata, + /*candidate*/ None, + CompactionFallback::Foreground, ) .await + .map(|_| ()) } +#[allow(clippy::too_many_arguments)] async fn run_remote_compact_task_inner( sess: &Arc, step_context: &Arc, @@ -154,7 +172,9 @@ async fn run_remote_compact_task_inner( client_session: Option<&mut ModelClientSession>, world_state: Arc, compaction_metadata: CompactionTurnMetadata, -) -> CodexResult<()> { + candidate: Option, + fallback: CompactionFallback, +) -> CodexResult { let turn_context = &step_context.turn; let trigger = compaction_metadata.trigger(); let reason = compaction_metadata.reason(); @@ -189,16 +209,38 @@ async fn run_remote_compact_task_inner( return Err(error); } } - let result = run_remote_compact_task_inner_impl( + let started_item = if candidate.is_none() { + let item = TurnItem::ContextCompaction(ContextCompactionItem::new()); + sess.emit_turn_item_started(turn_context, &item).await; + Some(item) + } else { + None + }; + let candidate = if candidate.is_none() && fallback == CompactionFallback::Foreground { + crate::session::async_compaction::wait_for_pending(sess, step_context).await + } else { + candidate + }; + // Box the combined request/install future to keep it off the turn stack. + let result = Box::pin(run_compaction( sess, step_context, replacement_step_context, client_session, - world_state, + &world_state, compaction_metadata, &mut analytics_details, - ) + started_item.as_ref(), + candidate, + fallback, + )) .await; + if matches!(result, Ok(false)) { + if let Some(item) = started_item { + sess.emit_turn_item_completed(turn_context, item).await; + } + return Ok(false); + } let status = compaction_status_from_result(&result); let codex_error = result.as_ref().err(); if result.is_ok() { @@ -214,7 +256,7 @@ async fn run_remote_compact_task_inner( .track(sess.as_ref(), status, codex_error, analytics_details) .await; match result { - Ok(()) => Ok(()), + Ok(installed) => Ok(installed), Err(err) if matches!(err.details(), CodexErrorDetails::TurnAborted) || matches!(phase, CompactionPhase::PostTurn) => @@ -235,161 +277,6 @@ async fn run_remote_compact_task_inner( } } -async fn run_remote_compact_task_inner_impl( - sess: &Arc, - step_context: &Arc, - replacement_step_context: &Arc, - mut client_session: Option<&mut ModelClientSession>, - world_state: Arc, - compaction_metadata: CompactionTurnMetadata, - analytics_details: &mut CompactionAnalyticsDetails, -) -> CodexResult<()> { - let turn_context = &step_context.turn; - let context_compaction_item = ContextCompactionItem::new(); - let compaction_id = context_compaction_item.id.clone(); - let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( - turn_context.sub_id.as_str(), - compaction_id.as_str(), - turn_context.model_info().slug.as_str(), - turn_context.provider.info().name.as_str(), - ); - let compaction_item = TurnItem::ContextCompaction(context_compaction_item); - sess.emit_turn_item_started(turn_context, &compaction_item) - .await; - - let attempt = run_remote_compact_v2_attempt( - sess, - step_context, - &compaction_trace, - compaction_metadata, - analytics_details, - RemoteCompactV2AttemptInput::Foreground(client_session.as_deref_mut()), - ) - .await; - let (attempt, compaction_turn_context) = match attempt { - Ok(attempt) => (attempt, turn_context), - Err(error) => { - let fallback_turn_context = &replacement_step_context.turn; - if !should_retry_with_current_model(&error, turn_context, fallback_turn_context) { - return Err(error); - } - sess.set_last_known_step_context(replacement_step_context) - .await; - let fallback_compaction_trace = - sess.services.rollout_thread_trace.compaction_trace_context( - fallback_turn_context.sub_id.as_str(), - compaction_id.as_str(), - fallback_turn_context.model_info().slug.as_str(), - fallback_turn_context.provider.info().name.as_str(), - ); - let fallback_result = run_remote_compact_v2_attempt( - sess, - replacement_step_context, - &fallback_compaction_trace, - compaction_metadata, - analytics_details, - RemoteCompactV2AttemptInput::Foreground(client_session), - ) - .await; - record_model_fallback( - &sess.services.session_telemetry, - turn_context.model_info().slug.as_str(), - fallback_turn_context.model_info().slug.as_str(), - compaction_metadata.reason(), - compaction_metadata.implementation(), - fallback_result.as_ref().err(), - ); - match fallback_result { - Ok(attempt) => (attempt, fallback_turn_context), - Err(_) => return Err(error), - } - } - }; - let RemoteCompactV2Attempt { - input_goal_ids, - trace_input_history, - prompt_input, - prompt_input_metadata, - compaction_output, - compaction_response_id, - token_usage, - owned_client_session: _owned_client_session, - } = attempt; - if let Some(token_usage) = token_usage { - sess.record_rollout_budget_usage(&token_usage).await?; - analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); - analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens); - 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, - sess.features().get(), - if sess.enabled(Feature::CompactionImageBudget) { - RetainedImageBudget::Enabled - } else { - RetainedImageBudget::Disabled - }, - ); - 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) = build_compaction_replacement_history( - sess.as_ref(), - replacement_step_context, - &world_state, - compacted_history, - ) - .await; - - if let Some(trace_input_history) = trace_input_history.as_deref() { - let replacement_history = new_history - .iter() - .map(|envelope| envelope.item.clone()) - .collect::>(); - compaction_trace.record_installed(&CompactionCheckpointTracePayload { - input_history: trace_input_history, - replacement_history: &replacement_history, - }); - } - let reviewer_compaction_hash = if crate::context::GuardianContextMode::from_history( - sess.conversation_history_snapshot().await.as_ref(), - ) == crate::context::GuardianContextMode::Legacy - && let Some(review_turn) = sess.turn_context_for_sub_id(&turn_context.sub_id).await - { - // Previous-model compaction must remain compatible with the continuing turn's - // reviewer, including model changes accepted while compaction was running. - let mut review_context = crate::guardian::GuardianReviewContext::from(&review_turn); - review_context.model_info = review_turn.capture_current_model_info(); - let (_, reviewer) = crate::guardian::resolve_review_model(sess, &review_context).await; - reviewer.comp_hash.clone() - } else { - None - }; - sess.replace_compacted_history( - new_history, - replacement_step_context.to_turn_context_item(), - world_state_baseline, - CompactedHistoryMetadata { - input_goal_ids, - message: String::new(), - window_number: new_window_number, - window_ids: new_window_ids, - compaction_response_id: Some(compaction_response_id), - compaction_model_hash: compaction_turn_context.model_info().comp_hash.clone(), - reviewer_compaction_hash, - }, - ) - .await; - sess.recompute_token_usage(&replacement_step_context.turn) - .await; - - sess.emit_turn_item_completed(compaction_turn_context, compaction_item) - .await; - Ok(()) -} - struct RemoteCompactionV2Output { compaction_output: ResponseItem, response_id: String, diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 0194f6526b..3f14ac5be7 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -48,9 +48,8 @@ pub(super) enum RemoteCompactV2AttemptInput<'a> { Background(BackgroundCompactionInput), } -#[allow(dead_code)] pub(crate) struct BackgroundCompactionCandidate { - pub(super) prefix: CompactionPrefix, + pub(crate) prefix: CompactionPrefix, pub(super) model_info: Arc, pub(super) cutoff_world_state: Arc, pub(super) attempt: RemoteCompactV2Attempt, @@ -255,6 +254,7 @@ pub(crate) fn capture_background_input( base_instructions }; let prefix = CompactionPrefix { + generation: snapshot.history_version(), len: snapshot.annotated_items().len(), auth_owner_generation, }; diff --git a/codex-rs/core/src/compact_remote_v2_installation.rs b/codex-rs/core/src/compact_remote_v2_installation.rs new file mode 100644 index 0000000000..36e11b7153 --- /dev/null +++ b/codex-rs/core/src/compact_remote_v2_installation.rs @@ -0,0 +1,283 @@ +//! Materializes and installs foreground or prepared V2 compaction through one lifecycle. +//! A rejected prepared result may fall back once; building and installation share one path. +//! Prepared prefixes are validated by the history replacement primitive before mutation. + +use super::*; +use crate::compact::build_compaction_replacement_history; +use crate::context::world_state::split_prefix_updates; +use crate::context_manager::updates::merge_world_state_updates; +use crate::session::CompactedHistoryInstallation; +use codex_rollout_trace::CompactionCheckpointTracePayload; + +#[allow(clippy::too_many_arguments)] +pub(super) async fn run_compaction( + sess: &Arc, + step_context: &Arc, + replacement_step_context: &Arc, + mut client_session: Option<&mut ModelClientSession>, + world_state: &WorldState, + compaction_metadata: CompactionTurnMetadata, + analytics_details: &mut CompactionAnalyticsDetails, + started_item: Option<&TurnItem>, + candidate: Option, + fallback: CompactionFallback, +) -> CodexResult { + if let Some(candidate) = candidate { + let installed = build_and_install_compaction( + sess, + step_context, + replacement_step_context, + client_session.as_deref_mut(), + world_state, + compaction_metadata, + analytics_details, + started_item, + Some(candidate), + ) + .await?; + if installed || fallback == CompactionFallback::Skip { + return Ok(installed); + } + } + build_and_install_compaction( + sess, + step_context, + replacement_step_context, + client_session, + world_state, + compaction_metadata, + analytics_details, + started_item, + /*candidate*/ None, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn build_and_install_compaction( + sess: &Arc, + step_context: &Arc, + replacement_step_context: &Arc, + mut client_session: Option<&mut ModelClientSession>, + world_state: &WorldState, + compaction_metadata: CompactionTurnMetadata, + analytics_details: &mut CompactionAnalyticsDetails, + started_item: Option<&TurnItem>, + candidate: Option, +) -> CodexResult { + let turn_context = &step_context.turn; + if let Some(candidate) = &candidate + && (candidate.model_info.slug != step_context.settings.model_info.slug + || candidate.model_info.comp_hash != step_context.settings.model_info.comp_hash) + { + return Ok(false); + } + let needs_start = started_item.is_none(); + let compaction_item = started_item + .cloned() + .unwrap_or_else(|| TurnItem::ContextCompaction(ContextCompactionItem::new())); + let compaction_id = compaction_item.id(); + let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( + turn_context.sub_id.as_str(), + compaction_id.as_str(), + candidate + .as_ref() + .map_or(turn_context.model_info().slug.as_str(), |candidate| { + candidate.model_info.slug.as_str() + }), + turn_context.provider.info().name.as_str(), + ); + if needs_start && candidate.is_none() { + sess.emit_turn_item_started(turn_context, &compaction_item) + .await; + } + let (attempt, compaction_turn_context, prefix) = if let Some(candidate) = candidate { + let mut attempt = candidate.attempt; + if compaction_trace.is_enabled() { + attempt.trace_input_history = + Some(sess.clone_history().await.raw_items().cloned().collect()); + } + ( + attempt, + turn_context, + Some(( + candidate.prefix, + candidate.model_info, + candidate.cutoff_world_state, + )), + ) + } else { + let attempt = run_remote_compact_v2_attempt( + sess, + step_context, + &compaction_trace, + compaction_metadata, + analytics_details, + RemoteCompactV2AttemptInput::Foreground(client_session.as_deref_mut()), + ) + .await; + let (attempt, compaction_turn_context) = match attempt { + Ok(attempt) => (attempt, turn_context), + Err(error) => { + let fallback_turn_context = &replacement_step_context.turn; + if !should_retry_with_current_model(&error, turn_context, fallback_turn_context) { + return Err(error); + } + sess.set_last_known_step_context(replacement_step_context) + .await; + let fallback_compaction_trace = + sess.services.rollout_thread_trace.compaction_trace_context( + fallback_turn_context.sub_id.as_str(), + compaction_id.as_str(), + fallback_turn_context.model_info().slug.as_str(), + fallback_turn_context.provider.info().name.as_str(), + ); + let fallback_result = run_remote_compact_v2_attempt( + sess, + replacement_step_context, + &fallback_compaction_trace, + compaction_metadata, + analytics_details, + RemoteCompactV2AttemptInput::Foreground(client_session), + ) + .await; + record_model_fallback( + &sess.services.session_telemetry, + turn_context.model_info().slug.as_str(), + fallback_turn_context.model_info().slug.as_str(), + compaction_metadata.reason(), + compaction_metadata.implementation(), + fallback_result.as_ref().err(), + ); + match fallback_result { + Ok(attempt) => (attempt, fallback_turn_context), + Err(_) => return Err(error), + } + } + }; + (attempt, compaction_turn_context, None) + }; + let RemoteCompactV2Attempt { + input_goal_ids, + trace_input_history, + prompt_input, + prompt_input_metadata, + compaction_output, + compaction_response_id, + token_usage, + owned_client_session: _owned_client_session, + } = attempt; + if prefix.is_none() + && let Some(token_usage) = token_usage.as_ref() + { + sess.record_rollout_budget_usage(token_usage).await?; + analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); + } + let (compacted_history, retained_images) = build_v2_compacted_history( + prompt_input, + prompt_input_metadata, + compaction_output, + sess.features().get(), + if sess.enabled(Feature::CompactionImageBudget) { + RetainedImageBudget::Enabled + } else { + RetainedImageBudget::Disabled + }, + ); + let reviewer_compaction_hash = reviewer_compaction_hash(sess, turn_context).await; + let installed = if let Some((prefix, model_info, cutoff_world_state)) = prefix { + let (initial_context, world_state_baseline) = sess + .build_initial_context_with_world_state(replacement_step_context, &cutoff_world_state) + .await; + let (prefix_context, initial_context) = split_prefix_updates(initial_context); + let initial_context = merge_world_state_updates(initial_context) + .into_iter() + .map(ResponseItemEnvelope::new) + .collect(); + sess.replace_compacted_history( + compacted_history, + replacement_step_context.to_turn_context_item(), + world_state_baseline, + CompactedHistoryInstallation::Prefix { + prefix, + prefix_context, + initial_context, + model_info: &model_info, + compaction_response_id, + reviewer_compaction_hash, + turn_context, + started_item: needs_start.then_some(&compaction_item), + }, + ) + .await + } else { + let (new_history, world_state_baseline) = build_compaction_replacement_history( + sess, + replacement_step_context, + world_state, + compacted_history, + ) + .await; + let (window_number, window_ids) = sess.advance_auto_compact_window().await; + let installation = CompactedHistoryMetadata { + input_goal_ids, + message: String::new(), + window_number, + window_ids, + compaction_response_id: Some(compaction_response_id), + compaction_model_hash: compaction_turn_context.model_info().comp_hash.clone(), + reviewer_compaction_hash, + }; + sess.replace_compacted_history( + new_history, + replacement_step_context.to_turn_context_item(), + world_state_baseline, + installation, + ) + .await + }; + if !installed { + return Ok(false); + } + // Rejected candidates must not supply analytics for a later foreground attempt. + if let Some(token_usage) = token_usage { + analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens); + analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens); + analytics_details.cache_write_input_tokens = Some(token_usage.cache_write_input_tokens); + } + analytics_details.retained_image_count = Some(retained_images); + if let Some(trace_input_history) = trace_input_history.as_deref() { + let replacement_history = sess + .clone_history() + .await + .raw_items() + .cloned() + .collect::>(); + compaction_trace.record_installed(&CompactionCheckpointTracePayload { + input_history: trace_input_history, + replacement_history: &replacement_history, + }); + } + sess.recompute_token_usage(&replacement_step_context.turn) + .await; + sess.emit_turn_item_completed(compaction_turn_context, compaction_item) + .await; + Ok(true) +} + +async fn reviewer_compaction_hash(sess: &Session, turn_context: &TurnContext) -> Option { + if crate::context::GuardianContextMode::from_history( + sess.conversation_history_snapshot().await.as_ref(), + ) == crate::context::GuardianContextMode::Legacy + && let Some(review_turn) = sess.turn_context_for_sub_id(&turn_context.sub_id).await + { + // Previous-model compaction must remain compatible with the continuing turn's + // reviewer, including model changes accepted while compaction was running. + let mut review_context = crate::guardian::GuardianReviewContext::from(&review_turn); + review_context.model_info = review_turn.capture_current_model_info(); + let (_, reviewer) = crate::guardian::resolve_review_model(sess, &review_context).await; + reviewer.comp_hash.clone() + } else { + None + } +} diff --git a/codex-rs/core/src/session/async_compaction.rs b/codex-rs/core/src/session/async_compaction.rs index 1dedc11b90..ea2f5a1d9b 100644 --- a/codex-rs/core/src/session/async_compaction.rs +++ b/codex-rs/core/src/session/async_compaction.rs @@ -1,4 +1,4 @@ -//! Holds the latest speculative request in memory until the foreground can use it. +//! Holds the current speculative request, polled at checkpoints or awaited when compaction is due. //! Discarded requests still charge completed usage; session shutdown cancels remaining network work. use std::sync::Arc; @@ -137,6 +137,44 @@ pub(super) async fn maybe_start( .spawn(storage_originator.scope(task)); } +pub(super) async fn take_ready( + sess: &Session, + turn_context: &TurnContext, +) -> Option { + if !enabled(turn_context) { + return None; + } + let mut state = sess.state.lock().await; + let candidate = state.async_compaction.result.as_mut()?.try_recv().ok()?; + state.async_compaction.result = None; + (candidate.prefix.generation == state.history.history_version()).then_some(candidate) +} + +pub(crate) async fn wait_for_pending( + sess: &Session, + step_context: &StepContext, +) -> Option { + if !enabled(&step_context.turn) { + return None; + } + let result = { + let mut state = sess.state.lock().await; + let result = state.async_compaction.result.take()?; + let model = &step_context.settings.model_info; + let attempt = state.async_compaction.last_attempt.as_ref()?; + if attempt.history_version != state.history.history_version() + || attempt.model_slug != model.slug + || attempt.compaction_hash != model.comp_hash + || attempt.auth_owner_generation != sess.services.model_client.auth_owner_generation() + { + return None; + } + result + }; + // Release session state while the request finishes and charges its usage. + result.await.ok() +} + pub(crate) async fn shutdown_cancellation(sess: &Session) -> CancellationToken { sess.state .lock() diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b5f12424b3..2fa6d80352 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -446,13 +446,13 @@ pub(crate) enum ForkPersistence { /// Normal appends keep this generation; compaction and rollback start a new one. #[derive(Clone, Copy)] pub(crate) struct CompactionPrefix { + pub(crate) generation: u64, pub(crate) len: usize, pub(crate) auth_owner_generation: Option, } pub(crate) enum CompactedHistoryInstallation<'a> { Complete(CompactedHistoryMetadata), - #[allow(dead_code)] Prefix { prefix: CompactionPrefix, prefix_context: Vec, @@ -460,7 +460,8 @@ pub(crate) enum CompactedHistoryInstallation<'a> { model_info: &'a ModelInfo, compaction_response_id: String, reviewer_compaction_hash: Option, - started_item: (&'a TurnContext, &'a TurnItem), + turn_context: &'a TurnContext, + started_item: Option<&'a TurnItem>, }, } @@ -4155,9 +4156,9 @@ impl Session { &self, mut items: Vec, reference_context_item: TurnContextItem, - world_state_baseline: WorldStateSnapshot, + mut world_state_baseline: WorldStateSnapshot, installation: impl Into>, - ) { + ) -> bool { let installation = installation.into(); let compaction_model_hash = match &installation { CompactedHistoryInstallation::Complete(metadata) => &metadata.compaction_model_hash, @@ -4187,9 +4188,10 @@ impl Session { // Wait for accepted updates to finish persisting, then keep later updates from // overtaking the current settings snapshot while its checkpoint is written. let _settings_guard = thread_settings::acquire_persistence_lock(self).await; - let world_state_item = WorldStateItem::full(world_state_baseline.clone().into_object()); let (compacted_item, started_item) = { let mut state = self.state.lock().await; + let complete_replacement = + matches!(&installation, CompactedHistoryInstallation::Complete(_)); let (metadata, started_item) = match installation { CompactedHistoryInstallation::Complete(metadata) => (metadata, None), CompactedHistoryInstallation::Prefix { @@ -4199,14 +4201,23 @@ impl Session { model_info, compaction_response_id, reviewer_compaction_hash, - started_item: (turn_context, started_item), + turn_context, + started_item, } => { + // Validate immediately before slicing and replacing the captured prefix. + if state.history.history_version() != prefix.generation + || state.history.annotated_items().len() < prefix.len + || self.services.model_client.auth_owner_generation() + != prefix.auth_owner_generation + { + return false; + } let current = state.history.annotated_items(); let suffix = ¤t[prefix.len..]; - if suffix - .iter() - .any(|envelope| compact::is_real_user_or_agent_message(&envelope.item)) - { + if suffix.iter().any(|envelope| { + crate::compact_retained_tool_outputs::is_retained_tool_output(envelope) + || compact::is_real_user_or_agent_message(&envelope.item) + }) { items.extend(initial_context); } else { items = compact::insert_initial_context_before_last_real_user_or_summary( @@ -4219,8 +4230,12 @@ impl Session { prefix_context.into_iter().map(ResponseItemEnvelope::new), ); items.extend_from_slice(suffix); - let input_goal_ids = - crate::context::UserGoalUpdate::message_ids(state.history.raw_items()); + // The unchanged suffix already contains updates through this baseline. + if let Some(checkpoint) = state.history.world_state_checkpoint() { + world_state_baseline = WorldStateSnapshot::from(&checkpoint.state); + } + // Prefix replacement already retains all edits in the live suffix. + let input_goal_ids = Default::default(); let (window_number, window_ids) = state.advance_auto_compact_window(); ( CompactedHistoryMetadata { @@ -4232,7 +4247,7 @@ impl Session { compaction_model_hash: model_info.comp_hash.clone(), reviewer_compaction_hash, }, - Some((turn_context, started_item)), + started_item.map(|item| (turn_context, item)), ) } }; @@ -4241,23 +4256,25 @@ impl Session { } // Goal edits are published outside the running task. Keep edits accepted after // the compaction input snapshot, in their original order, after its replacement. - let replacement_goal_ids = crate::context::UserGoalUpdate::message_ids( - items.iter().map(|envelope| &envelope.item), - ); - items.extend( - state - .history - .annotated_items() - .iter() - .filter(|envelope| { - crate::context::UserGoalUpdate::message_text(&envelope.item).is_some() - && envelope.item.id().is_some_and(|id| { - !metadata.input_goal_ids.contains(id) - && !replacement_goal_ids.contains(id) - }) - }) - .cloned(), - ); + if complete_replacement { + let replacement_goal_ids = crate::context::UserGoalUpdate::message_ids( + items.iter().map(|envelope| &envelope.item), + ); + items.extend( + state + .history + .annotated_items() + .iter() + .filter(|envelope| { + crate::context::UserGoalUpdate::message_text(&envelope.item).is_some() + && envelope.item.id().is_some_and(|id| { + !metadata.input_goal_ids.contains(id) + && !replacement_goal_ids.contains(id) + }) + }) + .cloned(), + ); + } let replacement_history = items.clone(); state.replace_annotated_history( items, @@ -4267,7 +4284,9 @@ impl Session { }, ); state.reasoning_effort_pin = ReasoningEffortPin::Compacted; - state.history.set_world_state_baseline(world_state_baseline); + state + .history + .set_world_state_baseline(world_state_baseline.clone()); let compacted_item = CompactedItem { message: metadata.message, replacement_history: Some(replacement_history), @@ -4297,6 +4316,7 @@ impl Session { // The prefix is now known to be valid; announce the install before checkpoint I/O. self.emit_turn_item_started(turn_context, item).await; } + let world_state_item = WorldStateItem::full(world_state_baseline.into_object()); let rollout_items = [ RolloutItem::Compacted(compacted_item), RolloutItem::WorldState(world_state_item), @@ -4315,6 +4335,7 @@ impl Session { let mut state = self.state.lock().await; state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); } + true } pub fn enabled(&self, feature: Feature) -> bool { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e77d26239b..967e8b1d92 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -12031,6 +12031,8 @@ async fn remote_compaction_v2_retains_only_the_selected_step(first_attempt: Firs ), CompactionReason::ModelDownshift, CompactionPhase::PreTurn, + /*candidate*/ None, + crate::compact_remote_v2::CompactionFallback::Foreground, ) .await .expect("compaction succeeds"); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 35cf4d6ef5..2895e42826 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -7,6 +7,7 @@ use crate::client::ModelClientSession; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::compact::run_inline_auto_compact_task; +use crate::compact_remote_v2::BackgroundCompactionCandidate; use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_remote_auto_compact_task_v2; use crate::connectors; use crate::context::ContextualUserFragment; @@ -351,6 +352,7 @@ pub(crate) async fn run_turn( &mut client_session, CompactionReason::ContextLimit, CompactionPhase::PreTurn, + /*candidate*/ None, ) .await?; sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref()) @@ -626,6 +628,7 @@ pub(crate) async fn run_turn( &mut client_session, CompactionReason::ContextLimit, CompactionPhase::MidTurn, + /*candidate*/ None, ) .await { @@ -734,6 +737,7 @@ pub(crate) async fn run_turn( &mut client_session, CompactionReason::ContextLimit, CompactionPhase::PostTurn, + /*candidate*/ None, ) .await { @@ -784,6 +788,7 @@ pub(crate) async fn run_turn( &mut client_session, CompactionReason::ContextLimit, CompactionPhase::MidTurn, + /*candidate*/ None, ) .await?; can_drain_pending_input = false; @@ -1325,6 +1330,7 @@ async fn run_pre_sampling_compact( client_session, CompactionReason::ContextLimit, CompactionPhase::PreTurn, + /*candidate*/ None, ) .await?; } @@ -1421,12 +1427,14 @@ async fn maybe_run_previous_model_inline_compact( client_session, reason, CompactionPhase::PreTurn, + /*candidate*/ None, ) .await?; Ok(()) } /// The replacement context may use a different model than the compactor. +#[allow(clippy::too_many_arguments)] #[instrument( level = "trace", skip_all, @@ -1439,13 +1447,14 @@ async fn run_auto_compact( client_session: &mut ModelClientSession, reason: CompactionReason, phase: CompactionPhase, -) -> CodexResult<()> { + candidate: Option, +) -> CodexResult { let turn_context = &step_context.turn; if matches!(phase, CompactionPhase::PreTurn) && crate::guardian::is_basic_session_source(&turn_context.session_source) && !crate::guardian::should_compact_guardian_input(sess)? { - return Ok(()); + return Ok(false); } let world_state = Arc::new( sess.build_world_state_for_step(&replacement_step_context, /*new_window*/ true) @@ -1467,17 +1476,37 @@ async fn run_auto_compact( world_state, ) .await?; - return Ok(()); + return Ok(true); } match turn_context.provider.capabilities().remote_compaction { RemoteCompactionSupport::V2 => { - emit_compact_metric( - &sess.services.session_telemetry, - "remote_v2", - /*manual*/ false, - ); - run_inline_remote_auto_compact_task_v2( + let has_candidate = candidate.is_some(); + if !has_candidate { + emit_compact_metric( + &sess.services.session_telemetry, + "remote_v2", + /*manual*/ false, + ); + } + let fallback = if has_candidate && matches!(reason, CompactionReason::ContextLimit) { + let status = + super::context_window::context_window_token_status(sess, turn_context).await; + let required = match phase { + CompactionPhase::PostTurn => status.turn_end_compaction_threshold_reached, + CompactionPhase::PreTurn + | CompactionPhase::MidTurn + | CompactionPhase::StandaloneTurn => status.token_limit_reached, + }; + if required { + crate::compact_remote_v2::CompactionFallback::Foreground + } else { + crate::compact_remote_v2::CompactionFallback::Skip + } + } else { + crate::compact_remote_v2::CompactionFallback::Foreground + }; + let installed = run_inline_remote_auto_compact_task_v2( Arc::clone(sess), step_context, replacement_step_context, @@ -1485,8 +1514,18 @@ async fn run_auto_compact( world_state, reason, phase, + candidate, + fallback, ) .await?; + if has_candidate && installed { + emit_compact_metric( + &sess.services.session_telemetry, + "remote_v2", + /*manual*/ false, + ); + } + Ok(installed) } RemoteCompactionSupport::Unsupported => { emit_compact_metric( @@ -1503,9 +1542,9 @@ async fn run_auto_compact( phase, ) .await?; + Ok(true) } } - Ok(()) } pub(super) fn collect_explicit_app_ids_from_skill_items(