optimize codex resume and fork history

This commit is contained in:
anais
2026-06-19 10:45:37 -07:00
parent a57b268d61
commit 1d5dbfd346
25 changed files with 1688 additions and 247 deletions
@@ -252,6 +252,14 @@ impl ThreadHistoryBuilder {
}
}
pub fn with_initial_rollout_index(initial_rollout_index: usize) -> Self {
Self {
current_rollout_index: initial_rollout_index,
next_rollout_index: initial_rollout_index,
..Self::new()
}
}
pub fn reset(&mut self) {
*self = Self::new();
}
@@ -214,6 +214,7 @@ impl ExternalAgentSessionImporter {
},
dynamic_tools: Vec::new(),
multi_agent_version: Some(MultiAgentVersion::V1),
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: Some(cwd.clone()),
model_provider: model_provider.clone(),
File diff suppressed because it is too large Load Diff
@@ -324,6 +324,160 @@ mod thread_processor_behavior_tests {
assert_eq!(turns.last(), Some(&active_turn));
}
fn compacted_resume_test_line(item: RolloutItem) -> codex_protocol::protocol::RolloutLine {
codex_protocol::protocol::RolloutLine {
timestamp: "2025-09-05T16:53:11.850Z".to_string(),
item,
}
}
fn compacted_resume_test_session(thread_id: ThreadId) -> RolloutItem {
RolloutItem::SessionMeta(codex_protocol::protocol::SessionMetaLine {
meta: codex_protocol::protocol::SessionMeta {
id: thread_id,
timestamp: "2025-09-05T16:53:11.850Z".to_string(),
..codex_protocol::protocol::SessionMeta::default()
},
git: None,
})
}
fn compacted_resume_test_user_event(text: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::UserMessage(
codex_protocol::protocol::UserMessageEvent {
message: text.to_string(),
..Default::default()
},
))
}
fn compacted_resume_test_agent_event(text: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::AgentMessage(
codex_protocol::protocol::AgentMessageEvent {
message: text.to_string(),
phase: None,
memory_citation: None,
},
))
}
fn compacted_resume_test_message(
role: &str,
text: &str,
) -> codex_protocol::models::ResponseItem {
let content = if role == "user" {
vec![codex_protocol::models::ContentItem::InputText {
text: text.to_string(),
}]
} else {
vec![codex_protocol::models::ContentItem::OutputText {
text: text.to_string(),
}]
};
codex_protocol::models::ResponseItem::Message {
id: None,
role: role.to_string(),
content,
phase: None,
}
}
fn compacted_resume_test_compaction(message: &str, window_id: u64) -> RolloutItem {
RolloutItem::Compacted(codex_protocol::protocol::CompactedItem {
message: message.to_string(),
replacement_history: Some(vec![compacted_resume_test_message("assistant", message)]),
window_id: Some(window_id),
})
}
fn compacted_resume_test_rollout_contents(items: Vec<RolloutItem>) -> Result<String> {
let mut contents = String::new();
for item in items {
contents.push_str(&serde_json::to_string(&compacted_resume_test_line(item))?);
contents.push('\n');
}
Ok(contents)
}
fn compacted_resume_test_replacement_text(item: &RolloutItem) -> Option<&str> {
let RolloutItem::Compacted(compacted) = item else {
return None;
};
let [codex_protocol::models::ResponseItem::Message { content, .. }] =
compacted.replacement_history.as_deref()?
else {
return None;
};
let [codex_protocol::models::ContentItem::OutputText { text }] = content.as_slice() else {
return None;
};
Some(text.as_str())
}
#[test]
fn compacted_resume_rollout_items_reads_latest_checkpoint_suffix_across_chunks() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let large_tail_message = "tail".repeat(20_000);
let contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
compacted_resume_test_user_event("first user event"),
RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")),
compacted_resume_test_compaction("older checkpoint", 1),
compacted_resume_test_user_event("middle user event"),
compacted_resume_test_compaction("latest checkpoint", 2),
compacted_resume_test_agent_event(&large_tail_message),
])?;
std::fs::write(&path, contents)?;
let items = read_compacted_resume_rollout_items(path.as_path(), thread_id)?
.expect("compacted resume items");
assert_eq!(items.len(), 5);
assert!(matches!(items[0], RolloutItem::SessionMeta(_)));
assert!(matches!(
items[1],
RolloutItem::EventMsg(EventMsg::UserMessage(_))
));
assert!(matches!(items[2], RolloutItem::ResponseItem(_)));
assert_eq!(
compacted_resume_test_replacement_text(&items[3]),
Some("latest checkpoint")
);
assert!(matches!(
&items[4],
RolloutItem::EventMsg(EventMsg::AgentMessage(event))
if event.message == large_tail_message
));
assert_eq!(
items
.iter()
.filter(|item| matches!(item, RolloutItem::Compacted(_)))
.count(),
1
);
Ok(())
}
#[test]
fn compacted_resume_rollout_items_falls_back_on_malformed_tail_line() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let mut contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
compacted_resume_test_user_event("first user event"),
RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")),
compacted_resume_test_compaction("latest checkpoint", 2),
])?;
contents.push_str("{not json}\n");
std::fs::write(&path, contents)?;
assert!(read_compacted_resume_rollout_items(path.as_path(), thread_id)?.is_none());
Ok(())
}
#[test]
fn validate_dynamic_tools_rejects_empty_namespace() {
let tools = vec![dynamic_tool(
@@ -130,6 +130,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() ->
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: None,
model_provider: "test-provider".to_string(),
@@ -156,6 +156,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: Some(codex_home.path().to_path_buf()),
model_provider: "mock_provider".to_string(),
@@ -1366,6 +1366,7 @@ async fn seed_pathless_store_thread(
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: None,
model_provider: "test-provider".to_string(),
@@ -217,6 +217,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> {
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: None,
model_provider: "test-provider".to_string(),
+1
View File
@@ -101,6 +101,7 @@ pub(crate) async fn run_codex_thread_interactive(
conversation_history,
session_source: SessionSource::SubAgent(subagent_source.clone()),
forked_from_thread_id,
initial_rollout_copy: None,
parent_thread_id: Some(parent_session.thread_id),
thread_source: Some(ThreadSource::Subagent),
agent_control: parent_session.services.agent_control.clone(),
+28 -6
View File
@@ -423,6 +423,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) parent_thread_id: Option<ThreadId>,
pub(crate) thread_source: Option<ThreadSource>,
pub(crate) agent_control: AgentControl,
@@ -512,6 +513,7 @@ impl Codex {
conversation_history,
session_source,
forked_from_thread_id,
initial_rollout_copy,
parent_thread_id,
thread_source,
agent_control,
@@ -644,6 +646,7 @@ impl Codex {
app_server_client_version: None,
session_source,
forked_from_thread_id,
initial_rollout_copy,
parent_thread_id,
thread_source,
dynamic_tools,
@@ -1259,12 +1262,15 @@ impl Session {
}
async fn record_initial_history(&self, conversation_history: InitialHistory) {
let is_subagent = {
let (is_subagent, skip_fork_history_persistence) = {
let state = self.state.lock().await;
state
.session_configuration
.session_source
.is_non_root_agent()
(
state
.session_configuration
.session_source
.is_non_root_agent(),
state.session_configuration.initial_rollout_copy.is_some(),
)
};
let has_prior_user_turns = initial_history_has_prior_user_turns(&conversation_history);
{
@@ -1331,7 +1337,12 @@ impl Session {
}
// If persisting, persist all rollout items as-is (the store filters).
if !rollout_items.is_empty() {
if skip_fork_history_persistence {
let suffix = Self::fork_rollout_copy_suffix(&rollout_items);
if !suffix.is_empty() {
self.persist_rollout_items(suffix).await;
}
} else if !rollout_items.is_empty() {
self.persist_rollout_items(&rollout_items).await;
}
@@ -1346,6 +1357,17 @@ 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,
+15 -2
View File
@@ -102,6 +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>,
/// 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.
@@ -557,6 +559,9 @@ impl Session {
},
dynamic_tools: session_configuration.dynamic_tools.clone(),
multi_agent_version: initial_multi_agent_version,
initial_rollout_copy: session_configuration
.initial_rollout_copy
.clone(),
metadata: ThreadPersistenceMetadata {
cwd: Some(config.cwd.to_path_buf()),
model_provider: config.model_provider_id.clone(),
@@ -570,10 +575,13 @@ impl Session {
LiveThread::create(Arc::clone(&thread_store), params).await?
}
InitialHistory::Resumed(resumed_history) => {
let local_resume_with_path = resumed_history.rollout_path.is_some()
&& thread_store.as_any().is::<LocalThreadStore>();
let params = ResumeThreadParams {
thread_id: resumed_history.conversation_id,
rollout_path: resumed_history.rollout_path.clone(),
history: Some(resumed_history.history.clone()),
history: (!local_resume_with_path)
.then(|| resumed_history.history.clone()),
include_archived: true,
metadata: ThreadPersistenceMetadata {
cwd: Some(config.cwd.to_path_buf()),
@@ -585,7 +593,12 @@ impl Session {
},
},
};
LiveThread::resume(Arc::clone(&thread_store), params).await?
LiveThread::resume_with_observed_history(
Arc::clone(&thread_store),
params,
Some(resumed_history.history.as_slice()),
)
.await?
}
};
Ok(Some(live_thread))
+10
View File
@@ -3378,6 +3378,7 @@ async fn set_rate_limits_retains_previous_credits() {
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -3485,6 +3486,7 @@ async fn set_rate_limits_updates_plan_type_when_present() {
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -3738,6 +3740,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf {
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: Some(config.cwd.to_path_buf()),
model_provider: config.model_provider_id.clone(),
@@ -4013,6 +4016,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -4880,6 +4884,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -4993,6 +4998,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -5239,6 +5245,7 @@ async fn make_session_with_config_and_rx(
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -5346,6 +5353,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
app_server_client_version: None,
session_source: session_source.clone(),
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools: Vec::new(),
@@ -6578,6 +6586,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() {
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: Some(config.cwd.to_path_buf()),
model_provider: config.model_provider_id.clone(),
@@ -7053,6 +7062,7 @@ where
app_server_client_version: None,
session_source: SessionSource::Exec,
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
dynamic_tools,
@@ -720,6 +720,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
GUARDIAN_REVIEWER_NAME.to_string(),
)),
forked_from_thread_id: None,
initial_rollout_copy: None,
parent_thread_id: None,
thread_source: None,
agent_control: AgentControl::default(),
+99 -17
View File
@@ -954,6 +954,32 @@ impl ThreadManager {
thread_source,
parent_trace,
supports_openai_form_elicitation,
/*initial_rollout_copy*/ None,
)
.await
}
pub async fn fork_thread_from_history_with_initial_rollout_copy<S>(
&self,
snapshot: S,
config: Config,
history: InitialHistory,
thread_source: Option<ThreadSource>,
parent_trace: Option<W3cTraceContext>,
supports_openai_form_elicitation: bool,
initial_rollout_copy: PathBuf,
) -> CodexResult<NewThread>
where
S: Into<ForkSnapshot>,
{
self.fork_thread_with_initial_history(
snapshot.into(),
config,
history,
thread_source,
parent_trace,
supports_openai_form_elicitation,
Some(initial_rollout_copy),
)
.await
}
@@ -966,6 +992,7 @@ impl ThreadManager {
thread_source: Option<ThreadSource>,
parent_trace: Option<W3cTraceContext>,
supports_openai_form_elicitation: bool,
initial_rollout_copy: Option<PathBuf>,
) -> CodexResult<NewThread> {
// `forked_from_id()` describes this history's existing lineage. When
// forking a resumed thread, the child copies the resumed thread itself.
@@ -999,23 +1026,30 @@ impl ThreadManager {
&config.cwd,
);
let agent_control = self.agent_control_for_config(&config);
Box::pin(self.state.spawn_thread(
config,
history,
Arc::clone(&self.state.auth_manager),
agent_control,
/*parent_thread_id*/ None,
source_thread_id,
thread_source,
Vec::new(),
/*metrics_service_name*/ None,
initial_multi_agent_mode,
parent_trace,
environments,
/*thread_extension_init*/ ExtensionDataInit::default(),
supports_openai_form_elicitation,
/*user_shell_override*/ None,
))
Box::pin(
self.state
.spawn_thread_with_source_and_initial_rollout_copy(
config,
history,
Arc::clone(&self.state.auth_manager),
agent_control,
self.state.session_source.clone(),
/*parent_thread_id*/ None,
source_thread_id,
thread_source,
Vec::new(),
/*metrics_service_name*/ None,
initial_multi_agent_mode,
/*inherited_environments*/ None,
/*inherited_exec_policy*/ None,
parent_trace,
environments,
/*thread_extension_init*/ ExtensionDataInit::default(),
supports_openai_form_elicitation,
/*user_shell_override*/ None,
initial_rollout_copy,
),
)
.await
}
@@ -1428,6 +1462,53 @@ impl ThreadManagerState {
thread_extension_init: ExtensionDataInit,
supports_openai_form_elicitation: bool,
user_shell_override: Option<crate::shell::Shell>,
) -> CodexResult<NewThread> {
Box::pin(self.spawn_thread_with_source_and_initial_rollout_copy(
config,
initial_history,
auth_manager,
agent_control,
session_source,
parent_thread_id,
forked_from_thread_id,
thread_source,
dynamic_tools,
metrics_service_name,
initial_multi_agent_mode,
inherited_environments,
inherited_exec_policy,
parent_trace,
environments,
thread_extension_init,
supports_openai_form_elicitation,
user_shell_override,
/*initial_rollout_copy*/ None,
))
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn spawn_thread_with_source_and_initial_rollout_copy(
&self,
config: Config,
initial_history: InitialHistory,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
session_source: SessionSource,
parent_thread_id: Option<ThreadId>,
forked_from_thread_id: Option<ThreadId>,
thread_source: Option<ThreadSource>,
dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
metrics_service_name: Option<String>,
initial_multi_agent_mode: Option<MultiAgentMode>,
inherited_environments: Option<TurnEnvironmentSnapshot>,
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
parent_trace: Option<W3cTraceContext>,
environments: Vec<TurnEnvironmentSelection>,
thread_extension_init: ExtensionDataInit,
supports_openai_form_elicitation: bool,
user_shell_override: Option<crate::shell::Shell>,
initial_rollout_copy: Option<PathBuf>,
) -> CodexResult<NewThread> {
let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
if let InitialHistory::Resumed(resumed) = &initial_history {
@@ -1482,6 +1563,7 @@ impl ThreadManagerState {
conversation_history: initial_history,
session_source,
forked_from_thread_id,
initial_rollout_copy,
parent_thread_id,
thread_source,
agent_control,
+2
View File
@@ -63,9 +63,11 @@ pub use metadata::builder_from_items;
pub use policy::is_persisted_rollout_item;
pub use policy::persisted_rollout_items;
pub use policy::should_persist_response_item_for_memories;
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;
+19
View File
@@ -1258,6 +1258,25 @@ fn event_msg_preview(event: &EventMsg) -> Option<String> {
/// Read the SessionMetaLine from the head of a rollout file for reuse by
/// callers that need the session metadata (e.g. to derive a cwd for config).
pub async fn read_session_meta_line(path: &Path) -> io::Result<SessionMetaLine> {
let mut lines = compression::open_rollout_line_reader(path).await?;
while let Some(line) = lines.next_line().await? {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
if let Ok(RolloutLine {
item: RolloutItem::SessionMeta(session_meta_line),
..
}) = serde_json::from_str::<RolloutLine>(trimmed)
{
return Ok(session_meta_line);
}
break;
}
// Preserve the legacy tolerant behavior for unusual files whose first
// parseable summary record is still session metadata.
let head = read_head_for_summary(path).await?;
let Some(first) = head.first() else {
return Err(io::Error::other(format!(
+319 -32
View File
@@ -50,6 +50,7 @@ use crate::state_db;
use crate::state_db::StateDbHandle;
use codex_git_utils::collect_git_info;
use codex_git_utils::get_git_repo_root;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::GitInfo as ProtocolGitInfo;
use codex_protocol::protocol::InitialHistory;
use codex_protocol::protocol::MultiAgentVersion;
@@ -63,6 +64,15 @@ use codex_protocol::protocol::ThreadSource;
use codex_state::StateRuntime;
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)]
pub struct ForkParentRolloutRef {
pub path: PathBuf,
pub byte_len: u64,
}
/// Writes canonical session rollout items to JSONL.
///
/// Rollouts are recorded as JSONL and can be inspected with tools such as:
@@ -89,6 +99,7 @@ pub enum RolloutRecorderParams {
base_instructions: BaseInstructions,
dynamic_tools: Vec<DynamicToolSpec>,
multi_agent_version: Option<MultiAgentVersion>,
initial_rollout_copy: Option<PathBuf>,
},
Resume {
path: PathBuf,
@@ -175,6 +186,7 @@ impl RolloutRecorderParams {
base_instructions,
dynamic_tools,
multi_agent_version: None,
initial_rollout_copy: None,
}
}
@@ -192,6 +204,17 @@ impl RolloutRecorderParams {
self
}
pub fn with_initial_rollout_copy(mut self, initial_rollout_copy: Option<PathBuf>) -> Self {
if let Self::Create {
initial_rollout_copy: copy,
..
} = &mut self
{
*copy = initial_rollout_copy;
}
self
}
pub fn resume(path: PathBuf) -> Self {
Self::Resume { path }
}
@@ -694,7 +717,8 @@ impl RolloutRecorder {
config: &impl RolloutConfigView,
params: RolloutRecorderParams,
) -> std::io::Result<Self> {
let (file, deferred_log_file_info, rollout_path, meta) = match params {
let (file, deferred_log_file_info, rollout_path, meta, initial_rollout_copy) = match params
{
RolloutRecorderParams::Create {
conversation_id,
forked_from_id,
@@ -704,6 +728,7 @@ impl RolloutRecorder {
base_instructions,
dynamic_tools,
multi_agent_version,
initial_rollout_copy,
} => {
let log_file_info = precompute_log_file_info(config, conversation_id)?;
let path = log_file_info.path.clone();
@@ -742,7 +767,13 @@ impl RolloutRecorder {
multi_agent_version,
};
(None, Some(log_file_info), path, Some(session_meta))
(
None,
Some(log_file_info),
path,
Some(session_meta),
initial_rollout_copy,
)
}
RolloutRecorderParams::Resume { path } => {
let path = compression::materialize_rollout_for_append(path.as_path()).await?;
@@ -756,6 +787,7 @@ impl RolloutRecorder {
None,
path,
None,
None,
)
}
};
@@ -779,6 +811,7 @@ impl RolloutRecorder {
deferred_log_file_info,
rx,
meta,
initial_rollout_copy,
cwd,
rollout_path_for_spawn.clone(),
)
@@ -878,35 +911,25 @@ impl RolloutRecorder {
continue;
}
saw_non_empty_line = true;
let mut v: Value = match serde_json::from_str(&line) {
Ok(v) => v,
Err(e) => {
warn!("failed to parse line as JSON: {line:?}, error: {e}");
parse_errors = parse_errors.saturating_add(1);
continue;
}
};
if strip_legacy_ghost_snapshot_rollout_line(&mut v) {
trace!("skipping legacy ghost_snapshot rollout line");
continue;
}
// Parse the rollout line structure
match serde_json::from_value::<RolloutLine>(v.clone()) {
Ok(rollout_line) => {
let item = rollout_line.item;
// Use the FIRST SessionMeta encountered in the file as the canonical
// thread id and main session information. Keep all items intact.
if thread_id.is_none()
&& let RolloutItem::SessionMeta(session_meta_line) = &item
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, None)
{
thread_id = Some(session_meta_line.meta.id);
push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item);
}
items.push(item);
}
Err(e) => {
trace!("failed to parse rollout line: {e}");
parse_errors = parse_errors.saturating_add(1);
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);
}
}
}
}
@@ -923,6 +946,65 @@ impl RolloutRecorder {
Ok((items, thread_id, parse_errors))
}
pub async fn load_rollout_items_prefix(
path: &Path,
byte_len: u64,
) -> std::io::Result<(Vec<RolloutItem>, Option<ThreadId>, usize)> {
trace!("Resuming rollout prefix from {path:?}, byte_len={byte_len}");
let mut items: Vec<RolloutItem> = Vec::new();
let mut thread_id: Option<ThreadId> = None;
let mut parse_errors = 0usize;
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"
)));
}
bytes_read = bytes_read.saturating_add(line_len);
if line.trim().is_empty() {
continue;
}
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, 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);
}
}
}
}
if !saw_non_empty_line {
return Err(IoError::other("empty session file prefix"));
}
Ok((items, thread_id, parse_errors))
}
pub async fn get_rollout_history(path: &Path) -> std::io::Result<InitialHistory> {
let (items, thread_id, _parse_errors) = Self::load_rollout_items(path).await?;
let conversation_id = thread_id
@@ -968,6 +1050,101 @@ 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));
}
Ok(None)
}
pub fn parse_fork_parent_rollout_ref_from_line(line: &str) -> Option<ForkParentRolloutRef> {
let value: Value = serde_json::from_str(line).ok()?;
let payload = value.get("payload").unwrap_or(&value);
let path = payload
.get(FORK_PARENT_ROLLOUT_PATH_FIELD)
.and_then(Value::as_str)
.filter(|path| !path.is_empty())?;
let byte_len = payload
.get(FORK_PARENT_ROLLOUT_BYTE_LEN_FIELD)
.and_then(Value::as_u64)
.filter(|byte_len| *byte_len > 0)?;
Some(ForkParentRolloutRef {
path: PathBuf::from(path),
byte_len,
})
}
fn push_loaded_rollout_item(
items: &mut Vec<RolloutItem>,
thread_id: &mut Option<ThreadId>,
item: RolloutItem,
) {
// Use the FIRST SessionMeta encountered in the file as the canonical
// thread id and main session information. Keep all items intact.
if thread_id.is_none()
&& let RolloutItem::SessionMeta(session_meta_line) = &item
{
*thread_id = Some(session_meta_line.meta.id);
}
items.push(item);
}
fn rollout_item_may_need_legacy_cleanup(item: &RolloutItem) -> bool {
match item {
RolloutItem::ResponseItem(ResponseItem::Other) => true,
RolloutItem::Compacted(compacted) => {
compacted
.replacement_history
.as_ref()
.is_some_and(|history| {
history
.iter()
.any(|item| matches!(item, ResponseItem::Other))
})
}
_ => false,
}
}
fn parse_rollout_line_after_legacy_cleanup(
line: &str,
parse_errors: &mut usize,
rollout_error: Option<&serde_json::Error>,
) -> Option<RolloutLine> {
let mut value: Value = match serde_json::from_str(line) {
Ok(value) => value,
Err(err) => {
warn!("failed to parse line as JSON: {line:?}, error: {err}");
*parse_errors = parse_errors.saturating_add(1);
return None;
}
};
if strip_legacy_ghost_snapshot_rollout_line(&mut value) {
trace!("skipping legacy ghost_snapshot rollout line");
return None;
}
match serde_json::from_value::<RolloutLine>(value) {
Ok(rollout_line) => Some(rollout_line),
Err(err) => {
if let Some(rollout_error) = rollout_error {
trace!("failed to parse rollout line: {rollout_error}; fallback error: {err}");
} else {
trace!("failed to parse rollout line after legacy cleanup: {err}");
}
*parse_errors = parse_errors.saturating_add(1);
None
}
}
}
fn strip_legacy_ghost_snapshot_rollout_line(value: &mut Value) -> bool {
match value.get("type").and_then(Value::as_str) {
Some("response_item") => value
@@ -1461,6 +1638,7 @@ struct RolloutWriterState {
deferred_log_file_info: Option<LogFileInfo>,
pending_items: Vec<RolloutItem>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
cwd: PathBuf,
rollout_path: PathBuf,
last_logged_error: Option<String>,
@@ -1471,6 +1649,7 @@ impl RolloutWriterState {
file: Option<tokio::fs::File>,
deferred_log_file_info: Option<LogFileInfo>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
cwd: PathBuf,
rollout_path: PathBuf,
) -> Self {
@@ -1479,6 +1658,7 @@ impl RolloutWriterState {
deferred_log_file_info,
pending_items: Vec::new(),
meta,
initial_rollout_copy,
cwd,
rollout_path,
last_logged_error: None,
@@ -1584,14 +1764,41 @@ impl RolloutWriterState {
let Some(session_meta) = self.meta.as_ref().cloned() else {
return Ok(());
};
write_session_meta(self.writer.as_mut(), session_meta, &self.cwd).await?;
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,
};
write_session_meta(
self.writer.as_mut(),
session_meta,
&self.cwd,
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;
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?;
@@ -1629,10 +1836,18 @@ async fn rollout_writer(
deferred_log_file_info: Option<LogFileInfo>,
mut rx: mpsc::Receiver<RolloutCmd>,
meta: Option<SessionMeta>,
initial_rollout_copy: Option<PathBuf>,
cwd: PathBuf,
rollout_path: PathBuf,
) -> std::io::Result<()> {
let mut state = RolloutWriterState::new(file, deferred_log_file_info, meta, cwd, rollout_path);
let mut state = RolloutWriterState::new(
file,
deferred_log_file_info,
meta,
initial_rollout_copy,
cwd,
rollout_path,
);
// Process rollout commands
while let Some(cmd) = rx.recv().await {
@@ -1662,10 +1877,36 @@ 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,
cwd: &Path,
fork_parent_rollout_ref: Option<&ForkParentRolloutRef>,
) -> std::io::Result<()> {
let git_info = if get_git_repo_root(cwd).is_some() {
collect_git_info(cwd).await.map(|info| ProtocolGitInfo {
@@ -1681,9 +1922,10 @@ async fn write_session_meta(
git: git_info,
};
let rollout_item = RolloutItem::SessionMeta(session_meta_line);
if let Some(writer) = writer.as_mut() {
writer.write_rollout_item(&rollout_item).await?;
writer
.write_session_meta_line(&session_meta_line, fork_parent_rollout_ref)
.await?;
}
Ok(())
}
@@ -1717,7 +1959,52 @@ struct RolloutLineRef<'a> {
item: &'a RolloutItem,
}
#[derive(serde::Serialize)]
struct SessionMetaLineRef<'a> {
timestamp: String,
#[serde(rename = "type")]
item_type: &'static str,
payload: SessionMetaPayloadRef<'a>,
}
#[derive(serde::Serialize)]
struct SessionMetaPayloadRef<'a> {
#[serde(flatten)]
meta: &'a SessionMeta,
#[serde(skip_serializing_if = "Option::is_none")]
git: Option<&'a ProtocolGitInfo>,
#[serde(skip_serializing_if = "Option::is_none")]
fork_parent_rollout_path: Option<&'a PathBuf>,
#[serde(skip_serializing_if = "Option::is_none")]
fork_parent_rollout_byte_len: Option<u64>,
}
impl JsonlWriter {
async fn write_session_meta_line(
&mut self,
session_meta_line: &SessionMetaLine,
fork_parent_rollout_ref: Option<&ForkParentRolloutRef>,
) -> std::io::Result<()> {
let timestamp_format: &[FormatItem] = format_description!(
"[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
);
let timestamp = OffsetDateTime::now_utc()
.format(timestamp_format)
.map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?;
let line = SessionMetaLineRef {
timestamp,
item_type: "session_meta",
payload: SessionMetaPayloadRef {
meta: &session_meta_line.meta,
git: session_meta_line.git.as_ref(),
fork_parent_rollout_path: fork_parent_rollout_ref.map(|reference| &reference.path),
fork_parent_rollout_byte_len: fork_parent_rollout_ref
.map(|reference| reference.byte_len),
},
};
self.write_line(&line).await
}
async fn write_rollout_item(&mut self, rollout_item: &RolloutItem) -> std::io::Result<()> {
let timestamp_format: &[FormatItem] = format_description!(
"[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
+1
View File
@@ -506,6 +506,7 @@ async fn writer_state_retries_write_error_before_reporting_flush_success() -> st
Some(tokio::fs::File::from_std(read_only_file)),
/*deferred_log_file_info*/ None,
/*meta*/ None,
/*initial_rollout_copy*/ None,
home.path().to_path_buf(),
rollout_path.clone(),
);
+1
View File
@@ -119,6 +119,7 @@ mod tests {
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: ThreadPersistenceMetadata {
cwd: None,
model_provider: "test-provider".to_string(),
+25 -7
View File
@@ -100,13 +100,28 @@ impl LiveThread {
pub async fn resume(
thread_store: Arc<dyn ThreadStore>,
mut params: ResumeThreadParams,
params: ResumeThreadParams,
) -> ThreadStoreResult<Self> {
Self::resume_with_observed_history(thread_store, params, None).await
}
pub async fn resume_with_observed_history(
thread_store: Arc<dyn ThreadStore>,
params: ResumeThreadParams,
observed_history: Option<&[RolloutItem]>,
) -> ThreadStoreResult<Self> {
let thread_id = params.thread_id;
let should_load_history = params.history.is_none();
let include_archived = params.include_archived;
thread_store.resume_thread(params.clone()).await?;
if should_load_history {
let metadata = params.metadata.clone();
let metadata_sync = observed_history
.or(params.history.as_deref())
.map(|history| {
ThreadMetadataSync::for_resume_parts(thread_id, &metadata, Some(history))
});
thread_store.resume_thread(params).await?;
let metadata_sync = if let Some(metadata_sync) = metadata_sync {
metadata_sync
} else {
match thread_store
.load_history(LoadThreadHistoryParams {
thread_id,
@@ -114,7 +129,11 @@ impl LiveThread {
})
.await
{
Ok(history) => params.history = Some(history.items),
Ok(history) => ThreadMetadataSync::for_resume_parts(
thread_id,
&metadata,
Some(history.items.as_slice()),
),
Err(err) => {
if let Err(discard_err) = thread_store.discard_thread(thread_id).await {
warn!(
@@ -124,8 +143,7 @@ impl LiveThread {
return Err(err);
}
}
}
let metadata_sync = ThreadMetadataSync::for_resume(&params);
};
Ok(Self {
thread_id,
thread_store,
@@ -36,7 +36,8 @@ pub(super) async fn create_thread(
params.base_instructions,
params.dynamic_tools,
)
.with_multi_agent_version(params.multi_agent_version),
.with_multi_agent_version(params.multi_agent_version)
.with_initial_rollout_copy(params.initial_rollout_copy),
)
.await
.map_err(|err| ThreadStoreError::Internal {
+1
View File
@@ -1132,6 +1132,7 @@ mod tests {
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
multi_agent_version: None,
initial_rollout_copy: None,
metadata: thread_metadata(),
}
}
+75 -10
View File
@@ -39,12 +39,8 @@ pub(super) async fn read_thread(
metadata.rollout_path.as_path(),
)))
&& (!params.include_history
|| sqlite_rollout_path_can_load_history_for_thread(
store,
&metadata.rollout_path,
thread_id,
)
.await)
|| sqlite_rollout_path_can_load_history_for_thread(&metadata.rollout_path, thread_id)
.await)
{
let metadata_sandbox_policy = metadata.sandbox_policy.clone();
let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await;
@@ -87,7 +83,6 @@ pub(super) async fn read_thread(
}
async fn sqlite_rollout_path_can_load_history_for_thread(
store: &LocalThreadStore,
path: &std::path::Path,
thread_id: codex_protocol::ThreadId,
) -> bool {
@@ -96,10 +91,11 @@ async fn sqlite_rollout_path_can_load_history_for_thread(
}
// SQLite metadata can outlive a moved/recreated rollout path. When history is
// requested, verify the path still resolves to the requested thread before
// trusting it as the source replay.
read_thread_from_rollout_path(store, path.to_path_buf())
// trusting it as the source replay. The session metadata line is enough for
// this identity check and avoids doing a summary scan before full replay.
read_session_meta_line(path)
.await
.is_ok_and(|thread| thread.thread_id == thread_id)
.is_ok_and(|meta_line| meta_line.meta.id == thread_id)
}
pub(super) async fn read_thread_by_rollout_path(
@@ -288,6 +284,75 @@ 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)
}
@@ -16,8 +16,10 @@ use codex_protocol::protocol::UserMessageEvent;
use crate::CreateThreadParams;
use crate::GitInfoPatch;
#[cfg(test)]
use crate::ResumeThreadParams;
use crate::ThreadMetadataPatch;
use crate::ThreadPersistenceMetadata;
const IMAGE_ONLY_USER_MESSAGE_PLACEHOLDER: &str = "[Image]";
#[cfg(not(test))]
@@ -90,11 +92,23 @@ impl ThreadMetadataSync {
}
}
#[cfg(test)]
pub(crate) fn for_resume(params: &ResumeThreadParams) -> Self {
Self::for_resume_parts(
params.thread_id,
&params.metadata,
params.history.as_deref(),
)
}
pub(crate) fn for_resume_parts(
thread_id: ThreadId,
metadata: &ThreadPersistenceMetadata,
history: Option<&[RolloutItem]>,
) -> Self {
let mut sync = Self {
thread_id: params.thread_id,
cwd_seen: params
.metadata
thread_id,
cwd_seen: metadata
.cwd
.as_ref()
.is_some_and(|cwd| !cwd.as_os_str().is_empty()),
@@ -107,7 +121,7 @@ impl ThreadMetadataSync {
defer_create_update_until_history_exists: false,
defer_resume_update_until_append: false,
};
if let Some(history) = params.history.as_deref() {
if let Some(history) = history {
let update = sync.observe_resume_history(history);
sync.merge_pending_update(update);
sync.defer_resume_update_until_append = sync.pending_update.is_some();
+4
View File
@@ -81,6 +81,10 @@ 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.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_rollout_copy: Option<PathBuf>,
/// Metadata captured for the newly created thread.
pub metadata: ThreadPersistenceMetadata,
}