speed up rollout history loading

2026-11-05 20:10:33 +00:00 · 2026-06-15 21:49:54 -07:00
parent 08901fc8e1
commit 7b2a18a6c7
2 changed files with 305 additions and 84 deletions
@@ -1,7 +1,12 @@
 #![allow(warnings, clippy::all)]

+use chrono::DateTime;
+use chrono::Datelike;
+use chrono::Timelike;
+use chrono::Utc;
 use codex_utils_path as path_utils;
 use std::cmp::Reverse;
+use std::collections::BinaryHeap;
 use std::ffi::OsStr;
 use std::io;
 use std::num::NonZero;
@@ -23,7 +28,6 @@ use crate::state_db;
 use codex_file_search as file_search;
 use codex_protocol::ThreadId;
 use codex_protocol::protocol::RolloutItem;
-use codex_protocol::protocol::RolloutLine;
 use codex_protocol::protocol::SessionMetaLine;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::USER_MESSAGE_BEGIN;
@@ -928,20 +932,62 @@ pub(crate) fn parse_timestamp_uuid_from_filename(name: &str) -> Option<(OffsetDa
    // Expected: rollout-YYYY-MM-DDThh-mm-ss-<uuid>.jsonl[.zst]
    let name = compression::parse_rollout_file_name(name)?;
    let core = name.strip_prefix("rollout-")?.strip_suffix(".jsonl")?;
-
-    // Scan from the right for a '-' such that the suffix parses as a UUID.
-    let (sep_idx, uuid) = core
-        .match_indices('-')
-        .rev()
-        .find_map(|(i, _)| Uuid::parse_str(&core[i + 1..]).ok().map(|u| (i, u)))?;
-
-    let ts_str = &core[..sep_idx];
-    let format: &[FormatItem] =
-        format_description!("[year]-[month]-[day]T[hour]-[minute]-[second]");
-    let ts = PrimitiveDateTime::parse(ts_str, format).ok()?.assume_utc();
+    let uuid_start = core.len().checked_sub(36)?;
+    let (ts_with_sep, uuid_str) = core.split_at(uuid_start);
+    let ts_str = ts_with_sep.strip_suffix('-')?;
+    let uuid = Uuid::parse_str(uuid_str).ok()?;
+    let ts = parse_rollout_filename_timestamp(ts_str)?;
    Some((ts, uuid))
 }

+fn parse_rollout_filename_timestamp(ts_str: &str) -> Option<OffsetDateTime> {
+    let bytes = ts_str.as_bytes();
+    if bytes.len() != 19
+        || bytes[4] != b'-'
+        || bytes[7] != b'-'
+        || bytes[10] != b'T'
+        || bytes[13] != b'-'
+        || bytes[16] != b'-'
+    {
+        return None;
+    }
+
+    let year = parse_decimal_i32(&bytes[0..4])?;
+    let month = parse_decimal_u8(&bytes[5..7])?;
+    let day = parse_decimal_u8(&bytes[8..10])?;
+    let hour = parse_decimal_u8(&bytes[11..13])?;
+    let minute = parse_decimal_u8(&bytes[14..16])?;
+    let second = parse_decimal_u8(&bytes[17..19])?;
+    let date =
+        time::Date::from_calendar_date(year, time::Month::try_from(month).ok()?, day).ok()?;
+    let time = time::Time::from_hms(hour, minute, second).ok()?;
+    Some(time::PrimitiveDateTime::new(date, time).assume_utc())
+}
+
+fn parse_decimal_i32(bytes: &[u8]) -> Option<i32> {
+    let mut value = 0i32;
+    for &byte in bytes {
+        let digit = byte.checked_sub(b'0')?;
+        if digit > 9 {
+            return None;
+        }
+        value = value.checked_mul(10)?.checked_add(i32::from(digit))?;
+    }
+    Some(value)
+}
+
+fn parse_decimal_u8(bytes: &[u8]) -> Option<u8> {
+    let mut value = 0u8;
+    for &byte in bytes {
+        let digit = byte.checked_sub(b'0')?;
+        if digit > 9 {
+            return None;
+        }
+        value = value.checked_mul(10)?.checked_add(digit)?;
+    }
+    Some(value)
+}
+
 struct ThreadCandidate {
    path: PathBuf,
    id: Uuid,
@@ -1087,10 +1133,10 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result<HeadTai
        }
        lines_scanned += 1;

