Add infrastructure for speculative remote V2 compaction (#53614)

## What changed

- Capture history, prompt metadata, routing state, authentication ownership, and context-window identity for background compaction requests.
- Send background requests over HTTP while preserving foreground routing, and retain the latest candidate in session memory.
- Add feature and threshold eligibility checks and deduplicate attempts by history version, model, compaction hash, and authentication ownership.
- Cancel and drain background tasks during session shutdown, preserving completed usage accounting.

The background launch helper is not yet connected to the turn loop.

GitOrigin-RevId: 0e5e6ab7dac4d98a836733668a775ac48b5a59f1

codex-rs/core/src/session/async_compaction.rs

@@ -0,0 +1,160 @@
+//! Holds the latest speculative request in memory until the foreground can use it.
+//! Discarded requests still charge completed usage; session shutdown cancels remaining network work.
+
+use std::sync::Arc;
+
+use crate::compact_remote_v2::BackgroundCompactionCandidate;
+use crate::compact_remote_v2::capture_background_input;
+use crate::compact_remote_v2::run_background_compaction;
+use crate::context::world_state::WorldState;
+use crate::session::context_window::ContextWindowTokenStatus;
+use crate::session::session::Session;
+use crate::session::step_context::StepContext;
+use crate::session::turn_context::TurnContext;
+use codex_analytics::CompactionPhase;
+use codex_features::Feature;
+use codex_model_provider::RemoteCompactionSupport;
+use codex_otel::auth_storage::AuthStorageOriginator;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
+use tracing::debug;
+
+#[derive(PartialEq, Eq)]
+struct CompactionAttempt {
+    history_version: u64,
+    model_slug: String,
+    compaction_hash: Option<String>,
+    auth_owner_generation: Option<u64>,
+}
+
+#[derive(Default)]
+pub(crate) struct AsyncCompactionState {
+    last_attempt: Option<CompactionAttempt>,
+    result: Option<oneshot::Receiver<BackgroundCompactionCandidate>>,
+    tasks: TaskTracker,
+    shutdown_cancellation: CancellationToken,
+}
+
+pub(super) fn enabled(turn_context: &TurnContext) -> bool {
+    turn_context
+        .config
+        .features
+        .enabled(Feature::AsyncCompaction)
+        && !turn_context.config.features.enabled(Feature::TokenBudget)
+        && matches!(
+            turn_context.provider.capabilities().remote_compaction,
+            RemoteCompactionSupport::V2
+        )
+}
+
+pub(super) async fn maybe_start(
+    sess: &Arc<Session>,
+    step_context: &Arc<StepContext>,
+    world_state: &Arc<WorldState>,
+    phase: CompactionPhase,
+    token_status: &ContextWindowTokenStatus,
+    (turn_state, auth_owner_generation): (Option<String>, Option<u64>),
+) {
+    let turn_context = &step_context.turn;
+    // Let the foreground establish the backend's model pin before starting a parallel request.
+    if matches!(phase, CompactionPhase::PreTurn) && turn_state.is_none() {
+        return;
+    }
+    if !enabled(turn_context)
+        || token_status.token_limit_reached
+        || (matches!(phase, CompactionPhase::PostTurn)
+            && token_status.turn_end_compaction_threshold_reached)
+        || step_context
+            .settings
+            .model_info
+            .resolved_context_window()
+            .is_none_or(|limit| {
+                i128::from(token_status.active_context_tokens) * 100
+                    < i128::from(limit)
+                        * i128::from(turn_context.config.async_compaction_threshold_percent)
+            })
+    {
+        return;
+    }
+
+    let base_instructions = sess.get_prompt_base_instructions().await;
+    let mut state = sess.state.lock().await;
+    let model_info = &step_context.settings.model_info;
+    let attempt = CompactionAttempt {
+        history_version: state.history.history_version(),
+        model_slug: model_info.slug.clone(),
+        compaction_hash: model_info.comp_hash.clone(),
+        auth_owner_generation,
+    };
+    if state.shutting_down
+        || state.async_compaction.last_attempt.as_ref() == Some(&attempt)
+        || state.history.annotated_items().is_empty()
+        || sess.services.model_client.auth_owner_generation() != auth_owner_generation
+    {
+        return;
+    }
+
+    // Capture and register without yielding, so shutdown or history replacement cannot
+    // overtake this request. Tool attribution still belongs to the foreground's current step.
+    let prepared_input = capture_background_input(
+        sess,
+        step_context,
+        state.history.clone(),
+        base_instructions,
+        turn_state,
+        auth_owner_generation,
+        sess.current_window_from_state(&state),
+    );
+    let task_session = Arc::clone(sess);
+    let task_step = Arc::clone(step_context);
+    let cutoff_world_state = Arc::clone(world_state);
+    let storage_originator = AuthStorageOriginator::from_client_name(&turn_context.originator);
+    state.async_compaction.last_attempt = Some(attempt);
+    let (send, result) = oneshot::channel();
+    state.async_compaction.result = Some(result);
+    let task = async move {
+        match run_background_compaction(
+            &task_session,
+            task_step,
+            cutoff_world_state,
+            phase,
+            prepared_input,
+        )
+        .await
+        {
+            Ok(candidate) => {
+                let _ = send.send(candidate);
+            }
+            Err(error) => {
+                debug!(%error, "discarding unsuccessful background compaction");
+            }
+        }
+    };
+    state
+        .async_compaction
+        .tasks
+        .spawn(storage_originator.scope(task));
+}
+
+pub(crate) async fn shutdown_cancellation(sess: &Session) -> CancellationToken {
+    sess.state
+        .lock()
+        .await
+        .async_compaction
+        .shutdown_cancellation
+        .clone()
+}
+
+pub(super) async fn shutdown(sess: &Session) {
+    let tasks = {
+        let mut state = sess.state.lock().await;
+        let async_compaction = &mut state.async_compaction;
+        async_compaction.result = None;
+        async_compaction.tasks.close();
+        async_compaction.shutdown_cancellation.cancel();
+        async_compaction.tasks.clone()
+    };
+    // Drain completed usage charging before shutdown finishes.
+    tasks.wait().await;
+}

codex-rs/core/src/session/handlers.rs

@@ -299,6 +299,7 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
     }
     let _ = sess.conversation.shutdown().await;
     sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
