Load paginated model context across rollout lineages (#34621)

## What changed

- Resolve the full rollout lineage when loading model context for paginated threads.
- Reverse-scan each lineage segment up to its recorded byte boundary, stopping at segment metadata and preserving the requested thread's canonical session metadata.
- Add a bounded reverse-scanner constructor for scanning a frozen JSONL prefix.

## Testing

- Cover frozen-prefix reverse scans and nested lineages whose root rollout has been archived.

GitOrigin-RevId: 8b97da43a263b7bd09cb7883450bb19eba5f195f

codex-rs/rollout/src/reverse_jsonl_scanner.rs

@@ -31,9 +31,24 @@ where
 {
     pub fn new(mut reader: R) -> io::Result<Self> {
         let next_chunk_end = reader.seek(SeekFrom::End(0))?;
+        Self::new_at(reader, next_chunk_end)
+    }
+
+    /// Creates a reverse scanner whose logical end is the given byte offset.
+    ///
+    /// This lets callers scan a frozen JSONL prefix without reading records appended after that
+    /// prefix was captured.
+    pub fn new_at(mut reader: R, end_byte_offset: u64) -> io::Result<Self> {
+        let file_len = reader.seek(SeekFrom::End(0))?;
+        if end_byte_offset > file_len {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "reverse JSONL scan end is past the file",
+            ));
+        }
         Ok(Self {
             reader,
-            next_chunk_end,
+            next_chunk_end: end_byte_offset,
             chunk_position: 0,
             chunk: vec![0; READ_CHUNK_SIZE],
             record_reversed: Vec::new(),

codex-rs/rollout/src/reverse_jsonl_scanner_tests.rs

@@ -78,6 +78,18 @@ fn accepts_valid_json_at_eof() -> std::io::Result<()> {
     )
 }
 
+#[test]
+fn scans_from_a_frozen_prefix_end() -> std::io::Result<()> {
+    let prefix = b"{\"value\":\"first\"}\n{\"value\":\"second\"}\n";
+    let mut input = prefix.to_vec();
+    input.extend_from_slice(b"{\"value\":\"later\"}\n");
+
+    assert_records(
+        &mut ReverseJsonlScanner::new_at(Cursor::new(input), prefix.len() as u64)?,
+        &["second", "first"],
+    )
+}
+
 #[test]
 fn rejects_invalid_json_at_eof_and_continues_scanning() -> std::io::Result<()> {
     let input = b"{\"value\":\"first\"}\n{\"value\":";

codex-rs/thread-store/src/local/model_context.rs

@@ -1,7 +1,5 @@
 use std::fs::File;
 use std::io;
-use std::path::Path;
-use std::path::PathBuf;
 
 use codex_protocol::protocol::RolloutItem;
 use codex_protocol::protocol::RolloutLine;
@@ -14,6 +12,7 @@ use codex_rollout::ScanOutcome;
 
 use super::LocalThreadStore;
 use super::read_thread;
+use super::rollout_lineage::RolloutLineage;
 use crate::LoadThreadHistoryParams;
 use crate::StoredModelContext;
 use crate::ThreadStoreError;
@@ -64,7 +63,8 @@ pub(super) async fn load_latest_model_context(
             .and_then(|file_name| file_name.to_str())
             .is_some_and(|file_name| file_name.ends_with(".jsonl.zst"))
     {
-        scan_model_context_from_end(path, session_meta).await?
+        let lineage = store.resolve_rollout_lineage(params.thread_id).await?;
+        scan_model_context_from_lineage(lineage, session_meta).await?
     } else {
         read_thread::load_history_items(path.as_path()).await?
     };
@@ -75,13 +75,12 @@ pub(super) async fn load_latest_model_context(
     })
 }
 
-async fn scan_model_context_from_end(
-    path: PathBuf,
+async fn scan_model_context_from_lineage(
+    lineage: RolloutLineage,
     session_meta: SessionMetaLine,
 ) -> ThreadStoreResult<Vec<RolloutItem>> {
-    let path_for_scan = path.clone();
     let scan = tokio::task::spawn_blocking(move || {
-        scan_model_context_from_end_blocking(&path_for_scan, session_meta)
+        scan_model_context_from_lineage_blocking(&lineage, session_meta)
     })
     .await
     .map_err(|err| ThreadStoreError::Internal {
@@ -89,33 +88,43 @@ async fn scan_model_context_from_end(
     })?;
     match scan {
         Ok(items) => Ok(items),
-        Err(err) if err.kind() == io::ErrorKind::NotFound => {
-            // Compression can replace the resolved plain rollout with its compressed sibling
-            // before the blocking reverse scanner opens it. The forward loader re-resolves that
-            // representation transition and already supports compressed rollouts.
-            read_thread::load_history_items(path.as_path()).await
-        }
         Err(err) => Err(ThreadStoreError::Internal {
-            message: format!("failed to scan model context {}: {err}", path.display()),
+            message: format!("failed to scan paginated model context lineage: {err}"),
         }),
     }
 }
 
-fn scan_model_context_from_end_blocking(
-    path: &Path,
+fn scan_model_context_from_lineage_blocking(
+    lineage: &RolloutLineage,
     session_meta: SessionMetaLine,
 ) -> io::Result<Vec<RolloutItem>> {
     let mut scan = ModelContextScan::default();
-    let mut scanner = ReverseJsonlScanner::new(File::open(path)?)?;
-    while let Some(outcome) = scanner.scan_next::<RolloutLine>()? {
-        let ScanOutcome::Parsed(line) = outcome else {
-            continue;
+    'segments: for segment in lineage.segments().iter().rev() {
+        let file = File::open(segment.rollout_path.as_path())?;
+        let mut scanner = match segment.end.map(|end| end.end_byte_offset) {
+            Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?,
+            None => ReverseJsonlScanner::new(file)?,
         };
-        match scan.push(line.item) {
-            ModelContextScanProgress::Continue => {}
-            ModelContextScanProgress::Complete => break,
+        while let Some(outcome) = scanner.scan_next::<RolloutLine>()? {
+            let ScanOutcome::Parsed(line) = outcome else {
+                continue;
+            };
+            // Each physical segment contributes only its local delta. Its head metadata is
+            // replaced with the requested thread's canonical SessionMeta after replay.
+            if matches!(&line.item, RolloutItem::SessionMeta(_)) {
+                break;
+            }
+            match scan.push(line.item) {
+                ModelContextScanProgress::Continue => {}
+                ModelContextScanProgress::Complete => break 'segments,
+            }
         }
     }
 
-    Ok(scan.finish(session_meta))
+    let canonical_meta = session_meta.clone();
+    let mut items = scan.finish(session_meta);
+    if !matches!(items.first(), Some(RolloutItem::SessionMeta(_))) {
+        items.insert(0, RolloutItem::SessionMeta(canonical_meta));
+    }
+    Ok(items)
 }

codex-rs/thread-store/src/local/model_context_tests.rs

@@ -3,6 +3,7 @@ use std::io::Write;
 use std::path::Path;
 use std::path::PathBuf;
 
+use codex_protocol::ThreadId;
 use codex_protocol::config_types::ReasoningSummary;
 use codex_protocol::items::TurnItem;
 use codex_protocol::items::UserMessageItem;
@@ -12,6 +13,7 @@ use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::CompactedItem;
 use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::HistoryPosition;
 use codex_protocol::protocol::ItemCompletedEvent;
 use codex_protocol::protocol::RolloutItem;
 use codex_protocol::protocol::RolloutLine;
@@ -145,7 +147,7 @@ async fn returns_scanned_full_history_for_unsupported_compaction() {
         ],
     );
 
-    assert_reverse_scan_matches_full_history(path.as_path()).await;
+    assert_reverse_scan_matches_full_history(home.path(), path.as_path()).await;
 }
 
 #[tokio::test]
@@ -165,7 +167,7 @@ async fn returns_scanned_full_history_at_bof_without_checkpoint() {
         ],
     );
 
-    assert_reverse_scan_matches_full_history(path.as_path()).await;
+    assert_reverse_scan_matches_full_history(home.path(), path.as_path()).await;
 }
 
 #[tokio::test]
@@ -245,6 +247,112 @@ async fn ignores_contextual_user_messages_when_selecting_turn_context() {
     }));
 }
 
