From 7e085334ad09c39267590b7401d1de161ac25dd8 Mon Sep 17 00:00:00 2001 From: anais Date: Wed, 17 Jun 2026 23:04:55 -0700 Subject: [PATCH] fix copy on write fork persistence --- .../request_processors/thread_processor.rs | 74 +++-- .../thread_processor_tests.rs | 8 +- codex-rs/core/src/session/mod.rs | 33 +- codex-rs/core/src/session/session.rs | 28 +- codex-rs/core/src/thread_manager.rs | 16 +- codex-rs/rollout/src/lib.rs | 1 - codex-rs/rollout/src/list.rs | 113 +++---- codex-rs/rollout/src/recorder.rs | 287 +++++++++++------- codex-rs/thread-store/src/live_thread.rs | 18 +- .../thread-store/src/local/read_thread.rs | 69 ----- .../thread-store/src/thread_metadata_sync.rs | 13 +- codex-rs/thread-store/src/types.rs | 6 +- 12 files changed, 366 insertions(+), 300 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 848afd1781..ae58cde722 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -3311,10 +3311,10 @@ impl ThreadRequestProcessor { ) .await? { - Some(items) => { + Some(compacted_history) => { stored_thread.history = Some(codex_thread_store::StoredThreadHistory { thread_id: stored_thread.thread_id, - items, + items: compacted_history.items, }); ResumeHistoryExtent::CompactedCheckpoint } @@ -3391,14 +3391,15 @@ impl ThreadRequestProcessor { &self, thread_id: &str, path: Option<&PathBuf>, - ) -> Result, JSONRPCErrorError> { + ) -> Result, JSONRPCErrorError> + { let mut stored_thread = self .read_stored_thread_for_resume(thread_id, path, /*include_history*/ false) .await?; let Some(rollout_path) = stored_thread.rollout_path.clone() else { return Ok(None); }; - let Some(items) = try_read_compacted_resume_rollout_items_from_rollout_path( + let Some(compacted_history) = try_read_compacted_resume_rollout_items_from_rollout_path( rollout_path.as_path(), stored_thread.thread_id, ) @@ -3408,9 +3409,9 @@ impl ThreadRequestProcessor { }; stored_thread.history = Some(codex_thread_store::StoredThreadHistory { thread_id: stored_thread.thread_id, - items, + items: compacted_history.items, }); - Ok(Some(stored_thread)) + Ok(Some((stored_thread, compacted_history.parent_ref))) } async fn load_full_resume_history_items( @@ -3623,21 +3624,13 @@ impl ThreadRequestProcessor { )); } let optimized_source = if !include_turns && !ephemeral { - match self - .read_stored_thread_for_resume_with_compacted_history(&thread_id, path.as_ref()) + self.read_stored_thread_for_resume_with_compacted_history(&thread_id, path.as_ref()) .await? - { - Some(stored_thread) => stored_thread - .rollout_path - .clone() - .map(|rollout_path| (stored_thread, rollout_path)), - None => None, - } } else { None }; let (mut source_thread, initial_rollout_copy) = match optimized_source { - Some((source_thread, rollout_path)) => (source_thread, Some(rollout_path)), + Some((source_thread, parent_ref)) => (source_thread, Some(parent_ref)), None => ( self.read_stored_thread_for_resume( &thread_id, @@ -4140,6 +4133,11 @@ struct RecentRolloutItems { items: Vec, } +struct CompactedResumeRollout { + items: Vec, + parent_ref: codex_rollout::ForkParentRolloutRef, +} + fn thread_resume_initial_page_is_recent_not_loaded( params: &ThreadResumeInitialTurnsPageParams, ) -> bool { @@ -4206,20 +4204,32 @@ fn rollout_line_has_fork_parent_rollout_ref(line: &str) -> bool { async fn try_read_compacted_resume_rollout_items_from_rollout_path( rollout_path: &Path, expected_thread_id: ThreadId, -) -> Result>, JSONRPCErrorError> { +) -> Result, JSONRPCErrorError> { let Some(existing_path) = codex_rollout::existing_rollout_path(rollout_path).await else { return Ok(None); }; if !is_plain_jsonl_rollout_path(existing_path.as_path()) { return Ok(None); } + let Ok(metadata) = tokio::fs::metadata(existing_path.as_path()).await else { + return Ok(None); + }; + if !metadata.is_file() || metadata.len() == 0 { + return Ok(None); + } + let byte_len = metadata.len(); + let parent_ref = codex_rollout::ForkParentRolloutRef { + path: existing_path.clone(), + byte_len, + }; let items = tokio::task::spawn_blocking(move || { - read_compacted_resume_rollout_items(existing_path.as_path(), expected_thread_id) + read_compacted_resume_rollout_items(existing_path.as_path(), expected_thread_id, byte_len) }) .await .map_err(|err| internal_error(format!("failed to join compacted rollout read: {err}")))?; match items { - Ok(items) => Ok(items), + Ok(Some(items)) => Ok(Some(CompactedResumeRollout { items, parent_ref })), + Ok(None) => Ok(None), Err(_) => Ok(None), } } @@ -4227,12 +4237,14 @@ async fn try_read_compacted_resume_rollout_items_from_rollout_path( fn read_compacted_resume_rollout_items( path: &Path, expected_thread_id: ThreadId, + byte_len: u64, ) -> std::io::Result>> { - let Some(mut prefix_items) = read_compacted_resume_prefix_items(path, expected_thread_id)? + let Some(mut prefix_items) = + read_compacted_resume_prefix_items(path, expected_thread_id, byte_len)? else { return Ok(None); }; - let Some(mut suffix_items) = read_compacted_resume_suffix_items(path)? else { + let Some(mut suffix_items) = read_compacted_resume_suffix_items(path, byte_len)? else { return Ok(None); }; @@ -4245,11 +4257,13 @@ fn read_compacted_resume_rollout_items( fn read_compacted_resume_prefix_items( path: &Path, expected_thread_id: ThreadId, + byte_len: u64, ) -> std::io::Result>> { use std::io::BufRead; + use std::io::Read; let file = std::fs::File::open(path)?; - let reader = std::io::BufReader::new(file); + let reader = std::io::BufReader::new(file.take(byte_len)); let mut prefix_items = Vec::with_capacity(3); let mut saw_session_meta = false; let mut saw_user_event = false; @@ -4301,12 +4315,15 @@ fn read_compacted_resume_prefix_items( Ok(saw_session_meta.then_some(prefix_items)) } -fn read_compacted_resume_suffix_items(path: &Path) -> std::io::Result>> { +fn read_compacted_resume_suffix_items( + path: &Path, + byte_len: u64, +) -> std::io::Result>> { let mut suffix_items = Vec::new(); let mut found_checkpoint = false; let mut should_fallback = false; - read_rollout_lines_from_end(path, |line| { + read_rollout_lines_from_end(path, byte_len, |line| { let trimmed = line.trim(); if trimmed.is_empty() { return Ok(true); @@ -4339,7 +4356,7 @@ fn read_compacted_resume_suffix_items(path: &Path) -> std::io::Result(path: &Path, mut visit: F) -> std::io::Result<()> +fn read_rollout_lines_from_end(path: &Path, byte_len: u64, mut visit: F) -> std::io::Result<()> where F: FnMut(&str) -> std::io::Result, { @@ -4350,7 +4367,12 @@ where const CHUNK_SIZE: u64 = 64 * 1024; let mut file = std::fs::File::open(path)?; - let mut position = file.metadata()?.len(); + if file.metadata()?.len() < byte_len { + return Err(IoError::other(format!( + "rollout ended before snapshot boundary {byte_len}" + ))); + } + let mut position = byte_len; let mut pending = Vec::new(); while position > 0 { diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 5179af11f5..00d2ea523b 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -431,7 +431,8 @@ mod thread_processor_behavior_tests { ])?; std::fs::write(&path, contents)?; - let items = read_compacted_resume_rollout_items(path.as_path(), thread_id)? + let byte_len = std::fs::metadata(path.as_path())?.len(); + let items = read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)? .expect("compacted resume items"); assert_eq!(items.len(), 5); @@ -474,7 +475,10 @@ mod thread_processor_behavior_tests { contents.push_str("{not json}\n"); std::fs::write(&path, contents)?; - assert!(read_compacted_resume_rollout_items(path.as_path(), thread_id)?.is_none()); + let byte_len = std::fs::metadata(path.as_path())?.len(); + assert!( + read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)?.is_none() + ); Ok(()) } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index bffd64d0cf..b3ed8619de 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -409,6 +409,12 @@ pub struct CodexSpawnOk { pub thread_id: ThreadId, } +#[derive(Clone)] +pub(crate) struct InitialRolloutCopy { + pub(crate) parent_ref: codex_rollout::ForkParentRolloutRef, + pub(crate) source_history_item_count: usize, +} + pub(crate) struct CodexSpawnArgs { pub(crate) config: Config, pub(crate) user_instructions: LoadedUserInstructions, @@ -423,7 +429,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, pub(crate) forked_from_thread_id: Option, - pub(crate) initial_rollout_copy: Option, + pub(crate) initial_rollout_copy: Option, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, pub(crate) agent_control: AgentControl, @@ -1262,14 +1268,18 @@ impl Session { } async fn record_initial_history(&self, conversation_history: InitialHistory) { - let (is_subagent, skip_fork_history_persistence) = { + let (is_subagent, fork_source_history_item_count) = { let state = self.state.lock().await; ( state .session_configuration .session_source .is_non_root_agent(), - state.session_configuration.initial_rollout_copy.is_some(), + state + .session_configuration + .initial_rollout_copy + .as_ref() + .map(|copy| copy.source_history_item_count), ) }; let has_prior_user_turns = initial_history_has_prior_user_turns(&conversation_history); @@ -1337,8 +1347,10 @@ impl Session { } // If persisting, persist all rollout items as-is (the store filters). - if skip_fork_history_persistence { - let suffix = Self::fork_rollout_copy_suffix(&rollout_items); + if let Some(source_history_item_count) = fork_source_history_item_count { + let suffix = rollout_items + .get(source_history_item_count..) + .unwrap_or_default(); if !suffix.is_empty() { self.persist_rollout_items(suffix).await; } @@ -1357,17 +1369,6 @@ impl Session { } } - fn fork_rollout_copy_suffix(rollout_items: &[RolloutItem]) -> &[RolloutItem] { - let Some(RolloutItem::EventMsg(EventMsg::TurnAborted(event))) = rollout_items.last() else { - return &[]; - }; - if event.reason == TurnAbortReason::Interrupted { - &rollout_items[rollout_items.len() - 1..] - } else { - &[] - } - } - #[instrument( level = "trace", skip_all, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 922af9c6d0..245c72c6e3 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -102,8 +102,8 @@ pub(crate) struct SessionConfiguration { pub(super) session_source: SessionSource, /// Immediate history source copied into this thread, when this thread was forked. pub(super) forked_from_thread_id: Option, - /// Optional rollout JSONL copied into this thread after its own session metadata. - pub(super) initial_rollout_copy: Option, + /// Optional immutable parent rollout snapshot referenced by this thread. + pub(super) initial_rollout_copy: Option, /// Immediate control/spawn parent for this thread, when it has one. pub(super) parent_thread_id: Option, /// Optional analytics source classification for this thread. @@ -541,6 +541,16 @@ impl Session { // - initialize thread persistence with new or resumed session info // - perform default shell discovery // - load history metadata (skipped for subagents) + let initial_rollout_copy_history = + session_configuration + .initial_rollout_copy + .as_ref() + .and(match &initial_history { + InitialHistory::Forked(items) => Some(items.as_slice()), + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Resumed(_) => { + None + } + }); let thread_persistence_fut = async { if config.ephemeral { Ok::<_, anyhow::Error>(None) @@ -561,7 +571,8 @@ impl Session { multi_agent_version: initial_multi_agent_version, initial_rollout_copy: session_configuration .initial_rollout_copy - .clone(), + .as_ref() + .map(|copy| copy.parent_ref.clone()), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -572,7 +583,16 @@ impl Session { }, }, }; - LiveThread::create(Arc::clone(&thread_store), params).await? + if let Some(observed_history) = initial_rollout_copy_history { + LiveThread::create_with_observed_history( + Arc::clone(&thread_store), + params, + observed_history, + ) + .await? + } else { + LiveThread::create(Arc::clone(&thread_store), params).await? + } } InitialHistory::Resumed(resumed_history) => { let local_resume_with_path = resumed_history.rollout_path.is_some() diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 131cb48cd9..7b5489b7a4 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -13,6 +13,7 @@ use crate::session::Codex; use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; use crate::session::INITIAL_SUBMIT_ID; +use crate::session::InitialRolloutCopy; use crate::session::resolve_multi_agent_version; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; @@ -55,6 +56,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; +use codex_rollout::ForkParentRolloutRef; use codex_rollout::state_db::StateDbHandle; use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::InMemoryThreadStore; @@ -967,7 +969,7 @@ impl ThreadManager { thread_source: Option, parent_trace: Option, supports_openai_form_elicitation: bool, - initial_rollout_copy: PathBuf, + initial_rollout_copy: ForkParentRolloutRef, ) -> CodexResult where S: Into, @@ -992,7 +994,7 @@ impl ThreadManager { thread_source: Option, parent_trace: Option, supports_openai_form_elicitation: bool, - initial_rollout_copy: Option, + initial_rollout_copy: Option, ) -> CodexResult { // `forked_from_id()` describes this history's existing lineage. When // forking a resumed thread, the child copies the resumed thread itself. @@ -1018,6 +1020,14 @@ impl ThreadManager { &config, ) .await; + let initial_rollout_copy = initial_rollout_copy.map(|parent_ref| InitialRolloutCopy { + parent_ref, + source_history_item_count: match &history { + InitialHistory::New | InitialHistory::Cleared => 0, + InitialHistory::Resumed(resumed) => resumed.history.len(), + InitialHistory::Forked(items) => items.len(), + }, + }); let interrupted_marker = InterruptedTurnHistoryMarker::from_config_and_version(&config, multi_agent_version); let history = fork_history_from_snapshot(snapshot, history, interrupted_marker); @@ -1508,7 +1518,7 @@ impl ThreadManagerState { thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, - initial_rollout_copy: Option, + initial_rollout_copy: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index d1b067233b..4c669bbbcc 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -67,7 +67,6 @@ pub use recorder::ForkParentRolloutRef; pub use recorder::RolloutRecorder; pub use recorder::RolloutRecorderParams; pub use recorder::append_rollout_item_to_path; -pub use recorder::read_fork_parent_rollout_ref; pub use search::first_rollout_content_match_snippet; pub use search::search_rollout_matches; pub use search::search_rollout_paths; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 49c27926ed..dd31ca2d17 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1391,62 +1391,8 @@ async fn find_thread_path_by_id_str_in_subdir( } } - let mut root = codex_home.to_path_buf(); - root.push(subdir); - if !root.exists() { - return Ok(unverified_db_path); - } - let (filename_match, filename_scan_error) = match find_rollout_path_by_id_from_filenames( - root.as_path(), - id_str, - ) - .await - { - Ok(path) => (path, None), - Err(err) => { - tracing::warn!( - "rollout filename lookup failed during find_thread_path_by_id_str_in_subdir: {err}" - ); - (None, Some(err)) - } - }; - - let found = match filename_match { - Some(path) => Some(path), - None => { - // This is safe because we know the values are valid. - #[allow(clippy::unwrap_used)] - let limit = NonZero::new(1).unwrap(); - let options = file_search::FileSearchOptions { - limit, - compute_indices: false, - respect_gitignore: false, - ..Default::default() - }; - - let results = file_search::run( - id_str, - vec![root.clone()], - options, - /*cancel_flag*/ None, - ) - .map_err(|e| io::Error::other(format!("file search failed: {e}")))?; - - let found = results - .matches - .into_iter() - .map(|m| m.full_path()) - .find_map(compression::RolloutFile::from_path) - .map(compression::RolloutFile::into_path); - - if found.is_none() - && let Some(err) = filename_scan_error - { - return Err(err); - } - found - } - }; + let found = + find_thread_path_by_id_str_in_subdir_from_filesystem(codex_home, subdir, id_str).await?; if let Some(found_path) = found.as_ref() { tracing::debug!("state db missing rollout path for thread {id_str}"); tracing::warn!( @@ -1471,6 +1417,61 @@ async fn find_thread_path_by_id_str_in_subdir( Ok(found.or(unverified_db_path)) } +pub(crate) async fn find_thread_path_by_id_str_in_subdir_from_filesystem( + codex_home: &Path, + subdir: &str, + id_str: &str, +) -> io::Result> { + if Uuid::parse_str(id_str).is_err() { + return Ok(None); + } + let root = codex_home.join(subdir); + if !root.exists() { + return Ok(None); + } + let (filename_match, filename_scan_error) = match find_rollout_path_by_id_from_filenames( + root.as_path(), + id_str, + ) + .await + { + Ok(path) => (path, None), + Err(err) => { + tracing::warn!( + "rollout filename lookup failed during find_thread_path_by_id_str_in_subdir: {err}" + ); + (None, Some(err)) + } + }; + if filename_match.is_some() { + return Ok(filename_match); + } + + // This is safe because we know the values are valid. + #[allow(clippy::unwrap_used)] + let limit = NonZero::new(1).unwrap(); + let options = file_search::FileSearchOptions { + limit, + compute_indices: false, + respect_gitignore: false, + ..Default::default() + }; + let results = file_search::run(id_str, vec![root], options, /*cancel_flag*/ None) + .map_err(|err| io::Error::other(format!("file search failed: {err}")))?; + let found = results + .matches + .into_iter() + .map(|matched| matched.full_path()) + .find_map(compression::RolloutFile::from_path) + .map(compression::RolloutFile::into_path); + if found.is_none() + && let Some(err) = filename_scan_error + { + return Err(err); + } + Ok(found) +} + async fn find_rollout_path_by_id_from_filenames( root: &Path, id_str: &str, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 7d7cca6334..c503f2a95d 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -13,6 +13,8 @@ use chrono::SecondsFormat; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; +use serde::Deserialize; +use serde::Serialize; use serde_json::Value; use time::OffsetDateTime; use time::format_description::FormatItem; @@ -38,6 +40,7 @@ use super::list::ThreadListConfig; use super::list::ThreadListLayout; use super::list::ThreadSortKey; use super::list::ThreadsPage; +use super::list::find_thread_path_by_id_str_in_subdir_from_filesystem; use super::list::get_threads; use super::list::get_threads_in_root; use super::list::parse_cursor; @@ -67,7 +70,7 @@ use codex_utils_path as path_utils; const FORK_PARENT_ROLLOUT_PATH_FIELD: &str = "fork_parent_rollout_path"; const FORK_PARENT_ROLLOUT_BYTE_LEN_FIELD: &str = "fork_parent_rollout_byte_len"; -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] pub struct ForkParentRolloutRef { pub path: PathBuf, pub byte_len: u64, @@ -100,7 +103,7 @@ pub enum RolloutRecorderParams { base_instructions: BaseInstructions, dynamic_tools: Vec, multi_agent_version: Option, - initial_rollout_copy: Option, + initial_rollout_copy: Option, }, Resume { path: PathBuf, @@ -205,7 +208,10 @@ impl RolloutRecorderParams { self } - pub fn with_initial_rollout_copy(mut self, initial_rollout_copy: Option) -> Self { + pub fn with_initial_rollout_copy( + mut self, + initial_rollout_copy: Option, + ) -> Self { if let Self::Create { initial_rollout_copy: copy, .. @@ -902,42 +908,47 @@ impl RolloutRecorder { path: &Path, ) -> std::io::Result<(Vec, Option, usize)> { trace!("Resuming rollout from {path:?}"); - let mut items: Vec = Vec::new(); - let mut thread_id: Option = 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; + let (mut items, thread_id, mut parse_errors, mut parent_ref) = + Self::load_rollout_items_unexpanded(path, /*byte_len*/ None).await?; + let mut current_path = path.to_path_buf(); + let mut child_segments = Vec::new(); + let mut visited_paths = HashSet::from([current_path.clone()]); + + while let Some(reference) = parent_ref { + let parent_thread_id = fork_parent_thread_id(items.as_slice(), current_path.as_path())?; + let parent_path = resolve_fork_parent_rollout_path( + current_path.as_path(), + &reference, + parent_thread_id, + ) + .await?; + if !visited_paths.insert(parent_path.clone()) { + return Err(IoError::other(format!( + "fork parent cycle detected at {}", + parent_path.display() + ))); } - saw_non_empty_line = true; - match serde_json::from_str::(&line) { - Ok(rollout_line) if rollout_item_may_need_legacy_cleanup(&rollout_line.item) => { - if let Some(rollout_line) = parse_rollout_line_after_legacy_cleanup( - &line, - &mut parse_errors, - /*rollout_error*/ None, - ) { - push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); - } - } - Ok(rollout_line) => { - push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); - } - Err(rollout_error) => { - if let Some(rollout_line) = parse_rollout_line_after_legacy_cleanup( - &line, - &mut parse_errors, - Some(&rollout_error), - ) { - push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); - } - } + let (parent_items, loaded_parent_thread_id, parent_parse_errors, next_parent_ref) = + Self::load_rollout_items_unexpanded( + parent_path.as_path(), + Some(reference.byte_len), + ) + .await?; + if loaded_parent_thread_id != Some(parent_thread_id) { + return Err(IoError::other(format!( + "fork parent identity mismatch for {}", + current_path.display() + ))); } + parse_errors = parse_errors.saturating_add(parent_parse_errors); + child_segments.push(items); + items = parent_items; + current_path = parent_path; + parent_ref = next_parent_ref; } - if !saw_non_empty_line { - return Err(IoError::other("empty session file")); + + for child_items in child_segments.into_iter().rev() { + items = prepend_parent_items(child_items, items)?; } tracing::debug!( @@ -949,35 +960,51 @@ impl RolloutRecorder { Ok((items, thread_id, parse_errors)) } - pub async fn load_rollout_items_prefix( + async fn load_rollout_items_unexpanded( path: &Path, - byte_len: u64, - ) -> std::io::Result<(Vec, Option, usize)> { - trace!("Resuming rollout prefix from {path:?}, byte_len={byte_len}"); + byte_len: Option, + ) -> std::io::Result<( + Vec, + Option, + usize, + Option, + )> { let mut items: Vec = Vec::new(); let mut thread_id: Option = None; let mut parse_errors = 0usize; + let mut parent_ref = None; let mut bytes_read = 0u64; let mut reader = compression::open_rollout_line_reader(path).await?; let mut saw_non_empty_line = false; - while bytes_read < byte_len { - let Some(line) = reader.next_line().await? else { - return Err(IoError::other(format!( - "rollout prefix ended before {byte_len} bytes" - ))); - }; - let line_len = u64::try_from(line.len()) - .unwrap_or(u64::MAX) - .saturating_add(1); - if bytes_read.saturating_add(line_len) > byte_len { - return Err(IoError::other(format!( - "rollout prefix boundary {byte_len} splits a JSONL line" - ))); + loop { + if byte_len.is_some_and(|byte_len| bytes_read >= byte_len) { + break; + } + let Some(line) = reader.next_line().await? else { + if let Some(byte_len) = byte_len { + return Err(IoError::other(format!( + "rollout prefix ended before {byte_len} bytes" + ))); + } + break; + }; + if let Some(byte_len) = byte_len { + let line_len = u64::try_from(line.len()) + .unwrap_or(u64::MAX) + .saturating_add(1); + if bytes_read.saturating_add(line_len) > byte_len { + return Err(IoError::other(format!( + "rollout prefix boundary {byte_len} splits a JSONL line" + ))); + } + bytes_read = bytes_read.saturating_add(line_len); } - bytes_read = bytes_read.saturating_add(line_len); if line.trim().is_empty() { continue; } + if !saw_non_empty_line { + parent_ref = parse_fork_parent_rollout_ref_from_line(line.as_str()); + } saw_non_empty_line = true; match serde_json::from_str::(&line) { Ok(rollout_line) if rollout_item_may_need_legacy_cleanup(&rollout_line.item) => { @@ -1004,10 +1031,13 @@ impl RolloutRecorder { } } if !saw_non_empty_line { - return Err(IoError::other("empty session file prefix")); + return Err(IoError::other(if byte_len.is_some() { + "empty session file prefix" + } else { + "empty session file" + })); } - - Ok((items, thread_id, parse_errors)) + Ok((items, thread_id, parse_errors, parent_ref)) } pub async fn get_rollout_history(path: &Path) -> std::io::Result { @@ -1055,18 +1085,85 @@ impl RolloutRecorder { } } -pub async fn read_fork_parent_rollout_ref( - path: &Path, -) -> std::io::Result> { - let mut reader = compression::open_rollout_line_reader(path).await?; - while let Some(line) = reader.next_line().await? { - let trimmed = line.trim(); - if trimmed.is_empty() { - continue; - } - return Ok(parse_fork_parent_rollout_ref_from_line(trimmed)); +fn fork_parent_thread_id(items: &[RolloutItem], path: &Path) -> std::io::Result { + items + .iter() + .find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.forked_from_id, + _ => None, + }) + .ok_or_else(|| { + IoError::other(format!( + "fork parent reference in {} is missing a parent thread id", + path.display() + )) + }) +} + +async fn resolve_fork_parent_rollout_path( + child_path: &Path, + parent_ref: &ForkParentRolloutRef, + parent_thread_id: ThreadId, +) -> std::io::Result { + if let Some(path) = compression::existing_rollout_path(parent_ref.path.as_path()).await { + return Ok(path); } - Ok(None) + + let codex_home = child_path + .ancestors() + .find(|ancestor| { + ancestor + .file_name() + .is_some_and(|name| name == SESSIONS_SUBDIR || name == ARCHIVED_SESSIONS_SUBDIR) + }) + .and_then(Path::parent) + .ok_or_else(|| { + IoError::other(format!( + "fork parent rollout {} does not exist", + parent_ref.path.display() + )) + })?; + let parent_thread_id = parent_thread_id.to_string(); + if let Some(path) = find_thread_path_by_id_str_in_subdir_from_filesystem( + codex_home, + SESSIONS_SUBDIR, + parent_thread_id.as_str(), + ) + .await? + { + return Ok(path); + } + if let Some(path) = find_thread_path_by_id_str_in_subdir_from_filesystem( + codex_home, + ARCHIVED_SESSIONS_SUBDIR, + parent_thread_id.as_str(), + ) + .await? + { + return Ok(path); + } + Err(IoError::other(format!( + "fork parent rollout {} does not exist and thread {parent_thread_id} was not found under {}", + parent_ref.path.display(), + codex_home.display() + ))) +} + +fn prepend_parent_items( + child_items: Vec, + parent_items: Vec, +) -> std::io::Result> { + let mut child_items = child_items.into_iter(); + let Some(RolloutItem::SessionMeta(child_meta)) = child_items.next() else { + return Err(IoError::other( + "fork child rollout does not start with session metadata", + )); + }; + let mut items = Vec::with_capacity(1 + parent_items.len() + child_items.size_hint().0); + items.push(RolloutItem::SessionMeta(child_meta)); + items.extend(parent_items); + items.extend(child_items); + Ok(items) } pub fn parse_fork_parent_rollout_ref_from_line(line: &str) -> Option { @@ -1643,7 +1740,7 @@ struct RolloutWriterState { deferred_log_file_info: Option, pending_items: Vec, meta: Option, - initial_rollout_copy: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, last_logged_error: Option, @@ -1654,7 +1751,7 @@ impl RolloutWriterState { file: Option, deferred_log_file_info: Option, meta: Option, - initial_rollout_copy: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, ) -> Self { @@ -1769,10 +1866,7 @@ impl RolloutWriterState { let Some(session_meta) = self.meta.as_ref().cloned() else { return Ok(()); }; - let fork_parent_rollout_ref = match self.initial_rollout_copy.as_deref() { - Some(source_path) => fork_parent_rollout_ref_for_initial_copy(source_path).await?, - None => None, - }; + let fork_parent_rollout_ref = self.initial_rollout_copy.clone(); write_session_meta( self.writer.as_mut(), session_meta, @@ -1780,30 +1874,14 @@ impl RolloutWriterState { fork_parent_rollout_ref.as_ref(), ) .await?; - if fork_parent_rollout_ref.is_some() { - self.initial_rollout_copy = None; - } - self.meta = None; - Ok(()) - } - - async fn copy_initial_rollout_if_needed(&mut self) -> std::io::Result<()> { - let Some(source_path) = self.initial_rollout_copy.as_ref().cloned() else { - return Ok(()); - }; - let Some(writer) = self.writer.as_mut() else { - return Err(IoError::other("rollout writer is not open")); - }; - let mut source = tokio::fs::File::open(source_path.as_path()).await?; - tokio::io::copy(&mut source, &mut writer.file).await?; self.initial_rollout_copy = None; + self.meta = None; Ok(()) } async fn write_pending_once(&mut self) -> std::io::Result<()> { self.ensure_writer_open().await?; self.write_session_meta_if_needed().await?; - self.copy_initial_rollout_if_needed().await?; self.write_pending_items_once().await?; @@ -1841,7 +1919,7 @@ async fn rollout_writer( deferred_log_file_info: Option, mut rx: mpsc::Receiver, meta: Option, - initial_rollout_copy: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, ) -> std::io::Result<()> { @@ -1882,31 +1960,6 @@ async fn rollout_writer( Ok(()) } -async fn fork_parent_rollout_ref_for_initial_copy( - source_path: &Path, -) -> std::io::Result> { - let Some(existing_path) = compression::existing_rollout_path(source_path).await else { - return Ok(None); - }; - if !is_plain_jsonl_rollout_path(existing_path.as_path()) { - return Ok(None); - } - let metadata = tokio::fs::metadata(existing_path.as_path()).await?; - if !metadata.is_file() || metadata.len() == 0 { - return Ok(None); - } - Ok(Some(ForkParentRolloutRef { - path: existing_path, - byte_len: metadata.len(), - })) -} - -fn is_plain_jsonl_rollout_path(path: &Path) -> bool { - path.file_name() - .and_then(|name| name.to_str()) - .is_some_and(|name| name.ends_with(".jsonl")) -} - async fn write_session_meta( mut writer: Option<&mut JsonlWriter>, session_meta: SessionMeta, diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index d24b3c2076..ba7319c609 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -87,9 +87,25 @@ impl LiveThread { pub async fn create( thread_store: Arc, params: CreateThreadParams, + ) -> ThreadStoreResult { + Self::create_inner(thread_store, params, /*observed_history*/ None).await + } + + pub async fn create_with_observed_history( + thread_store: Arc, + params: CreateThreadParams, + observed_history: &[RolloutItem], + ) -> ThreadStoreResult { + Self::create_inner(thread_store, params, Some(observed_history)).await + } + + async fn create_inner( + thread_store: Arc, + params: CreateThreadParams, + observed_history: Option<&[RolloutItem]>, ) -> ThreadStoreResult { let thread_id = params.thread_id; - let metadata_sync = ThreadMetadataSync::for_create(¶ms).await; + let metadata_sync = ThreadMetadataSync::for_create(¶ms, observed_history).await; thread_store.create_thread(params).await?; Ok(Self { thread_id, diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index f27ca0c688..6ae2fe6e36 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -284,75 +284,6 @@ async fn load_history_items( .map_err(|err| ThreadStoreError::Internal { message: format!("failed to load thread history {}: {err}", path.display()), })?; - expand_fork_parent_rollout_ref(path, items).await -} - -async fn expand_fork_parent_rollout_ref( - path: &std::path::Path, - child_items: Vec, -) -> ThreadStoreResult> { - let Some(parent_ref) = codex_rollout::read_fork_parent_rollout_ref(path) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to read fork parent reference {}: {err}", - path.display() - ), - })? - else { - return Ok(child_items); - }; - - let Some(child_forked_from_id) = child_items.iter().find_map(|item| match item { - codex_protocol::protocol::RolloutItem::SessionMeta(meta_line) => { - meta_line.meta.forked_from_id - } - _ => None, - }) else { - return Err(ThreadStoreError::Internal { - message: format!( - "failed to load fork parent history for {}: child fork parent missing", - path.display() - ), - }); - }; - let Some(parent_path) = codex_rollout::existing_rollout_path(parent_ref.path.as_path()).await - else { - return Err(ThreadStoreError::Internal { - message: format!( - "failed to load fork parent history for {}: parent rollout {} does not exist", - path.display(), - parent_ref.path.display() - ), - }); - }; - let (parent_items, parent_thread_id, _) = - RolloutRecorder::load_rollout_items_prefix(parent_path.as_path(), parent_ref.byte_len) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to load fork parent history {} for {}: {err}", - parent_path.display(), - path.display() - ), - })?; - if parent_thread_id != Some(child_forked_from_id) { - return Err(ThreadStoreError::Internal { - message: format!( - "failed to load fork parent history for {}: parent identity mismatch", - path.display() - ), - }); - } - - let mut child_items = child_items.into_iter(); - let Some(child_meta) = child_items.next() else { - return Ok(parent_items); - }; - let mut items = Vec::with_capacity(1 + parent_items.len() + child_items.size_hint().0); - items.push(child_meta); - items.extend(parent_items); - items.extend(child_items); Ok(items) } diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 94a1d2934d..fbe88c8318 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -51,7 +51,10 @@ pub(crate) struct PendingThreadMetadataPatch { } impl ThreadMetadataSync { - pub(crate) async fn for_create(params: &CreateThreadParams) -> Self { + pub(crate) async fn for_create( + params: &CreateThreadParams, + observed_history: Option<&[RolloutItem]>, + ) -> Self { let created_at = Utc::now(); let cwd = params.metadata.cwd.clone().unwrap_or_default(); let git_info = if get_git_repo_root(cwd.as_path()).is_some() { @@ -78,7 +81,7 @@ impl ThreadMetadataSync { memory_mode: Some(params.metadata.memory_mode), ..Default::default() }; - Self { + let mut sync = Self { thread_id: params.thread_id, cwd_seen: !cwd.as_os_str().is_empty(), preview_seen: false, @@ -89,7 +92,13 @@ impl ThreadMetadataSync { last_touch_persisted_at: None, defer_create_update_until_history_exists: true, defer_resume_update_until_append: false, + }; + if let Some(observed_history) = observed_history { + sync.defer_create_update_until_history_exists = false; + let update = sync.observe_resume_history(observed_history); + sync.merge_pending_update(update); } + sync } #[cfg(test)] diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 9b171176d2..9b904c5af3 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadMemoryMode as MemoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TokenUsage; +use codex_rollout::ForkParentRolloutRef; use serde::Deserialize; use serde::Deserializer; use serde::Serialize; @@ -81,10 +82,9 @@ pub struct CreateThreadParams { pub dynamic_tools: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, - /// Optional existing rollout JSONL to seed into the new rollout. Local - /// stores may persist this as a copy-on-write parent reference when safe. + /// Optional immutable parent rollout snapshot for copy-on-write history. #[serde(default, skip_serializing_if = "Option::is_none")] - pub initial_rollout_copy: Option, + pub initial_rollout_copy: Option, /// Metadata captured for the newly created thread. pub metadata: ThreadPersistenceMetadata, }