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
This commit is contained in:
rka-oai
2026-10-13 09:18:45 +00:00
committed by copyberry
parent 308ddf5e7c
commit 9167a0b2d4
7 changed files with 477 additions and 207 deletions
+50 -163
View File
@@ -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<String>, Option<u64>),
}
/// 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<Session>,
step_context: Arc<StepContext>,
@@ -107,7 +117,9 @@ pub(crate) async fn run_inline_remote_auto_compact_task(
world_state: Arc<WorldState>,
reason: CompactionReason,
phase: CompactionPhase,
) -> CodexResult<()> {
candidate: Option<BackgroundCompactionCandidate>,
fallback: CompactionFallback,
) -> CodexResult<bool> {
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<Session>,
step_context: &Arc<StepContext>,
@@ -154,7 +172,9 @@ async fn run_remote_compact_task_inner(
client_session: Option<&mut ModelClientSession>,
world_state: Arc<WorldState>,
compaction_metadata: CompactionTurnMetadata,
) -> CodexResult<()> {
candidate: Option<BackgroundCompactionCandidate>,
fallback: CompactionFallback,
) -> CodexResult<bool> {
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<Session>,
step_context: &Arc<StepContext>,
replacement_step_context: &Arc<StepContext>,
mut client_session: Option<&mut ModelClientSession>,
world_state: Arc<WorldState>,
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::<Vec<_>>();
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,
@@ -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<ModelInfo>,
pub(super) cutoff_world_state: Arc<WorldState>,
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,
};
@@ -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<Session>,
step_context: &Arc<StepContext>,
replacement_step_context: &Arc<StepContext>,
mut client_session: Option<&mut ModelClientSession>,
world_state: &WorldState,
compaction_metadata: CompactionTurnMetadata,
analytics_details: &mut CompactionAnalyticsDetails,
started_item: Option<&TurnItem>,
candidate: Option<BackgroundCompactionCandidate>,
fallback: CompactionFallback,
) -> CodexResult<bool> {
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<Session>,
step_context: &Arc<StepContext>,
replacement_step_context: &Arc<StepContext>,
mut client_session: Option<&mut ModelClientSession>,
world_state: &WorldState,
compaction_metadata: CompactionTurnMetadata,
analytics_details: &mut CompactionAnalyticsDetails,
started_item: Option<&TurnItem>,
candidate: Option<BackgroundCompactionCandidate>,
) -> CodexResult<bool> {
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::<Vec<_>>();
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<String> {
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
}
}
+39 -1
View File
@@ -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<BackgroundCompactionCandidate> {
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<BackgroundCompactionCandidate> {
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()
+52 -31
View File
@@ -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<u64>,
}
pub(crate) enum CompactedHistoryInstallation<'a> {
Complete(CompactedHistoryMetadata),
#[allow(dead_code)]
Prefix {
prefix: CompactionPrefix,
prefix_context: Vec<ResponseItem>,
@@ -460,7 +460,8 @@ pub(crate) enum CompactedHistoryInstallation<'a> {
model_info: &'a ModelInfo,
compaction_response_id: String,
reviewer_compaction_hash: Option<String>,
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<ResponseItemEnvelope>,
reference_context_item: TurnContextItem,
world_state_baseline: WorldStateSnapshot,
mut world_state_baseline: WorldStateSnapshot,
installation: impl Into<CompactedHistoryInstallation<'a>>,
) {
) -> 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 = &current[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 {
+2
View File
@@ -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");
+49 -10
View File
@@ -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<BackgroundCompactionCandidate>,
) -> CodexResult<bool> {
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(