fix copy on write fork persistence

This commit is contained in:
anais
2026-06-19 10:45:38 -07:00
parent 594fe659b6
commit 7e085334ad
12 changed files with 366 additions and 300 deletions
@@ -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<Option<StoredThread>, JSONRPCErrorError> {
) -> Result<Option<(StoredThread, codex_rollout::ForkParentRolloutRef)>, 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<RolloutItem>,
}
struct CompactedResumeRollout {
items: Vec<RolloutItem>,
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<Option<Vec<RolloutItem>>, JSONRPCErrorError> {
) -> Result<Option<CompactedResumeRollout>, 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<Option<Vec<RolloutItem>>> {
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<Option<Vec<RolloutItem>>> {
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<Option<Vec<RolloutItem>>> {
fn read_compacted_resume_suffix_items(
path: &Path,
byte_len: u64,
) -> std::io::Result<Option<Vec<RolloutItem>>> {
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<Option<Vec
Ok(Some(suffix_items))
}
fn read_rollout_lines_from_end<F>(path: &Path, mut visit: F) -> std::io::Result<()>
fn read_rollout_lines_from_end<F>(path: &Path, byte_len: u64, mut visit: F) -> std::io::Result<()>
where
F: FnMut(&str) -> std::io::Result<bool>,
{
@@ -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 {
@@ -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(())
}
+17 -16
View File
@@ -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<ThreadId>,
pub(crate) initial_rollout_copy: Option<PathBuf>,
pub(crate) initial_rollout_copy: Option<InitialRolloutCopy>,
pub(crate) parent_thread_id: Option<ThreadId>,
pub(crate) thread_source: Option<ThreadSource>,
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,
+24 -4
View File
@@ -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<ThreadId>,
/// Optional rollout JSONL copied into this thread after its own session metadata.
pub(super) initial_rollout_copy: Option<PathBuf>,
/// Optional immutable parent rollout snapshot referenced by this thread.
pub(super) initial_rollout_copy: Option<InitialRolloutCopy>,
/// Immediate control/spawn parent for this thread, when it has one.
pub(super) parent_thread_id: Option<ThreadId>,
/// 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()
+13 -3
View File
@@ -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<ThreadSource>,
parent_trace: Option<W3cTraceContext>,
supports_openai_form_elicitation: bool,
initial_rollout_copy: PathBuf,
initial_rollout_copy: ForkParentRolloutRef,
) -> CodexResult<NewThread>
where
S: Into<ForkSnapshot>,
@@ -992,7 +994,7 @@ impl ThreadManager {
thread_source: Option<ThreadSource>,
parent_trace: Option<W3cTraceContext>,
supports_openai_form_elicitation: bool,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<ForkParentRolloutRef>,
) -> CodexResult<NewThread> {
// `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<crate::shell::Shell>,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<InitialRolloutCopy>,
) -> CodexResult<NewThread> {
let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
if let InitialHistory::Resumed(resumed) = &initial_history {
-1
View File
@@ -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;
+57 -56
View File
@@ -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<Option<PathBuf>> {
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,
+170 -117
View File
@@ -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<DynamicToolSpec>,
multi_agent_version: Option<MultiAgentVersion>,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<ForkParentRolloutRef>,
},
Resume {
path: PathBuf,
@@ -205,7 +208,10 @@ impl RolloutRecorderParams {
self
}
pub fn with_initial_rollout_copy(mut self, initial_rollout_copy: Option<PathBuf>) -> Self {
pub fn with_initial_rollout_copy(
mut self,
initial_rollout_copy: Option<ForkParentRolloutRef>,
) -> Self {
if let Self::Create {
initial_rollout_copy: copy,
..
@@ -902,42 +908,47 @@ 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;
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::<RolloutLine>(&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<RolloutItem>, Option<ThreadId>, usize)> {
trace!("Resuming rollout prefix from {path:?}, byte_len={byte_len}");
byte_len: Option<u64>,
) -> std::io::Result<(
Vec<RolloutItem>,
Option<ThreadId>,
usize,
Option<ForkParentRolloutRef>,
)> {
let mut items: Vec<RolloutItem> = Vec::new();
let mut thread_id: Option<ThreadId> = 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::<RolloutLine>(&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<InitialHistory> {
@@ -1055,18 +1085,85 @@ impl RolloutRecorder {
}
}
pub async fn read_fork_parent_rollout_ref(
path: &Path,
) -> std::io::Result<Option<ForkParentRolloutRef>> {
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<ThreadId> {
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<PathBuf> {
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<RolloutItem>,
parent_items: Vec<RolloutItem>,
) -> std::io::Result<Vec<RolloutItem>> {
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<ForkParentRolloutRef> {
@@ -1643,7 +1740,7 @@ struct RolloutWriterState {
deferred_log_file_info: Option<LogFileInfo>,
pending_items: Vec<RolloutItem>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<ForkParentRolloutRef>,
cwd: PathBuf,
rollout_path: PathBuf,
last_logged_error: Option<String>,
@@ -1654,7 +1751,7 @@ impl RolloutWriterState {
file: Option<tokio::fs::File>,
deferred_log_file_info: Option<LogFileInfo>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<ForkParentRolloutRef>,
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<LogFileInfo>,
mut rx: mpsc::Receiver<RolloutCmd>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
initial_rollout_copy: Option<ForkParentRolloutRef>,
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<Option<ForkParentRolloutRef>> {
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,
+17 -1
View File
@@ -87,9 +87,25 @@ impl LiveThread {
pub async fn create(
thread_store: Arc<dyn ThreadStore>,
params: CreateThreadParams,
) -> ThreadStoreResult<Self> {
Self::create_inner(thread_store, params, /*observed_history*/ None).await
}
pub async fn create_with_observed_history(
thread_store: Arc<dyn ThreadStore>,
params: CreateThreadParams,
observed_history: &[RolloutItem],
) -> ThreadStoreResult<Self> {
Self::create_inner(thread_store, params, Some(observed_history)).await
}
async fn create_inner(
thread_store: Arc<dyn ThreadStore>,
params: CreateThreadParams,
observed_history: Option<&[RolloutItem]>,
) -> ThreadStoreResult<Self> {
let thread_id = params.thread_id;
let metadata_sync = ThreadMetadataSync::for_create(&params).await;
let metadata_sync = ThreadMetadataSync::for_create(&params, observed_history).await;
thread_store.create_thread(params).await?;
Ok(Self {
thread_id,
@@ -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<codex_protocol::protocol::RolloutItem>,
) -> ThreadStoreResult<Vec<codex_protocol::protocol::RolloutItem>> {
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)
}
@@ -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)]
+3 -3
View File
@@ -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<DynamicToolSpec>,
/// Multi-agent runtime selected when the thread was created.
pub multi_agent_version: Option<MultiAgentVersion>,
/// 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<PathBuf>,
pub initial_rollout_copy: Option<ForkParentRolloutRef>,
/// Metadata captured for the newly created thread.
pub metadata: ThreadPersistenceMetadata,
}