-        let parsed: Result<RolloutLine, _> = serde_json::from_str(trimmed);
-        let Ok(rollout_line) = parsed else { continue };
+        let parsed: Result<RolloutItem, _> = serde_json::from_str(trimmed);
+        let Ok(rollout_item) = parsed else { continue };

-        match rollout_line.item {
+        match rollout_item {
            RolloutItem::SessionMeta(session_meta_line) => {
                if !summary.saw_session_meta {
                    summary.source = Some(session_meta_line.meta.source.clone());
@@ -1118,9 +1164,8 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result<HeadTai
                }
            }
            RolloutItem::ResponseItem(_) | RolloutItem::InterAgentCommunication(_) => {
-                summary
-                    .created_at
-                    .get_or_insert_with(|| rollout_line.timestamp.clone());
+                // The listing path only returns files with session metadata, which carries
+                // the canonical creation timestamp.
            }
            RolloutItem::TurnContext(_) => {
                // Not included in `head`; skip.
@@ -1167,8 +1212,8 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result<Vec<serde_json::Va
        if trimmed.is_empty() {
            continue;
        }
-        if let Ok(rollout_line) = serde_json::from_str::<RolloutLine>(trimmed) {
-            match rollout_line.item {
+        if let Ok(rollout_item) = serde_json::from_str::<RolloutItem>(trimmed) {
+            match rollout_item {
                RolloutItem::SessionMeta(session_meta_line) => {
                    if let Ok(value) = serde_json::to_value(session_meta_line) {
                        head.push(value);
@@ -1229,19 +1274,32 @@ fn event_msg_preview(event: &EventMsg) -> Option<String> {
 /// Read the SessionMetaLine from the head of a rollout file for reuse by
 /// callers that need the session metadata (e.g. to derive a cwd for config).
 pub async fn read_session_meta_line(path: &Path) -> io::Result<SessionMetaLine> {
-    let head = read_head_for_summary(path).await?;
-    let Some(first) = head.first() else {
-        return Err(io::Error::other(format!(
-            "rollout at {} is empty",
-            path.display()
-        )));
-    };
-    serde_json::from_value::<SessionMetaLine>(first.clone()).map_err(|_| {
-        io::Error::other(format!(
-            "rollout at {} does not start with session metadata",
-            path.display()
-        ))
-    })
+    let mut lines = compression::open_rollout_line_reader(path).await?;
+    while let Some(line) = lines.next_line().await? {
+        let trimmed = line.trim();
+        if trimmed.is_empty() {
+            continue;
+        }
+        let Ok(rollout_item) = serde_json::from_str::<RolloutItem>(trimmed) else {
+            continue;
+        };
+        return match rollout_item {
+            RolloutItem::SessionMeta(session_meta_line) => Ok(session_meta_line),
+            RolloutItem::ResponseItem(_) | RolloutItem::InterAgentCommunication(_) => {
+                Err(io::Error::other(format!(
+                    "rollout at {} does not start with session metadata",
+                    path.display()
+                )))
+            }
+            RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => {
+                continue;
+            }
+        };
+    }
+    Err(io::Error::other(format!(
+        "rollout at {} is empty",
+        path.display()
+    )))
 }

 async fn file_modified_time(path: &Path) -> io::Result<Option<OffsetDateTime>> {
@@ -1,9 +1,13 @@
 //! Persist Codex session rollouts (.jsonl) so sessions can be replayed or inspected later.

 use std::collections::HashSet;
+use std::ffi::OsStr;
 use std::fs;
 use std::fs::File;
+use std::io::BufRead;
+use std::io::BufReader;
 use std::io::Error as IoError;
+use std::io::ErrorKind;
 use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -13,6 +17,7 @@ use chrono::SecondsFormat;
 use codex_protocol::ThreadId;
 use codex_protocol::dynamic_tools::DynamicToolSpec;
 use codex_protocol::models::BaseInstructions;
+use codex_protocol::models::ResponseItem;
 use serde_json::Value;
 use time::OffsetDateTime;
 use time::format_description::FormatItem;
@@ -55,7 +60,6 @@ use codex_protocol::protocol::InitialHistory;
 use codex_protocol::protocol::MultiAgentVersion;
 use codex_protocol::protocol::ResumedHistory;
 use codex_protocol::protocol::RolloutItem;
-use codex_protocol::protocol::RolloutLine;
 use codex_protocol::protocol::SessionMeta;
 use codex_protocol::protocol::SessionMetaLine;
 use codex_protocol::protocol::SessionSource;
@@ -63,6 +67,9 @@ use codex_protocol::protocol::ThreadSource;
 use codex_state::StateRuntime;
 use codex_utils_path as path_utils;

+const ROLLOUT_HISTORY_READ_BUFFER_SIZE: usize = 64 * 1024;
+const ROLLOUT_HISTORY_BYTES_PER_ITEM_HINT: u64 = 128;
+
 /// Writes canonical session rollout items to JSONL.
 ///
 /// Rollouts are recorded as JSONL and can be inspected with tools such as:
@@ -847,59 +854,22 @@ impl RolloutRecorder {
        path: &Path,
    ) -> std::io::Result<(Vec<RolloutItem>, Option<ThreadId>, usize)> {
        trace!("Resuming rollout from {path:?}");
-        let mut items: Vec<RolloutItem> = Vec::new();
-        let mut thread_id: Option<ThreadId> = None;
-        let mut parse_errors = 0usize;
-        let mut reader = compression::open_rollout_line_reader(path).await?;
-        let mut saw_non_empty_line = false;
-        while let Some(line) = reader.next_line().await? {
-            if line.trim().is_empty() {
-                continue;
-            }
-            saw_non_empty_line = true;
-            let mut v: Value = match serde_json::from_str(&line) {
-                Ok(v) => v,
-                Err(e) => {
-                    warn!("failed to parse line as JSON: {line:?}, error: {e}");
-                    parse_errors = parse_errors.saturating_add(1);
-                    continue;
-                }
-            };
-            if strip_legacy_ghost_snapshot_rollout_line(&mut v) {
-                trace!("skipping legacy ghost_snapshot rollout line");
-                continue;
-            }
-
-            // Parse the rollout line structure
-            match serde_json::from_value::<RolloutLine>(v.clone()) {
-                Ok(rollout_line) => {
-                    let item = rollout_line.item;
-                    // Use the FIRST SessionMeta encountered in the file as the canonical
-                    // thread id and main session information. Keep all items intact.
-                    if thread_id.is_none()
-                        && let RolloutItem::SessionMeta(session_meta_line) = &item
-                    {
-                        thread_id = Some(session_meta_line.meta.id);
-                    }
-                    items.push(item);
-                }
-                Err(e) => {
-                    trace!("failed to parse rollout line: {e}");
-                    parse_errors = parse_errors.saturating_add(1);
+        let requested_path = path.to_path_buf();
+        for _ in 0..3 {
+            let path = requested_path.clone();
+            let result = tokio::task::spawn_blocking(move || load_rollout_items_blocking(path))
+                .await
+                .map_err(IoError::other)?;
+            match result {
+                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+                    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
                }
+                other => return other,
            }
        }
-        if !saw_non_empty_line {
-            return Err(IoError::other("empty session file"));
-        }
-
-        tracing::debug!(
-            "Resumed rollout with {} items, thread ID: {:?}, parse errors: {}",
-            items.len(),
-            thread_id,
-            parse_errors,
-        );
-        Ok((items, thread_id, parse_errors))
+        tokio::task::spawn_blocking(move || load_rollout_items_blocking(requested_path))
+            .await
+            .map_err(IoError::other)?
    }

    pub async fn get_rollout_history(path: &Path) -> std::io::Result<InitialHistory> {
@@ -947,6 +917,199 @@ impl RolloutRecorder {
    }
 }

+fn load_rollout_items_blocking(
+    path: PathBuf,
+) -> std::io::Result<(Vec<RolloutItem>, Option<ThreadId>, usize)> {
+    let (file, path) = open_rollout_file_for_history(path)?;
+    let capacity_hint = file
+        .metadata()
+        .ok()
+        .map(|metadata| rollout_items_capacity_hint(metadata.len()))
+        .unwrap_or_default();
+    if is_compressed_rollout_path(path.as_path()) {
+        let decoder = zstd::stream::read::Decoder::new(file)?;
+        load_rollout_items_from_reader(
+            BufReader::with_capacity(ROLLOUT_HISTORY_READ_BUFFER_SIZE, decoder),
+            capacity_hint,
+        )
+    } else {
+        load_rollout_items_from_reader(
+            BufReader::with_capacity(ROLLOUT_HISTORY_READ_BUFFER_SIZE, file),
+            capacity_hint,
+        )
+    }
+}
+
+fn open_rollout_file_for_history(path: PathBuf) -> std::io::Result<(File, PathBuf)> {
+    let plain_path = compression::plain_rollout_path(path.as_path());
+    if let Some(file) = open_regular_file(plain_path.as_path())? {
+        return Ok((file, plain_path));
+    }
+
+    let compressed_path = compressed_rollout_path(plain_path.as_path());
+    if let Some(file) = open_regular_file(compressed_path.as_path())? {
+        return Ok((file, compressed_path));
+    }
+
+    Err(IoError::new(
+        ErrorKind::NotFound,
+        format!("rollout file not found: {}", path.display()),
+    ))
+}
+
+fn open_regular_file(path: &Path) -> std::io::Result<Option<File>> {
+    let file = match File::open(path) {
+        Ok(file) => file,
+        Err(err) if err.kind() == ErrorKind::NotFound => return Ok(None),
+        Err(err) => return Err(err),
+    };
+    if file.metadata().is_ok_and(|metadata| metadata.is_file()) {
+        return Ok(Some(file));
+    }
+    Ok(None)
+}
+
+fn compressed_rollout_path(path: &Path) -> PathBuf {
+    if is_compressed_rollout_path(path) {
+        return path.to_path_buf();
+    }
+    let mut file_name = path
+        .file_name()
+        .map(OsStr::to_os_string)
+        .unwrap_or_else(|| OsStr::new("rollout.jsonl").to_os_string());
+    file_name.push(".zst");
+    path.with_file_name(file_name)
+}
+
+fn is_compressed_rollout_path(path: &Path) -> bool {
+    path.file_name()
+        .and_then(|name| name.to_str())
+        .is_some_and(|name| name.ends_with(".jsonl.zst"))
+}
+
+fn load_rollout_items_from_reader<R: BufRead>(
+    mut reader: R,
+    capacity_hint: usize,
+) -> std::io::Result<(Vec<RolloutItem>, Option<ThreadId>, usize)> {
+    let mut items: Vec<RolloutItem> = Vec::with_capacity(capacity_hint);
+    let mut thread_id: Option<ThreadId> = None;
+    let mut parse_errors = 0usize;
+    let mut saw_non_empty_line = false;
+    let mut line = Vec::new();
+    loop {
+        line.clear();
+        if reader.read_until(b'\n', &mut line)? == 0 {
+            break;
+        }
+        let trimmed_line = trim_ascii_whitespace(line.as_slice());
+        if trimmed_line.is_empty() {
+            continue;
+        }
+        saw_non_empty_line = true;
+        // Parse the rollout line structure
+        match parse_rollout_item_preserving_legacy_ghost_snapshot_filter(trimmed_line) {
+            Ok(None) => continue,
+            Ok(Some(item)) => {
+                // Use the FIRST SessionMeta encountered in the file as the canonical
+                // thread id and main session information. Keep all items intact.
+                if thread_id.is_none()
+                    && let RolloutItem::SessionMeta(session_meta_line) = &item
+                {
+                    thread_id = Some(session_meta_line.meta.id);
+                }
+                items.push(item);
+            }
+            Err(e) => {
+                if matches!(
+                    e.classify(),
+                    serde_json::error::Category::Eof | serde_json::error::Category::Syntax
+                ) {
+                    warn!("failed to parse line as JSON: {line:?}, error: {e}");
+                } else {
+                    trace!("failed to parse rollout line: {e}");
+                }
+                parse_errors = parse_errors.saturating_add(1);
+            }
+        }
+    }
+    if !saw_non_empty_line {
+        return Err(IoError::other("empty session file"));
+    }
+
+    tracing::debug!(
+        "Resumed rollout with {} items, thread ID: {:?}, parse errors: {}",
+        items.len(),
+        thread_id,
+        parse_errors,
+    );
+    Ok((items, thread_id, parse_errors))
+}
+
+fn rollout_items_capacity_hint(file_size_bytes: u64) -> usize {
+    usize::try_from(file_size_bytes / ROLLOUT_HISTORY_BYTES_PER_ITEM_HINT)
+        .unwrap_or(usize::MAX)
+        .saturating_add(1)
+        .min(16_384)
+}
+
+fn trim_ascii_whitespace(bytes: &[u8]) -> &[u8] {
+    let mut start = 0usize;
+    let mut end = bytes.len();
+    while start < end && bytes[start].is_ascii_whitespace() {
+        start += 1;
+    }
+    while end > start && bytes[end - 1].is_ascii_whitespace() {
+        end -= 1;
+    }
+    &bytes[start..end]
+}
+
+fn parse_rollout_item_preserving_legacy_ghost_snapshot_filter(
+    line: &[u8],
+) -> Result<Option<RolloutItem>, serde_json::Error> {
+    match serde_json::from_slice::<RolloutItem>(line) {
+        Ok(item) => {
+            if rollout_item_needs_legacy_ghost_snapshot_filter(&item)
+                && contains_bytes(line, b"\"ghost_snapshot\"")
+            {
+                return parse_rollout_item_after_legacy_ghost_snapshot_strip(line);
+            }
+            Ok(Some(item))
+        }
+        Err(err) => {
+            if contains_bytes(line, b"\"ghost_snapshot\"") {
+                parse_rollout_item_after_legacy_ghost_snapshot_strip(line)
+            } else {
+                Err(err)
+            }
+        }
+    }
+}
+
+fn rollout_item_needs_legacy_ghost_snapshot_filter(item: &RolloutItem) -> bool {
+    matches!(
+        item,
+        RolloutItem::ResponseItem(ResponseItem::Other) | RolloutItem::Compacted(_)
+    )
+}
+
+fn parse_rollout_item_after_legacy_ghost_snapshot_strip(
+    line: &[u8],
+) -> Result<Option<RolloutItem>, serde_json::Error> {
+    let mut v: Value = serde_json::from_slice(line)?;
+    if strip_legacy_ghost_snapshot_rollout_line(&mut v) {
+        trace!("skipping legacy ghost_snapshot rollout line");
+        return Ok(None);
+    }
+    serde_json::from_value::<RolloutItem>(v).map(Some)
+}
+
+fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool {
+    haystack
+        .windows(needle.len())
+        .any(|window| window == needle)
+}
+
 fn strip_legacy_ghost_snapshot_rollout_line(value: &mut Value) -> bool {
    match value.get("type").and_then(Value::as_str) {
        Some("response_item") => value