From 0cf8861225d203ee1f7f2e4e983076a8dd7faef6 Mon Sep 17 00:00:00 2001 From: Friel Date: Fri, 12 Jun 2026 22:27:50 +0000 Subject: [PATCH] Materialize complete rollout reference histories Separate bounded model replay from fallible complete-history reconstruction, validate immutable segment resolution, and make app-server, search, memory, and feedback readers reference-aware. --- .../src/protocol/thread_history.rs | 2 + .../app-server/src/bespoke_event_handling.rs | 22 +- codex-rs/app-server/src/request_processors.rs | 3 + .../request_processors/feedback_processor.rs | 136 ++++- .../request_processors/thread_lifecycle.rs | 32 +- .../request_processors/thread_processor.rs | 131 ++++- .../thread_processor_tests.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 2 + .../app-server/tests/suite/v2/thread_list.rs | 65 +++ .../app-server/tests/suite/v2/thread_read.rs | 148 ++++++ .../tests/suite/v2/thread_resume.rs | 151 ++++++ codex-rs/core/src/agent/control/spawn.rs | 1 + codex-rs/core/src/lib.rs | 2 + .../core/src/personality_migration_tests.rs | 1 + codex-rs/core/src/rollout.rs | 1 + .../src/session/rollout_reconstruction.rs | 5 +- codex-rs/core/src/session/session.rs | 6 + codex-rs/core/src/thread_manager.rs | 10 + codex-rs/core/src/thread_manager_tests.rs | 62 +++ .../core/src/thread_rollout_truncation.rs | 284 ++++++++++ .../src/thread_rollout_truncation_tests.rs | 485 ++++++++++++++++++ .../core/tests/suite/personality_migration.rs | 2 + codex-rs/core/tests/suite/resume.rs | 92 ++++ codex-rs/core/tests/suite/sqlite_state.rs | 1 + codex-rs/memories/write/src/phase1.rs | 117 ++++- codex-rs/rollout/src/list.rs | 209 +++++--- codex-rs/rollout/src/metadata.rs | 21 +- codex-rs/rollout/src/metadata_tests.rs | 5 + codex-rs/rollout/src/search.rs | 247 ++++++++- codex-rs/rollout/src/tests.rs | 191 ++++++- codex-rs/thread-store/src/in_memory.rs | 1 + .../thread-store/src/local/read_thread.rs | 96 +++- .../thread-store/src/thread_metadata_sync.rs | 1 + 33 files changed, 2384 insertions(+), 149 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 09bfa10a5c..5c00cdf371 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -75,6 +75,8 @@ use codex_protocol::protocol::PatchApplyStatus as CorePatchApplyStatus; /// /// When available, this uses `TurnContext.turn_id` as the canonical turn id so /// resumed/rebuilt thread history preserves the original turn identifiers. +/// File-backed callers must resolve every [`RolloutItem::RolloutReference`] before calling this +/// function. The builder is synchronous and deliberately does not perform rollout file I/O. pub fn build_turns_from_rollout_items(items: &[RolloutItem]) -> Vec { let mut builder = ThreadHistoryBuilder::new(); for item in items { diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index d1099d3b19..402334fb99 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -86,6 +86,7 @@ use codex_app_server_protocol::guardian_auto_approval_review_notification; use codex_app_server_protocol::item_event_to_server_notification; use codex_core::CodexThread; use codex_core::ThreadManager; +use codex_core::materialize_rollout_items_for_complete_history; use codex_core::review_format::format_review_findings_block; use codex_core::review_prompts; use codex_protocol::ThreadId; @@ -1204,13 +1205,17 @@ pub(crate) async fn apply_bespoke_event_handling( let loaded_status = thread_watch_manager .loaded_status_for_thread(&conversation_id.to_string()) .await; + let config = conversation.config().await; let response = match thread_rollback_response_from_stored_thread( stored_thread, conversation.session_configured().session_id.to_string(), fallback_model_provider.as_str(), &fallback_cwd, loaded_status, - ) { + config.codex_home.as_path(), + ) + .await + { Ok(response) => response, Err(err) => { outgoing @@ -1567,12 +1572,13 @@ async fn handle_thread_rollback_failed( } } -fn thread_rollback_response_from_stored_thread( +async fn thread_rollback_response_from_stored_thread( stored_thread: codex_thread_store::StoredThread, session_id: String, fallback_model_provider: &str, fallback_cwd: &AbsolutePathBuf, loaded_status: ThreadStatus, + codex_home: &std::path::Path, ) -> std::result::Result { let thread_id = stored_thread.thread_id; let (mut thread, history) = @@ -1583,7 +1589,10 @@ fn thread_rollback_response_from_stored_thread( "thread {thread_id} did not include persisted history after rollback" )); }; - populate_thread_turns_from_history(&mut thread, &history.items, /*active_turn*/ None); + let history_items = materialize_rollout_items_for_complete_history(codex_home, &history.items) + .await + .map_err(|err| format!("failed to materialize thread {thread_id} history: {err}"))?; + populate_thread_turns_from_history(&mut thread, &history_items, /*active_turn*/ None); thread.status = loaded_status; Ok(ThreadRollbackResponse { thread }) } @@ -2182,8 +2191,8 @@ mod tests { } } - #[test] - fn rollback_response_rebuilds_pathless_thread_from_stored_history() -> Result<()> { + #[tokio::test] + async fn rollback_response_rebuilds_pathless_thread_from_stored_history() -> Result<()> { let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000789")?; let created_at = Utc::now(); let history_items = vec![ @@ -2233,6 +2242,7 @@ mod tests { }), }; let fallback_cwd = test_path_buf("/tmp").abs(); + let codex_home = tempfile::tempdir()?; let response = thread_rollback_response_from_stored_thread( stored_thread, @@ -2240,7 +2250,9 @@ mod tests { "fallback-provider", &fallback_cwd, ThreadStatus::NotLoaded, + codex_home.path(), ) + .await .expect("rollback response should rebuild from stored history"); assert_eq!(response.thread.id, thread_id.to_string()); diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 83e704cdfd..188761cbdb 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -290,6 +290,7 @@ use codex_core::CodexThreadSettingsOverrides; use codex_core::ForkSnapshot; use codex_core::McpManager; use codex_core::NewThread; +use codex_core::RolloutRecorder; #[cfg(test)] use codex_core::SessionMeta; use codex_core::StartThreadOptions; @@ -306,6 +307,7 @@ use codex_core::exec::ExecCapturePolicy; use codex_core::exec::ExecExpiration; use codex_core::exec::ExecParams; use codex_core::exec_env::create_env; +use codex_core::materialize_rollout_items_for_complete_history; use codex_core::path_utils; #[cfg(test)] use codex_core::read_head_for_summary; @@ -345,6 +347,7 @@ use codex_features::FEATURES; use codex_features::Feature; use codex_features::Stage; use codex_feedback::CodexFeedback; +use codex_feedback::FeedbackAttachment; use codex_feedback::FeedbackAttachmentPath; use codex_feedback::FeedbackUploadOptions; use codex_git_utils::git_diff_to_remote; diff --git a/codex-rs/app-server/src/request_processors/feedback_processor.rs b/codex-rs/app-server/src/request_processors/feedback_processor.rs index 73da67f735..f346062b53 100644 --- a/codex-rs/app-server/src/request_processors/feedback_processor.rs +++ b/codex-rs/app-server/src/request_processors/feedback_processor.rs @@ -158,6 +158,7 @@ impl FeedbackRequestProcessor { (Vec::new(), None, None) }; + let mut extra_attachments = Vec::new(); let mut attachment_paths = Vec::new(); let mut seen_attachment_paths = HashSet::new(); if include_logs { @@ -169,6 +170,19 @@ impl FeedbackRequestProcessor { continue; }; if seen_attachment_paths.insert(rollout_path.clone()) { + match materialized_rollout_feedback_attachment( + self.config.codex_home.as_path(), + rollout_path.as_path(), + ) + .await + { + Ok(Some(attachment)) => extra_attachments.push(attachment), + Ok(None) => {} + Err(err) => warn!( + "failed to materialize feedback rollout {}: {err}", + rollout_path.display() + ), + } attachment_paths.push(FeedbackAttachmentPath { path: rollout_path, attachment_filename_override: None, @@ -206,7 +220,6 @@ impl FeedbackRequestProcessor { } } - let mut extra_attachments = Vec::new(); if include_logs && let Some(doctor_report) = super::feedback_doctor_report::doctor_feedback_report(&self.config).await @@ -269,6 +282,30 @@ impl FeedbackRequestProcessor { } } +async fn materialized_rollout_feedback_attachment( + codex_home: &Path, + rollout_path: &Path, +) -> anyhow::Result> { + let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; + if !rollout_items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))) + { + return Ok(None); + } + let rollout_items = + materialize_rollout_items_for_complete_history(codex_home, &rollout_items).await?; + let file_stem = rollout_path + .file_stem() + .and_then(|value| value.to_str()) + .unwrap_or("rollout"); + Ok(Some(FeedbackAttachment { + filename: format!("{file_stem}-complete-history.json"), + content_type: Some("application/json".to_string()), + buffer: serde_json::to_vec_pretty(&rollout_items)?, + })) +} + fn auto_review_rollout_filename(thread_id: ThreadId) -> String { format!("auto-review-rollout-{thread_id}.jsonl") } @@ -289,11 +326,17 @@ fn windows_sandbox_log_attachment(_codex_home: &Path) -> Option>() + .join("\n"); + tokio::fs::write(&predecessor_path, format!("{predecessor_lines}\n")) + .await + .expect("write predecessor"); + let current = RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "after rotation".to_string(), + }], + phase: None, + }); + let current_lines = [ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: predecessor_path, + thread_id: None, + rollout_timestamp: None, + segment_id: None, + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + current.clone(), + ] + .into_iter() + .map(|item| { + serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item, + }) + .expect("serialize current rollout line") + }) + .collect::>() + .join("\n"); + tokio::fs::write(¤t_path, format!("{current_lines}\n")) + .await + .expect("write current rollout"); + + let attachment = + materialized_rollout_feedback_attachment(codex_home.path(), current_path.as_path()) + .await + .expect("materialize attachment") + .expect("reference-backed rollout should produce an attachment"); + let items: Vec = + serde_json::from_slice(&attachment.buffer).expect("parse attachment"); + + assert_eq!( + serde_json::to_value(items).expect("serialize attachment items"), + serde_json::to_value(vec![predecessor_meta, predecessor, current]) + .expect("serialize expected items") + ); + assert_eq!( + attachment.filename, + "current-complete-history.json".to_string() + ); + } } diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index fd9e93e189..e89ba91d81 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -529,7 +529,7 @@ pub(super) async fn handle_thread_listener_command( pub(super) async fn handle_pending_thread_resume_request( conversation_id: ThreadId, conversation: &Arc, - _codex_home: &Path, + codex_home: &Path, thread_state_manager: &ThreadStateManager, thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, @@ -557,13 +557,29 @@ pub(super) async fn handle_pending_thread_resume_request( let request_id = pending.request_id; let connection_id = request_id.connection_id; + let history_items = if pending.include_turns || pending.initial_turns_page.is_some() { + match materialize_rollout_items_for_complete_history(codex_home, &pending.history_items) + .await + { + Ok(items) => items, + Err(err) => { + outgoing + .send_error( + request_id, + internal_error(format!( + "failed to materialize thread {conversation_id} history: {err}" + )), + ) + .await; + return; + } + } + } else { + Vec::new() + }; let mut thread = pending.thread_summary; if pending.include_turns { - populate_thread_turns_from_history( - &mut thread, - &pending.history_items, - active_turn.as_ref(), - ); + populate_thread_turns_from_history(&mut thread, &history_items, active_turn.as_ref()); } let thread_status = thread_watch_manager @@ -578,7 +594,7 @@ pub(super) async fn handle_pending_thread_resume_request( let token_usage_thread = pending.include_turns.then(|| thread.clone()); let mut initial_turns_page = if let Some(params) = pending.initial_turns_page.as_ref() { match super::thread_processor::build_thread_resume_initial_turns_page( - &pending.history_items, + &history_items, thread.status.clone(), has_live_in_progress_turn, active_turn, @@ -668,7 +684,7 @@ pub(super) async fn handle_pending_thread_resume_request( // paying the cost of turn reconstruction for historical usage replay. if let Some(token_usage_thread) = token_usage_thread { let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( - &pending.history_items, + &history_items, token_usage_thread.turns.as_slice(), ); // Rejoining a loaded thread has the same UI contract as a cold resume, but 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 813a85fb4d..cf23661812 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2224,7 +2224,17 @@ impl ThreadRequestProcessor { let (mut thread, history) = thread_from_stored_thread(stored_thread, fallback_provider, &self.config.cwd); if include_turns && let Some(history) = history { - thread.turns = build_api_turns_from_rollout_items(&history.items); + let history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history.items, + ) + .await + .map_err(|err| { + ThreadReadViewError::Internal(format!( + "failed to materialize thread {thread_id} history: {err}" + )) + })?; + thread.turns = build_api_turns_from_rollout_items(&history_items); } Ok(Some(thread)) } @@ -2290,7 +2300,17 @@ impl ThreadRequestProcessor { .load_history(/*include_archived*/ true) .await .map_err(|err| thread_read_history_load_error(thread_id, err))?; - thread.turns = build_api_turns_from_rollout_items(&history.items); + let history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history.items, + ) + .await + .map_err(|err| { + ThreadReadViewError::Internal(format!( + "failed to materialize thread {thread_id} history: {err}" + )) + })?; + thread.turns = build_api_turns_from_rollout_items(&history_items); } Ok(()) @@ -2369,7 +2389,16 @@ impl ThreadRequestProcessor { "thread store did not return history for thread {thread_id}" )) })?; - return Ok(history.items); + return materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history.items, + ) + .await + .map_err(|err| { + ThreadReadViewError::Internal(format!( + "failed to materialize thread {thread_id} history: {err}" + )) + }); } Err(ThreadStoreError::InvalidRequest { message }) if message == format!("no rollout found for thread id {thread_id}") => {} @@ -2400,11 +2429,21 @@ impl ThreadRequestProcessor { )); } - thread + let history_items = thread .load_history(/*include_archived*/ true) .await .map(|history| history.items) - .map_err(|err| thread_turns_list_history_load_error(thread_id, err)) + .map_err(|err| thread_turns_list_history_load_error(thread_id, err))?; + materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history_items, + ) + .await + .map_err(|err| { + ThreadReadViewError::Internal(format!( + "failed to materialize thread {thread_id} history: {err}" + )) + }) } pub(crate) fn thread_created_receiver(&self) -> broadcast::Receiver { @@ -2620,6 +2659,20 @@ impl ThreadRequestProcessor { }; let response_history = thread_history.clone(); + let response_history_items = if include_turns || initial_turns_page.is_some() { + materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &response_history.get_rollout_items(), + ) + .await + .map_err(|err| { + internal_error(format!( + "failed to materialize thread {thread_id} history: {err}" + )) + })? + } else { + Vec::new() + }; match self .thread_manager @@ -2718,7 +2771,7 @@ impl ThreadRequestProcessor { let token_usage_thread = include_turns.then(|| thread.clone()); let mut initial_turns_page = if let Some(params) = initial_turns_page.as_ref() { match build_thread_resume_initial_turns_page( - &response_history.get_rollout_items(), + &response_history_items, thread.status.clone(), /*has_live_running_thread*/ false, /*active_turn*/ None, @@ -2762,7 +2815,7 @@ impl ThreadRequestProcessor { // rebuilding history only to attribute a replayed usage update. if let Some(token_usage_thread) = token_usage_thread { let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( - &response_history.get_rollout_items(), + &response_history_items, token_usage_thread.turns.as_slice(), ); // The client needs restored usage before it starts another turn. @@ -2928,6 +2981,16 @@ impl ThreadRequestProcessor { "thread {existing_thread_id} did not include persisted history" )) })?; + let history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history_items, + ) + .await + .map_err(|err| { + internal_error(format!( + "failed to materialize thread {existing_thread_id} history: {err}" + )) + })?; let thread_state = self .thread_state_manager @@ -2948,11 +3011,13 @@ impl ThreadRequestProcessor { let mut summary_source_thread = source_thread; summary_source_thread.history = None; - let mut thread_summary = self.stored_thread_to_api_thread( - summary_source_thread, - config_snapshot.model_provider_id.as_str(), - /*include_turns*/ false, - ); + let mut thread_summary = self + .stored_thread_to_api_thread( + summary_source_thread, + config_snapshot.model_provider_id.as_str(), + /*include_turns*/ false, + ) + .await?; thread_summary.session_id = existing_thread.session_configured().session_id.to_string(); let instruction_sources = existing_thread.instruction_sources().await; @@ -3086,22 +3151,33 @@ impl ThreadRequestProcessor { })) } - fn stored_thread_to_api_thread( + async fn stored_thread_to_api_thread( &self, stored_thread: StoredThread, fallback_provider: &str, include_turns: bool, - ) -> Thread { + ) -> Result { let (mut thread, history) = thread_from_stored_thread(stored_thread, fallback_provider, &self.config.cwd); if include_turns && let Some(history) = history { + let history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history.items, + ) + .await + .map_err(|err| { + internal_error(format!( + "failed to materialize thread {} history: {err}", + thread.id + )) + })?; populate_thread_turns_from_history( &mut thread, - &history.items, + &history_items, /*active_turn*/ None, ); } - thread + Ok(thread) } async fn read_stored_thread_for_new_fork( @@ -3208,6 +3284,12 @@ impl ThreadRequestProcessor { thread.path = Some(rollout_path.to_path_buf()); if include_turns { let history_items = thread_history.get_rollout_items(); + let history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history_items, + ) + .await + .map_err(|err| format!("failed to materialize thread {thread_id} history: {err}"))?; populate_thread_turns_from_history( &mut thread, &history_items, @@ -3284,6 +3366,16 @@ impl ThreadRequestProcessor { "thread {source_thread_id} did not include persisted history" )) })?; + let materialized_history_items = materialize_rollout_items_for_complete_history( + self.config.codex_home.as_path(), + &history_items, + ) + .await + .map_err(|err| { + internal_error(format!( + "failed to materialize thread {source_thread_id} history: {err}" + )) + })?; let history_cwd = Some(source_thread.cwd.clone()); // Persist Windows sandbox mode. @@ -3408,6 +3500,7 @@ impl ThreadRequestProcessor { fallback_model_provider.as_str(), include_turns, ) + .await? } else { let config_snapshot = forked_thread.config_snapshot().await; let mut thread = build_thread_from_snapshot( @@ -3416,12 +3509,11 @@ impl ThreadRequestProcessor { &config_snapshot, /*path*/ None, ); - thread.preview = preview_from_rollout_items(&history_items); thread.forked_from_id = Some(source_thread_id.to_string()); if include_turns { populate_thread_turns_from_history( &mut thread, - &history_items, + &materialized_history_items, /*active_turn*/ None, ); } @@ -3430,6 +3522,7 @@ impl ThreadRequestProcessor { if let Some(name) = source_thread_name { set_thread_name_from_title(&mut thread, name); } + thread.preview = preview_from_rollout_items(&materialized_history_items); thread.session_id = session_configured.session_id.to_string(); thread.thread_source = forked_thread .config_snapshot() @@ -3478,7 +3571,7 @@ impl ThreadRequestProcessor { // instead of rebuilding history only to attribute a historical update. if let Some(token_usage_thread) = token_usage_thread { let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( - &history_items, + &materialized_history_items, token_usage_thread.turns.as_slice(), ); // Mirror the resume contract for forks: the new thread is usable as soon 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 cc93600f2d..472ce49ae2 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 @@ -1095,6 +1095,7 @@ mod thread_processor_behavior_tests { let session_meta = SessionMeta { id: conversation_id, + segment_id: None, forked_from_id: Some(forked_from_id), timestamp: timestamp.clone(), model_provider: Some("test-provider".to_string()), diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 6d446e1c42..84c926da4a 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -179,6 +179,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( // Build JSONL lines let meta = SessionMeta { id: conversation_id, + segment_id: None, forked_from_id: None, parent_thread_id, timestamp: meta_rfc3339.to_string(), @@ -265,6 +266,7 @@ pub fn create_fake_rollout_with_text_elements( // Build JSONL lines let meta = SessionMeta { id: conversation_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: meta_rfc3339.to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_list.rs b/codex-rs/app-server/tests/suite/v2/thread_list.rs index 0b49647128..5f03fae2af 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_list.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_list.rs @@ -1890,6 +1890,71 @@ async fn thread_list_archived_filter() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_list_keeps_rotated_live_thread_in_recent_results() -> Result<()> { + let codex_home = TempDir::new()?; + create_minimal_config(codex_home.path())?; + + let thread_id = create_fake_rollout( + codex_home.path(), + "2025-03-01T10-00-00", + "2025-03-01T10:00:00Z", + "Active after rotation", + Some("mock_provider"), + /*git_info*/ None, + )?; + let live_rollout_path = rollout_path(codex_home.path(), "2025-03-01T10-00-00", &thread_id); + let legacy_rotated_path = codex_home.path().join(ARCHIVED_SESSIONS_SUBDIR).join( + live_rollout_path + .file_name() + .expect("live rollout should have a file name"), + ); + fs::create_dir_all( + legacy_rotated_path + .parent() + .expect("legacy rotated rollout should have a parent"), + )?; + fs::copy(live_rollout_path.as_path(), legacy_rotated_path.as_path())?; + + let mut mcp = init_mcp(codex_home.path()).await?; + let state_db = + codex_state::StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()) + .await?; + let thread_id = ThreadId::from_string(&thread_id)?; + let mut metadata = state_db + .get_thread(thread_id) + .await? + .expect("thread should be backfilled into sqlite"); + metadata.rollout_path = legacy_rotated_path; + metadata.archived_at = Some(Utc::now()); + state_db.upsert_thread(&metadata).await?; + + let ThreadListResponse { data, .. } = list_threads( + &mut mcp, + /*cursor*/ None, + Some(10), + Some(vec!["mock_provider".to_string()]), + /*source_kinds*/ None, + /*archived*/ None, + ) + .await?; + assert_eq!(data.len(), 1); + assert_eq!(data[0].id, thread_id.to_string()); + + let ThreadListResponse { data, .. } = list_threads( + &mut mcp, + /*cursor*/ None, + Some(10), + Some(vec!["mock_provider".to_string()]), + /*source_kinds*/ None, + Some(true), + ) + .await?; + assert!(data.is_empty()); + + Ok(()) +} + #[tokio::test] async fn thread_list_invalid_cursor_returns_error() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 272ecac79f..9595319125 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::TestAppServer; +use app_test_support::create_fake_rollout; use app_test_support::create_fake_rollout_with_text_elements; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::rollout_path; @@ -51,6 +52,7 @@ use codex_protocol::models::BaseInstructions; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SessionSource as ProtocolSessionSource; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::UserMessageEvent; @@ -505,6 +507,152 @@ async fn thread_read_loaded_include_turns_reads_store_history_without_rollout_pa Ok(()) } +#[tokio::test] +async fn thread_read_unloaded_include_turns_materializes_rollout_reference() -> Result<()> { + let codex_home = TempDir::new()?; + let referenced_thread_id = create_fake_rollout( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "history before segment rotation", + Some("mock_provider"), + /*git_info*/ None, + )?; + let referenced_rollout_path = rollout_path( + codex_home.path(), + "2025-01-05T12-00-00", + &referenced_thread_id, + ); + + let thread_id = codex_protocol::ThreadId::from_string("00000000-0000-4000-8000-000000000125")?; + let store_id = Uuid::new_v4().to_string(); + create_config_toml_with_thread_store(codex_home.path(), &store_id)?; + let store = InMemoryThreadStore::for_id(store_id.clone()); + let _in_memory_store = InMemoryThreadStoreId { store_id }; + store + .create_thread(CreateThreadParams { + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: ProtocolSessionSource::Cli, + thread_source: None, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + multi_agent_version: None, + metadata: ThreadPersistenceMetadata { + cwd: None, + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Disabled, + }, + }) + .await?; + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: referenced_rollout_path, + thread_id: Some(codex_protocol::ThreadId::from_string( + &referenced_thread_id, + )?), + rollout_timestamp: None, + segment_id: None, + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "history after segment rotation".to_string(), + ..Default::default() + })), + ], + }) + .await?; + + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .loader_overrides(loader_overrides.clone()) + .build() + .await?; + let client = in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + config: Arc::new(config), + cli_overrides: Vec::new(), + loader_overrides, + strict_config: false, + cloud_config_bundle: CloudConfigBundleLoader::default(), + thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + config_warnings: Vec::new(), + session_source: SessionSource::Cli.into(), + enable_codex_api_key_env: false, + initialize: InitializeParams { + client_info: ClientInfo { + name: "codex-app-server-tests".to_string(), + title: None, + version: "0.1.0".to_string(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + }, + channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await?; + + let result = client + .request(ClientRequest::ThreadRead { + request_id: RequestId::Integer(1), + params: ThreadReadParams { + thread_id: thread_id.to_string(), + include_turns: true, + }, + }) + .await? + .expect("thread/read should succeed"); + let ThreadReadResponse { thread, .. } = serde_json::from_value(result)?; + + assert_eq!( + turn_user_texts(&thread.turns), + vec![ + "history before segment rotation", + "history after segment rotation", + ] + ); + + let result = client + .request(ClientRequest::ThreadTurnsList { + request_id: RequestId::Integer(2), + params: ThreadTurnsListParams { + thread_id: thread_id.to_string(), + cursor: None, + limit: Some(10), + sort_direction: Some(SortDirection::Asc), + items_view: None, + }, + }) + .await? + .expect("thread/turns/list should succeed"); + let ThreadTurnsListResponse { data, .. } = serde_json::from_value(result)?; + assert_eq!( + turn_user_texts(&data), + vec![ + "history before segment rotation", + "history after segment rotation", + ] + ); + + client.shutdown().await?; + Ok(()) +} + #[tokio::test] async fn thread_list_includes_store_thread_without_rollout_path() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index cf797ff6be..7a9b9da5bf 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -53,18 +53,22 @@ use codex_app_server_protocol::UserInput; use codex_config::types::AuthCredentialsStoreMode; use codex_core::ARCHIVED_SESSIONS_SUBDIR; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; +use codex_protocol::SegmentId; use codex_protocol::ThreadId; use codex_protocol::config_types::Personality; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::ContextCompactedEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ImageGenerationEndEvent; use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::McpToolCallEndEvent; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource as RolloutSessionSource; @@ -74,6 +78,7 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::UserMessageEvent; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; use codex_rollout::append_rollout_item_to_path; @@ -115,6 +120,50 @@ fn normalized_existing_path(path: impl AsRef) -> Result { Ok(AbsolutePathBuf::from_absolute_path(path.as_ref().canonicalize()?)?.into_path_buf()) } +fn session_meta_rollout_item( + thread_id: ThreadId, + segment_id: Option, + timestamp: &str, +) -> RolloutItem { + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + segment_id, + timestamp: timestamp.to_string(), + cwd: PathBuf::from("/"), + originator: "codex".to_string(), + cli_version: "0.0.0".to_string(), + source: RolloutSessionSource::Cli, + model_provider: Some("mock_provider".to_string()), + ..SessionMeta::default() + }, + git: None, + }) +} + +fn user_message_item(text: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: text.to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + ..Default::default() + })) +} + +fn write_rollout_items(path: &Path, timestamp: &str, items: &[RolloutItem]) -> Result<()> { + let mut jsonl = String::new(); + for item in items { + jsonl.push_str(&serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item: item.clone(), + })?); + jsonl.push('\n'); + } + std::fs::write(path, jsonl)?; + Ok(()) +} + async fn wait_for_responses_request_count( server: &wiremock::MockServer, expected_count: usize, @@ -826,6 +875,107 @@ fn append_resume_redaction_history( Ok(()) } +#[tokio::test] +async fn thread_resume_materializes_rollout_references_when_session_segmentation_is_disabled() +-> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let thread_id = ThreadId::new(); + let thread_id_str = thread_id.to_string(); + let old_ts = "2025-01-05T12-00-00"; + let current_ts = "2025-01-05T12-05-00"; + let old_rfc3339 = "2025-01-05T12:00:00Z"; + let current_rfc3339 = "2025-01-05T12:05:00Z"; + let active_old_path = rollout_path(codex_home.path(), old_ts, &thread_id_str); + let archived_old_path = codex_home + .path() + .join("archived_sessions") + .join(active_old_path.file_name().expect("old rollout file name")); + std::fs::create_dir_all(archived_old_path.parent().expect("archived rollout parent"))?; + let current_path = rollout_path(codex_home.path(), current_ts, &thread_id_str); + std::fs::create_dir_all(current_path.parent().expect("current rollout parent"))?; + + let old_items = vec![ + session_meta_rollout_item(thread_id, Some(SegmentId::new()), old_rfc3339), + user_message_item("before compaction"), + ]; + write_rollout_items(archived_old_path.as_path(), old_rfc3339, &old_items)?; + + let current_items = vec![ + session_meta_rollout_item(thread_id, Some(SegmentId::new()), current_rfc3339), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: active_old_path, + thread_id: Some(thread_id), + rollout_timestamp: Some(old_ts.to_string()), + segment_id: None, + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::EventMsg(EventMsg::ContextCompacted(ContextCompactedEvent {})), + user_message_item("after compaction"), + ]; + write_rollout_items(current_path.as_path(), current_rfc3339, ¤t_items)?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread_id_str, + ..Default::default() + }) + .await?; + let resume_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + + let user_messages: Vec<&str> = thread + .turns + .iter() + .flat_map(|turn| turn.items.iter()) + .filter_map(|item| match item { + ThreadItem::UserMessage { content, .. } => { + content.iter().find_map(|input| match input { + UserInput::Text { text, .. } => Some(text.as_str()), + UserInput::Image { .. } + | UserInput::LocalImage { .. } + | UserInput::Skill { .. } + | UserInput::Mention { .. } => None, + }) + } + ThreadItem::AgentMessage { .. } + | ThreadItem::Reasoning { .. } + | ThreadItem::CommandExecution { .. } + | ThreadItem::FileChange { .. } + | ThreadItem::McpToolCall { .. } + | ThreadItem::WebSearch { .. } + | ThreadItem::ImageView { .. } + | ThreadItem::ImageGeneration { .. } + | ThreadItem::EnteredReviewMode { .. } + | ThreadItem::ExitedReviewMode { .. } + | ThreadItem::ContextCompaction { .. } + | ThreadItem::HookPrompt { .. } + | ThreadItem::CollabAgentToolCall { .. } + | ThreadItem::SubAgentActivity { .. } + | ThreadItem::DynamicToolCall { .. } + | ThreadItem::Plan { .. } => None, + }) + .collect(); + assert_eq!(user_messages, vec!["before compaction", "after compaction"]); + assert!(thread.turns.iter().any(|turn| { + turn.items + .iter() + .any(|item| matches!(item, ThreadItem::ContextCompaction { .. })) + })); + Ok(()) +} + #[tokio::test] async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; @@ -1940,6 +2090,7 @@ stream_max_retries = 0 std::fs::create_dir_all(rollout_dir)?; let session_meta = SessionMeta { id: conversation_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: "2025-01-05T12:00:00Z".to_string(), diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index dcfa1b83a9..add512cec3 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -56,6 +56,7 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: | ResponseItem::Other, ) => false, RolloutItem::InterAgentCommunication(_) => false, + RolloutItem::RolloutReference(_) => true, // Full-history forks preserve the cached prompt prefix and can keep diffing // from the parent's durable baseline. Truncated forks drop part of that prompt, // so they must rebuild context on their first child turn. diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 073c06cede..c809694851 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -119,6 +119,8 @@ pub use thread_manager::ThreadManager; pub use thread_manager::ThreadShutdownReport; pub use thread_manager::build_models_manager; pub use thread_manager::thread_store_from_config; +pub use thread_rollout_truncation::materialize_rollout_items_for_complete_history; +pub use thread_rollout_truncation::materialize_rollout_items_for_model_replay; pub use web_search::web_search_action_detail; pub use web_search::web_search_detail; pub use windows_sandbox_read_grants::grant_read_root_non_elevated; diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 4cf8ab9d8f..f05afa354b 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -44,6 +44,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R let session_meta = SessionMetaLine { meta: SessionMeta { id: thread_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: TEST_TIMESTAMP.to_string(), diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs index 73fbe5b5e4..7a84b17c8a 100644 --- a/codex-rs/core/src/rollout.rs +++ b/codex-rs/core/src/rollout.rs @@ -21,6 +21,7 @@ pub use codex_rollout::find_thread_path_by_id_str; pub use codex_rollout::parse_cursor; pub use codex_rollout::read_head_for_summary; pub use codex_rollout::read_session_meta_line; +pub use codex_rollout::resolve_rollout_reference_rollout_path; pub use codex_rollout::rollout_date_parts; impl codex_rollout::RolloutConfigView for Config { diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index c844a630fc..117885ebf8 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -225,7 +225,9 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); active_segment.counts_as_user_turn = true; } - RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::EventMsg(_) + | RolloutItem::RolloutReference(_) + | RolloutItem::SessionMeta(_) => {} } if base_replacement_history.is_some() @@ -309,6 +311,7 @@ impl Session { history.drop_last_n_user_turns(rollback.num_turns); } RolloutItem::EventMsg(_) + | RolloutItem::RolloutReference(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 1601a66543..158dd299e2 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -4,6 +4,7 @@ use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; use crate::skills::SkillError; use crate::state::ActiveTurn; +use crate::thread_rollout_truncation::materialize_initial_history_for_model_replay; use codex_extension_api::ExtensionDataInit; use codex_protocol::SessionId; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; @@ -499,6 +500,11 @@ impl Session { session_configuration.collaboration_mode.model(), session_configuration.provider ); + let initial_history = materialize_initial_history_for_model_replay( + config.codex_home.as_path(), + initial_history, + ) + .await; let forked_from_id = session_configuration .forked_from_thread_id .or_else(|| initial_history.forked_from_id()); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c332622acb..e9b00cbb08 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -673,6 +673,11 @@ impl ThreadManager { let inherited_multi_agent_version = fork_source .multi_agent_version() .unwrap_or(MultiAgentVersion::V1); + let history = truncation::materialize_initial_history_for_model_replay( + options.config.codex_home.as_path(), + history, + ) + .await; options.initial_history = fork_history_from_snapshot( ForkSnapshot::Interrupted, history, @@ -946,6 +951,11 @@ impl ThreadManager { .await; let interrupted_marker = InterruptedTurnHistoryMarker::from_config_and_version(&config, multi_agent_version); + let history = truncation::materialize_initial_history_for_model_replay( + config.codex_home.as_path(), + history, + ) + .await; let history = fork_history_from_snapshot(snapshot, history, interrupted_marker); let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 790f5ab8c8..1b1bfce44d 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -16,9 +16,12 @@ use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::ResumedHistory; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; @@ -160,6 +163,65 @@ fn out_of_range_truncation_drops_only_unfinished_suffix_mid_turn() { ); } +#[tokio::test] +async fn interrupted_snapshot_detects_turn_started_in_referenced_segment() { + let codex_home = tempdir().expect("codex home"); + let predecessor_path = codex_home.path().join("predecessor.jsonl"); + let timestamp = "2026-06-12T00:00:00Z"; + let predecessor_items = [ + RolloutItem::ResponseItem(user_msg("turn started before rotation")), + RolloutItem::ResponseItem(assistant_msg("partial answer")), + ]; + let predecessor_json = predecessor_items + .iter() + .cloned() + .map(|item| { + serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item, + }) + .expect("serialize predecessor item") + }) + .collect::>() + .join("\n"); + tokio::fs::write(&predecessor_path, format!("{predecessor_json}\n")) + .await + .expect("write predecessor"); + let raw_history = InitialHistory::Forked(vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: predecessor_path, + thread_id: None, + rollout_timestamp: None, + segment_id: None, + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::Compacted(CompactedItem { + message: "compacted during active turn".to_string(), + replacement_history: None, + window_id: None, + }), + ]); + let materialized = + crate::thread_rollout_truncation::materialize_initial_history_for_model_replay( + codex_home.path(), + raw_history, + ) + .await; + assert!(snapshot_turn_state(&materialized).ends_mid_turn); + let interrupted = fork_history_from_snapshot( + ForkSnapshot::Interrupted, + materialized, + InterruptedTurnHistoryMarker::Disabled, + ); + + assert!(matches!( + interrupted.get_rollout_items().last(), + Some(RolloutItem::EventMsg(EventMsg::TurnAborted(_))) + )); +} + #[test] fn fork_thread_accepts_legacy_usize_snapshot_argument() { fn assert_legacy_snapshot_callsite( diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index 407586715a..8d5e893be4 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -5,12 +5,20 @@ use crate::context_manager::is_user_turn_boundary; use crate::event_mapping; +use crate::rollout::RolloutRecorder; +use crate::rollout::resolve_rollout_reference_rollout_path; use codex_protocol::items::TurnItem; +use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::RolloutItem; +use std::collections::HashSet; +use std::io; +use std::path::Path; +use std::path::PathBuf; +use tracing::warn; pub(crate) fn initial_history_has_prior_user_turns(conversation_history: &InitialHistory) -> bool { conversation_history.scan_rollout_items(rollout_item_is_user_turn_boundary) @@ -160,6 +168,282 @@ pub(crate) fn truncate_rollout_to_last_n_fork_turns( items[keep_idx..].to_vec() } +#[derive(Clone, Copy)] +enum RolloutMaterialization { + ModelReplay, + CompleteHistory, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum RolloutReferenceIdentity { + Segment(codex_protocol::SegmentId), + Path(PathBuf), +} + +pub async fn materialize_rollout_items_for_model_replay( + codex_home: &Path, + rollout_items: &[RolloutItem], +) -> Vec { + match materialize_rollout_items( + codex_home, + rollout_items, + RolloutMaterialization::ModelReplay, + ) + .await + { + Ok(items) => items, + Err(err) => { + warn!("failed to materialize rollout references for model replay: {err}"); + Vec::new() + } + } +} + +pub async fn materialize_rollout_items_for_complete_history( + codex_home: &Path, + rollout_items: &[RolloutItem], +) -> io::Result> { + materialize_rollout_items( + codex_home, + rollout_items, + RolloutMaterialization::CompleteHistory, + ) + .await +} + +pub(crate) async fn materialize_initial_history_for_model_replay( + codex_home: &Path, + initial_history: InitialHistory, +) -> InitialHistory { + match initial_history { + InitialHistory::New => InitialHistory::New, + InitialHistory::Cleared => InitialHistory::Cleared, + InitialHistory::Resumed(mut resumed) => { + resumed.history = + materialize_rollout_items_for_model_replay(codex_home, &resumed.history).await; + InitialHistory::Resumed(resumed) + } + InitialHistory::Forked(items) => InitialHistory::Forked( + materialize_rollout_items_for_model_replay(codex_home, &items).await, + ), + } +} + +async fn materialize_rollout_items( + codex_home: &Path, + rollout_items: &[RolloutItem], + materialization: RolloutMaterialization, +) -> io::Result> { + enum Work { + Items { + rollout_items: Vec, + remaining_segment_depth: Option, + }, + Item { + item: Box, + remaining_segment_depth: Option, + }, + TruncateSuffix { + start: usize, + nth_user_message: usize, + }, + LeaveReference(RolloutReferenceIdentity), + } + + let mut materialized = Vec::new(); + let mut active_references = HashSet::new(); + let mut work = vec![Work::Items { + rollout_items: rollout_items.to_vec(), + remaining_segment_depth: None, + }]; + while let Some(next) = work.pop() { + match next { + Work::Items { + rollout_items, + remaining_segment_depth, + } => { + for item in rollout_items.into_iter().rev() { + work.push(Work::Item { + item: Box::new(item), + remaining_segment_depth, + }); + } + } + Work::Item { + item, + remaining_segment_depth, + } => { + let reference = match *item { + RolloutItem::RolloutReference(reference) => reference, + item => { + materialized.push(item); + continue; + } + }; + let has_prefix_truncation = reference.nth_user_message.is_some(); + let next_remaining_segment_depth = match materialization { + RolloutMaterialization::CompleteHistory => None, + RolloutMaterialization::ModelReplay if has_prefix_truncation => { + remaining_segment_depth + } + RolloutMaterialization::ModelReplay => { + let available_depth = remaining_segment_depth + .map_or(reference.max_depth, |remaining| { + remaining.min(reference.max_depth) + }); + let Some(next_remaining_depth) = available_depth.checked_sub(1) else { + warn!("rollout reference materialization reached max depth"); + continue; + }; + Some(next_remaining_depth) + } + }; + let resolved_path = + match resolve_rollout_reference_rollout_path(codex_home, &reference).await { + Ok(path) => path, + Err(err) => { + if matches!(materialization, RolloutMaterialization::CompleteHistory) { + return Err(io::Error::new( + err.kind(), + format!( + "failed to resolve rollout reference {}: {err}", + reference.rollout_path.display() + ), + )); + } + warn!( + "failed to resolve rollout reference {}: {err}", + reference.rollout_path.display() + ); + continue; + } + }; + let identity = reference + .segment_id + .map(RolloutReferenceIdentity::Segment) + .unwrap_or_else(|| RolloutReferenceIdentity::Path(resolved_path.clone())); + if !active_references.insert(identity.clone()) { + if matches!(materialization, RolloutMaterialization::CompleteHistory) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "rollout reference cycle detected at {}", + resolved_path.display() + ), + )); + } + warn!( + "rollout reference cycle detected at {}", + resolved_path.display() + ); + continue; + } + match RolloutRecorder::load_rollout_items(&resolved_path).await { + Ok((mut reference_items, _, _)) => { + if let Some(filter_texts) = reference + .compacted_replacement_history_filter_texts + .as_deref() + { + apply_compacted_replacement_history_filter( + &mut reference_items, + filter_texts, + ); + } + work.push(Work::LeaveReference(identity)); + if let Some(nth_user_message) = reference.nth_user_message { + work.push(Work::TruncateSuffix { + start: materialized.len(), + nth_user_message, + }); + } + work.push(Work::Items { + rollout_items: reference_items, + remaining_segment_depth: next_remaining_segment_depth, + }); + } + Err(err) => { + active_references.remove(&identity); + if matches!(materialization, RolloutMaterialization::CompleteHistory) { + return Err(io::Error::new( + err.kind(), + format!( + "failed to load rollout reference {}: {err}", + resolved_path.display() + ), + )); + } + warn!( + "failed to load rollout reference {}: {err}", + resolved_path.display() + ); + } + } + } + Work::TruncateSuffix { + start, + nth_user_message, + } => { + let suffix = truncate_rollout_before_nth_user_message_from_start( + &materialized[start..], + nth_user_message, + ); + materialized.truncate(start); + materialized.extend(suffix); + } + Work::LeaveReference(identity) => { + active_references.remove(&identity); + } + } + } + Ok(materialized) +} + +fn apply_compacted_replacement_history_filter( + rollout_items: &mut [RolloutItem], + filter_texts: &[String], +) { + for item in rollout_items { + match item { + RolloutItem::Compacted(compacted) => { + if let Some(replacement_history) = compacted.replacement_history.as_mut() { + replacement_history.retain(|response_item| { + !matches_filtered_developer_message(response_item, filter_texts) + }); + } + } + RolloutItem::RolloutReference(reference) => { + let nested_filter_texts = reference + .compacted_replacement_history_filter_texts + .get_or_insert_default(); + for filter_text in filter_texts { + if !nested_filter_texts.contains(filter_text) { + nested_filter_texts.push(filter_text.clone()); + } + } + } + RolloutItem::SessionMeta(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::InterAgentCommunication(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => {} + } + } +} + +fn matches_filtered_developer_message(item: &ResponseItem, filter_texts: &[String]) -> bool { + let ResponseItem::Message { role, content, .. } = item else { + return false; + }; + if role != "developer" { + return false; + } + let [ContentItem::InputText { text }] = content.as_slice() else { + return false; + }; + + filter_texts.iter().any(|filter_text| filter_text == text) +} + fn is_real_user_message_boundary(item: &ResponseItem) -> bool { matches!( event_mapping::parse_turn_item(item), diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index 6526827769..bc1d5b6b48 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -1,11 +1,20 @@ use super::*; use crate::session::tests::make_session_and_context; use codex_protocol::AgentPath; +use codex_protocol::SegmentId; +use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; +use codex_protocol::protocol::CompactedItem; +use codex_protocol::protocol::DEFAULT_ROLLOUT_REFERENCE_DEPTH; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::RolloutReferenceItem; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::ThreadRolledBackEvent; use pretty_assertions::assert_eq; +use std::path::PathBuf; fn user_msg(text: &str) -> ResponseItem { ResponseItem::Message { @@ -61,6 +70,34 @@ fn inter_agent_communication(text: &str, trigger_turn: bool) -> RolloutItem { )) } +async fn write_rollout(path: &std::path::Path, items: &[RolloutItem]) { + let mut jsonl = String::new(); + for item in items { + let line = RolloutLine { + timestamp: "2026-04-30T00:00:00.000Z".to_string(), + item: item.clone(), + }; + jsonl.push_str(&serde_json::to_string(&line).expect("serialize rollout line")); + jsonl.push('\n'); + } + tokio::fs::write(path, jsonl).await.expect("write rollout"); +} + +fn session_meta_item(thread_id: ThreadId, segment_id: SegmentId) -> RolloutItem { + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + segment_id: Some(segment_id), + timestamp: "2026-04-30T00:00:00.000Z".to_string(), + cwd: PathBuf::from("/tmp"), + originator: "test".to_string(), + cli_version: "0.0.0".to_string(), + ..SessionMeta::default() + }, + git: None, + }) +} + #[test] fn truncates_rollout_from_start_before_nth_user_only() { let items = [ @@ -129,6 +166,454 @@ fn truncation_max_keeps_full_rollout() { ); } +#[test] +fn legacy_fork_reference_decodes_as_rollout_reference() { + let item: RolloutItem = serde_json::from_value(serde_json::json!({ + "type": "fork_reference", + "payload": { + "rollout_path": "/tmp/legacy.jsonl", + "nth_user_message": 3 + } + })) + .expect("deserialize legacy fork reference"); + + let RolloutItem::RolloutReference(reference) = item else { + panic!("legacy fork reference should decode as rollout reference"); + }; + assert_eq!(reference.rollout_path, PathBuf::from("/tmp/legacy.jsonl")); + assert_eq!(reference.max_depth, DEFAULT_ROLLOUT_REFERENCE_DEPTH); + assert_eq!(reference.nth_user_message, Some(3)); + + let value = serde_json::to_value(RolloutItem::RolloutReference(reference)) + .expect("serialize rollout reference"); + assert_eq!(value["type"], "rollout_reference"); +} + +#[tokio::test] +async fn materializes_prefix_truncated_rollout_reference_before_replay() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let segment_id = SegmentId::new(); + let source_path = temp + .path() + .join("rollout-2026-04-30T00-00-00-00000000-0000-0000-0000-000000000001.jsonl"); + let source_items = vec![ + session_meta_item(thread_id, segment_id), + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + RolloutItem::ResponseItem(user_msg("u2")), + RolloutItem::ResponseItem(assistant_msg("a2")), + ]; + write_rollout(&source_path, &source_items).await; + + let compact_reference = vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: source_path.clone(), + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: Some(1), + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("child request")), + ]; + + let materialized = + materialize_rollout_items_for_model_replay(temp.path(), &compact_reference).await; + + let expected = vec![ + session_meta_item(thread_id, segment_id), + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + RolloutItem::ResponseItem(user_msg("child request")), + ]; + assert_eq!( + serde_json::to_value(&materialized).unwrap(), + serde_json::to_value(&expected).unwrap() + ); +} + +#[tokio::test] +async fn materializes_prefix_truncated_rollout_reference_by_segment_id_after_source_rollover() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let old_segment_id = SegmentId::new(); + let new_segment_id = SegmentId::new(); + + let old_active_path = temp + .path() + .join("sessions/2026/04/30") + .join(format!("rollout-2026-04-30T00-00-00-{thread_id}.jsonl")); + let old_archived_path = temp + .path() + .join("archived_sessions/2026/04/30") + .join(format!("rollout-2026-04-30T00-00-00-{thread_id}.jsonl")); + let new_active_path = temp + .path() + .join("sessions/2026/05/01") + .join(format!("rollout-2026-05-01T00-00-00-{thread_id}.jsonl")); + + tokio::fs::create_dir_all(old_archived_path.parent().expect("archived parent")) + .await + .expect("create archived parent"); + tokio::fs::create_dir_all(new_active_path.parent().expect("active parent")) + .await + .expect("create active parent"); + + write_rollout( + &old_archived_path, + &[ + session_meta_item(thread_id, old_segment_id), + RolloutItem::ResponseItem(user_msg("old segment request")), + RolloutItem::ResponseItem(assistant_msg("old segment answer")), + ], + ) + .await; + write_rollout( + &new_active_path, + &[ + session_meta_item(thread_id, new_segment_id), + RolloutItem::ResponseItem(user_msg("new segment request")), + RolloutItem::ResponseItem(assistant_msg("new segment answer")), + ], + ) + .await; + + let compact_reference = vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: old_active_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(old_segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: Some(usize::MAX), + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("child request")), + ]; + + let materialized = + materialize_rollout_items_for_model_replay(temp.path(), &compact_reference).await; + let text = serde_json::to_string(&materialized).expect("serialize materialized rollout"); + + assert!(text.contains("old segment request")); + assert!(!text.contains("new segment request")); + assert!(text.contains("child request")); +} + +#[tokio::test] +async fn materializes_prefix_truncated_reference_before_bounded_rollout_references() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let old_segment_id = SegmentId::new(); + let current_segment_id = SegmentId::new(); + + let old_path = temp.path().join("old.jsonl"); + let current_path = temp.path().join("current.jsonl"); + + write_rollout( + &old_path, + &[ + session_meta_item(thread_id, old_segment_id), + RolloutItem::ResponseItem(user_msg("u1")), + RolloutItem::ResponseItem(assistant_msg("a1")), + ], + ) + .await; + write_rollout( + ¤t_path, + &[ + session_meta_item(thread_id, current_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: old_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(old_segment_id), + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("u2")), + RolloutItem::ResponseItem(assistant_msg("a2")), + RolloutItem::ResponseItem(user_msg("u3")), + ], + ) + .await; + + let reference_items = vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: current_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(current_segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: Some(2), + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("child request")), + ]; + + let materialized = + materialize_rollout_items_for_model_replay(temp.path(), &reference_items).await; + let text = serde_json::to_string(&materialized).expect("serialize materialized rollout"); + + assert!(text.contains("u1")); + assert!(text.contains("u2")); + assert!(!text.contains("u3")); + assert!(text.contains("child request")); +} + +#[tokio::test] +async fn materializes_rollout_reference_with_bounded_depth() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let oldest_segment_id = SegmentId::new(); + let old_segment_id = SegmentId::new(); + let middle_segment_id = SegmentId::new(); + let current_segment_id = SegmentId::new(); + + let oldest_path = temp.path().join("oldest.jsonl"); + let old_path = temp.path().join("old.jsonl"); + let middle_path = temp.path().join("middle.jsonl"); + let current_path = temp.path().join("current.jsonl"); + + write_rollout( + &oldest_path, + &[ + session_meta_item(thread_id, oldest_segment_id), + RolloutItem::ResponseItem(user_msg("oldest segment request")), + ], + ) + .await; + write_rollout( + &old_path, + &[ + session_meta_item(thread_id, old_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: oldest_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(oldest_segment_id), + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("old segment request")), + ], + ) + .await; + write_rollout( + &middle_path, + &[ + session_meta_item(thread_id, middle_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: old_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(old_segment_id), + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("middle segment request")), + ], + ) + .await; + + let current_items = vec![ + session_meta_item(thread_id, current_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: middle_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(middle_segment_id), + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("current segment request")), + ]; + write_rollout(¤t_path, ¤t_items).await; + let materialized = + materialize_rollout_items_for_model_replay(temp.path(), ¤t_items).await; + let text = serde_json::to_string(&materialized).expect("serialize materialized rollout"); + + assert!(!text.contains("oldest segment request")); + assert!(text.contains("old segment request")); + assert!(text.contains("middle segment request")); + assert!(text.contains("current segment request")); + + let complete_history = + materialize_rollout_items_for_complete_history(temp.path(), ¤t_items) + .await + .expect("materialize complete history"); + let complete_text = + serde_json::to_string(&complete_history).expect("serialize complete rollout"); + assert!(complete_text.contains("oldest segment request")); + assert!(complete_text.contains("old segment request")); + assert!(complete_text.contains("middle segment request")); + assert!(complete_text.contains("current segment request")); + + let fork_items = vec![RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: current_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(current_segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: Some(usize::MAX), + compacted_replacement_history_filter_texts: None, + })]; + let fork_materialized = + materialize_rollout_items_for_model_replay(temp.path(), &fork_items).await; + assert_eq!( + serde_json::to_value(&fork_materialized).unwrap(), + serde_json::to_value(&materialized).unwrap() + ); +} + +#[tokio::test] +async fn composes_compacted_history_filters_across_nested_references() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let oldest_segment_id = SegmentId::new(); + let middle_segment_id = SegmentId::new(); + let oldest_path = temp.path().join("oldest.jsonl"); + let middle_path = temp.path().join("middle.jsonl"); + let oldest_filter = developer_msg("oldest injected instructions"); + let outer_filter = developer_msg("outer injected instructions"); + write_rollout( + &oldest_path, + &[ + session_meta_item(thread_id, oldest_segment_id), + RolloutItem::Compacted(CompactedItem { + message: "summary".to_string(), + replacement_history: Some(vec![ + oldest_filter.clone(), + outer_filter.clone(), + developer_msg("retained instructions"), + ]), + window_id: None, + }), + ], + ) + .await; + write_rollout( + &middle_path, + &[ + session_meta_item(thread_id, middle_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: oldest_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(oldest_segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: None, + compacted_replacement_history_filter_texts: Some(vec![ + "oldest injected instructions".to_string(), + ]), + }), + ], + ) + .await; + + let materialized = materialize_rollout_items_for_complete_history( + temp.path(), + &[RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: middle_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(middle_segment_id), + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: None, + compacted_replacement_history_filter_texts: Some(vec![ + "outer injected instructions".to_string(), + ]), + })], + ) + .await + .expect("materialize complete history"); + let serialized = serde_json::to_string(&materialized).expect("serialize materialized rollout"); + + assert!(!serialized.contains("oldest injected instructions")); + assert!(!serialized.contains("outer injected instructions")); + assert!(serialized.contains("retained instructions")); +} + +#[tokio::test] +async fn outer_reference_depth_bounds_descendant_references() { + let temp = tempfile::tempdir().expect("tempdir"); + let thread_id = ThreadId::new(); + let oldest_segment_id = SegmentId::new(); + let middle_segment_id = SegmentId::new(); + let oldest_path = temp.path().join("oldest.jsonl"); + let middle_path = temp.path().join("middle.jsonl"); + write_rollout( + &oldest_path, + &[ + session_meta_item(thread_id, oldest_segment_id), + RolloutItem::ResponseItem(user_msg("oldest request")), + ], + ) + .await; + write_rollout( + &middle_path, + &[ + session_meta_item(thread_id, middle_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: oldest_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(oldest_segment_id), + max_depth: usize::MAX, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + RolloutItem::ResponseItem(user_msg("middle request")), + ], + ) + .await; + + let materialized = materialize_rollout_items_for_model_replay( + temp.path(), + &[RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: middle_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: Some(middle_segment_id), + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + })], + ) + .await; + let serialized = serde_json::to_string(&materialized).expect("serialize materialized rollout"); + + assert!(!serialized.contains("oldest request")); + assert!(serialized.contains("middle request")); +} + +#[tokio::test] +async fn complete_history_reports_unresolvable_references() { + let temp = tempfile::tempdir().expect("tempdir"); + let missing_path = temp.path().join("missing.jsonl"); + let items = [RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: missing_path, + thread_id: None, + rollout_timestamp: None, + segment_id: None, + max_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + })]; + + let err = materialize_rollout_items_for_complete_history(temp.path(), &items) + .await + .expect_err("complete history must not silently drop a missing predecessor"); + + assert_eq!(err.kind(), std::io::ErrorKind::NotFound); +} + #[test] fn truncates_rollout_from_start_applies_thread_rollback_markers() { let rollout_items = vec![ diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 5302274153..bfadbb2d97 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -60,6 +60,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R let session_meta = SessionMetaLine { meta: SessionMeta { id: thread_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: TEST_TIMESTAMP.to_string(), @@ -110,6 +111,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re let session_meta = SessionMetaLine { meta: SessionMeta { id: thread_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: TEST_TIMESTAMP.to_string(), diff --git a/codex-rs/core/tests/suite/resume.rs b/codex-rs/core/tests/suite/resume.rs index 8f760db417..01faaf8fd3 100644 --- a/codex-rs/core/tests/suite/resume.rs +++ b/codex-rs/core/tests/suite/resume.rs @@ -1,6 +1,9 @@ use anyhow::Result; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; use codex_protocol::user_input::UserInput; @@ -145,6 +148,95 @@ async fn resume_includes_initial_messages_from_rollout_events() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resume_materializes_references_before_session_configured() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut builder = test_codex(); + let initial = builder.build(&server).await?; + let codex = Arc::clone(&initial.codex); + let home = initial.home.clone(); + let thread_id = initial.session_configured.thread_id; + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-initial"), + ev_assistant_message("msg-1", "Referenced answer"), + ev_completed("resp-initial"), + ]), + ) + .await; + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "Referenced question".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + codex.shutdown_and_wait().await?; + + let predecessor_path = rollout_path.with_file_name("predecessor.jsonl"); + let persisted = std::fs::read_to_string(&rollout_path)?; + let session_meta_line = persisted.lines().next().expect("session meta line"); + std::fs::rename(&rollout_path, &predecessor_path)?; + let reference_line = serde_json::to_string(&RolloutLine { + timestamp: "2026-06-12T00:00:00Z".to_string(), + item: RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: predecessor_path, + thread_id: Some(thread_id), + rollout_timestamp: None, + segment_id: None, + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + })?; + std::fs::write( + &rollout_path, + format!("{session_meta_line}\n{reference_line}\n"), + )?; + + let resumed = resume_until_initial_messages( + &mut builder, + &server, + home, + rollout_path, + |initial_messages| { + initial_messages.iter().any(|event| { + matches!( + event, + EventMsg::UserMessage(message) if message.message == "Referenced question" + ) + }) + }, + ) + .await?; + let initial_messages = resumed + .session_configured + .initial_messages + .expect("initial messages"); + + assert!(initial_messages.iter().any(|event| { + matches!( + event, + EventMsg::AgentMessage(message) if message.message == "Referenced answer" + ) + })); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resume_includes_initial_messages_from_reasoning_events() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index b4809e8a39..ea72d763dd 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -217,6 +217,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { let session_meta_line = SessionMetaLine { meta: SessionMeta { id: thread_id, + segment_id: None, forked_from_id: None, parent_thread_id: None, timestamp: "2026-01-27T12:00:00Z".to_string(), diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 3edad7ff75..9fbfc92c43 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -9,6 +9,7 @@ use codex_config::types::MemoriesConfig; use codex_core::Prompt; use codex_core::RolloutRecorder; use codex_core::config::Config; +use codex_core::materialize_rollout_items_for_complete_history; use codex_protocol::error::CodexErr; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; @@ -287,8 +288,9 @@ mod job { rollout_cwd: &Path, stage_one_context: &StageOneRequestContext, ) -> anyhow::Result<(StageOneOutput, Option)> { - let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; - let rollout_contents = serialize_filtered_rollout_response_items(&rollout_items)?; + let rollout_contents = + serialize_complete_rollout_for_memories(config.codex_home.as_path(), rollout_path) + .await?; let mut prompt = Prompt::default(); prompt.input = vec![ResponseItem::Message { @@ -322,6 +324,16 @@ mod job { Ok((output, token_usage)) } + async fn serialize_complete_rollout_for_memories( + codex_home: &Path, + rollout_path: &Path, + ) -> anyhow::Result { + let (rollout_items, _, _) = RolloutRecorder::load_rollout_items(rollout_path).await?; + let rollout_items = + materialize_rollout_items_for_complete_history(codex_home, &rollout_items).await?; + Ok(serialize_filtered_rollout_response_items(&rollout_items)?) + } + mod result { use super::*; @@ -411,6 +423,7 @@ mod job { Some(communication.to_model_input_item()) } RolloutItem::SessionMeta(_) + | RolloutItem::RolloutReference(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -482,6 +495,12 @@ mod job { #[cfg(test)] mod tests { use super::*; + use codex_protocol::ThreadId; + use codex_protocol::protocol::RolloutLine; + use codex_protocol::protocol::RolloutReferenceItem; + use codex_protocol::protocol::SessionMeta; + use codex_protocol::protocol::SessionMetaLine; + use tempfile::TempDir; #[test] fn classifies_memory_excluded_fragments() { @@ -560,6 +579,100 @@ mod job { ); assert_eq!(rollout_slug_types, vec!["null", "string"]); } + + #[tokio::test] + async fn serializes_referenced_predecessors_for_memory_extraction() { + let codex_home = TempDir::new().expect("create codex home"); + let predecessor_path = codex_home.path().join("predecessor.jsonl"); + let current_path = codex_home.path().join("current.jsonl"); + let predecessor = RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "before rotation".to_string(), + }], + phase: None, + }); + let current = RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "after rotation".to_string(), + }], + phase: None, + }); + let timestamp = "2026-06-12T00:00:00Z"; + let predecessor_json = [ + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: ThreadId::new(), + timestamp: timestamp.to_string(), + ..SessionMeta::default() + }, + git: None, + }), + predecessor.clone(), + ] + .into_iter() + .map(|item| { + serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item, + }) + .expect("serialize predecessor") + }) + .collect::>() + .join("\n"); + tokio::fs::write(&predecessor_path, format!("{predecessor_json}\n")) + .await + .expect("write predecessor"); + let current_json = [ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: predecessor_path, + thread_id: None, + rollout_timestamp: None, + segment_id: None, + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }), + current.clone(), + ] + .into_iter() + .map(|item| { + serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item, + }) + .expect("serialize current rollout line") + }) + .collect::>() + .join("\n"); + tokio::fs::write(¤t_path, format!("{current_json}\n")) + .await + .expect("write current rollout"); + + let serialized = + serialize_complete_rollout_for_memories(codex_home.path(), current_path.as_path()) + .await + .expect("serialize complete rollout"); + let response_items: Vec = + serde_json::from_str(&serialized).expect("parse serialized response items"); + + assert_eq!( + response_items, + vec![ + match predecessor { + RolloutItem::ResponseItem(item) => item, + _ => unreachable!(), + }, + match current { + RolloutItem::ResponseItem(item) => item, + _ => unreachable!(), + }, + ] + ); + } } } diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 6f0792a19b..1786552966 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -2,6 +2,7 @@ use codex_utils_path as path_utils; use std::cmp::Reverse; +use std::collections::HashSet; use std::ffi::OsStr; use std::io; use std::num::NonZero; @@ -1095,19 +1096,26 @@ impl<'a> ProviderMatcher<'a> { } async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { - const MAX_SUMMARY_REFERENCE_DEPTH: usize = 8; + const MAX_SUMMARY_REFERENCE_FILES: usize = 128; - let (mut summary, mut reference) = read_head_summary_file(path, head_limit).await?; - let mut referenced_summaries = Vec::new(); - let mut remaining_depth = MAX_SUMMARY_REFERENCE_DEPTH; - while remaining_depth > 0 { + let (summary, mut reference) = read_head_summary_file(path, head_limit).await?; + let mut summaries = vec![summary]; + let mut segment_edges = Vec::new(); + let mut visited_references = HashSet::new(); + loop { let Some(next_reference) = reference.take() else { break; }; - if next_reference.max_depth == 0 || next_reference.nth_user_message.is_some() { + if next_reference.nth_user_message == Some(0) { break; } - let reference_depth = next_reference.max_depth; + if visited_references.len() >= MAX_SUMMARY_REFERENCE_FILES { + tracing::warn!( + "stopped reading rollout reference metadata after {MAX_SUMMARY_REFERENCE_FILES} files" + ); + break; + } + let is_segment_edge = next_reference.nth_user_message.is_none(); let reference_path = if let Some(codex_home) = codex_home_from_rollout_path(path) { match Box::pin(resolve_rollout_reference_rollout_path( codex_home, @@ -1125,33 +1133,34 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result io::Result> { + let Some(path) = compression::existing_rollout_path(path).await else { + return Ok(None); }; - let Some((rollout_timestamp, uuid)) = parse_timestamp_string_uuid_from_filename(file_name) - else { - return Ok(rollout_path.to_path_buf()); + let Ok(meta) = read_session_meta_line(path.as_path()).await else { + return Ok(None); }; - let archived_path = codex_home.join(ARCHIVED_SESSIONS_SUBDIR).join(file_name); - if tokio::fs::try_exists(archived_path.as_path()) - .await - .unwrap_or(false) + if (meta.meta.segment_id.is_some() && path.starts_with(codex_home.join(SESSIONS_SUBDIR))) + || reference + .thread_id + .is_some_and(|thread_id| meta.meta.id != thread_id) { - return Ok(archived_path); + return Ok(None); } - if let Some(active_path) = - rollout_path_for_timestamp_file(codex_home, rollout_timestamp, file_name) - && tokio::fs::try_exists(active_path.as_path()) - .await - .unwrap_or(false) - { - return Ok(active_path); + if let Some(expected_timestamp) = reference.rollout_timestamp.as_deref() { + let plain_path = compression::plain_rollout_path(path.as_path()); + let Some((actual_timestamp, _)) = plain_path + .file_name() + .and_then(|file_name| file_name.to_str()) + .and_then(parse_timestamp_string_uuid_from_filename) + else { + return Ok(None); + }; + if actual_timestamp != expected_timestamp { + return Ok(None); + } } - let id = uuid.to_string(); - if let Some(path) = - find_thread_path_by_id_str(codex_home, id.as_str(), /*state_db_ctx*/ None).await? - { - return Ok(path); - } - if let Some(path) = - find_archived_thread_path_by_id_str(codex_home, id.as_str(), /*state_db_ctx*/ None).await? - { - return Ok(path); - } - Ok(rollout_path.to_path_buf()) + Ok(Some(path)) } fn rollout_path_for_timestamp_file( diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index f85d803ca3..a8dee286da 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -115,17 +115,20 @@ pub async fn extract_metadata_from_rollout( for item in &items { apply_rollout_item(&mut metadata, item, default_provider); } - let has_segment_reference = items.iter().any(|item| { - matches!( - item, - RolloutItem::RolloutReference(reference) if reference.nth_user_message.is_none() - ) - }); - if has_segment_reference + let has_rollout_reference = items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))); + if has_rollout_reference && let Ok((first_user_message, preview)) = read_preview_metadata(rollout_path).await { - if first_user_message.is_some() { - metadata.first_user_message = first_user_message; + if let Some(first_user_message) = first_user_message { + let current_first_user_message = metadata.first_user_message.as_deref().map(str::trim); + if metadata.title.trim().is_empty() + || current_first_user_message == Some(metadata.title.trim()) + { + metadata.title = first_user_message.clone(); + } + metadata.first_user_message = Some(first_user_message); } if preview.is_some() { metadata.preview = preview; diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 8748d4a8ab..b2cb5d1a49 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -224,6 +224,10 @@ async fn extract_metadata_from_rollout_reads_referenced_preview() { nth_user_message: None, compacted_replacement_history_filter_texts: None, }), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "Current segment preview".to_string(), + ..Default::default() + })), ] { let line = RolloutLine { timestamp: "2026-01-27T12:34:56Z".to_string(), @@ -250,6 +254,7 @@ async fn extract_metadata_from_rollout_reads_referenced_preview() { outcome.metadata.preview.as_deref(), Some("Referenced preview") ); + assert_eq!(outcome.metadata.title, "Referenced preview"); } #[test] diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs index 67b79cef30..a63132428e 100644 --- a/codex-rs/rollout/src/search.rs +++ b/codex-rs/rollout/src/search.rs @@ -17,6 +17,8 @@ use tokio::process::Command; use super::ARCHIVED_SESSIONS_SUBDIR; use super::SESSIONS_SUBDIR; use super::compression; +use super::list::resolve_rollout_reference_rollout_path; +use super::recorder::RolloutRecorder; const MATCH_CONTEXT_BEFORE_CHARS: usize = 48; const MATCH_CONTEXT_AFTER_CHARS: usize = 96; @@ -50,14 +52,157 @@ pub async fn search_rollout_matches( SESSIONS_SUBDIR }); let json_search_term = json_escaped_search_term(search_term)?; - let Some(plain_matches) = - ripgrep_rollout_paths(rg_command, root.as_path(), json_search_term.as_str()).await? - else { - return scan_rollout_matches(root.as_path(), json_search_term.as_str(), search_term).await; - }; - let mut matches: RolloutSearchMatches = - plain_matches.into_iter().map(|path| (path, None)).collect(); - matches.extend(scan_compressed_rollout_matches(root.as_path(), search_term).await?); + let mut matches = + match ripgrep_rollout_paths(rg_command, root.as_path(), json_search_term.as_str()).await? { + Some(plain_matches) => { + let mut matches: RolloutSearchMatches = + plain_matches.into_iter().map(|path| (path, None)).collect(); + matches.extend(scan_compressed_rollout_matches(root.as_path(), search_term).await?); + matches + } + None => { + scan_rollout_matches(root.as_path(), json_search_term.as_str(), search_term).await? + } + }; + matches.extend( + search_reference_backed_rollouts(rg_command, codex_home, root.as_path(), search_term) + .await?, + ); + Ok(matches) +} + +async fn search_reference_backed_rollouts( + rg_command: &Path, + codex_home: &Path, + root: &Path, + search_term: &str, +) -> io::Result { + let mut reference_paths = HashSet::new(); + for reference_marker in ["\"rollout_reference\"", "\"fork_reference\""] { + match ripgrep_rollout_paths(rg_command, root, reference_marker).await? { + Some(paths) => { + reference_paths.extend(paths); + reference_paths + .extend(scan_compressed_rollout_reference_paths(root, reference_marker).await?); + } + None => { + reference_paths.extend(scan_rollout_reference_paths(root, reference_marker).await?); + } + } + } + let mut matches = HashMap::new(); + for path in reference_paths { + let Ok((items, _, _)) = RolloutRecorder::load_rollout_items(path.as_path()).await else { + continue; + }; + if let Some(snippet) = + first_segment_reference_content_match_snippet(codex_home, items.as_slice(), search_term) + .await? + { + matches.insert( + compression::plain_rollout_path(path.as_path()), + Some(snippet), + ); + } + } + Ok(matches) +} + +async fn first_segment_reference_content_match_snippet( + codex_home: &Path, + items: &[RolloutItem], + search_term: &str, +) -> io::Result> { + let search_term = case_insensitive_literal_regex(search_term)?; + let mut references = items + .iter() + .filter_map(|item| match item { + // Segment predecessors belong to the current thread. Prefix-truncated fork + // references belong to the source thread and should not duplicate that thread's + // search matches on every child. + RolloutItem::RolloutReference(reference) if reference.nth_user_message.is_none() => { + Some(reference.clone()) + } + _ => None, + }) + .collect::>(); + let mut visited_paths = HashSet::new(); + while let Some(reference) = references.pop() { + let Ok(path) = resolve_rollout_reference_rollout_path(codex_home, &reference).await else { + continue; + }; + if !visited_paths.insert(path.clone()) { + continue; + } + let Ok((referenced_items, _, _)) = RolloutRecorder::load_rollout_items(&path).await else { + continue; + }; + for item in &referenced_items { + if let Some(text) = conversation_text_from_item(item) + && let Some(snippet) = excerpt_around_match(text.as_str(), &search_term) + { + return Ok(Some(snippet)); + } + } + references.extend(referenced_items.into_iter().filter_map(|item| match item { + RolloutItem::RolloutReference(reference) if reference.nth_user_message.is_none() => { + Some(reference) + } + _ => None, + })); + } + Ok(None) +} + +async fn scan_rollout_reference_paths( + root: &Path, + reference_marker: &str, +) -> io::Result> { + scan_rollout_reference_paths_by_compression(root, reference_marker, /*compressed*/ None).await +} + +async fn scan_compressed_rollout_reference_paths( + root: &Path, + reference_marker: &str, +) -> io::Result> { + scan_rollout_reference_paths_by_compression(root, reference_marker, Some(true)).await +} + +async fn scan_rollout_reference_paths_by_compression( + root: &Path, + reference_marker: &str, + compressed: Option, +) -> io::Result> { + let mut matches = HashSet::new(); + let mut dirs = vec![root.to_path_buf()]; + let reference_marker = case_insensitive_literal_regex(reference_marker)?; + while let Some(dir) = dirs.pop() { + let mut entries = match tokio::fs::read_dir(dir).await { + Ok(entries) => entries, + Err(err) if err.kind() == io::ErrorKind::NotFound => continue, + Err(err) => return Err(err), + }; + while let Some(entry) = entries.next_entry().await? { + let path = entry.path(); + let file_type = entry.file_type().await?; + if file_type.is_dir() { + dirs.push(path); + continue; + } + if !file_type.is_file() { + continue; + } + let Some(rollout_file) = compression::RolloutFile::from_path(path) else { + continue; + }; + if compressed.is_some_and(|compressed| compressed != rollout_file.is_compressed()) { + continue; + } + if rollout_contains(rollout_file.path(), &reference_marker).await? { + matches.insert(rollout_file.into_path()); + } + } + } Ok(matches) } @@ -347,3 +492,89 @@ fn char_end_after(text: &str, byte_index: usize, chars_after: usize) -> usize { .map(|(offset, _)| byte_index.saturating_add(offset)) .unwrap_or(text.len()) } + +#[cfg(test)] +mod tests { + use super::*; + use codex_protocol::ThreadId; + use codex_protocol::protocol::RolloutReferenceItem; + use codex_protocol::protocol::SessionMeta; + use codex_protocol::protocol::SessionMetaLine; + use codex_protocol::protocol::UserMessageEvent; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + async fn write_rollout(path: &Path, items: Vec) { + tokio::fs::create_dir_all(path.parent().expect("rollout parent")) + .await + .expect("create rollout parent"); + let contents = items + .into_iter() + .map(|item| { + serde_json::to_string(&RolloutLine { + timestamp: "2026-06-12T00:00:00Z".to_string(), + item, + }) + .expect("serialize rollout line") + }) + .collect::>() + .join("\n"); + tokio::fs::write(path, format!("{contents}\n")) + .await + .expect("write rollout"); + } + + #[tokio::test] + async fn search_maps_predecessor_segment_matches_to_live_rollout() { + let codex_home = TempDir::new().expect("create codex home"); + let predecessor_path = codex_home.path().join("segments/predecessor.jsonl"); + let current_path = codex_home.path().join( + "sessions/2026/06/12/rollout-2026-06-12T00-00-00-00000000-0000-4000-8000-000000000001.jsonl", + ); + write_rollout( + predecessor_path.as_path(), + vec![ + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: ThreadId::new(), + timestamp: "2026-06-12T00:00:00Z".to_string(), + ..SessionMeta::default() + }, + git: None, + }), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "needle in predecessor".to_string(), + ..Default::default() + })), + ], + ) + .await; + write_rollout( + current_path.as_path(), + vec![RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: predecessor_path, + thread_id: None, + rollout_timestamp: None, + segment_id: None, + max_depth: 1, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + })], + ) + .await; + + let matches = search_rollout_matches( + Path::new("/missing/rg"), + codex_home.path(), + /*archived*/ false, + "needle", + ) + .await + .expect("search rollouts"); + + assert_eq!( + matches, + HashMap::from([(current_path, Some("needle in predecessor".to_string()))]) + ); + } +} diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 22d6d86850..d5b6784316 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -251,8 +251,7 @@ async fn rollout_reference_resolves_archived_file_by_stable_thread_and_timestamp let ts = "2025-01-03T13-00-00"; let active_path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl")); let archived_path = home.join(format!("archived_sessions/rollout-{ts}-{uuid}.jsonl")); - fs::create_dir_all(archived_path.parent().expect("archived parent")).unwrap(); - fs::write(&archived_path, "").unwrap(); + write_legacy_session_meta(archived_path.as_path(), thread_id, ts); let resolved = crate::resolve_rollout_reference_rollout_path( home, @@ -314,6 +313,36 @@ async fn rollout_reference_prefers_segment_id_over_live_path_with_same_thread_ti assert_eq!(resolved, archived_path); } +#[tokio::test] +async fn rollout_reference_rejects_live_path_with_mismatched_segment_id() { + let temp = TempDir::new().expect("tempdir"); + let home = temp.path(); + let uuid = Uuid::new_v4(); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let referenced_segment_id = SegmentId::new(); + let live_segment_id = SegmentId::new(); + let ts = "2025-01-03T13-00-00"; + let active_path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl")); + write_session_meta(active_path.as_path(), thread_id, live_segment_id, ts); + + let err = crate::resolve_rollout_reference_rollout_path( + home, + &RolloutReferenceItem { + rollout_path: active_path, + thread_id: Some(thread_id), + rollout_timestamp: Some(ts.to_string()), + segment_id: Some(referenced_segment_id), + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }, + ) + .await + .expect_err("a segment reference must not fall back to different live content"); + + assert_eq!(err.kind(), std::io::ErrorKind::NotFound); +} + #[tokio::test] async fn rollout_reference_prefers_existing_prior_segment_path_without_segment_id() { let temp = TempDir::new().expect("tempdir"); @@ -326,12 +355,7 @@ async fn rollout_reference_prefers_existing_prior_segment_path_without_segment_i let active_path = home.join(format!("sessions/2025/01/03/{file_name}")); let legacy_rotated_path = home.join("archived_sessions").join(file_name); write_session_meta(active_path.as_path(), thread_id, live_segment_id, ts); - write_session_meta( - legacy_rotated_path.as_path(), - thread_id, - SegmentId::new(), - ts, - ); + write_legacy_session_meta(legacy_rotated_path.as_path(), thread_id, ts); let resolved = crate::resolve_rollout_reference_rollout_path( home, @@ -351,6 +375,66 @@ async fn rollout_reference_prefers_existing_prior_segment_path_without_segment_i assert_eq!(resolved, legacy_rotated_path); } +#[tokio::test] +async fn legacy_reference_prefers_compressed_archived_predecessor_over_new_live_segment() { + let temp = TempDir::new().expect("tempdir"); + let home = temp.path(); + let uuid = Uuid::new_v4(); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let ts = "2025-01-03T13-00-00"; + let file_name = format!("rollout-{ts}-{uuid}.jsonl"); + let active_path = home.join(format!("sessions/2025/01/03/{file_name}")); + let archived_path = home.join("archived_sessions").join(file_name); + write_session_meta(active_path.as_path(), thread_id, SegmentId::new(), ts); + write_legacy_session_meta(archived_path.as_path(), thread_id, ts); + let compressed_archived_path = compress_rollout(archived_path.as_path()); + + let resolved = crate::resolve_rollout_reference_rollout_path( + home, + &RolloutReferenceItem { + rollout_path: active_path, + thread_id: Some(thread_id), + rollout_timestamp: Some(ts.to_string()), + segment_id: None, + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }, + ) + .await + .expect("resolve legacy rollout reference"); + + assert_eq!(resolved, compressed_archived_path); +} + +#[tokio::test] +async fn legacy_reference_rejects_new_live_segment_without_predecessor_snapshot() { + let temp = TempDir::new().expect("tempdir"); + let home = temp.path(); + let uuid = Uuid::new_v4(); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let ts = "2025-01-03T13-00-00"; + let active_path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl")); + write_session_meta(active_path.as_path(), thread_id, SegmentId::new(), ts); + + let err = crate::resolve_rollout_reference_rollout_path( + home, + &RolloutReferenceItem { + rollout_path: active_path, + thread_id: Some(thread_id), + rollout_timestamp: Some(ts.to_string()), + segment_id: None, + max_depth: 2, + nth_user_message: None, + compacted_replacement_history_filter_texts: None, + }, + ) + .await + .expect_err("legacy reference must not resolve to a newer identified segment"); + + assert_eq!(err.kind(), std::io::ErrorKind::NotFound); +} + #[tokio::test] async fn rollout_reference_resolves_rotated_segment_file_by_segment_id() { let temp = TempDir::new().expect("tempdir"); @@ -437,13 +521,26 @@ async fn rollout_reference_resolves_compressed_segment_file_by_segment_id() { } fn write_session_meta(path: &Path, thread_id: ThreadId, segment_id: SegmentId, timestamp: &str) { + write_session_meta_with_segment_id(path, thread_id, Some(segment_id), timestamp); +} + +fn write_legacy_session_meta(path: &Path, thread_id: ThreadId, timestamp: &str) { + write_session_meta_with_segment_id(path, thread_id, /*segment_id*/ None, timestamp); +} + +fn write_session_meta_with_segment_id( + path: &Path, + thread_id: ThreadId, + segment_id: Option, + timestamp: &str, +) { fs::create_dir_all(path.parent().expect("rollout parent")).expect("create rollout parent"); let line = RolloutLine { timestamp: timestamp.to_string(), item: RolloutItem::SessionMeta(SessionMetaLine { meta: SessionMeta { id: thread_id, - segment_id: Some(segment_id), + segment_id, forked_from_id: None, parent_thread_id: None, timestamp: timestamp.to_string(), @@ -1310,7 +1407,7 @@ async fn test_list_threads_reads_preview_from_referenced_segment() { } #[tokio::test] -async fn test_list_threads_does_not_read_preview_from_fork_reference() { +async fn test_list_threads_prefers_child_preview_over_fork_reference() { let temp = TempDir::new().unwrap(); let home = temp.path(); @@ -1388,6 +1485,80 @@ async fn test_list_threads_does_not_read_preview_from_fork_reference() { ); } +#[tokio::test] +async fn test_list_threads_reads_inherited_preview_from_fresh_fork_reference() { + let temp = TempDir::new().unwrap(); + let home = temp.path(); + + let parent_uuid = Uuid::from_u128(106); + let child_uuid = Uuid::from_u128(107); + let parent_thread_id = thread_id_from_uuid(parent_uuid); + let child_thread_id = thread_id_from_uuid(child_uuid); + let parent_segment_id = SegmentId::new(); + let child_segment_id = SegmentId::new(); + let ts = "2025-05-04T10-30-00"; + let parent_path = home + .join(crate::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(parent_thread_id.to_string()) + .join(parent_segment_id.to_string()) + .join(format!("rollout-{ts}-{parent_uuid}.jsonl")); + let child_path = home + .join("sessions/2025/05/04") + .join(format!("rollout-{ts}-{child_uuid}.jsonl")); + write_session_meta( + parent_path.as_path(), + parent_thread_id, + parent_segment_id, + ts, + ); + append_rollout_item( + parent_path.as_path(), + ts, + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "Inherited parent preview".to_string(), + ..Default::default() + })), + ); + write_session_meta(child_path.as_path(), child_thread_id, child_segment_id, ts); + append_rollout_item( + child_path.as_path(), + ts, + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: parent_path, + thread_id: Some(parent_thread_id), + rollout_timestamp: Some(ts.to_string()), + segment_id: Some(parent_segment_id), + max_depth: 2, + nth_user_message: Some(usize::MAX), + compacted_replacement_history_filter_texts: None, + }), + ); + + let provider_filter = provider_vec(&[TEST_PROVIDER]); + let page = get_threads( + home, + /*page_size*/ 10, + /*cursor*/ None, + ThreadSortKey::CreatedAt, + INTERACTIVE_SESSION_SOURCES.as_slice(), + Some(provider_filter.as_slice()), + /*cwd_filters*/ None, + TEST_PROVIDER, + ) + .await + .unwrap(); + + assert_eq!(page.items.len(), 1); + assert_eq!( + page.items[0].preview.as_deref(), + Some("Inherited parent preview") + ); + assert_eq!( + page.items[0].first_user_message.as_deref(), + Some("Inherited parent preview") + ); +} + #[tokio::test] async fn test_list_threads_uses_goal_objective_as_preview() { let temp = TempDir::new().unwrap(); diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 472e3d35e3..1ce3f14e3b 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -166,6 +166,7 @@ impl InMemoryThreadStore { state.calls.create_thread += 1; let session_meta = SessionMeta { id: params.thread_id, + segment_id: None, forked_from_id: params.forked_from_id, parent_thread_id: params.parent_thread_id, cwd: params.metadata.cwd.clone().unwrap_or_default(), diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index 24bc32993c..902837b7a9 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -4,7 +4,9 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_rollout::ArchivedThreadRolloutDisposition; use codex_rollout::RolloutRecorder; +use codex_rollout::classify_archived_thread_rollout; use codex_rollout::find_archived_thread_path_by_id_str; use codex_rollout::find_thread_name_by_id; use codex_rollout::find_thread_path_by_id_str; @@ -32,12 +34,12 @@ pub(super) async fn read_thread( ) -> ThreadStoreResult { let thread_id = params.thread_id; if let Some(metadata) = read_sqlite_metadata(store, thread_id).await - && (params.include_archived - || (metadata.archived_at.is_none() - && !rollout_path_is_archived( - store.config.codex_home.as_path(), - metadata.rollout_path.as_path(), - ))) + && sqlite_metadata_matches_requested_archive_state( + store, + &metadata, + params.include_archived, + ) + .await && (!params.include_history || sqlite_rollout_path_can_load_history_for_thread( store, @@ -85,6 +87,35 @@ pub(super) async fn read_thread( Ok(thread) } +async fn sqlite_metadata_matches_requested_archive_state( + store: &LocalThreadStore, + metadata: &ThreadMetadata, + include_archived: bool, +) -> bool { + let rollout_path_is_archived = rollout_path_is_archived( + store.config.codex_home.as_path(), + metadata.rollout_path.as_path(), + ); + if !include_archived { + return metadata.archived_at.is_none() && !rollout_path_is_archived; + } + if metadata.archived_at.is_none() { + return !rollout_path_is_archived; + } + if !rollout_path_is_archived { + return true; + } + matches!( + classify_archived_thread_rollout( + store.config.codex_home.as_path(), + metadata.rollout_path.as_path(), + store.state_db().await.as_deref(), + ) + .await, + Ok(ArchivedThreadRolloutDisposition::CanonicalArchivedThread) + ) +} + async fn sqlite_rollout_path_can_load_history_for_thread( store: &LocalThreadStore, path: &std::path::Path, @@ -1067,6 +1098,59 @@ mod tests { assert_eq!(history.items.len(), 2); } + #[tokio::test] + async fn read_repair_prefers_live_rollout_over_legacy_rotated_segment_when_including_archived() + { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(223); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let live_rollout_path = + write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); + let legacy_rotated_path = + write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid) + .expect("legacy rotated session file"); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + legacy_rotated_path, + Utc::now(), + SessionSource::Cli, + ); + builder.archived_at = Some(Utc::now()); + let mut metadata = builder.build(config.default_model_provider_id.as_str()); + metadata.first_user_message = Some("legacy rotated sqlite preview".to_string()); + runtime + .upsert_thread(&metadata) + .await + .expect("state db upsert should succeed"); + + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: true, + }) + .await + .expect("read active thread"); + + assert_eq!(thread.rollout_path, Some(live_rollout_path.clone())); + assert_eq!(thread.archived_at, None); + let repaired = runtime + .get_thread(thread_id) + .await + .expect("state db get thread") + .expect("repaired metadata"); + assert_eq!(repaired.rollout_path, live_rollout_path); + assert_eq!(repaired.archived_at, None); + } + #[tokio::test] async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 8eb23e3632..5d37e733eb 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -270,6 +270,7 @@ impl ThreadMetadataSync { } } RolloutItem::SessionMeta(_) + | RolloutItem::RolloutReference(_) | RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::InterAgentCommunication(_)