+    super::async_compaction::shutdown(sess).await;
     let shell_snapshot_prewarm = sess.state.lock().await.shell_snapshot_prewarm.take();
     if let Some(shell_snapshot_prewarm) = shell_snapshot_prewarm {
         shell_snapshot_prewarm.abort();

codex-rs/core/src/session/mod.rs

@@ -237,6 +237,8 @@ use codex_protocol::error::Result as CodexResult;
 #[cfg(test)]
 use codex_protocol::exec_output::StreamOutput;
 
+#[allow(dead_code)]
+pub(crate) mod async_compaction;
 mod code_mode_warning;
 mod config_refresh;
 pub(crate) mod context_window;
@@ -445,6 +447,7 @@ pub(crate) enum ForkPersistence {
 #[derive(Clone, Copy)]
 pub(crate) struct CompactionPrefix {
     pub(crate) len: usize,
+    pub(crate) auth_owner_generation: Option<u64>,
 }
 
 pub(crate) enum CompactedHistoryInstallation<'a> {
@@ -3745,7 +3748,7 @@ impl Session {
     pub(crate) async fn record_step_world_state_if_changed(
         &self,
         step_context: &step_context::StepContext,
-    ) -> CodexResult<()> {
+    ) -> CodexResult<Arc<WorldState>> {
         let turn_context = step_context.turn.as_ref();
         // Render model-visible state from the same step used to build and run tools.
         let world_state = self
@@ -3773,7 +3776,7 @@ impl Session {
             self.persist_rollout_items(&[RolloutItem::WorldState(world_state_item)])
                 .await;
         }
-        Ok(())
+        Ok(Arc::new(world_state))
     }
 
     /// Retains the step captured for execution.
@@ -4822,7 +4825,7 @@ impl Session {
     pub(crate) async fn record_context_updates_and_set_reference_context_item(
         &self,
         step_context: &StepContext,
-    ) -> CodexResult<()> {
+    ) -> CodexResult<Arc<WorldState>> {
         let turn_context = step_context.turn.as_ref();
         let reference_context_item = {
             let state = self.state.lock().await;
@@ -4870,7 +4873,7 @@ impl Session {
         // A snapshot can change without producing model-visible or TurnContext updates.
         let only_world_state_changed = !turn_context_changed && context_items.is_empty();
         if only_world_state_changed && world_state_item.is_none() {
-            return Ok(());
+            return Ok(Arc::new(world_state));
         }
         if !context_items.is_empty() {
             self.record_conversation_items(
@@ -4887,7 +4890,7 @@ impl Session {
         }
         // A snapshot-only change does not require a duplicate TurnContext record.
         if only_world_state_changed {
-            return Ok(());
+            return Ok(Arc::new(world_state));
         }
         // Persist one `TurnContextItem` per real user turn so resume/lazy replay can recover the
         // latest durable baseline even when this turn emitted no model-visible context diffs.
@@ -4911,7 +4914,7 @@ impl Session {
         // context items.
         let mut state = self.state.lock().await;
         state.set_reference_context_item(Some(turn_context_item));
-        Ok(())
+        Ok(Arc::new(world_state))
     }
 
     pub(crate) async fn update_token_usage_info(