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

use super::*;
use crate::context::world_state::WorldStateSnapshot;
use crate::context_manager::is_user_turn_boundary;
use codex_history::ResponseItemEnvelope;
use codex_history::TurnAttribution;
use codex_protocol::ContextWindowId;
use codex_protocol::protocol::SessionContextWindow;
use codex_protocol::protocol::ThreadHistoryMode;
use codex_utils_output_truncation::TruncationPolicy;
use std::collections::BTreeSet;

// Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with
// the resume/fork hydration metadata derived from the same replay.
#[derive(Debug, PartialEq)]
pub(crate) struct RolloutReconstruction {
    pub(super) history: Vec<ResponseItemEnvelope>,
    pub(super) rejected_audio_file_ids: BTreeSet<String>,
    pub(crate) retained_context: codex_history::RetainedContext,
    pub(super) guardian_history: Option<codex_history::GuardianHistoryCheckpoint>,
    pub(super) last_started_turn_id: Option<String>,
    pub(super) previous_turn_settings: Option<PreviousTurnSettings>,
    pub(super) reference_context_item: Option<TurnContextItem>,
    pub(super) turn_attribution: Option<TurnAttribution>,
    pub(super) world_state_baseline: Option<WorldStateSnapshot>,
    pub(super) window_number: u64,
    pub(super) first_window_id: Option<ContextWindowId>,
    pub(super) previous_window_id: Option<ContextWindowId>,
    pub(super) window_id: Option<ContextWindowId>,
}

#[derive(Debug, Clone)]
struct ReconstructedWindow {
    number: u64,
    first_id: Option<ContextWindowId>,
    previous_id: Option<ContextWindowId>,
    id: Option<ContextWindowId>,
}

#[derive(Debug, Default)]
enum TurnReferenceContextItem {
    /// No `TurnContextItem` has been seen for this replay span yet.
    ///
    /// This differs from `Cleared`: `NeverSet` means there is no evidence this turn ever
    /// established a baseline, while `Cleared` means a baseline existed and a later compaction
    /// invalidated it. Only the latter must emit an explicit clearing segment for resume/fork
    /// hydration.
    #[default]
    NeverSet,
    /// A previously established baseline was invalidated by later compaction.
    Cleared,
    /// The latest baseline established by this replay span.
    Latest(Box<TurnContextItem>),
}

#[derive(Debug, Clone, Copy)]
// The selected compaction and its replay tail must belong to the same surviving segment.
struct ReplayCheckpoint<'a> {
    compacted: &'a CompactedItem,
    suffix: &'a [RolloutItem],
}

/// Selects the newest compaction that can safely bound replay.
///
/// Returns `None` when reconstruction must replay all supplied items, either because there is no
/// compaction or the newest compaction cannot bound replay.
fn select_input_compaction(
    rollout_items: &[RolloutItem],
    history_mode: ThreadHistoryMode,
) -> Option<ReplayCheckpoint<'_>> {
    // Only the newest compaction can bound replay. If it is incomplete, an older compaction
    // cannot replace the history or window state that the newer one may have changed.
    let (index, compacted) = rollout_items
        .iter()
        .enumerate()
        .rev()
        .find_map(|(index, item)| match item {
            RolloutItem::Compacted(compacted) => Some((index, compacted)),
            _ => None,
        })?;
    // Paginated histories always honor this boundary. Other histories only do so when resume
    // metadata identifies a compaction written under the newer resume contract.
    if compacted.replacement_history.is_none()
        || compacted.window_number.is_none()
        || (compacted.resume_metadata.is_none()
            && !matches!(history_mode, ThreadHistoryMode::Paginated))
    {
        return None;
    }
    Some(ReplayCheckpoint {
        compacted,
        suffix: &rollout_items[index + 1..],
    })
}

#[derive(Debug, Default)]
struct ActiveReplaySegment<'a> {
    turn_id: Option<String>,
    turn_completed: bool,
    counts_as_user_turn: bool,
    previous_turn_settings: Option<PreviousTurnSettings>,
    reference_context_item: TurnReferenceContextItem,
    turn_attribution: Option<TurnAttribution>,
    world_state_replay: Vec<&'a RolloutItem>,
    history_checkpoint: Option<ReplayCheckpoint<'a>>,
    window: Option<ReconstructedWindow>,
}