+#[tokio::test]
+async fn replays_nested_archived_lineage_from_frozen_prefix() {
+    let home = TempDir::new().expect("temp dir");
+    let root_uuid = Uuid::from_u128(/*v*/ 2001);
+    let root_id = ThreadId::from_string(&root_uuid.to_string()).expect("root id");
+    let root_path = write_ordinaled_paginated_rollout(
+        home.path(),
+        "2025-01-03T13-01-00",
+        root_uuid,
+        [
+            user_message("root before checkpoint"),
+            compacted("root checkpoint", Some(Vec::new())),
+            turn_started("root-excluded"),
+            user_message("root after cutoff"),
+        ],
+    );
+    let archived_root = home
+        .path()
+        .join("archived_sessions")
+        .join(root_path.file_name().expect("root filename"));
+    std::fs::create_dir_all(archived_root.parent().expect("archive parent"))
+        .expect("create archive directory");
+    std::fs::rename(root_path, &archived_root).expect("archive root rollout");
+
+    let middle_uuid = Uuid::from_u128(/*v*/ 2002);
+    let middle_id = ThreadId::from_string(&middle_uuid.to_string()).expect("middle id");
+    let middle_path = write_ordinaled_paginated_rollout(
+        home.path(),
+        "2025-01-03T13-01-01",
+        middle_uuid,
+        [
+            turn_started("middle-turn"),
+            user_message("middle inherited"),
+            completed_user_message("middle-turn", "middle inherited"),
+            turn_context(home.path(), "middle-turn"),
+            turn_complete("middle-turn"),
+        ],
+    );
+    set_history_base(
+        middle_path.as_path(),
+        history_position(
+            archived_root.as_path(),
+            root_id,
+            /*end_ordinal_exclusive*/ 3,
+        ),
+    );
+
+    let child_uuid = Uuid::from_u128(/*v*/ 2003);
+    let child_id = ThreadId::from_string(&child_uuid.to_string()).expect("child id");
+    let child_path = write_ordinaled_paginated_rollout(
+        home.path(),
+        "2025-01-03T13-01-02",
+        child_uuid,
+        [
+            turn_started("child-turn"),
+            user_message("child local"),
+            completed_user_message("child-turn", "child local"),
+            turn_context(home.path(), "child-turn"),
+            turn_complete("child-turn"),
+        ],
+    );
+    set_history_base(
+        child_path.as_path(),
+        history_position(
+            middle_path.as_path(),
+            middle_id,
+            /*end_ordinal_exclusive*/ 6,
+        ),
+    );
+    let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
+
+    let context = store
+        .load_latest_model_context(LoadThreadHistoryParams {
+            thread_id: child_id,
+            include_archived: false,
+        })
+        .await
+        .expect("load lineage model context");
+
+    assert!(matches!(
+        context.items.first(),
+        Some(RolloutItem::SessionMeta(meta)) if meta.meta.id == child_id
+    ));
+    let child_meta = codex_rollout::read_session_meta_line(child_path.as_path())
+        .await
+        .expect("read child metadata");
+    let expected = vec![
+        RolloutItem::SessionMeta(child_meta),
+        compacted("root checkpoint", Some(Vec::new())),
+        turn_started("middle-turn"),
+        user_message("middle inherited"),
+        completed_user_message("middle-turn", "middle inherited"),
+        turn_context(home.path(), "middle-turn"),
+        turn_complete("middle-turn"),
+        turn_started("child-turn"),
+        user_message("child local"),
+        completed_user_message("child-turn", "child local"),
+        turn_context(home.path(), "child-turn"),
+        turn_complete("child-turn"),
+    ];
+    assert_eq!(
+        serde_json::to_value(context.items).expect("serialize context"),
+        serde_json::to_value(expected).expect("serialize expected context")
+    );
+}
+
 fn write_paginated_rollout<const N: usize>(
     home: &Path,
     timestamp: &str,
@@ -258,12 +366,90 @@ fn write_paginated_rollout<const N: usize>(
     path
 }
 
-async fn assert_reverse_scan_matches_full_history(path: &Path) {
+fn write_ordinaled_paginated_rollout<const N: usize>(
+    home: &Path,
+    timestamp: &str,
+    uuid: Uuid,
+    items: [RolloutItem; N],
+) -> PathBuf {
+    let path =
+        write_session_file_with_history_mode(home, timestamp, uuid, ThreadHistoryMode::Paginated)
+            .expect("write session file");
+    let mut file = OpenOptions::new()
+        .append(true)
+        .open(path.as_path())
+        .expect("open session file");
+    for (index, item) in items.into_iter().enumerate() {
+        let line = RolloutLine {
+            timestamp: "2025-01-03T13:00:01Z".to_string(),
+            ordinal: Some(u64::try_from(index).expect("fixture index fits u64") + 1),
+            item,
+        };
+        writeln!(
+            file,
+            "{}",
+            serde_json::to_string(&line).expect("serialize line")
+        )
+        .expect("append rollout line");
+    }
+    path
+}
+
+fn set_history_base(path: &Path, history_base: HistoryPosition) {
+    let contents = std::fs::read_to_string(path).expect("read rollout");
+    let mut lines = contents.lines();
+    let mut head: serde_json::Value =
+        serde_json::from_str(lines.next().expect("session meta line")).expect("parse head");
+    head["payload"]["history_base"] =
+        serde_json::to_value(history_base).expect("serialize history base");
+    let mut updated = serde_json::to_string(&head).expect("serialize head");
+    for line in lines {
+        updated.push('\n');
+        updated.push_str(line);
+    }
+    updated.push('\n');
+    std::fs::write(path, updated).expect("write history base");
+}
+
+fn history_position(
+    path: &Path,
+    thread_id: ThreadId,
+    end_ordinal_exclusive: u64,
+) -> HistoryPosition {
+    HistoryPosition {
+        thread_id,
+        end_ordinal_exclusive,
+        end_byte_offset: rollout_end_byte_offset(path, end_ordinal_exclusive),
+    }
+}
+
+fn rollout_end_byte_offset(path: &Path, end_ordinal_exclusive: u64) -> u64 {
+    let contents = std::fs::read(path).expect("read rollout");
+    let mut byte_offset = 0_u64;
+    for line in contents.split_inclusive(|byte| *byte == b'\n') {
+        let parsed: RolloutLine =
+            serde_json::from_slice(line).expect("parse rollout line for byte offset");
+        if parsed.ordinal == Some(end_ordinal_exclusive) {
+            return byte_offset;
+        }
+        byte_offset += u64::try_from(line.len()).expect("line length fits u64");
+    }
+    byte_offset
+}
+
+async fn assert_reverse_scan_matches_full_history(home: &Path, path: &Path) {
     let session_meta = codex_rollout::read_session_meta_line(path)
         .await
         .expect("read session metadata");
-    let items =
-        scan_model_context_from_end_blocking(path, session_meta).expect("scan model context");
+    let store = LocalThreadStore::new(test_config(home), /*state_db*/ None);
+    let items = store
+        .load_latest_model_context(LoadThreadHistoryParams {
+            thread_id: session_meta.meta.id,
+            include_archived: false,
+        })
+        .await
+        .expect("scan model context")
+        .items;
     let full_items = read_thread::load_history_items(path)
         .await
         .expect("load full history");