diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 23be07996c..218e92cd1f 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -7,6 +7,7 @@ use codex_protocol::protocol::SessionSource; pub(crate) mod compression; pub(crate) mod config; pub(crate) mod list; +mod load_error; pub(crate) mod metadata; pub(crate) mod policy; pub(crate) mod recorder; @@ -59,6 +60,7 @@ pub use list::read_head_for_summary; pub use list::read_session_meta_line; pub use list::read_thread_item_from_rollout; pub use list::rollout_date_parts; +pub use load_error::LoadRolloutItemsForThreadError; pub use metadata::builder_from_items; pub use policy::is_persisted_rollout_item; pub use policy::persisted_rollout_items; diff --git a/codex-rs/rollout/src/load_error.rs b/codex-rs/rollout/src/load_error.rs new file mode 100644 index 0000000000..d4401729bf --- /dev/null +++ b/codex-rs/rollout/src/load_error.rs @@ -0,0 +1,58 @@ +use codex_protocol::ThreadId; +use std::io::Error as IoError; +use std::io::ErrorKind; + +/// Failure modes from loading a rollout for one expected thread. +#[derive(Debug)] +pub enum LoadRolloutItemsForThreadError { + /// The rollout could not be read. + Io(std::io::Error), + /// The first session metadata record belongs to another thread. + ThreadIdMismatch { + /// Thread found in the rollout. + actual_thread_id: ThreadId, + }, + /// The rollout contained no session metadata record. + MissingSessionMeta, +} + +impl LoadRolloutItemsForThreadError { + pub(crate) fn into_io_error(self) -> std::io::Error { + match self { + Self::Io(err) => err, + err @ (Self::ThreadIdMismatch { .. } | Self::MissingSessionMeta) => { + IoError::new(ErrorKind::InvalidData, err) + } + } + } +} + +impl std::fmt::Display for LoadRolloutItemsForThreadError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Io(err) => err.fmt(formatter), + Self::ThreadIdMismatch { actual_thread_id } => { + write!( + formatter, + "rollout contains history for thread {actual_thread_id}" + ) + } + Self::MissingSessionMeta => formatter.write_str("rollout contains no session metadata"), + } + } +} + +impl std::error::Error for LoadRolloutItemsForThreadError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Io(err) => Some(err), + Self::ThreadIdMismatch { .. } | Self::MissingSessionMeta => None, + } + } +} + +impl From for LoadRolloutItemsForThreadError { + fn from(err: std::io::Error) -> Self { + Self::Io(err) + } +} diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 7ba111e944..cbbc528abb 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -46,6 +46,7 @@ use super::metadata; use super::session_index::find_thread_names_by_ids; use crate::config::RolloutConfigView; use crate::default_client::originator; +use crate::load_error::LoadRolloutItemsForThreadError; use crate::state_db; use crate::state_db::StateDbHandle; use codex_git_utils::collect_git_info; @@ -867,6 +868,26 @@ impl RolloutRecorder { pub async fn load_rollout_items( path: &Path, ) -> std::io::Result<(Vec, Option, usize)> { + Self::load_rollout_items_inner(path, /*expected_thread_id*/ None) + .await + .map_err(LoadRolloutItemsForThreadError::into_io_error) + } + + /// Loads rollout items and rejects a mismatched first `SessionMeta` before reading the rest of + /// the rollout. + pub async fn load_rollout_items_for_thread( + path: &Path, + expected_thread_id: ThreadId, + ) -> Result<(Vec, usize), LoadRolloutItemsForThreadError> { + let (items, _, parse_errors) = + Self::load_rollout_items_inner(path, Some(expected_thread_id)).await?; + Ok((items, parse_errors)) + } + + async fn load_rollout_items_inner( + path: &Path, + expected_thread_id: Option, + ) -> Result<(Vec, Option, usize), LoadRolloutItemsForThreadError> { trace!("Resuming rollout from {path:?}"); let mut items: Vec = Vec::new(); let mut thread_id: Option = None; @@ -900,7 +921,15 @@ impl RolloutRecorder { if thread_id.is_none() && let RolloutItem::SessionMeta(session_meta_line) = &item { - thread_id = Some(session_meta_line.meta.id); + let actual_thread_id = session_meta_line.meta.id; + if let Some(expected_thread_id) = expected_thread_id.as_ref() + && *expected_thread_id != actual_thread_id + { + return Err(LoadRolloutItemsForThreadError::ThreadIdMismatch { + actual_thread_id, + }); + } + thread_id = Some(actual_thread_id); } items.push(item); } @@ -911,7 +940,10 @@ impl RolloutRecorder { } } if !saw_non_empty_line { - return Err(IoError::other("empty session file")); + return Err(IoError::other("empty session file").into()); + } + if expected_thread_id.is_some() && thread_id.is_none() { + return Err(LoadRolloutItemsForThreadError::MissingSessionMeta); } tracing::debug!( diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 8701578ebf..02112850ff 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -68,6 +68,33 @@ fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let actual_uuid = Uuid::from_u128(1); + let actual_thread_id = + ThreadId::from_string(&actual_uuid.to_string()).expect("valid actual thread id"); + let expected_thread_id = + ThreadId::from_string(&Uuid::from_u128(2).to_string()).expect("valid expected thread id"); + let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", actual_uuid)?; + let mut file = fs::OpenOptions::new().append(true).open(&rollout_path)?; + // Invalid UTF-8 after SessionMeta proves the mismatch returns before reading the tail. + file.write_all(&[0xff, b'\n'])?; + + let err = RolloutRecorder::load_rollout_items_for_thread(&rollout_path, expected_thread_id) + .await + .expect_err("mismatched thread id should fail"); + + let LoadRolloutItemsForThreadError::ThreadIdMismatch { + actual_thread_id: error_actual_thread_id, + } = err + else { + panic!("expected thread id mismatch, got {err}"); + }; + assert_eq!(error_actual_thread_id, actual_thread_id); + Ok(()) +} + #[tokio::test] async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/local/benchmark.rs b/codex-rs/thread-store/src/local/benchmark.rs deleted file mode 100644 index 9b446da8cc..0000000000 --- a/codex-rs/thread-store/src/local/benchmark.rs +++ /dev/null @@ -1,184 +0,0 @@ -use std::io::Write; -use std::time::Duration; -use std::time::Instant; - -use codex_protocol::ThreadId; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::ThreadMemoryMode; -use tempfile::TempDir; - -use super::LoadThreadHistoryParams; -use super::LocalThreadStore; -use super::ReadThreadParams; -use super::ResumeThreadParams; -use super::read_thread; -use super::test_support::test_config; -use super::test_support::write_session_file; -use crate::ThreadPersistenceMetadata; -use crate::ThreadStore; - -const BENCHMARK_HISTORY_ITEMS: usize = 10_000; -const BENCHMARK_RUNS: usize = 20; - -fn write_thread_history_benchmark_fixture( - root: &std::path::Path, - uuid: uuid::Uuid, -) -> std::path::PathBuf { - let timestamp = "2025-01-04T10-00-00"; - let rollout_path = write_session_file(root, timestamp, uuid).expect("write benchmark session"); - let message = "x".repeat(1_024); - let event = serde_json::json!({ - "timestamp": timestamp, - "type": "event_msg", - "payload": { - "type": "user_message", - "message": message, - "kind": "plain", - }, - }); - let mut file = std::fs::OpenOptions::new() - .append(true) - .open(&rollout_path) - .expect("open benchmark session"); - for _ in 0..BENCHMARK_HISTORY_ITEMS { - writeln!(file, "{event}").expect("append benchmark history item"); - } - rollout_path -} - -fn print_thread_history_benchmark( - label: &str, - mut durations: Vec, - rollout_bytes: u64, - read_work: &read_thread::read_work::ReadWork, -) { - durations.sort_unstable(); - let p50 = durations[durations.len() / 2]; - let p95 = durations[durations.len() * 95 / 100]; - let max = *durations.last().expect("benchmark duration"); - eprintln!( - "{BENCHMARK_HISTORY_ITEMS}-item thread history {label} benchmark: rollout_bytes={rollout_bytes} p50={p50:?} p95={p95:?} max={max:?} work={read_work:?}" - ); -} - -fn thread_metadata() -> ThreadPersistenceMetadata { - ThreadPersistenceMetadata { - cwd: Some(std::env::current_dir().expect("cwd")), - model_provider: "test-provider".to_string(), - memory_mode: ThreadMemoryMode::Enabled, - } -} - -#[tokio::test] -#[ignore = "release benchmark"] -async fn thread_history_benchmark_10_000_items() { - let home = TempDir::new().expect("temp dir"); - let uuid = uuid::Uuid::from_u128(408); - let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); - let rollout_path = write_thread_history_benchmark_fixture(home.path(), uuid); - let rollout_bytes = std::fs::metadata(&rollout_path) - .expect("benchmark rollout metadata") - .len(); - - let config = test_config(home.path()); - let runtime = codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.default_model_provider_id.clone(), - ) - .await - .expect("state db should initialize"); - let mut builder = codex_state::ThreadMetadataBuilder::new( - thread_id, - rollout_path.clone(), - chrono::Utc::now(), - SessionSource::Cli, - ); - builder.model_provider = Some(config.default_model_provider_id.clone()); - builder.cwd = home.path().to_path_buf(); - runtime - .upsert_thread(&builder.build(config.default_model_provider_id.as_str())) - .await - .expect("state db upsert should succeed"); - let sqlite_store = LocalThreadStore::new(config.clone(), Some(runtime)); - sqlite_store - .read_thread(ReadThreadParams { - thread_id, - include_archived: false, - include_history: true, - }) - .await - .expect("warm SQLite history read"); - - let mut sqlite_durations = Vec::with_capacity(BENCHMARK_RUNS); - let mut sqlite_work = None; - for _ in 0..BENCHMARK_RUNS { - let started = Instant::now(); - let (thread, read_work) = - read_thread::read_work::measure(sqlite_store.read_thread(ReadThreadParams { - thread_id, - include_archived: false, - include_history: true, - })) - .await; - assert_eq!( - thread - .expect("read SQLite-backed thread") - .history - .expect("SQLite-backed history") - .items - .len(), - BENCHMARK_HISTORY_ITEMS + 2 - ); - sqlite_durations.push(started.elapsed()); - sqlite_work = Some(read_work); - } - print_thread_history_benchmark( - "SQLite", - sqlite_durations, - rollout_bytes, - &sqlite_work.expect("SQLite read work"), - ); - - let live_store = LocalThreadStore::new(config, /*state_db*/ None); - live_store - .resume_thread(ResumeThreadParams { - thread_id, - rollout_path: Some(rollout_path), - history: None, - include_archived: true, - metadata: thread_metadata(), - }) - .await - .expect("resume benchmark thread"); - live_store - .load_history(LoadThreadHistoryParams { - thread_id, - include_archived: false, - }) - .await - .expect("warm active-writer history read"); - - let mut live_durations = Vec::with_capacity(BENCHMARK_RUNS); - let mut live_work = None; - for _ in 0..BENCHMARK_RUNS { - let started = Instant::now(); - let (history, read_work) = - read_thread::read_work::measure(live_store.load_history(LoadThreadHistoryParams { - thread_id, - include_archived: false, - })) - .await; - assert_eq!( - history.expect("load active-writer history").items.len(), - BENCHMARK_HISTORY_ITEMS + 2 - ); - live_durations.push(started.elapsed()); - live_work = Some(read_work); - } - print_thread_history_benchmark( - "active writer", - live_durations, - rollout_bytes, - &live_work.expect("active-writer read work"), - ); -} diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index bc4c1c3c9d..7d9f31688b 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1,6 +1,4 @@ mod archive_thread; -#[cfg(test)] -mod benchmark; mod create_thread; mod delete_thread; mod helpers; @@ -950,11 +948,11 @@ mod tests { .await .expect("flush live thread"); - let (history, read_work) = - read_thread::read_work::measure(store.load_history(LoadThreadHistoryParams { + let history = store + .load_history(LoadThreadHistoryParams { thread_id, include_archived: false, - })) + }) .await; let history = history.expect("load external live history"); @@ -964,13 +962,6 @@ mod tests { RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "external history item" ) })); - assert_eq!( - read_work, - read_thread::read_work::ReadWork { - summary_reads: 0, - history_reads: 1, - } - ); } #[tokio::test] diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index 8a662b7945..7e7ab8a736 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -5,6 +5,7 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_rollout::LoadRolloutItemsForThreadError; use codex_rollout::RolloutRecorder; use codex_rollout::find_archived_thread_path_by_id_str; use codex_rollout::find_thread_name_by_id; @@ -27,54 +28,6 @@ use crate::StoredThreadHistory; use crate::ThreadStoreError; use crate::ThreadStoreResult; -#[cfg(test)] -pub(crate) mod read_work { - use std::future::Future; - use std::sync::Arc; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - #[derive(Clone, Default)] - struct ReadWorkRecorder { - summary_reads: Arc, - history_reads: Arc, - } - - #[derive(Debug, PartialEq, Eq)] - pub(crate) struct ReadWork { - pub(crate) summary_reads: usize, - pub(crate) history_reads: usize, - } - - tokio::task_local! { - static READ_WORK: ReadWorkRecorder; - } - - pub(super) fn record_summary_read() { - let _ = READ_WORK.try_with(|work| { - work.summary_reads.fetch_add(1, Ordering::Relaxed); - }); - } - - pub(super) fn record_history_read() { - let _ = READ_WORK.try_with(|work| { - work.history_reads.fetch_add(1, Ordering::Relaxed); - }); - } - - pub(crate) async fn measure(future: impl Future) -> (T, ReadWork) { - let recorder = ReadWorkRecorder::default(); - let result = READ_WORK.scope(recorder.clone(), future).await; - ( - result, - ReadWork { - summary_reads: recorder.summary_reads.load(Ordering::Relaxed), - history_reads: recorder.history_reads.load(Ordering::Relaxed), - }, - ) - } -} - pub(super) async fn read_thread( store: &LocalThreadStore, params: ReadThreadParams, @@ -91,9 +44,12 @@ pub(super) async fn read_thread( let history = if params.include_history { // The full parse also validates the first SessionMeta thread ID, so do not precede it // with a separate summary read on the common SQLite path. - load_history_from_rollout_path(metadata.rollout_path.as_path(), thread_id) - .await - .ok() + match try_load_history_from_rollout_path(metadata.rollout_path.as_path(), thread_id) + .await? + { + HistoryLoadOutcome::Loaded(history) => Some(history), + HistoryLoadOutcome::MissingPath | HistoryLoadOutcome::MismatchedPath(_) => None, + } } else { None }; @@ -292,8 +248,6 @@ async fn read_thread_from_rollout_path( store: &LocalThreadStore, path: std::path::PathBuf, ) -> ThreadStoreResult { - #[cfg(test)] - read_work::record_summary_read(); let Some(item) = read_thread_item_from_rollout(path.clone()).await else { return stored_thread_from_session_meta(store, path).await; }; @@ -332,33 +286,57 @@ pub(super) async fn load_history_from_rollout_path( path: &std::path::Path, expected_thread_id: codex_protocol::ThreadId, ) -> ThreadStoreResult { - #[cfg(test)] - read_work::record_history_read(); - let Some(path) = codex_rollout::existing_rollout_path(path).await else { - return Err(ThreadStoreError::InvalidRequest { + match try_load_history_from_rollout_path(path, expected_thread_id).await? { + HistoryLoadOutcome::Loaded(history) => Ok(history), + HistoryLoadOutcome::MissingPath => Err(ThreadStoreError::InvalidRequest { message: format!( "failed to resolve rollout path `{}`: file does not exist", path.display() ), - }); - }; - let (items, thread_id, _) = RolloutRecorder::load_rollout_items(path.as_path()) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to load thread history {}: {err}", path.display()), - })?; - if thread_id != Some(expected_thread_id) { - return Err(ThreadStoreError::InvalidRequest { - message: format!( - "rollout {} does not contain history for thread {expected_thread_id}", - path.display() - ), - }); + }), + HistoryLoadOutcome::MismatchedPath(actual_thread_id) => { + Err(ThreadStoreError::InvalidRequest { + message: format!( + "rollout {} contains history for thread {actual_thread_id}, not {expected_thread_id}", + path.display() + ), + }) + } + } +} + +/// Result of loading history from a path that may have become stale in SQLite. +enum HistoryLoadOutcome { + /// History was loaded and belongs to the expected thread. + Loaded(StoredThreadHistory), + /// The path does not exist. + MissingPath, + /// The path belongs to another thread. + MismatchedPath(codex_protocol::ThreadId), +} + +async fn try_load_history_from_rollout_path( + path: &std::path::Path, + expected_thread_id: codex_protocol::ThreadId, +) -> ThreadStoreResult { + let Some(path) = codex_rollout::existing_rollout_path(path).await else { + return Ok(HistoryLoadOutcome::MissingPath); + }; + match RolloutRecorder::load_rollout_items_for_thread(path.as_path(), expected_thread_id).await { + Ok((items, _)) => Ok(HistoryLoadOutcome::Loaded(StoredThreadHistory { + thread_id: expected_thread_id, + items, + })), + Err(LoadRolloutItemsForThreadError::ThreadIdMismatch { actual_thread_id }) => { + Ok(HistoryLoadOutcome::MismatchedPath(actual_thread_id)) + } + Err( + err @ (LoadRolloutItemsForThreadError::Io(_) + | LoadRolloutItemsForThreadError::MissingSessionMeta), + ) => Err(ThreadStoreError::Internal { + message: format!("failed to load thread history {}: {err}", path.display()), + }), } - Ok(StoredThreadHistory { - thread_id: expected_thread_id, - items, - }) } async fn read_sqlite_metadata( @@ -1034,13 +1012,14 @@ mod tests { .await .expect("state db upsert should succeed"); - let (thread, read_work) = read_work::measure(store.read_thread(ReadThreadParams { - thread_id, - include_archived: false, - include_history: true, - })) - .await; - let thread = thread.expect("read thread"); + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + }) + .await + .expect("read thread"); assert_eq!(thread.thread_id, thread_id); assert_eq!(thread.rollout_path, Some(rollout_path)); @@ -1052,13 +1031,6 @@ mod tests { let history = thread.history.expect("history should load"); assert_eq!(history.thread_id, thread_id); assert_eq!(history.items.len(), 1); - assert_eq!( - read_work, - read_work::ReadWork { - summary_reads: 0, - history_reads: 1, - } - ); } #[tokio::test] @@ -1157,6 +1129,59 @@ mod tests { assert_eq!(history.items.len(), 2); } + #[tokio::test] + async fn read_thread_propagates_sqlite_rollout_read_failure() { + let home = TempDir::new().expect("temp dir"); + let external = TempDir::new().expect("external temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(223); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + write_session_file(home.path(), "2025-01-03T12-00-00", uuid) + .expect("fallback session file"); + let corrupt_path = write_session_file(external.path(), "2025-01-04T12-00-00", uuid) + .expect("sqlite session file"); + let mut corrupt_file = std::fs::OpenOptions::new() + .append(true) + .open(corrupt_path.as_path()) + .expect("open sqlite session file"); + corrupt_file + .write_all(&[0xff, b'\n']) + .expect("corrupt sqlite session file"); + + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let metadata = ThreadMetadataBuilder::new( + thread_id, + corrupt_path.clone(), + Utc::now(), + SessionSource::Cli, + ) + .build(config.default_model_provider_id.as_str()); + runtime + .upsert_thread(&metadata) + .await + .expect("state db upsert should succeed"); + + let err = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: true, + }) + .await + .expect_err("corrupt sqlite rollout should not trigger filesystem fallback"); + + let ThreadStoreError::Internal { message } = err else { + panic!("expected internal history read error, got {err}"); + }; + assert!(message.contains(corrupt_path.to_string_lossy().as_ref())); + } + #[tokio::test] async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() { let home = TempDir::new().expect("temp dir");