impl ActiveReplaySegment<'_> {
    fn take_checkpoint_attribution(&mut self, checkpoint: &mut Option<TurnAttribution>) {
        if checkpoint.as_ref().is_some_and(|attribution| {
            self.turn_id
                .as_ref()
                .is_none_or(|id| id == &attribution.turn_id)
        }) {
            // Consume a matching checkpoint even when rollback discards this segment.
            self.turn_attribution = self.turn_attribution.take().or(checkpoint.take());
        }
    }
}

fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool {
    active_turn_id
        .is_none_or(|turn_id| item_turn_id.is_none_or(|item_turn_id| item_turn_id == turn_id))
}

fn finalize_active_segment<'a>(
    active_segment: ActiveReplaySegment<'a>,
    history_checkpoint: &mut Option<ReplayCheckpoint<'a>>,
    previous_turn_settings: &mut Option<PreviousTurnSettings>,
    reference_context_item: &mut TurnReferenceContextItem,
    world_state_replay: &mut Vec<&'a RolloutItem>,
    window: &mut Option<ReconstructedWindow>,
    pending_rollback_turns: &mut usize,
) -> Option<TurnAttribution> {
    // Thread rollback drops the newest surviving real user-message boundaries. In replay, that
    // means skipping the next finalized segments that contain a non-contextual
    // `EventMsg::UserMessage`.
    if *pending_rollback_turns > 0 {
        if active_segment.counts_as_user_turn {
            *pending_rollback_turns -= 1;
        }
        return None;
    }

    // Full world-state snapshots are persisted after installing initial context. They still
    // establish a baseline when a child fork removes the parent turn's agent message. Do not
    // count these context-only segments as user turns for rollback, or use a snapshot from
    // before the segment's latest compaction.
    let has_context_baseline = active_segment.counts_as_user_turn
        || active_segment
            .world_state_replay
            .iter()
            .take_while(|item| !matches!(item, RolloutItem::Compacted(_)))
            .any(|item| matches!(item, RolloutItem::WorldState(state) if state.full));
    world_state_replay.extend(active_segment.world_state_replay);

    // A surviving replacement-history compaction is a complete history base. Once we
    // know the newest surviving one, older rollout items do not affect rebuilt history.
    if history_checkpoint.is_none()
        && let Some(segment_history_checkpoint) = active_segment.history_checkpoint
    {
        *history_checkpoint = Some(segment_history_checkpoint);
    }

    if window.is_none() {
        *window = active_segment.window;
    }

    // Restore settings from the newest surviving context baseline.
    if previous_turn_settings.is_none() && has_context_baseline {
        *previous_turn_settings = active_segment.previous_turn_settings;
    }

    // `reference_context_item` comes from the newest surviving context baseline, or
    // from a surviving compaction that explicitly cleared that baseline.
    if matches!(reference_context_item, TurnReferenceContextItem::NeverSet)
        && (has_context_baseline
            || matches!(
                active_segment.reference_context_item,
                TurnReferenceContextItem::Cleared
            ))
    {
        *reference_context_item = active_segment.reference_context_item;
    }
    // Automatic turns can persist attribution without establishing a model-context baseline.
    // Select recovery metadata from the same surviving segments, independently of that baseline.
    active_segment.turn_attribution
}

impl Session {
    pub(super) async fn reconstruct_history_from_rollout(
        &self,
        turn_context: &TurnContext,
        rollout_items: &[RolloutItem],
    ) -> RolloutReconstruction {
        ContextManager::reconstruct_rollout(
            rollout_items,
            turn_context.history_mode,
            ContextManager::for_session(
                turn_context.root_capabilities_restricted,
                &turn_context.config.features,
                turn_context.config.guardian_context_budgeting(),
                // Reconstructing retained evidence may mint temporary revisions before restoring
                // the durable originals. Replay must not consume the live thread's item sequence.
                Self::new_item_id_generator(),
            ),
            turn_context.model_info().truncation_policy.into(),
        )
    }
}

