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 e33d9eb0b0..16c36610c7 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -233,6 +233,7 @@ impl ThreadHistoryBuilder { RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0b49c789f6..95fbb99a51 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -4039,7 +4039,12 @@ impl CodexMessageProcessor { 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_turns_from_rollout_items(&history.items); + let items = codex_core::materialize_rollout_items_for_replay( + &self.config.codex_home, + &history.items, + ) + .await; + thread.turns = build_turns_from_rollout_items(&items); } Ok(Some(thread)) } 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 589c7c330a..a5682010df 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::McpProcess; +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; @@ -47,6 +48,7 @@ use codex_feedback::CodexFeedback; use codex_protocol::models::BaseInstructions; 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; @@ -421,6 +423,124 @@ 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, + forked_from_id: None, + source: ProtocolSessionSource::Cli, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + metadata: ThreadPersistenceMetadata { + cwd: None, + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Disabled, + }, + event_persistence_mode: ThreadEventPersistenceMode::default(), + }) + .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, + )?), + segment_id: None, + max_depth: 2, + }), + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: "history after segment rotation".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + })), + ], + }) + .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, + cloud_requirements: CloudRequirementsLoader::default(), + thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_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", + ] + ); + + 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/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1a5cacab22..e111a8841f 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -164,6 +164,7 @@ fn keep_forked_rollout_item(item: &RolloutItem) -> bool { RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::SessionMeta(_) => true, } } @@ -177,6 +178,7 @@ fn full_history_fork_reference_items( RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::TurnContext(_) => None, }); @@ -187,6 +189,7 @@ fn full_history_fork_reference_items( RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::TurnContext(_) => None, }) @@ -1841,6 +1844,7 @@ fn previous_response_fork_rollout_items( let source_session_meta = source_items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta) => Some(meta.clone()), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) @@ -1849,6 +1853,7 @@ fn previous_response_fork_rollout_items( let latest_turn_context = source_items.iter().rev().find_map(|item| match item { RolloutItem::TurnContext(turn_context) => Some(turn_context.clone()), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::SessionMeta(_) diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs index aa1bcc8d57..bfd4845dce 100644 --- a/codex-rs/core/src/rollout.rs +++ b/codex-rs/core/src/rollout.rs @@ -23,6 +23,7 @@ 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_fork_reference_rollout_path; +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/mod.rs b/codex-rs/core/src/session/mod.rs index 644fd99eb6..c96a7ad49b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -101,6 +101,7 @@ use codex_protocol::models::format_allow_prefixes; use codex_protocol::openai_models::ModelInfo; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::DEFAULT_ROLLOUT_REFERENCE_DEPTH; use codex_protocol::protocol::FileChange; use codex_protocol::protocol::HasLegacyEvent; use codex_protocol::protocol::InterAgentCommunication; @@ -135,6 +136,7 @@ use codex_thread_store::LiveThread; use codex_thread_store::LiveThreadInitGuard; use codex_thread_store::LocalThreadStore; use codex_thread_store::ResumeThreadParams; +use codex_thread_store::RotateThreadSegmentParams; use codex_thread_store::ThreadEventPersistenceMode; use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; @@ -1245,9 +1247,12 @@ impl Session { let state = self.state.lock().await; state.session_configuration.codex_home().clone() }; - let replay_rollout_items = if conversation_history - .scan_rollout_items(|item| matches!(item, RolloutItem::ForkReference(_))) - { + let replay_rollout_items = if conversation_history.scan_rollout_items(|item| { + matches!( + item, + RolloutItem::ForkReference(_) | RolloutItem::RolloutReference(_) + ) + }) { Some( materialize_rollout_items_for_replay( codex_home.as_path(), @@ -2599,15 +2604,70 @@ impl Session { self.replace_history(items, reference_context_item.clone()) .await; - self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)]) - .await; + let mut rollout_items = vec![RolloutItem::Compacted(compacted_item)]; if let Some(turn_context_item) = reference_context_item { - self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item)]) - .await; + rollout_items.push(RolloutItem::TurnContext(turn_context_item)); + } + if !self + .rotate_rollout_segment_after_compaction(rollout_items.clone()) + .await + { + self.persist_rollout_items(&rollout_items).await; } self.services.model_client.advance_window_generation(); } + async fn rotate_rollout_segment_after_compaction( + &self, + initial_items: Vec, + ) -> bool { + let Some(live_thread) = self.live_thread() else { + return false; + }; + let params = { + let state = self.state.lock().await; + let session_configuration = &state.session_configuration; + let event_persistence_mode = if session_configuration.persist_extended_history { + ThreadEventPersistenceMode::Extended + } else { + ThreadEventPersistenceMode::Limited + }; + RotateThreadSegmentParams { + source: session_configuration.session_source.clone(), + base_instructions: BaseInstructions { + text: session_configuration.base_instructions.clone(), + }, + dynamic_tools: session_configuration.dynamic_tools.clone(), + metadata: ThreadPersistenceMetadata { + cwd: Some(session_configuration.cwd.to_path_buf()), + model_provider: session_configuration + .original_config_do_not_use + .model_provider_id + .clone(), + memory_mode: if session_configuration + .original_config_do_not_use + .memories + .generate_memories + { + ThreadMemoryMode::Enabled + } else { + ThreadMemoryMode::Disabled + }, + }, + event_persistence_mode, + initial_items, + previous_segment_reference_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + } + }; + match live_thread.rotate_local_segment(params).await { + Ok(rotated) => rotated, + Err(err) => { + warn!("failed to rotate rollout segment after compaction: {err:#}"); + false + } + } + } + async fn persist_rollout_response_items(&self, items: &[ResponseItem]) { let rollout_items: Vec = items .iter() diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 3ca4098219..8ab63bd3ca 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -209,6 +209,7 @@ impl Session { } RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::SessionMeta(_) => {} } @@ -278,6 +279,7 @@ impl Session { } RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 99abeb80ac..aef299cd39 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2879,6 +2879,72 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { .expect("thread should have rollout path") } +#[tokio::test] +async fn replace_compacted_history_rotates_local_rollout_segment() { + let (mut sess, tc, _) = make_session_and_context_with_rx().await; + let sess = Arc::get_mut(&mut sess).expect("session should not have additional references"); + let old_rollout_path = attach_thread_persistence(sess).await; + sess.persist_rollout_items(&[ + RolloutItem::ResponseItem(user_message("before compaction")), + RolloutItem::ResponseItem(assistant_message("before compaction answer")), + ]) + .await; + sess.flush_rollout() + .await + .expect("pre-compaction rollout should flush"); + let replacement_history = vec![user_message("compacted summary")]; + + sess.replace_compacted_history( + replacement_history.clone(), + Some(tc.to_turn_context_item()), + CompactedItem { + message: "compacted summary".to_string(), + replacement_history: Some(replacement_history.clone()), + }, + ) + .await; + + let new_rollout_path = sess + .current_rollout_path() + .await + .expect("load current rollout path") + .expect("rollout path after compaction"); + assert_ne!(new_rollout_path, old_rollout_path); + let (new_items, new_thread_id, _) = + RolloutRecorder::load_rollout_items(new_rollout_path.as_path()) + .await + .expect("load new rollout segment"); + assert_eq!(new_thread_id, Some(sess.conversation_id)); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.rollout_path == old_rollout_path + && reference.thread_id == Some(sess.conversation_id) + ) + })); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::Compacted(CompactedItem { + replacement_history: Some(history), + .. + }) if *history == replacement_history + ) + })); + + let config = sess.get_config().await; + let replay_items = crate::thread_rollout_truncation::materialize_rollout_items_for_replay( + config.codex_home.as_path(), + &new_items, + ) + .await; + let reconstructed = sess + .reconstruct_history_from_rollout(tc.as_ref(), &replay_items) + .await; + assert_eq!(reconstructed.history, replacement_history); +} + fn text_block(s: &str) -> serde_json::Value { json!({ "type": "text", diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a3ca5a7b76..4ee40e84cb 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1416,6 +1416,7 @@ fn fork_reference_history_for_snapshot( RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::TurnContext(_) => None, }); @@ -1427,6 +1428,7 @@ fn fork_reference_history_for_snapshot( RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::TurnContext(_) => None, }) diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index bbe3de38a1..65b7d7fa9b 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -7,6 +7,7 @@ use crate::context_manager::is_user_turn_boundary; use crate::event_mapping; use crate::rollout::RolloutRecorder; use crate::rollout::resolve_fork_reference_rollout_path; +use crate::rollout::resolve_rollout_reference_rollout_path; use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; @@ -208,6 +209,41 @@ async fn materialize_rollout_items_for_replay_at_depth( } } } + RolloutItem::RolloutReference(reference) => { + if depth >= reference.max_depth { + warn!("rollout reference materialization reached max depth"); + continue; + } + let resolved_path = + match resolve_rollout_reference_rollout_path(codex_home, reference).await { + Ok(path) => path, + Err(err) => { + warn!( + "failed to resolve rollout reference {}: {err}", + reference.rollout_path.display() + ); + reference.rollout_path.clone() + } + }; + match RolloutRecorder::load_rollout_items(&resolved_path).await { + Ok((reference_items, _, _)) => { + let reference_materialized = + Box::pin(materialize_rollout_items_for_replay_at_depth( + codex_home, + &reference_items, + depth + 1, + )) + .await; + materialized.extend(reference_materialized); + } + Err(err) => { + warn!( + "failed to load rollout reference {}: {err}", + resolved_path.display() + ); + } + } + } other => materialized.push(other.clone()), } } diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index a26c895384..a2414242ba 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::protocol::ForkReferenceItem; 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; @@ -244,6 +245,73 @@ async fn materializes_fork_reference_by_segment_id_after_source_rollover() { 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 oldest_path = temp.path().join("oldest.jsonl"); + let old_path = temp.path().join("old.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 segment request")), + ], + ) + .await; + write_rollout( + &old_path, + &[ + session_meta_item(thread_id, old_segment_id), + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: oldest_path, + thread_id: None, + segment_id: None, + max_depth: 2, + }), + 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: None, + segment_id: None, + max_depth: 2, + }), + RolloutItem::ResponseItem(user_msg("middle segment request")), + ], + ) + .await; + + let current_items = vec![ + RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: middle_path, + thread_id: None, + segment_id: None, + max_depth: 2, + }), + RolloutItem::ResponseItem(user_msg("current segment request")), + ]; + let materialized = materialize_rollout_items_for_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")); +} + #[test] fn truncates_rollout_from_start_applies_thread_rollback_markers() { let rollout_items = vec![ diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 844cb1109a..f344222f8a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2784,11 +2784,29 @@ pub struct ForkReferenceItem { pub nth_user_message: usize, } +pub const DEFAULT_ROLLOUT_REFERENCE_DEPTH: usize = 2; + +#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema, TS)] +pub struct RolloutReferenceItem { + pub rollout_path: PathBuf, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub thread_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub segment_id: Option, + #[serde(default = "default_rollout_reference_depth")] + pub max_depth: usize, +} + +fn default_rollout_reference_depth() -> usize { + DEFAULT_ROLLOUT_REFERENCE_DEPTH +} + #[derive(Serialize, Deserialize, Debug, Clone, JsonSchema, TS)] #[serde(tag = "type", content = "payload", rename_all = "snake_case")] pub enum RolloutItem { SessionMeta(SessionMetaLine), ForkReference(ForkReferenceItem), + RolloutReference(RolloutReferenceItem), ResponseItem(ResponseItem), Compacted(CompactedItem), TurnContext(TurnContextItem), diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index f026bce563..8a12a69b55 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -52,6 +52,7 @@ pub use list::read_head_for_summary; pub use list::read_session_meta_line; pub use list::read_thread_item_from_rollout; pub use list::resolve_fork_reference_rollout_path; +pub use list::resolve_rollout_reference_rollout_path; pub use list::rollout_date_parts; pub use metadata::builder_from_items; pub use policy::EventPersistenceMode; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 00f9621490..159f652fbd 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1119,7 +1119,7 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { + RolloutItem::ForkReference(_) | RolloutItem::RolloutReference(_) => { // Not included in summaries; skip. } RolloutItem::ResponseItem(_) => { @@ -1186,6 +1186,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result {} @@ -1454,6 +1455,41 @@ pub async fn resolve_fork_reference_rollout_path( Ok(rollout_path.to_path_buf()) } +pub async fn resolve_rollout_reference_rollout_path( + codex_home: &Path, + reference: &codex_protocol::protocol::RolloutReferenceItem, +) -> io::Result { + if let (Some(thread_id), Some(segment_id)) = (reference.thread_id, reference.segment_id) + && let Some(path) = + find_rollout_path_by_segment_id(codex_home, thread_id, segment_id).await? + { + return Ok(path); + } + + let rollout_path = reference.rollout_path.as_path(); + if tokio::fs::try_exists(rollout_path).await.unwrap_or(false) { + return Ok(rollout_path.to_path_buf()); + } + + let Some(file_name) = rollout_path + .file_name() + .and_then(|file_name| file_name.to_str()) + else { + return Ok(rollout_path.to_path_buf()); + }; + let Some((_, uuid)) = parse_timestamp_uuid_from_filename(file_name) else { + return Ok(rollout_path.to_path_buf()); + }; + let id = uuid.to_string(); + if let Some(path) = find_thread_path_by_id_str(codex_home, id.as_str()).await? { + return Ok(path); + } + if let Some(path) = find_archived_thread_path_by_id_str(codex_home, id.as_str()).await? { + return Ok(path); + } + Ok(rollout_path.to_path_buf()) +} + /// Extract the `YYYY/MM/DD` directory components from a rollout filename. pub fn rollout_date_parts(file_name: &OsStr) -> Option<(String, String, String)> { let name = file_name.to_string_lossy(); diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index b832814249..7e42d9ca04 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -70,6 +70,7 @@ pub fn builder_from_items( if let Some(session_meta) = items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) @@ -125,6 +126,7 @@ pub async fn extract_metadata_from_rollout( memory_mode: items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index a06e3b987e..0ae70b8606 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -18,6 +18,7 @@ pub fn is_persisted_response_item(item: &RolloutItem, mode: EventPersistenceMode // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). RolloutItem::Compacted(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => true, } diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 32b52ffb15..ce5980f389 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -900,6 +900,9 @@ impl RolloutRecorder { RolloutItem::ForkReference(item) => { items.push(RolloutItem::ForkReference(item)); } + RolloutItem::RolloutReference(item) => { + items.push(RolloutItem::RolloutReference(item)); + } RolloutItem::Compacted(item) => { items.push(RolloutItem::Compacted(item)); } @@ -1373,28 +1376,43 @@ fn precompute_log_file_info( // Resolve ~/.codex/sessions/YYYY/MM/DD path. let timestamp = OffsetDateTime::now_local() .map_err(|e| IoError::other(format!("failed to get local time: {e}")))?; - let mut dir = config.codex_home().to_path_buf(); - dir.push(SESSIONS_SUBDIR); - dir.push(timestamp.year().to_string()); - dir.push(format!("{:02}", u8::from(timestamp.month()))); - dir.push(format!("{:02}", timestamp.day())); // Custom format for YYYY-MM-DDThh-mm-ss. Use `-` instead of `:` for // compatibility with filesystems that do not allow colons in filenames. let format: &[FormatItem] = format_description!("[year]-[month]-[day]T[hour]-[minute]-[second]"); - let date_str = timestamp - .format(format) - .map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?; - - let filename = format!("rollout-{date_str}-{conversation_id}.jsonl"); - - let path = dir.join(filename); + let mut selected_timestamp = timestamp; + let mut path = None; + for offset_seconds in 0..60 { + let candidate_timestamp = timestamp + .checked_add(time::Duration::seconds(offset_seconds)) + .ok_or_else(|| IoError::other("failed to compute rollout timestamp"))?; + let mut dir = config.codex_home().to_path_buf(); + dir.push(SESSIONS_SUBDIR); + dir.push(candidate_timestamp.year().to_string()); + dir.push(format!("{:02}", u8::from(candidate_timestamp.month()))); + dir.push(format!("{:02}", candidate_timestamp.day())); + let date_str = candidate_timestamp + .format(format) + .map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?; + let filename = format!("rollout-{date_str}-{conversation_id}.jsonl"); + let candidate_path = dir.join(filename); + if !candidate_path.exists() { + selected_timestamp = candidate_timestamp; + path = Some(candidate_path); + break; + } + } + let path = path.ok_or_else(|| { + IoError::other(format!( + "failed to find an unused rollout path for thread {conversation_id}" + )) + })?; Ok(LogFileInfo { path, conversation_id, - timestamp, + timestamp: selected_timestamp, }) } @@ -1891,6 +1909,7 @@ async fn resume_candidate_matches_cwd( RolloutItem::TurnContext(turn_context) => Some(turn_context.cwd.as_path()), RolloutItem::SessionMeta(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 8f301c53e3..32b6c18c34 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -287,6 +287,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< let segment_id = items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta) => meta.meta.segment_id, RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index ff55e18d1a..e893f0c811 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -22,7 +22,9 @@ pub fn apply_rollout_item( RolloutItem::TurnContext(turn_ctx) => apply_turn_context(metadata, turn_ctx), RolloutItem::EventMsg(event) => apply_event_msg(metadata, event), RolloutItem::ResponseItem(item) => apply_response_item(metadata, item), - RolloutItem::Compacted(_) | RolloutItem::ForkReference(_) => {} + RolloutItem::Compacted(_) + | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) => {} } if metadata.model_provider.is_empty() { metadata.model_provider = default_provider.to_string(); @@ -38,6 +40,7 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { ) => true, RolloutItem::EventMsg(_) | RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => false, } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 5c6e4d923d..beb53c0bfa 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -980,6 +980,7 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option Some(meta_line.meta.dynamic_tools.clone()), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) @@ -991,6 +992,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option { items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 52b7f5ea1f..5bc57f6f4a 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -32,6 +32,7 @@ pub use types::OptionalStringPatch; pub use types::ReadThreadByRolloutPathParams; pub use types::ReadThreadParams; pub use types::ResumeThreadParams; +pub use types::RotateThreadSegmentParams; pub use types::SortDirection; pub use types::StoredThread; pub use types::StoredThreadHistory; diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index bcce1c7645..d92c38a37e 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -11,6 +11,7 @@ use crate::CreateThreadParams; use crate::LoadThreadHistoryParams; use crate::LocalThreadStore; use crate::ResumeThreadParams; +use crate::RotateThreadSegmentParams; use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; use crate::ThreadStore; @@ -111,6 +112,23 @@ impl LiveThread { .await } + pub async fn rotate_local_segment( + &self, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult { + let Some(local_store) = self + .thread_store + .as_any() + .downcast_ref::() + else { + return Ok(false); + }; + local_store + .rotate_thread_segment(self.thread_id, params) + .await?; + Ok(true) + } + pub async fn persist(&self) -> ThreadStoreResult<()> { self.thread_store.persist_thread(self.thread_id).await } diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 643207b59d..939812d7ab 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -1,11 +1,15 @@ use std::path::PathBuf; use codex_protocol::ThreadId; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::RolloutConfig; use codex_rollout::RolloutRecorder; use codex_rollout::RolloutRecorderParams; use codex_rollout::builder_from_items; +use codex_rollout::read_session_meta_line; +use tracing::warn; use super::LocalThreadStore; use super::create_thread; @@ -13,6 +17,7 @@ use crate::AppendThreadItemsParams; use crate::CreateThreadParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::RotateThreadSegmentParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; @@ -161,6 +166,101 @@ pub(super) async fn rollout_path( .to_path_buf()) } +pub(super) async fn rotate_thread_segment( + store: &LocalThreadStore, + thread_id: ThreadId, + params: RotateThreadSegmentParams, +) -> ThreadStoreResult<()> { + let old_recorder = store.live_recorder(thread_id).await?; + old_recorder.flush().await.map_err(thread_store_io_error)?; + let old_rollout_path = old_recorder.rollout_path().to_path_buf(); + let old_meta = read_session_meta_line(old_rollout_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to read current rollout metadata from {}: {err}", + old_rollout_path.display() + ), + })?; + if old_meta.meta.id != thread_id { + return Err(ThreadStoreError::Internal { + message: format!( + "live rollout {} belongs to thread {} instead of {thread_id}", + old_rollout_path.display(), + old_meta.meta.id + ), + }); + } + + let cwd = params + .metadata + .cwd + .clone() + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: "local thread store requires a cwd".to_string(), + })?; + let config = RolloutConfig { + codex_home: store.config.codex_home.clone(), + sqlite_home: store.config.sqlite_home.clone(), + cwd, + model_provider_id: params.metadata.model_provider.clone(), + generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled), + }; + let mut initial_items = Vec::with_capacity(params.initial_items.len() + 1); + initial_items.push(RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: old_rollout_path.clone(), + thread_id: Some(thread_id), + segment_id: old_meta.meta.segment_id, + max_depth: params.previous_segment_reference_depth, + })); + initial_items.extend(params.initial_items); + + let state_db_ctx = store.state_db().await; + let new_recorder = RolloutRecorder::new( + &config, + RolloutRecorderParams::new( + thread_id, + old_meta.meta.forked_from_id, + params.source, + params.base_instructions, + params.dynamic_tools, + create_thread::event_persistence_mode(params.event_persistence_mode), + ), + state_db_ctx, + /*state_builder*/ None, + ) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to initialize rotated local thread recorder: {err}"), + })?; + new_recorder + .record_items(initial_items.as_slice()) + .await + .map_err(thread_store_io_error)?; + new_recorder.flush().await.map_err(thread_store_io_error)?; + + if let Err(err) = old_recorder.shutdown().await { + warn!( + "failed to close previous rollout segment {} for thread {thread_id}: {err}", + old_rollout_path.display() + ); + } + + let mut live_recorders = store.live_recorders.lock().await; + let current_path = live_recorders + .get(&thread_id) + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })? + .rollout_path() + .to_path_buf(); + if current_path != old_rollout_path { + return Err(ThreadStoreError::Conflict { + message: format!("live writer for thread {thread_id} changed during segment rotation"), + }); + } + live_recorders.insert(thread_id, new_recorder); + Ok(()) +} + fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError { ThreadStoreError::Internal { message: err.to_string(), diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 04dd8b2490..825881ea7d 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -29,6 +29,7 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::RotateThreadSegmentParams; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadPage; @@ -123,6 +124,14 @@ impl LocalThreadStore { live_writer::rollout_path(self, thread_id).await } + pub async fn rotate_thread_segment( + &self, + thread_id: ThreadId, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult<()> { + live_writer::rotate_thread_segment(self, thread_id, params).await + } + pub(super) async fn live_recorder( &self, thread_id: ThreadId, @@ -293,6 +302,7 @@ mod tests { use tempfile::TempDir; use super::*; + use crate::RotateThreadSegmentParams; use crate::ThreadEventPersistenceMode; use crate::ThreadPersistenceMetadata; use crate::local::test_support::test_config; @@ -348,6 +358,109 @@ mod tests { ); } + #[tokio::test] + async fn rotate_thread_segment_keeps_thread_id_and_references_previous_segment() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path())); + let thread_id = ThreadId::default(); + + store + .create_thread(create_thread_params(thread_id)) + .await + .expect("create live thread"); + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![user_message_item("before rotation")], + }) + .await + .expect("append pre-rotation item"); + store + .flush_thread(thread_id) + .await + .expect("flush pre-rotation item"); + let old_rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("old rollout path"); + let (old_items, _, _) = RolloutRecorder::load_rollout_items(old_rollout_path.as_path()) + .await + .expect("old rollout items"); + let old_segment_id = old_items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta) => meta.meta.segment_id, + RolloutItem::ForkReference(_) + | RolloutItem::RolloutReference(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => None, + }); + + store + .rotate_thread_segment( + thread_id, + RotateThreadSegmentParams { + source: SessionSource::Exec, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + metadata: thread_metadata(), + event_persistence_mode: ThreadEventPersistenceMode::Limited, + initial_items: vec![user_message_item("rotation checkpoint")], + previous_segment_reference_depth: 2, + }, + ) + .await + .expect("rotate segment"); + let new_rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("new rollout path"); + assert_ne!(new_rollout_path, old_rollout_path); + + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![user_message_item("after rotation")], + }) + .await + .expect("append post-rotation item"); + store + .flush_thread(thread_id) + .await + .expect("flush post-rotation item"); + + let (new_items, new_thread_id, _) = + RolloutRecorder::load_rollout_items(new_rollout_path.as_path()) + .await + .expect("new rollout items"); + assert_eq!(new_thread_id, Some(thread_id)); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.rollout_path == old_rollout_path + && reference.thread_id == Some(thread_id) + && reference.segment_id == old_segment_id + && reference.max_depth == 2 + ) + })); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::UserMessage(event)) + if event.message == "rotation checkpoint" + ) + })); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::UserMessage(event)) + if event.message == "after rotation" + ) + })); + assert_rollout_contains_message(old_rollout_path.as_path(), "before rotation").await; + } + #[tokio::test] async fn create_thread_rejects_missing_cwd() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 85bde023bd..7e1766f758 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -84,6 +84,25 @@ pub struct AppendThreadItemsParams { pub items: Vec, } +/// Parameters for replacing a live local rollout segment without changing the thread id. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct RotateThreadSegmentParams { + /// Runtime source for the thread. + pub source: SessionSource, + /// Base instructions persisted in the new session metadata. + pub base_instructions: BaseInstructions, + /// Dynamic tools available to the thread. + pub dynamic_tools: Vec, + /// Metadata captured for the new segment. + pub metadata: ThreadPersistenceMetadata, + /// Whether persistence should include the extended event surface. + pub event_persistence_mode: ThreadEventPersistenceMode, + /// Items written immediately after the new segment metadata. + pub initial_items: Vec, + /// Maximum number of rollout-reference segments materialized through the new segment. + pub previous_segment_reference_depth: usize, +} + /// Parameters for loading persisted history for resume, fork, rollback, and memory jobs. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct LoadThreadHistoryParams {