From f88ff940c0f5b20628e94abc27d04196a14c5b94 Mon Sep 17 00:00:00 2001 From: Bryan Ashley Date: Mon, 31 Aug 2026 18:52:54 +0000 Subject: [PATCH] Load bounded context after empty wake turns (#41901) ## Why Paginated threads can accumulate wake turns with no user-message boundary. The reverse context scan could not use those turns as a bounded replay cutoff, even when a compaction and its resume state provided everything needed to reconstruct the latest model context. ## What changed Treat a full world-state snapshot that survives the latest compaction as a durable baseline when it is paired with a compatible turn context. This lets the scan stop at empty wake turns while still scanning further back for patches, snapshots before compaction, or missing and mismatched turn contexts. ## Testing Add coverage for repeated empty wake turns and verify that bounded replay matches full rollout reconstruction across successive compaction windows. GitOrigin-RevId: ae90c5f939c660ea539fbfc11128cc7259efaa9c --- .../session/rollout_reconstruction_tests.rs | 112 ++++++++++++++++++ codex-rs/rollout/src/model_context.rs | 33 ++++-- .../src/local/model_context_tests.rs | 111 ++++++++++++++++- 3 files changed, 243 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 4816b92d99..3f0b410b9e 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -19,6 +19,8 @@ use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::WorldStateItem; use codex_protocol::security_risk::SecurityRiskScore; +use codex_rollout::ModelContextScan; +use codex_rollout::ModelContextScanProgress; use core_test_support::responses::strip_metadata_from_items; use pretty_assertions::assert_eq; use serde_json::json; @@ -1283,6 +1285,116 @@ async fn reconstruct_history_replays_world_state_from_latest_compaction_window() ); } +#[tokio::test] +async fn bounded_replay_matches_full_replay_after_empty_turn_compactions() { + let (session, turn_context) = make_session_and_context().await; + let session_meta = SessionMetaLine { + meta: SessionMeta { + history_mode: ThreadHistoryMode::Paginated, + ..Default::default() + }, + git: None, + }; + let initial_context = turn_context.to_turn_context_item(); + let mut rollout_items = vec![RolloutItem::SessionMeta(session_meta.clone())]; + rollout_items.extend(completed_user_turn_rollout( + initial_context.clone(), + vec![RolloutItem::ResponseItem( + user_message("original task").into(), + )], + )); + let window_ids = [Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()]; + let mut latest_wake_start = 0; + for window_number in 1..=2 { + let mut context = initial_context.clone(); + context.turn_id = Some(format!("wake-{window_number}")); + context.model = format!("model-{window_number}"); + context.comp_hash = Some(format!("hash-{window_number}")); + context.realtime_active = Some(window_number == 2); + let mut wake_items = completed_user_turn_rollout( + context.clone(), + vec![ + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(annotated(vec![assistant_message(&format!( + "summary-{window_number}" + ))])), + mcp_resource_origins: None, + window_number: Some(window_number as u64), + first_window_id: Some(window_ids[0].to_string()), + previous_window_id: Some(window_ids[window_number - 1].to_string()), + window_id: Some(window_ids[window_number].to_string()), + }), + RolloutItem::WorldState(WorldStateItem::full(object!({ + "environment": {"window": window_number, "status": "starting"} + }))), + RolloutItem::TurnContext(context), + RolloutItem::ResponseItem(assistant_message("continued working").into()), + RolloutItem::WorldState(WorldStateItem::patch(object!({ + "environment": {"status": "ready"} + }))), + ], + ); + // Clock wakes have no user-message boundary, and the next wait interrupts the turn. + wake_items.retain(|item| { + !matches!( + item, + RolloutItem::EventMsg(EventMsg::UserMessage(_) | EventMsg::TurnComplete(_)) + ) + }); + wake_items.push(RolloutItem::EventMsg(EventMsg::TurnAborted( + codex_protocol::protocol::TurnAbortedEvent { + turn_id: Some(format!("wake-{window_number}")), + reason: TurnAbortReason::Interrupted, + started_at: None, + completed_at: None, + duration_ms: None, + }, + ))); + latest_wake_start = rollout_items.len(); + rollout_items.extend(wake_items); + } + + let mut scan = ModelContextScan::default(); + let cutoff = rollout_items + .iter() + .enumerate() + .rev() + .find_map(|(index, item)| { + matches!(scan.push(item.clone()), ModelContextScanProgress::Complete).then_some(index) + }); + assert_eq!(cutoff, Some(latest_wake_start)); + let bounded_items = scan.finish(session_meta); + let full = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + let bounded = session + .reconstruct_history_from_rollout(&turn_context, &bounded_items) + .await; + assert_eq!( + ( + bounded.history, + bounded.previous_turn_settings, + bounded.reference_context_item, + bounded.world_state_baseline, + bounded.window_number, + bounded.first_window_id, + bounded.previous_window_id, + bounded.window_id, + ), + ( + full.history, + full.previous_turn_settings, + full.reference_context_item, + full.world_state_baseline, + full.window_number, + full.first_window_id, + full.previous_window_id, + full.window_id, + ), + ); +} + #[tokio::test] async fn reconstruct_history_preserves_legacy_compaction_count_with_session_meta_window() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/rollout/src/model_context.rs b/codex-rs/rollout/src/model_context.rs index 6a19e5f24e..f41021442e 100644 --- a/codex-rs/rollout/src/model_context.rs +++ b/codex-rs/rollout/src/model_context.rs @@ -25,19 +25,20 @@ pub enum ModelContextScanProgress { /// The scan stops once it has both: /// /// - `saw_compaction`: a `CompactedItem` with `replacement_history` and `window_number`; -/// - `saw_completed_turn_context`: a completed user turn with a compatible `TurnContextItem`. +/// - `saw_completed_turn_context`: a completed turn with a compatible `TurnContextItem` and a +/// durable context baseline. /// /// If the scan reaches the beginning before finding a bounded cutoff, it has already collected /// the complete replay and so we can return that directly. /// -/// `TurnContextItem` does not identify whether it came from a user turn, so one only counts after -/// the same turn also proves a user-turn boundary: a paginated -/// `ItemCompleted(UserMessage)` marker, agent message, or inter-agent message. Paginated writers -/// persist that marker for real user turns; older rollouts without it conservatively scan to the -/// beginning. A raw `role=user` response item is not sufficient because contextual user fragments -/// use that role but do not count as turn boundaries during reconstruction. The compaction restores -/// model-visible items; the turn context restores previous settings (`model`, `comp_hash`, and -/// `realtime_active`) and the reference baseline. +/// A turn establishes a context baseline with a user-turn boundary (a paginated +/// `ItemCompleted(UserMessage)` marker, agent message, or inter-agent message), or a full +/// `WorldState` snapshot newer than that turn's latest compaction. The snapshot also lets turns +/// with empty input supply resume metadata, matching rollout reconstruction. Without either, +/// the scan continues to older turns. A raw `role=user` response item is not sufficient because +/// contextual user fragments use that role but do not count as turn boundaries during reconstruction. +/// The compaction restores model-visible items; the turn context restores previous settings +/// (`model`, `comp_hash`, and `realtime_active`) and the reference baseline. /// /// These paginated shapes disable the bounded cutoff: /// @@ -91,6 +92,7 @@ impl ModelContextScan { } RolloutItem::Compacted(_) => { self.saw_compaction = true; + self.active_segment.saw_compaction = true; } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) => { // Paginated threads reject rollback. Keep old rollouts correct rather than @@ -150,12 +152,15 @@ impl ModelContextScan { RolloutItem::EventMsg(EventMsg::UserMessage(_)) => { self.active_segment.has_user_turn = true; } + RolloutItem::WorldState(state) => { + self.active_segment.has_full_world_state |= + state.full && !self.active_segment.saw_compaction; + } RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::RealtimeItem(_) - | RolloutItem::SecurityRiskScore(_) - | RolloutItem::WorldState(_) => {} + | RolloutItem::SecurityRiskScore(_) => {} } if self.has_bounded_cutoff() { @@ -166,7 +171,9 @@ impl ModelContextScan { } fn finalize_active_segment(&mut self) { - if self.active_segment.has_user_turn && self.active_segment.has_turn_context { + if self.active_segment.has_turn_context + && (self.active_segment.has_user_turn || self.active_segment.has_full_world_state) + { self.saw_completed_turn_context = true; } self.active_segment = ActiveTurnSegment::default(); @@ -182,6 +189,8 @@ struct ActiveTurnSegment { turn_id: Option, has_user_turn: bool, has_turn_context: bool, + has_full_world_state: bool, + saw_compaction: bool, } fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool { diff --git a/codex-rs/thread-store/src/local/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs index af8e75ae53..a6ec02169d 100644 --- a/codex-rs/thread-store/src/local/model_context_tests.rs +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -19,6 +19,7 @@ use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::WorldStateItem; use codex_protocol::user_input::UserInput; use codex_rollout::CompactedItem; use codex_rollout::RolloutItem; @@ -81,6 +82,114 @@ async fn loads_latest_checkpoint_with_required_turn_metadata() { })); } +#[tokio::test] +async fn loads_recent_context_after_many_empty_wake_turns() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1008); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = write_paginated_rollout( + home.path(), + "2025-01-03T13-00-07", + uuid, + [ + turn_started("user-turn"), + user_message("keep working"), + completed_user_message("user-turn", "keep working"), + turn_context(home.path(), "user-turn"), + turn_complete("user-turn"), + ], + ); + let mut expected_suffix = Vec::new(); + for index in 0..32 { + let turn_id = format!("wake-{index}"); + let mut items = vec![turn_started(&turn_id)]; + if index % 8 == 0 { + expected_suffix.clear(); + items.extend([ + compacted(&format!("checkpoint-{index}"), Some(Vec::new())), + RolloutItem::WorldState(WorldStateItem::full(Default::default())), + ]); + } + items.extend([ + turn_context(home.path(), &turn_id), + contextual_user_message(), + turn_complete(&turn_id), + ]); + append_items(&path, items.clone()); + expected_suffix.extend(items); + } + let session_meta = codex_rollout::read_session_meta_line(&path) + .await + .expect("read session metadata"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + expected_suffix.insert(0, RolloutItem::SessionMeta(session_meta)); + + assert_eq!( + serde_json::to_value(context.items).expect("serialize context"), + serde_json::to_value(expected_suffix).expect("serialize expected context") + ); +} + +#[tokio::test] +async fn empty_wake_requires_surviving_full_world_state_and_matching_context() { + enum MissingBaseline { + SnapshotBeforeCompaction, + PatchOnly, + MissingContext, + IncompatibleContext, + } + for baseline in [ + MissingBaseline::SnapshotBeforeCompaction, + MissingBaseline::PatchOnly, + MissingBaseline::MissingContext, + MissingBaseline::IncompatibleContext, + ] { + let home = TempDir::new().expect("temp dir"); + let path = write_paginated_rollout( + home.path(), + "2025-01-03T13-00-08", + Uuid::from_u128(/*v*/ 1009), + [ + turn_started("user-turn"), + completed_user_message("user-turn", "keep working"), + turn_context(home.path(), "user-turn"), + turn_complete("user-turn"), + turn_started("wake"), + ], + ); + let full = RolloutItem::WorldState(WorldStateItem::full(Default::default())); + let checkpoint = compacted("checkpoint", Some(Vec::new())); + let context = turn_context(home.path(), "wake"); + let items = match baseline { + MissingBaseline::SnapshotBeforeCompaction => vec![full, checkpoint, context], + MissingBaseline::PatchOnly => vec![ + checkpoint, + RolloutItem::WorldState(WorldStateItem::patch(Default::default())), + context, + ], + MissingBaseline::MissingContext => vec![checkpoint, full], + MissingBaseline::IncompatibleContext => { + vec![ + checkpoint, + full, + turn_context(home.path(), "different-turn"), + ] + } + }; + append_items(&path, items); + append_items(&path, [turn_complete("wake")]); + + assert_reverse_scan_matches_full_history(home.path(), &path).await; + } +} + #[tokio::test] async fn fork_context_excludes_items_after_frozen_cutoff() { let home = TempDir::new().expect("temp dir"); @@ -523,7 +632,7 @@ async fn assert_reverse_scan_matches_full_history(home: &Path, path: &Path) { ); } -fn append_items(path: &Path, items: [RolloutItem; N]) { +fn append_items(path: &Path, items: impl IntoIterator) { let mut file = OpenOptions::new() .append(true) .open(path)