impl ContextManager {
    /// Replays stored history without starting a session or acquiring its writer.
    pub(crate) fn reconstruct_rollout(
        rollout_items: &[RolloutItem],
        history_mode: ThreadHistoryMode,
        mut history: Self,
        truncation_policy: TruncationPolicy,
    ) -> RolloutReconstruction {
        // Rejection facts are durable across compaction boundaries. File IDs are immutable,
        // so retaining a fact after rollback is conservative and cannot suppress a different
        // future upload.
        let mut rejected_audio_file_ids = BTreeSet::new();
        for item in rollout_items {
            match item {
                RolloutItem::Compacted(compacted) => {
                    if let Some(metadata) = &compacted.resume_metadata {
                        rejected_audio_file_ids
                            .extend(metadata.rejected_audio_file_ids.iter().cloned());
                    }
                }
                RolloutItem::EventMsg(EventMsg::TurnComplete(completed)) => {
                    if let Some(error) = &completed.error {
                        rejected_audio_file_ids
                            .extend(error.rejected_audio_file_ids.iter().cloned());
                    }
                }
                _ => {}
            }
        }
        // Select the compaction and suffix that can affect reconstruction.
        let has_legacy_compaction_without_window_number =
            rollout_items.iter().any(|item| {
                matches!(item, RolloutItem::Compacted(compacted) if compacted.window_number.is_none())
            });
        let initial_window = if has_legacy_compaction_without_window_number {
            None
        } else {
            rollout_items.iter().find_map(|item| match item {
                RolloutItem::SessionMeta(session_meta) => session_meta
                    .meta
                    .context_window
                    .as_ref()
                    .and_then(reconstructed_window_from_session_context_window),
                _ => None,
            })
        };
        let input_checkpoint = select_input_compaction(rollout_items, history_mode);
        let replay_items = input_checkpoint.map_or(rollout_items, |checkpoint| checkpoint.suffix);
        let resume_metadata =
            input_checkpoint.and_then(|checkpoint| checkpoint.compacted.resume_metadata.as_ref());
        let mut history_checkpoint = input_checkpoint;
        let mut window = input_checkpoint
            .and_then(|checkpoint| reconstructed_window_from_compaction(checkpoint.compacted));

        // Scan the selected items backward to find the newest surviving turn state.
        let mut started_turn_ids = replay_items
            .iter()
            .rev()
            .filter_map(|item| match item {
                RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => Some(event.turn_id.as_str()),
                RolloutItem::DynamicToolsUpdated { .. }
                | RolloutItem::SessionMeta(_)
                | RolloutItem::ResponseItem(_)
                | RolloutItem::InterAgentCommunication(_)
                | RolloutItem::InterAgentCommunicationMetadata { .. }
                | RolloutItem::TurnContext(_)
                | RolloutItem::WorldState(_)
                | RolloutItem::RetainedContext(_)
                | RolloutItem::SecurityRiskScore(_)
                | RolloutItem::TokenUsageRecord(_)
                | RolloutItem::RealtimeItem(_)
                | RolloutItem::Compacted(_)
                | RolloutItem::EventMsg(_) => None,
            })
            .peekable();
        let last_started_turn_id = started_turn_ids
            .peek()
            .map(|id| (*id).to_owned())
            .or_else(|| resume_metadata.and_then(|metadata| metadata.last_started_turn_id.clone()));
        // Saved attribution still identifies the turn when continuation eligibility was cleared.
        let checkpoint_turn_id = resume_metadata.and_then(|metadata| {
            metadata.last_started_turn_id.as_deref().or_else(|| {
                metadata
                    .turn_attribution
                    .as_ref()
                    .map(|attribution| attribution.turn_id.as_str())
            })
        });

        let mut previous_turn_settings = None;
        let mut reference_context_item = TurnReferenceContextItem::NeverSet;
        let mut turn_attribution = None;
        let mut world_state_replay = Vec::new();
        // Rollback is "drop the newest N user turns". While scanning in reverse, that becomes
        // "skip the next N user-turn segments we finalize".
        let mut pending_rollback_turns = 0usize;
        // Reverse replay accumulates rollout items into the newest in-progress turn segment until
        // we hit its matching `TurnStarted`, at which point the segment can be finalized.
        let mut checkpoint_attribution =
            resume_metadata.and_then(|metadata| metadata.turn_attribution.clone());
        let mut late_completed_turn_ids = HashSet::new();
        let mut active_segment: Option<ActiveReplaySegment<'_>> = None;

        for (index, item) in replay_items.iter().enumerate().rev() {
            let preceding_turn_id = started_turn_ids.peek().copied().or(checkpoint_turn_id);
            match item {
                RolloutItem::Compacted(compacted) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.world_state_replay.push(item);
                    if active_segment.window.is_none()
                        && let Some(compaction_window) =
                            reconstructed_window_from_compaction(compacted)
                    {
                        active_segment.window = Some(compaction_window);
                    }
                    // Looking backward, compaction clears any older baseline unless a newer
                    // `TurnContextItem` in this same segment has already re-established it.
                    if matches!(
                        active_segment.reference_context_item,
                        TurnReferenceContextItem::NeverSet
                    ) {
                        active_segment.reference_context_item = TurnReferenceContextItem::Cleared;
                    }
                    if active_segment.history_checkpoint.is_none()
                        && compacted.replacement_history.is_some()
                    {
                        active_segment.history_checkpoint = Some(ReplayCheckpoint {
                            compacted,
                            suffix: &replay_items[index + 1..],
                        });
                    }
                }
                RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => {
                    pending_rollback_turns = pending_rollback_turns
                        .saturating_add(usize::try_from(rollback.num_turns).unwrap_or(usize::MAX));
                }
                RolloutItem::EventMsg(EventMsg::TurnComplete(event)) => {
                    // Stop hooks can finish after a newer turn has started.
                    if !turn_ids_are_compatible(preceding_turn_id, Some(event.turn_id.as_str())) {
                        late_completed_turn_ids.insert(event.turn_id.as_str());
                        continue;
                    }
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.turn_completed = true;
                    // Reverse replay often sees `TurnComplete` before any turn-scoped metadata.
                    // Capture the turn id early so later `TurnContext` / abort items can match it.
                    if active_segment.turn_id.is_none() {
                        active_segment.turn_id = Some(event.turn_id.clone());
                    }
                }
                RolloutItem::EventMsg(EventMsg::TurnAborted(event)) => {
                    if !turn_ids_are_compatible(preceding_turn_id, event.turn_id.as_deref()) {
                        continue;
                    }
                    if let Some(active_segment) = active_segment.as_mut() {
                        if active_segment.turn_id.is_none()
                            && let Some(turn_id) = &event.turn_id
                        {
                            active_segment.turn_id = Some(turn_id.clone());
                        }
                    } else if let Some(turn_id) = &event.turn_id {
                        active_segment = Some(ActiveReplaySegment {
                            turn_id: Some(turn_id.clone()),
                            ..Default::default()
                        });
                    }
                }
                RolloutItem::EventMsg(EventMsg::UserMessage(_)) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.counts_as_user_turn = true;
                }
                RolloutItem::TurnContext(ctx) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    // `TurnContextItem` can attach metadata to an existing segment, but only a
                    // real `UserMessage` event should make the segment count as a user turn.
                    if active_segment.turn_id.is_none() {
                        active_segment.turn_id = ctx.turn_id.clone();
                    }
                    if turn_ids_are_compatible(
                        active_segment.turn_id.as_deref(),
                        ctx.turn_id.as_deref(),
                    ) {
                        active_segment.previous_turn_settings = Some(PreviousTurnSettings {
                            model: ctx.model.clone(),
                            cyber_access_program: ctx.cyber_access_program,
                            comp_hash: ctx.comp_hash.clone(),
                            realtime_active: ctx.realtime_active,
                        });
                        if matches!(
                            active_segment.reference_context_item,
                            TurnReferenceContextItem::NeverSet
                        ) {
                            active_segment.reference_context_item =
                                TurnReferenceContextItem::Latest(Box::new(ctx.clone()));
                        }
                    }
                }
                RolloutItem::WorldState(_) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.world_state_replay.push(item);
                }
                RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => {
                    started_turn_ids.next();
                    // Startup can be suspended before any input or context is written.
                    // Explicit start attribution is authoritative, including unknown fields.
                    if let Some(attribution) = &event.turn_attribution
                        && attribution.turn_id == event.turn_id
                    {
                        let segment =
                            active_segment.get_or_insert_with(ActiveReplaySegment::default);
                        if turn_ids_are_compatible(
                            segment.turn_id.as_deref(),
                            Some(event.turn_id.as_str()),
                        ) {
                            segment.turn_attribution = Some(attribution.clone());
                        }
                    }
                    // `TurnStarted` is the oldest boundary of the active reverse segment.
                    if active_segment.as_ref().is_some_and(|active_segment| {
                        turn_ids_are_compatible(
                            active_segment.turn_id.as_deref(),
                            Some(event.turn_id.as_str()),
                        )
                    }) && let Some(mut active_segment) = active_segment.take()
                    {
                        active_segment
                            .turn_id
                            .get_or_insert_with(|| event.turn_id.clone());
                        active_segment.take_checkpoint_attribution(&mut checkpoint_attribution);
                        let segment_context = finalize_active_segment(
                            active_segment,
                            &mut history_checkpoint,
                            &mut previous_turn_settings,
                            &mut reference_context_item,
                            &mut world_state_replay,
                            &mut window,
                            &mut pending_rollback_turns,
                        );
                        turn_attribution = turn_attribution.or(segment_context);
                    }
                }
                RolloutItem::ResponseItem(response_item) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.counts_as_user_turn |=
                        is_user_turn_boundary(&response_item.item);
                }
                RolloutItem::InterAgentCommunication(_) => {
                    let active_segment =
                        active_segment.get_or_insert_with(ActiveReplaySegment::default);
                    active_segment.counts_as_user_turn = true;
                }
                RolloutItem::EventMsg(_)
                | RolloutItem::DynamicToolsUpdated { .. }
                | RolloutItem::SessionMeta(_)
                | RolloutItem::RealtimeItem(_)
                | RolloutItem::RetainedContext(_)
                | RolloutItem::SecurityRiskScore(_)
                | RolloutItem::TokenUsageRecord(_)
                | RolloutItem::InterAgentCommunicationMetadata { .. } => {}
            }
        }

        if let Some(mut active_segment) = active_segment.take() {
            active_segment.take_checkpoint_attribution(&mut checkpoint_attribution);
            active_segment.turn_completed |= active_segment
                .turn_id
                .as_deref()
                .or(checkpoint_turn_id)
                .is_some_and(|turn_id| late_completed_turn_ids.contains(turn_id));
            // A companion turn context only restores the context baseline. Once that turn
            // completes, its settings are newer than the compaction metadata.
            if resume_metadata.is_some() && !active_segment.turn_completed {
                active_segment.previous_turn_settings = None;
            }
            let segment_context = finalize_active_segment(
                active_segment,
                &mut history_checkpoint,
                &mut previous_turn_settings,
                &mut reference_context_item,
                &mut world_state_replay,
                &mut window,
                &mut pending_rollback_turns,
            );
            turn_attribution = turn_attribution.or(segment_context);
        }

        if pending_rollback_turns == 0 {
            turn_attribution = turn_attribution.or(checkpoint_attribution);
        }

        if previous_turn_settings.is_none() {
            previous_turn_settings =
                resume_metadata.and_then(|metadata| metadata.previous_turn_settings.clone());
        }

        let fallback_window_number = u64::try_from(
            rollout_items
                .iter()
                .filter(|item| matches!(item, RolloutItem::Compacted(_)))
                .count(),
        )
        .unwrap_or(u64::MAX);

        // Build model-visible history from the selected compaction and its newer suffix.
        let mut saw_legacy_compaction_without_replacement_history = false;
        if let Some(checkpoint) = history_checkpoint
            && let Some(items) = &checkpoint.compacted.replacement_history
        {
            history.replace_annotated(items.clone());
            history.restore_review_context(
                checkpoint.compacted.retained_context.as_ref(),
                checkpoint.compacted.guardian_history.as_ref(),
                // Keep the backup during replay; the installing session resolves its reviewer.
                /*reviewer_compaction_hash*/
                None,
            );
        }
        let rollout_suffix =
            history_checkpoint.map_or(rollout_items, |checkpoint| checkpoint.suffix);
        for item in rollout_suffix {
            match item {
                RolloutItem::RetainedContext(event) => {
                    history.record_retained_context(event);
                }
                RolloutItem::ResponseItem(response_item) => {
                    history.replay_annotated_item(response_item, truncation_policy);
                }
                RolloutItem::InterAgentCommunication(communication) => {
                    let response_item = communication.to_model_input_item();
                    history.record_items(std::iter::once(&response_item), truncation_policy);
                }
                RolloutItem::InterAgentCommunicationMetadata { .. } => {}
                RolloutItem::Compacted(compacted) => {
                    // Reverse replay already chose the newest surviving compaction. Any newer
                    // replacement compaction belongs to a rolled-back turn; replay its original
                    // items so the rollback can still find the removed user boundary.
                    if compacted.replacement_history.is_none() {
                        saw_legacy_compaction_without_replacement_history = true;
                        // Legacy rollouts without `replacement_history` should rebuild the
                        // historical TurnContext at the correct insertion point from persisted
                        // `TurnContextItem`s. These are rare enough that we currently just clear
                        // `reference_context_item`, reinject canonical context at the end of the
                        // resumed conversation, and accept the temporary out-of-distribution
                        // prompt shape.
                        // TODO(ccunningham): if we drop support for None replacement_history compaction items,
                        // we can get rid of this second loop entirely and just build `history` directly in the first loop.
                        let rebuilt =
                            compact::build_compacted_history(&history, &compacted.message);
                        let retained_context = history.retained_context().clone();
                        history.replace_annotated(rebuilt);
                        history.restore_retained_context(Some(&retained_context));
                    }
                }
                RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => {
                    history.drop_last_n_user_turns(rollback.num_turns);
                }
                RolloutItem::EventMsg(_)
                | RolloutItem::TurnContext(_)
                | RolloutItem::RealtimeItem(_)
                | RolloutItem::WorldState(_)
                | RolloutItem::SecurityRiskScore(_)
                | RolloutItem::TokenUsageRecord(_)
                | RolloutItem::DynamicToolsUpdated { .. }
                | RolloutItem::SessionMeta(_) => {}
            }
        }

        let reference_context_item = match reference_context_item {
            TurnReferenceContextItem::NeverSet | TurnReferenceContextItem::Cleared => None,
            TurnReferenceContextItem::Latest(turn_reference_context_item) => {
                Some(*turn_reference_context_item)
            }
        };
        let reference_context_item = if saw_legacy_compaction_without_replacement_history {
            None
        } else {
            reference_context_item
        };

        // Replay the collected world-state records chronologically so compaction resets and merge
        // patches keep their original meaning.
        world_state_replay.reverse();
        let mut world_state_baseline: Option<WorldStateSnapshot> = None;
        for item in world_state_replay {
            match item {
                RolloutItem::Compacted(_) => world_state_baseline = None,
                RolloutItem::WorldState(world_state) if world_state.full => {
                    world_state_baseline = Some(WorldStateSnapshot::from(&world_state.state));
                }
                RolloutItem::WorldState(world_state) => {
                    let Some(baseline) = world_state_baseline.as_mut() else {
                        tracing::warn!("ignored world-state patch without a full snapshot");
                        continue;
                    };
                    baseline.apply_merge_patch(&world_state.state);
                }
                RolloutItem::DynamicToolsUpdated { .. }
                | RolloutItem::SessionMeta(_)
                | RolloutItem::ResponseItem(_)
                | RolloutItem::InterAgentCommunication(_)
                | RolloutItem::InterAgentCommunicationMetadata { .. }
                | RolloutItem::TurnContext(_)
                | RolloutItem::RealtimeItem(_)
                | RolloutItem::TokenUsageRecord(_)
                | RolloutItem::RetainedContext(_)
                | RolloutItem::SecurityRiskScore(_)
                | RolloutItem::EventMsg(_) => {
                    unreachable!("only world-state replay items are collected")
                }
            }
        }

        let window = window.or(initial_window).unwrap_or(ReconstructedWindow {
            number: fallback_window_number,
            first_id: None,
            previous_id: None,
            id: None,
        });
        RolloutReconstruction {
            retained_context: history.retained_context().clone(),
            guardian_history: history.guardian_history_checkpoint(),
            last_started_turn_id,
            history: history.into_annotated_items(),
            rejected_audio_file_ids,
            previous_turn_settings,
            reference_context_item,
            turn_attribution,
            world_state_baseline,
            window_number: window.number,
            first_window_id: window.first_id,
            previous_window_id: window.previous_id,
            window_id: window.id,
        }
    }
}

fn reconstructed_window_from_compaction(compacted: &CompactedItem) -> Option<ReconstructedWindow> {
    Some(ReconstructedWindow {
        number: compacted.window_number?,
        first_id: compacted
            .first_window_id
            .as_ref()
            .and_then(ContextWindowId::validated),
        previous_id: compacted
            .previous_window_id
            .as_ref()
            .and_then(ContextWindowId::validated),
        id: compacted
            .window_id
            .as_ref()
            .and_then(ContextWindowId::validated),
    })
}

fn reconstructed_window_from_session_context_window(
    context_window: &SessionContextWindow,
) -> Option<ReconstructedWindow> {
    let id = context_window.window_id.validated()?;
    Some(ReconstructedWindow {
        number: 0,
        first_id: Some(id.clone()),
        previous_id: None,
        id: Some(id),
    })
}