From 848575dc0d9feddd91b6656d8b4cadfb18d2d73c Mon Sep 17 00:00:00 2001 From: Friel Date: Sat, 13 Jun 2026 00:09:20 +0000 Subject: [PATCH] Add feature-gated session segmentation writers --- codex-rs/core/config.schema.json | 8 +- codex-rs/core/src/agent/control/spawn.rs | 69 ++- codex-rs/core/src/agent/control_tests.rs | 193 +++++- codex-rs/core/src/codex_thread.rs | 10 + codex-rs/core/src/session/mod.rs | 90 ++- codex-rs/core/src/session/session.rs | 17 +- codex-rs/core/src/session/tests.rs | 194 ++++++ codex-rs/core/src/thread_manager.rs | 130 +++- codex-rs/core/src/thread_manager_tests.rs | 46 +- codex-rs/core/tests/suite/fork_thread.rs | 61 +- codex-rs/core/tests/suite/resume.rs | 17 +- codex-rs/features/src/lib.rs | 12 + codex-rs/features/src/tests.rs | 17 + codex-rs/rollout/src/recorder.rs | 157 +++-- codex-rs/rollout/src/recorder_tests.rs | 14 + codex-rs/thread-store/Cargo.toml | 1 + codex-rs/thread-store/src/lib.rs | 1 + codex-rs/thread-store/src/live_thread.rs | 38 ++ .../thread-store/src/local/archive_thread.rs | 1 + .../thread-store/src/local/delete_thread.rs | 21 +- .../thread-store/src/local/live_writer.rs | 135 +++-- codex-rs/thread-store/src/local/mod.rs | 64 +- codex-rs/thread-store/src/local/segment_gc.rs | 136 +++++ .../thread-store/src/local/segment_writer.rs | 563 ++++++++++++++++++ .../src/local/segment_writer_tests.rs | 459 ++++++++++++++ .../src/local/unarchive_thread.rs | 1 + codex-rs/thread-store/src/types.rs | 9 + 27 files changed, 2320 insertions(+), 144 deletions(-) create mode 100644 codex-rs/thread-store/src/local/segment_gc.rs create mode 100644 codex-rs/thread-store/src/local/segment_writer.rs create mode 100644 codex-rs/thread-store/src/local/segment_writer_tests.rs diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 31daf2feb4..df460645f6 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -602,6 +602,9 @@ "secret_auth_storage": { "type": "boolean" }, + "session_segmentation": { + "type": "boolean" + }, "shell_snapshot": { "type": "boolean" }, @@ -4757,6 +4760,9 @@ "secret_auth_storage": { "type": "boolean" }, + "session_segmentation": { + "type": "boolean" + }, "shell_snapshot": { "type": "boolean" }, @@ -5274,4 +5280,4 @@ }, "title": "ConfigToml", "type": "object" -} +} \ No newline at end of file diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index add512cec3..a245c76aa1 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,5 +1,6 @@ use super::residency::is_v2_resident_session_source; use super::*; +use codex_features::Feature; const AGENT_NAMES: &str = include_str!("../agent_names.txt"); @@ -411,28 +412,56 @@ impl AgentControl { let parent_thread_id = *parent_thread_id; let parent_thread = state.get_thread(parent_thread_id).await.ok(); - if let Some(parent_thread) = parent_thread.as_ref() { + let sealed_reference = if matches!(fork_mode, SpawnAgentForkMode::FullHistory) + && config.features.enabled(Feature::SessionSegmentation) + && let Some(parent_thread) = parent_thread.as_ref() + { + match parent_thread.seal_rollout_segment_for_reference().await { + Ok(reference) => reference, + Err(err) => { + warn!( + "failed to seal parent rollout segment for full-history fork; copying history instead: {err}" + ); + None + } + } + } else { + None + }; + if sealed_reference.is_none() + && let Some(parent_thread) = parent_thread.as_ref() + { // `record_conversation_items` only queues persistence writes asynchronously. - // Flush before snapshotting store history for a fork. + // Flush before snapshotting store history for a copied-history fork. parent_thread.ensure_rollout_materialized().await; parent_thread.flush_rollout().await?; } - let parent_history = state - .read_stored_thread(ReadThreadParams { - thread_id: parent_thread_id, - include_archived: true, - include_history: true, - }) - .await? - .history - .ok_or_else(|| { - CodexErr::Fatal(format!( - "parent thread history unavailable for fork: {parent_thread_id}" - )) - })?; - - let mut forked_rollout_items = parent_history.items; + let mut forked_rollout_items = if let Some(reference) = sealed_reference { + // A full-history fork intentionally preserves the parent's exact model prefix, + // including reasoning and tool items that copied-history forks sanitize. Rotation + // makes that prefix immutable before this reference is persisted. + vec![RolloutItem::RolloutReference(reference)] + } else { + let parent_history = state + .read_stored_thread(ReadThreadParams { + thread_id: parent_thread_id, + include_archived: true, + include_history: true, + }) + .await? + .history + .ok_or_else(|| { + CodexErr::Fatal(format!( + "parent thread history unavailable for fork: {parent_thread_id}" + )) + })?; + crate::thread_rollout_truncation::materialize_rollout_items_for_model_replay( + &config.codex_home, + &parent_history.items, + ) + .await + }; if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { forked_rollout_items = truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns); @@ -469,6 +498,12 @@ impl AgentControl { Vec::new() }; let preserve_reference_context_item = matches!(fork_mode, SpawnAgentForkMode::FullHistory); + for item in &mut forked_rollout_items { + if let RolloutItem::RolloutReference(reference) = item { + reference.compacted_replacement_history_filter_texts = + Some(multi_agent_v2_usage_hint_texts_to_filter.clone()); + } + } forked_rollout_items.retain(|item| { keep_forked_rollout_item(item, preserve_reference_context_item) && !matches!( diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index ac42984b58..ed5098aaba 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -9,6 +9,7 @@ use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; use crate::init_state_db; +use crate::rollout::RolloutRecorder; use assert_matches::assert_matches; use codex_features::Feature; use codex_login::CodexAuth; @@ -18,9 +19,13 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::CompactedItem; +use codex_protocol::protocol::DEFAULT_ROLLOUT_REFERENCE_DEPTH; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; @@ -168,6 +173,36 @@ fn history_contains_text(history_items: &[ResponseItem], needle: &str) -> bool { }) } +fn write_rollout_items(path: &std::path::Path, items: impl IntoIterator) { + let timestamp = "2026-06-10T00:00:00Z"; + let mut jsonl = serde_json::json!({ + "timestamp": timestamp, + "type": "session_meta", + "payload": { + "id": ThreadId::default(), + "timestamp": timestamp, + "cwd": path.parent().expect("rollout parent"), + "originator": "test", + "cli_version": "test", + "source": "exec" + } + }) + .to_string(); + jsonl.push('\n'); + for item in items { + jsonl.push_str( + serde_json::to_string(&RolloutLine { + timestamp: timestamp.to_string(), + item, + }) + .expect("serialize rollout") + .as_str(), + ); + jsonl.push('\n'); + } + std::fs::write(path, jsonl).expect("write rollout"); +} + fn history_contains_assistant_inter_agent_communication( history_items: &[ResponseItem], expected: &InterAgentCommunication, @@ -815,7 +850,7 @@ async fn spawn_agent_creates_thread_and_sends_prompt() { } #[tokio::test] -async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { +async fn spawn_agent_full_history_fork_preserves_parent_items() { let harness = AgentControlHarness::new().await; let mut parent_config = harness.config.clone(); let _ = parent_config.features.enable(Feature::MultiAgentV2); @@ -825,6 +860,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { Some("Parent subagent guidance.".to_string()); let mut child_config = harness.config.clone(); let _ = child_config.features.enable(Feature::MultiAgentV2); + let _ = child_config.features.enable(Feature::SessionSegmentation); child_config.multi_agent_v2.root_agent_usage_hint_text = Some("Child root guidance.".to_string()); child_config.multi_agent_v2.subagent_usage_hint_text = @@ -932,6 +968,23 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .await .expect("child thread should be registered"); assert_ne!(child_thread_id, parent_thread_id); + let child_rollout_path = child_thread + .rollout_path() + .expect("child rollout path should be present"); + let child_rollout = RolloutRecorder::get_rollout_history(&child_rollout_path) + .await + .expect("child rollout should be readable"); + assert!( + child_rollout + .get_rollout_items() + .iter() + .any(|item| matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.nth_user_message.is_some() + )), + "full-history forks should store a compact RolloutReference so fork rollout files do not copy parent rollout history" + ); let history = child_thread.codex.session.clone_history().await; let expected_history = [ ResponseItem::Message { @@ -942,7 +995,33 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { }], phase: None, }, + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "Parent root guidance.".to_string(), + }], + phase: None, + }, + ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "Parent subagent guidance.".to_string(), + }], + phase: None, + }, + assistant_message("parent commentary", Some(MessagePhase::Commentary)), assistant_message("parent final answer", Some(MessagePhase::FinalAnswer)), + assistant_message("parent unknown phase", /*phase*/ None), + ResponseItem::Reasoning { + id: String::new(), + summary: Vec::new(), + content: None, + encrypted_content: None, + }, + trigger_message.to_response_input_item().into(), + spawn_agent_call(&parent_spawn_call_id), ResponseItem::Message { id: None, role: "developer".to_string(), @@ -955,7 +1034,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { assert_eq!( history.raw_items(), &expected_history, - "full-history forked child history should replace parent usage hints with the child subagent hint while filtering non-final assistant/tool chatter" + "full-history forked child history should preserve the parent transcript up to the fork point before appending child context" ); assert_eq!( serde_json::to_value(child_thread.codex.session.reference_context_item().await) @@ -1368,6 +1447,116 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .expect("parent shutdown should submit"); } +#[tokio::test] +async fn spawn_agent_fork_last_n_turns_materializes_referenced_segments() { + Box::pin(spawn_agent_fork_last_n_turns_materializes_referenced_segments_impl()).await; +} + +async fn spawn_agent_fork_last_n_turns_materializes_referenced_segments_impl() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + let referenced_path = harness.config.codex_home.join("referenced-parent.jsonl"); + write_rollout_items( + referenced_path.as_path(), + [ + RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "previous segment task".to_string(), + }], + phase: None, + }), + RolloutItem::ResponseItem(assistant_message( + "previous segment answer", + Some(MessagePhase::FinalAnswer), + )), + ], + ); + parent_thread + .codex + .session + .persist_rollout_items(&[RolloutItem::RolloutReference(RolloutReferenceItem { + rollout_path: referenced_path.to_path_buf(), + 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, + })]) + .await; + parent_thread + .inject_user_message_without_turn("current segment task".to_string()) + .await; + let spawn_turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-last-n-referenced".to_string(); + parent_thread + .codex + .session + .record_conversation_items( + spawn_turn_context.as_ref(), + &[spawn_agent_call(&parent_spawn_call_id)], + ) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread + .codex + .session + .flush_rollout() + .await + .expect("parent rollout should flush"); + + let child_thread_id = Box::pin(harness.control.spawn_agent_with_metadata( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id), + fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), + ..Default::default() + }, + )) + .await + .expect("forked spawn should materialize referenced parent turns") + .thread_id; + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let history = child_thread.codex.session.clone_history().await; + assert!( + history_contains_text(history.raw_items(), "previous segment task"), + "last-N fork should include turns from referenced segments" + ); + assert!( + history_contains_text(history.raw_items(), "current segment task"), + "last-N fork should include turns from the current segment" + ); + + let _ = harness + .control + .shutdown_live_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); +} + #[tokio::test] async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_limit() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 3b8d972d76..9d0ac498b6 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -24,6 +24,7 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; @@ -237,6 +238,15 @@ impl CodexThread { self.codex.session.flush_rollout().await } + pub(crate) async fn seal_rollout_segment_for_reference( + &self, + ) -> ThreadStoreResult> { + self.codex + .session + .seal_rollout_segment_for_reference() + .await + } + pub async fn submit_with_trace( &self, op: Op, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b1013965d5..3e6407a332 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -106,6 +106,7 @@ use codex_protocol::openai_models::ModelPreset; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AdditionalContextEntry; +use codex_protocol::protocol::DEFAULT_ROLLOUT_REFERENCE_DEPTH; use codex_protocol::protocol::FileChange; use codex_protocol::protocol::HasLegacyEvent; use codex_protocol::protocol::InterAgentCommunication; @@ -115,6 +116,7 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RawResponseItemEvent; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSource; @@ -145,8 +147,10 @@ use codex_thread_store::LiveThreadInitGuard; use codex_thread_store::LocalThreadStore; use codex_thread_store::ReadThreadParams; use codex_thread_store::ResumeThreadParams; +use codex_thread_store::RotateThreadSegmentParams; use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; +use codex_thread_store::ThreadStoreResult; use codex_utils_output_truncation::TruncationPolicy; use futures::future::BoxFuture; use futures::future::Shared; @@ -1219,7 +1223,20 @@ impl Session { state.clear_connector_selection(); } + #[cfg(test)] async fn record_initial_history(&self, conversation_history: InitialHistory) { + self.record_initial_history_with_persisted_fork_history( + conversation_history, + /*persisted_fork_history*/ None, + ) + .await; + } + + async fn record_initial_history_with_persisted_fork_history( + &self, + conversation_history: InitialHistory, + persisted_fork_history: Option>, + ) { let is_subagent = { let state = self.state.lock().await; state @@ -1280,6 +1297,8 @@ impl Session { } } InitialHistory::Forked(rollout_items) => { + let rollout_items_to_persist = + persisted_fork_history.unwrap_or_else(|| rollout_items.clone()); let turn_context = self.new_default_turn().await; self.apply_rollout_reconstruction(&turn_context, &rollout_items) .await; @@ -1292,8 +1311,8 @@ impl Session { } // If persisting, persist all rollout items as-is (the store filters). - if !rollout_items.is_empty() { - self.persist_rollout_items(&rollout_items).await; + if !rollout_items_to_persist.is_empty() { + self.persist_rollout_items(&rollout_items_to_persist).await; } // Forked threads should remain file-backed immediately after startup. @@ -2760,18 +2779,65 @@ impl Session { state.replace_history(items, reference_context_item.clone()); } - self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)]) + self.persist_compaction_checkpoint(compacted_item, reference_context_item) .await; - if let Some(turn_context_item) = reference_context_item { - self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item)]) - .await; - } { let mut state = self.state.lock().await; state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); } } + async fn persist_compaction_checkpoint( + &self, + compacted_item: CompactedItem, + turn_context_item: Option, + ) { + let mut rollout_items = vec![RolloutItem::Compacted(compacted_item)]; + if let Some(turn_context_item) = turn_context_item { + rollout_items.push(RolloutItem::TurnContext(turn_context_item)); + } + if !self.features.enabled(Feature::SessionSegmentation) { + self.persist_rollout_items(&rollout_items).await; + return; + } + match self.rotate_rollout_segment(rollout_items.clone()).await { + Ok(Some(_reference)) => {} + Ok(None) => self.persist_rollout_items(&rollout_items).await, + Err(err) => { + warn!("failed to rotate rollout segment after compaction: {err:#}"); + self.persist_rollout_items(&rollout_items).await; + } + } + } + + pub(crate) async fn seal_rollout_segment_for_reference( + &self, + ) -> ThreadStoreResult> { + let Some(live_thread) = self.live_thread() else { + return Ok(None); + }; + live_thread + .seal_local_segment(RotateThreadSegmentParams { + initial_items: Vec::new(), + previous_segment_reference_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + }) + .await + } + + async fn rotate_rollout_segment( + &self, + initial_items: Vec, + ) -> ThreadStoreResult> { + let Some(live_thread) = self.live_thread() else { + return Ok(None); + }; + let params = RotateThreadSegmentParams { + initial_items, + previous_segment_reference_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + }; + live_thread.rotate_local_segment(params).await + } + async fn persist_rollout_response_items(&self, items: &[ResponseItem]) { let rollout_items: Vec = items .iter() @@ -3113,14 +3179,14 @@ impl Session { let mut state = self.state.lock().await; state.replace_history(replacement_history.clone(), Some(turn_context_item.clone())); }; - self.persist_rollout_items(&[ - RolloutItem::Compacted(CompactedItem { + self.persist_compaction_checkpoint( + CompactedItem { message: String::new(), replacement_history: Some(replacement_history), window_id: Some(window_id), - }), - RolloutItem::TurnContext(turn_context_item), - ]) + }, + Some(turn_context_item), + ) .await; { let mut state = self.state.lock().await; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 158dd299e2..c7ce54c41c 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -500,6 +500,17 @@ impl Session { session_configuration.collaboration_mode.model(), session_configuration.provider ); + let persisted_fork_history = match &initial_history { + InitialHistory::Forked(items) + if items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))) => + { + Some(items.clone()) + } + InitialHistory::Forked(_) => None, + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Resumed(_) => None, + }; let initial_history = materialize_initial_history_for_model_replay( config.codex_home.as_path(), initial_history, @@ -1186,7 +1197,11 @@ impl Session { }; // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. - Box::pin(sess.record_initial_history(initial_history)).await; + Box::pin(sess.record_initial_history_with_persisted_fork_history( + initial_history, + persisted_fork_history, + )) + .await; { let mut state = sess.state.lock().await; state.queue_pending_session_start_source(session_start_source); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 01135e50a1..7814d6969b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3777,6 +3777,200 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { .expect("thread should have rollout path") } +#[tokio::test] +async fn replace_compacted_history_rolls_over_local_segment_at_stable_path() { + 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"); + sess.features + .enable(Feature::SessionSegmentation) + .expect("enable session segmentation"); + 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 (old_items, _, _) = RolloutRecorder::load_rollout_items(old_rollout_path.as_path()) + .await + .expect("load pre-compaction rollout segment"); + let old_segment_id = old_items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta) => meta.meta.segment_id, + RolloutItem::InterAgentCommunication(_) + | RolloutItem::RolloutReference(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => None, + }); + let old_segment_id = old_segment_id.expect("pre-compaction rollout should have segment id"); + 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()), + window_id: None, + }, + ) + .await; + + let new_rollout_path = sess + .current_rollout_path() + .await + .expect("load current rollout path") + .expect("rollout path after compaction"); + assert_eq!(new_rollout_path, old_rollout_path); + let config = sess.get_config().await; + let rotated_old_rollout_path = config + .codex_home + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(sess.thread_id.to_string()) + .join(old_segment_id.to_string()) + .join( + old_rollout_path + .file_name() + .expect("old rollout path should have a file name"), + ); + assert!(old_rollout_path.exists()); + assert!(rotated_old_rollout_path.exists()); + let old_rollout_timestamp = old_rollout_path + .file_name() + .and_then(|file_name| file_name.to_str()) + .and_then(|file_name| file_name.strip_prefix("rollout-")) + .and_then(|file_name| file_name.strip_suffix(&format!("-{}.jsonl", sess.thread_id))) + .expect("old rollout timestamp"); + 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.thread_id)); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.rollout_path.as_path() == rotated_old_rollout_path.as_path() + && reference.thread_id == Some(sess.thread_id) + && reference.rollout_timestamp.as_deref() == Some(old_rollout_timestamp) + ) + })); + assert!(new_items.iter().any(|item| { + matches!( + item, + RolloutItem::Compacted(CompactedItem { + replacement_history: Some(history), + .. + }) if *history == replacement_history + ) + })); + + let replay_items = + crate::thread_rollout_truncation::materialize_rollout_items_for_model_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); +} + +#[tokio::test] +async fn replace_compacted_history_does_not_segment_when_feature_is_disabled() { + 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"); + assert!(!sess.features.enabled(Feature::SessionSegmentation)); + let 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), + window_id: None, + }, + ) + .await; + sess.flush_rollout() + .await + .expect("post-compaction rollout should flush"); + + assert_eq!( + sess.current_rollout_path() + .await + .expect("load current rollout path"), + Some(rollout_path.clone()) + ); + let (items, _, _) = RolloutRecorder::load_rollout_items(rollout_path.as_path()) + .await + .expect("load rollout after compaction"); + assert!( + !items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))) + ); + let config = sess.get_config().await; + assert!( + !config + .codex_home + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .exists() + ); +} + +#[tokio::test] +async fn requested_context_window_rolls_over_local_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"); + sess.features + .enable(Feature::SessionSegmentation) + .expect("enable session segmentation"); + let rollout_path = attach_thread_persistence(sess).await; + sess.persist_rollout_items(&[RolloutItem::ResponseItem(user_message( + "before requested context window", + ))]) + .await; + sess.request_new_context_window().await; + + let window_id = sess + .maybe_start_new_context_window(tc.as_ref()) + .await + .expect("start requested context window"); + + let (items, _, _) = RolloutRecorder::load_rollout_items(rollout_path.as_path()) + .await + .expect("load rollout after requested context window"); + assert!( + items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))) + ); + assert!(items.iter().any(|item| { + matches!( + item, + RolloutItem::Compacted(CompactedItem { + window_id: Some(compacted_window_id), + .. + }) if *compacted_window_id == window_id + ) + })); +} + 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 e9b00cbb08..c3ab8359f7 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -39,6 +39,7 @@ use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::openai_models::ModelPreset; +use codex_protocol::protocol::DEFAULT_ROLLOUT_REFERENCE_DEPTH; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; @@ -61,6 +62,7 @@ use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; use codex_thread_store::ReadThreadByRolloutPathParams; use codex_thread_store::ReadThreadParams; +use codex_thread_store::RotateThreadSegmentParams; use codex_thread_store::StoredThread; use codex_thread_store::ThreadMetadataPatch; use codex_thread_store::ThreadStore; @@ -673,19 +675,16 @@ 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( + options.initial_history = Box::pin(self.prepare_fork_history( ForkSnapshot::Interrupted, + &options.config, history, InterruptedTurnHistoryMarker::from_config_and_version( &options.config, inherited_multi_agent_version, ), - ); + )) + .await; self.start_thread_with_options_and_fork_source(options, Some(forked_from_thread_id)) .await } @@ -951,12 +950,9 @@ 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 history = + Box::pin(self.prepare_fork_history(snapshot, &config, history, interrupted_marker)) + .await; let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), &config.cwd, @@ -979,6 +975,114 @@ impl ThreadManager { .await } + async fn prepare_fork_history( + &self, + snapshot: ForkSnapshot, + config: &Config, + history: InitialHistory, + interrupted_marker: InterruptedTurnHistoryMarker, + ) -> InitialHistory { + let reference_history = self + .rollout_reference_history_for_snapshot(snapshot, config, &history, interrupted_marker) + .await; + match reference_history { + Some(history) => history, + None => { + let history = truncation::materialize_initial_history_for_model_replay( + config.codex_home.as_path(), + history, + ) + .await; + fork_history_from_snapshot(snapshot, history, interrupted_marker) + } + } + } + + async fn rollout_reference_history_for_snapshot( + &self, + snapshot: ForkSnapshot, + config: &Config, + history: &InitialHistory, + interrupted_marker: InterruptedTurnHistoryMarker, + ) -> Option { + if !config.features.enabled(Feature::SessionSegmentation) { + return None; + } + let InitialHistory::Resumed(resumed) = history else { + return None; + }; + let rollout_path = resumed.rollout_path.clone()?; + let local_store = self + .state + .thread_store + .as_any() + .downcast_ref::()?; + let mut reference = match local_store + .seal_thread_segment( + resumed.conversation_id, + rollout_path, + RotateThreadSegmentParams { + initial_items: Vec::new(), + previous_segment_reference_depth: DEFAULT_ROLLOUT_REFERENCE_DEPTH, + }, + ) + .await + { + Ok(reference) => reference, + Err(err) => { + warn!("failed to seal rollout segment for fork; copying history instead: {err}"); + return None; + } + }; + + let materialized = + match crate::thread_rollout_truncation::materialize_rollout_items_for_complete_history( + config.codex_home.as_path(), + &[RolloutItem::RolloutReference(reference.clone())], + ) + .await + { + Ok(materialized) => materialized, + Err(err) => { + warn!( + "failed to inspect immutable rollout snapshot for fork; copying history instead: {err}" + ); + return None; + } + }; + let materialized_history = InitialHistory::Forked(materialized); + let snapshot_state = snapshot_turn_state(&materialized_history); + reference.nth_user_message = match snapshot { + ForkSnapshot::TruncateBeforeNthUserMessage(nth_user_message) => { + let user_message_count = + crate::thread_rollout_truncation::user_message_positions_in_rollout( + &materialized_history.get_rollout_items(), + ) + .len(); + if snapshot_state.ends_mid_turn && nth_user_message >= user_message_count { + return None; + } + Some(if nth_user_message < user_message_count { + nth_user_message + } else { + usize::MAX + }) + } + ForkSnapshot::Interrupted => Some(usize::MAX), + }; + + let reference_history = + InitialHistory::Forked(vec![RolloutItem::RolloutReference(reference)]); + if snapshot == ForkSnapshot::Interrupted && snapshot_state.ends_mid_turn { + return Some(append_interrupted_boundary( + reference_history, + snapshot_state.active_turn_id, + interrupted_marker, + )); + } + Some(reference_history) + } + pub(crate) fn agent_control(&self) -> AgentControl { AgentControl::new(Arc::downgrade(&self.state)) } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 41ad9b104d..8f4e527e38 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1552,6 +1552,10 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let mut config = test_config().await; config.codex_home = temp_dir.path().join("codex-home").abs(); config.cwd = config.codex_home.abs(); + config + .features + .enable(Feature::SessionSegmentation) + .expect("enable session segmentation"); std::fs::create_dir_all(&config.codex_home).expect("create codex home"); let auth_manager = @@ -1591,6 +1595,11 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ .await .expect("read source rollout history"); assert!(snapshot_turn_state(&source_history).ends_mid_turn); + source + .thread + .shutdown_and_wait() + .await + .expect("shutdown source thread"); manager.remove_thread(&source.thread_id).await; let forked = manager @@ -1612,8 +1621,20 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ .expect("read forked rollout history"); assert!(!snapshot_turn_state(&history).ends_mid_turn); - let forked_rollout_items: Vec<_> = history - .get_rollout_items() + let forked_raw_items = history.get_rollout_items(); + assert!( + forked_raw_items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))), + "segmented fork should persist an immutable rollout reference" + ); + let forked_rollout_items: Vec<_> = + crate::thread_rollout_truncation::materialize_rollout_items_for_complete_history( + config.codex_home.as_path(), + &forked_raw_items, + ) + .await + .expect("materialize forked history") .into_iter() .filter(|item| !matches!(item, RolloutItem::SessionMeta(_))) .collect(); @@ -1632,6 +1653,11 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ 1, ); + forked + .thread + .shutdown_and_wait() + .await + .expect("shutdown forked thread"); manager.remove_thread(&forked.thread_id).await; let reforked = manager .fork_thread( @@ -1650,8 +1676,20 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let reforked_history = RolloutRecorder::get_rollout_history(&reforked_path) .await .expect("read re-forked rollout history"); - let reforked_rollout_items: Vec<_> = reforked_history - .get_rollout_items() + let reforked_raw_items = reforked_history.get_rollout_items(); + assert!( + reforked_raw_items + .iter() + .any(|item| matches!(item, RolloutItem::RolloutReference(_))), + "segmented re-fork should persist an immutable rollout reference" + ); + let reforked_rollout_items: Vec<_> = + crate::thread_rollout_truncation::materialize_rollout_items_for_complete_history( + config.codex_home.as_path(), + &reforked_raw_items, + ) + .await + .expect("materialize re-forked history") .into_iter() .filter(|item| !matches!(item, RolloutItem::SessionMeta(_))) .collect(); diff --git a/codex-rs/core/tests/suite/fork_thread.rs b/codex-rs/core/tests/suite/fork_thread.rs index b4898ea827..a5c2a10257 100644 --- a/codex-rs/core/tests/suite/fork_thread.rs +++ b/codex-rs/core/tests/suite/fork_thread.rs @@ -1,6 +1,8 @@ use codex_core::ForkSnapshot; use codex_core::NewThread; +use codex_core::materialize_rollout_items_for_model_replay; use codex_core::parse_turn_item; +use codex_features::Feature; use codex_protocol::items::TurnItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; @@ -40,7 +42,12 @@ async fn fork_thread_twice_drops_to_first_message() { .mount(&server) .await; - let mut builder = test_codex(); + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::SessionSegmentation) + .expect("enable session segmentation"); + }); let test = builder.build(&server).await.expect("create conversation"); let codex = test.codex.clone(); let thread_manager = test.thread_manager.clone(); @@ -111,13 +118,29 @@ async fn fork_thread_twice_drops_to_first_message() { let fork1_path = codex_fork1.rollout_path().expect("rollout path"); // GetHistory on fork1 flushed; the file is ready. - let fork1_items = read_rollout_items(&fork1_path); + let fork1_raw_items = read_rollout_items(&fork1_path); + assert!( + fork1_raw_items.iter().any(|item| matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.nth_user_message.is_some() + )), + "forked rollout should keep a compact RolloutReference instead of copying parent history" + ); + let fork1_items = materialize_rollout_items_for_model_replay( + test.config.codex_home.as_path(), + &fork1_raw_items, + ) + .await; + let fork1_items = without_session_meta(fork1_items); pretty_assertions::assert_eq!( serde_json::to_value(&fork1_items).unwrap(), serde_json::to_value(&expected_after_first).unwrap() ); - // Fork again with n=0 → drops the (new) last user message, leaving only the first. + // Fork again with n=0. The first fork's raw rollout only contains a + // RolloutReference, but truncation still applies to the materialized history + // referenced by that item. let NewThread { thread: codex_fork2, .. @@ -134,14 +157,27 @@ async fn fork_thread_twice_drops_to_first_message() { let fork2_path = codex_fork2.rollout_path().expect("rollout path"); // GetHistory on fork2 flushed; the file is ready. - let fork1_items = read_rollout_items(&fork1_path); let fork1_user_inputs = find_user_input_positions(&fork1_items); - let cut_last_on_fork1 = fork1_user_inputs - .get(fork1_user_inputs.len().saturating_sub(1)) + let cut2 = fork1_user_inputs + .first() .copied() - .unwrap_or(0); - let expected_after_second: Vec = fork1_items[..cut_last_on_fork1].to_vec(); - let fork2_items = read_rollout_items(&fork2_path); + .unwrap_or(fork1_items.len()); + let expected_after_second = fork1_items[..cut2].to_vec(); + let fork2_raw_items = read_rollout_items(&fork2_path); + assert!( + fork2_raw_items.iter().any(|item| matches!( + item, + RolloutItem::RolloutReference(reference) + if reference.nth_user_message.is_some() + )), + "re-forked rollout should keep a compact RolloutReference instead of copying parent history" + ); + let fork2_items = materialize_rollout_items_for_model_replay( + test.config.codex_home.as_path(), + &fork2_raw_items, + ) + .await; + let fork2_items = without_session_meta(fork2_items); pretty_assertions::assert_eq!( serde_json::to_value(&fork2_items).unwrap(), serde_json::to_value(&expected_after_second).unwrap() @@ -246,3 +282,10 @@ fn read_rollout_items(path: &std::path::Path) -> Vec { } items } + +fn without_session_meta(items: Vec) -> Vec { + items + .into_iter() + .filter(|item| !matches!(item, RolloutItem::SessionMeta(_))) + .collect() +} diff --git a/codex-rs/core/tests/suite/resume.rs b/codex-rs/core/tests/suite/resume.rs index 01faaf8fd3..59be4e92be 100644 --- a/codex-rs/core/tests/suite/resume.rs +++ b/codex-rs/core/tests/suite/resume.rs @@ -190,6 +190,19 @@ async fn resume_materializes_references_before_session_configured() -> Result<() 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"); + let mut successor_session_meta = serde_json::from_str::(session_meta_line)?; + let predecessor_segment_id = match &mut successor_session_meta.item { + RolloutItem::SessionMeta(session_meta) => { + let predecessor_segment_id = session_meta + .meta + .segment_id + .expect("new rollout has segment id"); + session_meta.meta.segment_id = Some(codex_protocol::SegmentId::new()); + predecessor_segment_id + } + item => panic!("expected session meta, got {item:?}"), + }; + let successor_session_meta_line = serde_json::to_string(&successor_session_meta)?; std::fs::rename(&rollout_path, &predecessor_path)?; let reference_line = serde_json::to_string(&RolloutLine { timestamp: "2026-06-12T00:00:00Z".to_string(), @@ -197,7 +210,7 @@ async fn resume_materializes_references_before_session_configured() -> Result<() rollout_path: predecessor_path, thread_id: Some(thread_id), rollout_timestamp: None, - segment_id: None, + segment_id: Some(predecessor_segment_id), max_depth: 1, nth_user_message: None, compacted_replacement_history_filter_texts: None, @@ -205,7 +218,7 @@ async fn resume_materializes_references_before_session_configured() -> Result<() })?; std::fs::write( &rollout_path, - format!("{session_meta_line}\n{reference_line}\n"), + format!("{successor_session_meta_line}\n{reference_line}\n"), )?; let resumed = resume_until_initial_messages( diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 65c1c3b28b..3b80ba471e 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -127,6 +127,8 @@ pub enum Feature { MemoryTool, /// Compress cold local thread-store rollout files. LocalThreadStoreCompression, + /// Split long-lived sessions into rollout segments connected by references. + SessionSegmentation, /// Enable the Chronicle sidecar for passive screen-context memories. Chronicle, /// Append additional AGENTS.md guidance to user instructions. @@ -872,6 +874,16 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::SessionSegmentation, + key: "session_segmentation", + stage: Stage::Experimental { + name: "Session segmentation", + menu_description: "Split long-lived sessions into linked rollout segments after compaction.", + announcement: "", + }, + default_enabled: false, + }, FeatureSpec { id: Feature::Chronicle, key: "chronicle", diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index db56819889..cc37eb286c 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -355,6 +355,23 @@ fn telepathy_is_legacy_alias_for_chronicle() { assert_eq!(feature_for_key("telepathy"), Some(Feature::Chronicle)); } +#[test] +fn session_segmentation_is_experimental_and_disabled_by_default() { + assert_eq!( + Feature::SessionSegmentation.stage(), + Stage::Experimental { + name: "Session segmentation", + menu_description: "Split long-lived sessions into linked rollout segments after compaction.", + announcement: "", + } + ); + assert!(!Feature::SessionSegmentation.default_enabled()); + assert_eq!( + feature_for_key("session_segmentation"), + Some(Feature::SessionSegmentation) + ); +} + #[test] fn collab_is_legacy_alias_for_multi_agent() { assert_eq!(feature_for_key("multi_agent"), Some(Feature::Collab)); diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index fe766060c0..9edee3ae19 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use std::sync::Mutex; use chrono::SecondsFormat; +use codex_protocol::SegmentId; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; @@ -91,6 +92,18 @@ pub enum RolloutRecorderParams { dynamic_tools: Vec, multi_agent_version: Option, }, + CreateAtPath { + path: PathBuf, + conversation_id: ThreadId, + forked_from_id: Option, + parent_thread_id: Option, + source: SessionSource, + thread_source: Option, + base_instructions: BaseInstructions, + dynamic_tools: Vec, + multi_agent_version: Option, + session_timestamp: Option, + }, Resume { path: PathBuf, }, @@ -715,44 +728,56 @@ impl RolloutRecorder { } => { let log_file_info = precompute_log_file_info(config, conversation_id)?; let path = log_file_info.path.clone(); - let session_id = log_file_info.conversation_id; - let started_at = log_file_info.timestamp; - - let timestamp_format: &[FormatItem] = format_description!( - "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" - ); - let timestamp = started_at - .to_offset(time::UtcOffset::UTC) - .format(timestamp_format) - .map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?; - - let session_meta = SessionMeta { - id: session_id, - segment_id: None, - forked_from_id, - parent_thread_id, - timestamp, - cwd: config.cwd().to_path_buf(), - originator: originator().value, - cli_version: env!("CARGO_PKG_VERSION").to_string(), - agent_nickname: source.get_nickname(), - agent_role: source.get_agent_role(), - agent_path: source.get_agent_path().map(Into::into), - source, - thread_source, - model_provider: Some(config.model_provider_id().to_string()), - base_instructions: Some(base_instructions), - dynamic_tools: if dynamic_tools.is_empty() { - None - } else { - Some(dynamic_tools) + let session_meta = create_session_meta( + config, + &log_file_info, + SessionMetaParams { + forked_from_id, + parent_thread_id, + source, + thread_source, + base_instructions, + dynamic_tools, + multi_agent_version, + timestamp_override: None, }, - memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), - multi_agent_version, - }; + )?; (None, Some(log_file_info), path, Some(session_meta)) } + RolloutRecorderParams::CreateAtPath { + path, + conversation_id, + forked_from_id, + parent_thread_id, + source, + thread_source, + base_instructions, + dynamic_tools, + multi_agent_version, + session_timestamp, + } => { + let log_file_info = LogFileInfo { + path: path.clone(), + conversation_id, + timestamp: OffsetDateTime::now_utc(), + }; + let session_meta = create_session_meta( + config, + &log_file_info, + SessionMetaParams { + forked_from_id, + parent_thread_id, + source, + thread_source, + base_instructions, + dynamic_tools, + multi_agent_version, + timestamp_override: session_timestamp, + }, + )?; + (None, Some(log_file_info), path, Some(session_meta)) + } RolloutRecorderParams::Resume { path } => { let path = compression::materialize_rollout_for_append(path.as_path()).await?; ( @@ -1453,6 +1478,17 @@ struct LogFileInfo { timestamp: OffsetDateTime, } +struct SessionMetaParams { + forked_from_id: Option, + parent_thread_id: Option, + source: SessionSource, + thread_source: Option, + base_instructions: BaseInstructions, + dynamic_tools: Vec, + multi_agent_version: Option, + timestamp_override: Option, +} + fn precompute_log_file_info( config: &impl RolloutConfigView, conversation_id: ThreadId, @@ -1475,7 +1511,6 @@ fn precompute_log_file_info( .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); Ok(LogFileInfo { @@ -1500,6 +1535,58 @@ fn open_log_file(path: &Path) -> std::io::Result { .open(path) } +fn create_session_meta( + config: &impl RolloutConfigView, + log_file_info: &LogFileInfo, + params: SessionMetaParams, +) -> std::io::Result { + let SessionMetaParams { + forked_from_id, + parent_thread_id, + source, + thread_source, + base_instructions, + dynamic_tools, + multi_agent_version, + timestamp_override, + } = params; + let timestamp_format: &[FormatItem] = + format_description!("[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"); + let timestamp = match timestamp_override { + Some(timestamp) => timestamp, + None => log_file_info + .timestamp + .to_offset(time::UtcOffset::UTC) + .format(timestamp_format) + .map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?, + }; + + Ok(SessionMeta { + id: log_file_info.conversation_id, + segment_id: Some(SegmentId::new()), + forked_from_id, + parent_thread_id, + timestamp, + cwd: config.cwd().to_path_buf(), + originator: originator().value, + cli_version: env!("CARGO_PKG_VERSION").to_string(), + agent_nickname: source.get_nickname(), + agent_role: source.get_agent_role(), + agent_path: source.get_agent_path().map(Into::into), + source, + thread_source, + model_provider: Some(config.model_provider_id().to_string()), + base_instructions: Some(base_instructions), + dynamic_tools: if dynamic_tools.is_empty() { + None + } else { + Some(dynamic_tools) + }, + memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), + multi_agent_version, + }) +} + /// Mutable state owned by the background rollout writer. /// /// Items are first appended to `pending_items`; persist/flush/shutdown remove each item from that diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index b37ef88267..f80fbcc34b 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -426,6 +426,20 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< text.contains("\"type\":\"session_meta\""), "expected session metadata in rollout" ); + let (items, _, _) = RolloutRecorder::load_rollout_items(&rollout_path).await?; + let segment_id = items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta) => meta.meta.segment_id, + RolloutItem::InterAgentCommunication(_) + | RolloutItem::RolloutReference(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => None, + }); + assert!( + segment_id.is_some(), + "new rollout metadata should include a segment_id for durable references" + ); let buffered_idx = text .find("buffered-event") .expect("buffered event in rollout"); diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 7c59cf72f3..5e8b63cbb9 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -22,6 +22,7 @@ codex-state = { workspace = true } codex-utils-path = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tempfile = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index d0ab3805c9..fe08fdef97 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -37,6 +37,7 @@ pub use types::LoadThreadHistoryParams; pub use types::ReadThreadByRolloutPathParams; pub use types::ReadThreadParams; pub use types::ResumeThreadParams; +pub use types::RotateThreadSegmentParams; pub use types::SearchThreadsParams; pub use types::SortDirection; pub use types::StoredThread; diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 88d347f57b..36d6908227 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use codex_protocol::ThreadId; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutReferenceItem; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::persisted_rollout_items; use tokio::sync::Mutex; @@ -14,6 +15,7 @@ use crate::LoadThreadHistoryParams; use crate::LocalThreadStore; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::RotateThreadSegmentParams; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; @@ -168,6 +170,42 @@ impl LiveThread { Ok(()) } + /// Replace the current local rollout with a successor segment. + pub async fn rotate_local_segment( + &self, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult> { + let Some(local_store) = self + .thread_store + .as_any() + .downcast_ref::() + else { + return Ok(None); + }; + local_store + .rotate_thread_segment(self.thread_id, params) + .await + .map(Some) + } + + /// Seal the current local segment for a fork without consuming model-replay segment depth. + pub async fn seal_local_segment( + &self, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult> { + let Some(local_store) = self + .thread_store + .as_any() + .downcast_ref::() + else { + return Ok(None); + }; + local_store + .seal_live_thread_segment(self.thread_id, params) + .await + .map(Some) + } + pub async fn persist(&self) -> ThreadStoreResult<()> { self.thread_store.persist_thread(self.thread_id).await?; self.flush_pending_metadata_update().await diff --git a/codex-rs/thread-store/src/local/archive_thread.rs b/codex-rs/thread-store/src/local/archive_thread.rs index 8fb214e98c..cbbcffae36 100644 --- a/codex-rs/thread-store/src/local/archive_thread.rs +++ b/codex-rs/thread-store/src/local/archive_thread.rs @@ -13,6 +13,7 @@ pub(super) async fn archive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; + let _segment_storage_guard = store.segment_storage_lock.clone().lock_owned().await; let state_db_ctx = store.state_db().await; let rollout_path = find_thread_path_by_id_str( store.config.codex_home.as_path(), diff --git a/codex-rs/thread-store/src/local/delete_thread.rs b/codex-rs/thread-store/src/local/delete_thread.rs index ecd962f647..b648fbfc0e 100644 --- a/codex-rs/thread-store/src/local/delete_thread.rs +++ b/codex-rs/thread-store/src/local/delete_thread.rs @@ -25,6 +25,8 @@ pub(super) async fn delete_thread( params: DeleteThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; + super::live_writer::close_thread_for_delete(store, thread_id).await?; + let _segment_storage_guard = store.segment_storage_lock.clone().lock_owned().await; let thread_id_str = thread_id.to_string(); let state_db_ctx = store.state_db().await; let mut rollout_paths = Vec::new(); @@ -66,6 +68,15 @@ pub(super) async fn delete_thread( } let found_rollout_path = !rollout_paths.is_empty(); + let found_segment_tree = store + .config + .codex_home + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(thread_id.to_string()) + .try_exists() + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to inspect rollout segments for thread {thread_id}: {err}"), + })?; for rollout_path in rollout_paths { delete_rollout_file(store, rollout_path.as_path(), thread_id)?; } @@ -75,11 +86,17 @@ pub(super) async fn delete_thread( message: format!("failed to delete thread name index entries for {thread_id}: {err}"), })?; - if !found_rollout_path { + if !found_rollout_path && !found_segment_tree { return Err(ThreadStoreError::ThreadNotFound { thread_id }); } - store.live_recorders.lock().await.remove(&thread_id); + super::segment_gc::collect_unreferenced_segments(store.config.codex_home.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to collect rollout segments after deleting {thread_id}: {err}" + ), + })?; Ok(()) } diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 61c15a43d9..822c27ed2c 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -1,4 +1,6 @@ +use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadMemoryMode; @@ -8,6 +10,7 @@ use codex_rollout::RolloutRecorderParams; use codex_rollout::persisted_rollout_items; use tracing::warn; +use super::LiveRecorderHandle; use super::LocalThreadStore; use super::create_thread; use crate::AppendThreadItemsParams; @@ -82,7 +85,9 @@ pub(super) async fn append_items( if canonical_items.is_empty() { return Ok(()); } - let recorder = store.live_recorder(params.thread_id).await?; + let handle = store.live_recorder_handle(params.thread_id).await?; + let state = Arc::clone(&handle).lock_owned().await; + let recorder = live_recorder(&state.recorder, params.thread_id)?; recorder .record_canonical_items(canonical_items.as_slice()) .await @@ -96,36 +101,47 @@ pub(super) async fn persist_thread( store: &LocalThreadStore, thread_id: ThreadId, ) -> ThreadStoreResult<()> { - store - .live_recorder(thread_id) - .await? - .persist() - .await - .map_err(thread_store_io_error)?; - sync_materialized_rollout_path(store, thread_id).await + let handle = store.live_recorder_handle(thread_id).await?; + let rollout_path = { + let state = Arc::clone(&handle).lock_owned().await; + let recorder = live_recorder(&state.recorder, thread_id)?; + recorder.persist().await.map_err(thread_store_io_error)?; + recorder.rollout_path().to_path_buf() + }; + sync_materialized_rollout_path(store, thread_id, rollout_path.as_path()).await } pub(super) async fn flush_thread( store: &LocalThreadStore, thread_id: ThreadId, ) -> ThreadStoreResult<()> { - store - .live_recorder(thread_id) - .await? - .flush() - .await - .map_err(thread_store_io_error)?; - sync_materialized_rollout_path(store, thread_id).await + let handle = store.live_recorder_handle(thread_id).await?; + let rollout_path = { + let state = Arc::clone(&handle).lock_owned().await; + let recorder = live_recorder(&state.recorder, thread_id)?; + recorder.flush().await.map_err(thread_store_io_error)?; + recorder.rollout_path().to_path_buf() + }; + sync_materialized_rollout_path(store, thread_id, rollout_path.as_path()).await } pub(super) async fn shutdown_thread( store: &LocalThreadStore, thread_id: ThreadId, ) -> ThreadStoreResult<()> { - let recorder = store.live_recorder(thread_id).await?; - recorder.shutdown().await.map_err(thread_store_io_error)?; - sync_materialized_rollout_path(store, thread_id).await?; - store.live_recorders.lock().await.remove(&thread_id); + let handle = store.live_recorder_handle(thread_id).await?; + let rollout_path = { + let mut state = Arc::clone(&handle).lock_owned().await; + let recorder = state + .recorder + .take() + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })?; + let rollout_path = recorder.rollout_path().to_path_buf(); + recorder.shutdown().await.map_err(thread_store_io_error)?; + rollout_path + }; + sync_materialized_rollout_path(store, thread_id, rollout_path.as_path()).await?; + remove_live_recorder_handle(store, thread_id, &handle).await; Ok(()) } @@ -133,35 +149,76 @@ pub(super) async fn discard_thread( store: &LocalThreadStore, thread_id: ThreadId, ) -> ThreadStoreResult<()> { - store - .live_recorders - .lock() - .await - .remove(&thread_id) - .map(|_| ()) - .ok_or(ThreadStoreError::ThreadNotFound { thread_id }) + let handle = store.live_recorder_handle(thread_id).await?; + { + let mut state = handle.lock().await; + state.recorder.take(); + } + remove_live_recorder_handle(store, thread_id, &handle).await; + Ok(()) +} + +pub(super) async fn close_thread_for_delete( + store: &LocalThreadStore, + thread_id: ThreadId, +) -> ThreadStoreResult<()> { + let Some(handle) = store.live_recorders.lock().await.get(&thread_id).cloned() else { + return Ok(()); + }; + let recorder = handle.lock().await.recorder.take(); + if let Some(recorder) = recorder { + recorder.shutdown().await.map_err(thread_store_io_error)?; + } + remove_live_recorder_handle(store, thread_id, &handle).await; + Ok(()) } pub(super) async fn rollout_path( store: &LocalThreadStore, thread_id: ThreadId, ) -> ThreadStoreResult { - Ok(store - .live_recorders - .lock() - .await - .get(&thread_id) - .ok_or(ThreadStoreError::ThreadNotFound { thread_id })? + let handle = store.live_recorder_handle(thread_id).await?; + let state = handle.lock().await; + Ok(live_recorder(&state.recorder, thread_id)? .rollout_path() .to_path_buf()) } -async fn sync_materialized_rollout_path( +pub(super) fn live_recorder( + recorder: &Option, + thread_id: ThreadId, +) -> ThreadStoreResult<&RolloutRecorder> { + recorder + .as_ref() + .ok_or(ThreadStoreError::ThreadNotFound { thread_id }) +} + +pub(super) fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError { + ThreadStoreError::Internal { + message: err.to_string(), + } +} + +pub(super) async fn remove_live_recorder_handle( store: &LocalThreadStore, thread_id: ThreadId, + handle: &LiveRecorderHandle, +) { + let mut live_recorders = store.live_recorders.lock().await; + if live_recorders + .get(&thread_id) + .is_some_and(|current| Arc::ptr_eq(current, handle)) + { + live_recorders.remove(&thread_id); + } +} + +pub(super) async fn sync_materialized_rollout_path( + store: &LocalThreadStore, + thread_id: ThreadId, + rollout_path: &Path, ) -> ThreadStoreResult<()> { - let rollout_path = rollout_path(store, thread_id).await?; - if codex_rollout::existing_rollout_path(rollout_path.as_path()) + if codex_rollout::existing_rollout_path(rollout_path) .await .is_none() { @@ -182,7 +239,7 @@ async fn sync_materialized_rollout_path( return Ok(()); }; if metadata.rollout_path != rollout_path { - metadata.rollout_path = rollout_path; + metadata.rollout_path = rollout_path.to_path_buf(); state_db .upsert_thread(&metadata) .await @@ -198,9 +255,3 @@ async fn sync_materialized_rollout_path( } 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 dd49fd2a59..9eecf96d49 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -6,6 +6,8 @@ mod list_threads; mod live_writer; mod read_thread; mod search_threads; +mod segment_gc; +mod segment_writer; mod unarchive_thread; mod update_thread_metadata; @@ -30,6 +32,7 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::RotateThreadSegmentParams; use crate::SearchThreadsParams; use crate::StoredThread; use crate::StoredThreadHistory; @@ -57,10 +60,27 @@ use crate::UpdateThreadMetadataParams; #[derive(Clone)] pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, - live_recorders: Arc>>, + live_recorders: Arc>>, + segment_storage_lock: Arc>, + #[cfg(test)] + rotation_test_hook: Arc>>, state_db: Option, } +pub(super) type LiveRecorderHandle = Arc>; + +pub(super) struct LiveRecorderState { + pub(super) recorder: Option, +} + +#[cfg(test)] +#[derive(Clone)] +pub(super) struct RotationTestHook { + pub(super) reached: Arc, + pub(super) release: Arc, + pub(super) fail_before_install: bool, +} + /// Process-scoped configuration for local thread storage. /// /// This describes where local storage lives. New-thread rollout metadata such @@ -97,6 +117,9 @@ impl LocalThreadStore { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), + segment_storage_lock: Arc::new(Mutex::new(())), + #[cfg(test)] + rotation_test_hook: Arc::new(Mutex::new(None)), state_db, } } @@ -127,10 +150,41 @@ impl LocalThreadStore { live_writer::rollout_path(self, thread_id).await } - pub(super) async fn live_recorder( + /// Replace a running local thread's canonical rollout with a successor segment. + pub async fn rotate_thread_segment( &self, thread_id: ThreadId, - ) -> ThreadStoreResult { + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult { + segment_writer::rotate_thread_segment(self, thread_id, params).await + } + + /// Seal a running local thread at a non-consuming snapshot boundary. + pub async fn seal_live_thread_segment( + &self, + thread_id: ThreadId, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult { + segment_writer::seal_live_thread_segment(self, thread_id, params).await + } + + /// Seal the current local segment and return an immutable reference to it. + /// + /// If the thread is not running, `rollout_path` is opened only for this transaction and the + /// successor writer is closed before this method returns. + pub async fn seal_thread_segment( + &self, + thread_id: ThreadId, + rollout_path: PathBuf, + params: RotateThreadSegmentParams, + ) -> ThreadStoreResult { + segment_writer::seal_thread_segment(self, thread_id, rollout_path, params).await + } + + pub(super) async fn live_recorder_handle( + &self, + thread_id: ThreadId, + ) -> ThreadStoreResult { self.live_recorders .lock() .await @@ -161,7 +215,9 @@ impl LocalThreadStore { message: format!("thread {} already has a live local writer", entry.key()), }), Entry::Vacant(entry) => { - entry.insert(recorder); + entry.insert(Arc::new(Mutex::new(LiveRecorderState { + recorder: Some(recorder), + }))); Ok(()) } } diff --git a/codex-rs/thread-store/src/local/segment_gc.rs b/codex-rs/thread-store/src/local/segment_gc.rs new file mode 100644 index 0000000000..6b295b6d00 --- /dev/null +++ b/codex-rs/thread-store/src/local/segment_gc.rs @@ -0,0 +1,136 @@ +use std::collections::HashSet; +use std::io; +use std::path::Path; +use std::path::PathBuf; + +use codex_protocol::protocol::RolloutItem; +use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR; +use codex_rollout::RolloutRecorder; +use codex_rollout::SESSIONS_SUBDIR; +use tokio::fs; +use tracing::warn; + +/// Remove immutable segments that are unreachable from every canonical active or archived +/// rollout. Forks may reference another thread's segment, so segment lifetime is determined by +/// reachability rather than by the thread id encoded in the storage path. +pub(super) async fn collect_unreferenced_segments(codex_home: &Path) -> io::Result<()> { + let segment_root = codex_home.join(ROTATED_ROLLOUT_SEGMENTS_SUBDIR); + if !fs::try_exists(segment_root.as_path()).await? { + return Ok(()); + } + let canonical_segment_root = fs::canonicalize(segment_root.as_path()).await?; + let mut pending = Vec::new(); + for root in [SESSIONS_SUBDIR, ARCHIVED_SESSIONS_SUBDIR] { + pending.extend(rollout_files_under(codex_home.join(root).as_path())?); + } + + let mut visited = HashSet::new(); + let mut reachable_segments = HashSet::new(); + while let Some(path) = pending.pop() { + let canonical_path = fs::canonicalize(path.as_path()).await?; + if !visited.insert(canonical_path) { + continue; + } + let (items, _, _) = RolloutRecorder::load_rollout_items(path.as_path()).await?; + for reference in items.iter().filter_map(|item| match item { + RolloutItem::RolloutReference(reference) => Some(reference), + RolloutItem::Compacted(_) + | RolloutItem::EventMsg(_) + | RolloutItem::InterAgentCommunication(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::TurnContext(_) => None, + }) { + let referenced_path = + codex_rollout::resolve_rollout_reference_rollout_path(codex_home, reference) + .await?; + let canonical_referenced_path = fs::canonicalize(referenced_path.as_path()).await?; + if canonical_referenced_path.starts_with(canonical_segment_root.as_path()) + && reachable_segments.insert(canonical_referenced_path) + { + pending.push(referenced_path); + } + } + } + + let segment_files = rollout_files_under(segment_root.as_path())?; + for segment_file in segment_files { + let canonical_path = fs::canonicalize(segment_file.as_path()).await?; + if !reachable_segments.contains(&canonical_path) { + fs::remove_file(segment_file.as_path()).await?; + } + } + remove_empty_directories(segment_root.as_path()); + Ok(()) +} + +fn rollout_files_under(root: &Path) -> io::Result> { + if !root.try_exists()? { + return Ok(Vec::new()); + } + let mut files = Vec::new(); + let mut pending = vec![root.to_path_buf()]; + while let Some(directory) = pending.pop() { + for entry in std::fs::read_dir(directory)? { + let entry = entry?; + let file_type = entry.file_type()?; + if file_type.is_symlink() { + continue; + } + let path = entry.path(); + if file_type.is_dir() { + pending.push(path); + } else if file_type.is_file() && is_rollout_file(path.as_path()) { + files.push(path); + } + } + } + files.sort(); + Ok(files) +} + +fn is_rollout_file(path: &Path) -> bool { + let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else { + return false; + }; + file_name.starts_with("rollout-") + && (file_name.ends_with(".jsonl") || file_name.ends_with(".jsonl.zst")) +} + +fn remove_empty_directories(root: &Path) { + let mut directories = Vec::new(); + let mut pending = vec![root.to_path_buf()]; + while let Some(directory) = pending.pop() { + let entries = match std::fs::read_dir(directory.as_path()) { + Ok(entries) => entries, + Err(err) => { + warn!( + "failed to inspect rollout segment directory {} during cleanup: {err}", + directory.display() + ); + continue; + } + }; + directories.push(directory); + for entry in entries.flatten() { + if entry.file_type().is_ok_and(|file_type| file_type.is_dir()) { + pending.push(entry.path()); + } + } + } + directories.sort_by_key(|path| std::cmp::Reverse(path.components().count())); + for directory in directories { + if let Err(err) = std::fs::remove_dir(directory.as_path()) + && !matches!( + err.kind(), + io::ErrorKind::NotFound | io::ErrorKind::DirectoryNotEmpty + ) + { + warn!( + "failed to remove empty rollout segment directory {}: {err}", + directory.display() + ); + } + } +} diff --git a/codex-rs/thread-store/src/local/segment_writer.rs b/codex-rs/thread-store/src/local/segment_writer.rs new file mode 100644 index 0000000000..2302ffb2c1 --- /dev/null +++ b/codex-rs/thread-store/src/local/segment_writer.rs @@ -0,0 +1,563 @@ +//! Atomic local rollout-segment rotation. +//! +//! The live recorder mutex serializes appends with rotation. Rotation publishes the predecessor +//! at a new immutable path, writes the complete successor to a temporary file in the canonical +//! file's directory, and atomically replaces the canonical file. The canonical path therefore +//! always names either the complete predecessor or the complete successor. Immutable segment +//! lifetime is managed by `segment_gc` from references in canonical rollouts. + +use std::collections::hash_map::Entry; +use std::io; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use codex_protocol::SegmentId; +use codex_protocol::ThreadId; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutReferenceItem; +use codex_rollout::RolloutConfig; +use codex_rollout::RolloutRecorder; +use codex_rollout::RolloutRecorderParams; +use codex_rollout::read_session_meta_line; +use tempfile::TempPath; +use tokio::fs; +use tracing::warn; + +use super::LiveRecorderHandle; +use super::LiveRecorderState; +use super::LocalThreadStore; +use super::live_writer::live_recorder; +use super::live_writer::remove_live_recorder_handle; +use super::live_writer::sync_materialized_rollout_path; +use super::live_writer::thread_store_io_error; +use crate::RotateThreadSegmentParams; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +#[derive(Clone, Copy)] +enum ReferenceBoundary { + Segment, + Snapshot, +} + +pub(super) async fn rotate_thread_segment( + store: &LocalThreadStore, + thread_id: ThreadId, + params: RotateThreadSegmentParams, +) -> ThreadStoreResult { + let handle = store.live_recorder_handle(thread_id).await?; + Box::pin(rotate_thread_segment_with_handle( + store, + thread_id, + params, + handle, + ReferenceBoundary::Segment, + )) + .await +} + +pub(super) async fn seal_live_thread_segment( + store: &LocalThreadStore, + thread_id: ThreadId, + params: RotateThreadSegmentParams, +) -> ThreadStoreResult { + let handle = store.live_recorder_handle(thread_id).await?; + Box::pin(rotate_thread_segment_with_handle( + store, + thread_id, + params, + handle, + ReferenceBoundary::Snapshot, + )) + .await +} + +pub(super) async fn seal_thread_segment( + store: &LocalThreadStore, + thread_id: ThreadId, + rollout_path: PathBuf, + params: RotateThreadSegmentParams, +) -> ThreadStoreResult { + let existing_handle = store.live_recorders.lock().await.get(&thread_id).cloned(); + if let Some(handle) = existing_handle { + ensure_requested_rollout(&handle, thread_id, rollout_path.as_path()).await?; + let reference = Box::pin(rotate_thread_segment_with_handle( + store, + thread_id, + params, + handle, + ReferenceBoundary::Snapshot, + )) + .await?; + let canonical_path = codex_rollout::plain_rollout_path(rollout_path.as_path()); + sync_materialized_rollout_path(store, thread_id, canonical_path.as_path()).await?; + return Ok(reference); + } + + let source_meta = read_session_meta_line(rollout_path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to read rollout metadata from {} before snapshot: {err}", + rollout_path.display() + ), + })?; + let config = rollout_config(store, &source_meta.meta); + let candidate = + RolloutRecorder::new(&config, RolloutRecorderParams::resume(rollout_path.clone())) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to open rollout {} for snapshot: {err}", + rollout_path.display() + ), + })?; + let candidate_handle = Arc::new(tokio::sync::Mutex::new(LiveRecorderState { + recorder: Some(candidate.clone()), + })); + let (handle, owns_handle) = { + let mut live_recorders = store.live_recorders.lock().await; + match live_recorders.entry(thread_id) { + Entry::Occupied(entry) => (Arc::clone(entry.get()), false), + Entry::Vacant(entry) => { + entry.insert(Arc::clone(&candidate_handle)); + (candidate_handle, true) + } + } + }; + if !owns_handle { + let _ = candidate.shutdown().await; + ensure_requested_rollout(&handle, thread_id, rollout_path.as_path()).await?; + } + + let result = Box::pin(rotate_thread_segment_with_handle( + store, + thread_id, + params, + Arc::clone(&handle), + ReferenceBoundary::Snapshot, + )) + .await; + if owns_handle { + let recorder = handle.lock().await.recorder.take(); + if let Some(recorder) = recorder + && let Err(err) = recorder.shutdown().await + { + warn!("failed to close snapshot-only writer for thread {thread_id}: {err}"); + } + remove_live_recorder_handle(store, thread_id, &handle).await; + } + if result.is_ok() { + let canonical_path = codex_rollout::plain_rollout_path(rollout_path.as_path()); + sync_materialized_rollout_path(store, thread_id, canonical_path.as_path()).await?; + } + result +} + +async fn ensure_requested_rollout( + handle: &LiveRecorderHandle, + thread_id: ThreadId, + requested_path: &Path, +) -> ThreadStoreResult<()> { + let current_path = { + let state = handle.lock().await; + live_recorder(&state.recorder, thread_id)? + .rollout_path() + .to_path_buf() + }; + if codex_rollout::plain_rollout_path(current_path.as_path()) + == codex_rollout::plain_rollout_path(requested_path) + { + return Ok(()); + } + Err(ThreadStoreError::Conflict { + message: format!( + "thread {thread_id} is running from {} instead of requested rollout {}", + current_path.display(), + requested_path.display() + ), + }) +} + +async fn rotate_thread_segment_with_handle( + store: &LocalThreadStore, + thread_id: ThreadId, + params: RotateThreadSegmentParams, + handle: LiveRecorderHandle, + reference_boundary: ReferenceBoundary, +) -> ThreadStoreResult { + let mut state = handle.lock_owned().await; + let _segment_storage_guard = Arc::clone(&store.segment_storage_lock).lock_owned().await; + let old_recorder = live_recorder(&state.recorder, thread_id)?.clone(); + 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 config = rollout_config(store, &old_meta.meta); + let immutable_path = rotated_segment_path( + store.config.codex_home.as_path(), + thread_id, + old_meta.meta.segment_id, + old_rollout_path.as_path(), + )?; + remove_unreferenced_immutable_collision(store, immutable_path.as_path()).await?; + copy_rollout_to_immutable_path(old_rollout_path.as_path(), immutable_path.as_path()).await?; + + let reference = RolloutReferenceItem { + rollout_path: immutable_path.clone(), + thread_id: Some(thread_id), + rollout_timestamp: rollout_timestamp_from_path(old_rollout_path.as_path()), + segment_id: old_meta.meta.segment_id, + max_depth: params.previous_segment_reference_depth, + nth_user_message: match reference_boundary { + ReferenceBoundary::Segment => None, + ReferenceBoundary::Snapshot => Some(usize::MAX), + }, + compacted_replacement_history_filter_texts: None, + }; + let mut initial_items = Vec::with_capacity(params.initial_items.len() + 1); + initial_items.push(RolloutItem::RolloutReference(reference.clone())); + initial_items.extend(params.initial_items); + + let staged_rollout_path = + create_staged_rollout_path(old_rollout_path.as_path()).map_err(thread_store_io_error)?; + let staged_recorder = match RolloutRecorder::new( + &config, + RolloutRecorderParams::CreateAtPath { + path: staged_rollout_path.to_path_buf(), + conversation_id: thread_id, + forked_from_id: old_meta.meta.forked_from_id, + parent_thread_id: old_meta.meta.parent_thread_id, + source: old_meta.meta.source.clone(), + thread_source: old_meta.meta.thread_source, + base_instructions: old_meta.meta.base_instructions.clone().unwrap_or_default(), + dynamic_tools: old_meta.meta.dynamic_tools.clone().unwrap_or_default(), + multi_agent_version: old_meta.meta.multi_agent_version, + session_timestamp: Some(old_meta.meta.timestamp.clone()), + }, + ) + .await + { + Ok(staged_recorder) => staged_recorder, + Err(err) => { + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + return Err(ThreadStoreError::Internal { + message: format!("failed to initialize rotated local thread recorder: {err}"), + }); + } + }; + if let Err(err) = write_staged_rollout(&staged_recorder, initial_items.as_slice()).await { + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + return Err(err); + } + + if let Err(err) = old_recorder.shutdown().await { + restore_recorder_at_path( + &mut state.recorder, + &config, + thread_id, + old_rollout_path.as_path(), + ) + .await; + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + return Err(ThreadStoreError::Internal { + message: format!( + "failed to close previous rollout segment {} for thread {thread_id}: {err}", + old_rollout_path.display() + ), + }); + } + + #[cfg(test)] + { + let hook = store.rotation_test_hook.lock().await.clone(); + if let Some(hook) = hook { + hook.reached.wait().await; + hook.release.wait().await; + if hook.fail_before_install { + restore_recorder_at_path( + &mut state.recorder, + &config, + thread_id, + old_rollout_path.as_path(), + ) + .await; + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + return Err(ThreadStoreError::Internal { + message: "injected rollout rotation failure before atomic install".to_string(), + }); + } + } + } + + if let Err(err) = staged_rollout_path.persist(old_rollout_path.as_path()) { + restore_recorder_at_path( + &mut state.recorder, + &config, + thread_id, + old_rollout_path.as_path(), + ) + .await; + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + return Err(ThreadStoreError::Internal { + message: format!( + "failed to atomically replace live rollout {}: {}", + old_rollout_path.display(), + err.error + ), + }); + } + + match RolloutRecorder::new( + &config, + RolloutRecorderParams::resume(old_rollout_path.clone()), + ) + .await + { + Ok(new_recorder) => { + state.recorder = Some(new_recorder); + Ok(reference) + } + Err(resume_err) => { + let restore_result = restore_previous_rollout_after_install( + &config, + old_rollout_path.as_path(), + immutable_path.as_path(), + ) + .await; + match restore_result { + Ok(restored_recorder) => { + state.recorder = Some(restored_recorder); + remove_immutable_after_failed_rotation(immutable_path.as_path()).await; + Err(ThreadStoreError::Internal { + message: format!( + "failed to resume rotated local thread recorder; restored the previous segment: {resume_err}" + ), + }) + } + Err(restore_err) => { + state.recorder = None; + Err(ThreadStoreError::Internal { + message: format!( + "failed to resume rotated local thread recorder: {resume_err}; failed to restore previous segment: {restore_err}" + ), + }) + } + } + } + } +} + +fn rollout_config(store: &LocalThreadStore, meta: &codex_rollout::SessionMeta) -> RolloutConfig { + RolloutConfig { + codex_home: store.config.codex_home.clone(), + sqlite_home: store.config.sqlite_home.clone(), + cwd: meta.cwd.clone(), + model_provider_id: meta + .model_provider + .clone() + .unwrap_or_else(|| store.config.default_model_provider_id.clone()), + generate_memories: meta.memory_mode.as_deref() != Some("disabled"), + } +} + +async fn write_staged_rollout( + recorder: &RolloutRecorder, + initial_items: &[RolloutItem], +) -> ThreadStoreResult<()> { + if let Err(err) = recorder.record_canonical_items(initial_items).await { + let _ = recorder.shutdown().await; + return Err(thread_store_io_error(err)); + } + if let Err(err) = recorder.flush().await { + let _ = recorder.shutdown().await; + return Err(thread_store_io_error(err)); + } + recorder.shutdown().await.map_err(thread_store_io_error) +} + +async fn restore_recorder_at_path( + recorder: &mut Option, + config: &RolloutConfig, + thread_id: ThreadId, + rollout_path: &Path, +) { + match RolloutRecorder::new( + config, + RolloutRecorderParams::resume(rollout_path.to_path_buf()), + ) + .await + { + Ok(restored) => *recorder = Some(restored), + Err(err) => { + *recorder = None; + warn!( + "failed to restore live rollout recorder {} for thread {thread_id}: {err}", + rollout_path.display() + ); + } + } +} + +async fn restore_previous_rollout_after_install( + config: &RolloutConfig, + live_rollout_path: &Path, + immutable_path: &Path, +) -> io::Result { + let rollback_path = create_staged_rollout_path(live_rollout_path)?; + copy_file_contents(immutable_path, rollback_path.as_ref()).await?; + rollback_path.persist(live_rollout_path)?; + RolloutRecorder::new( + config, + RolloutRecorderParams::resume(live_rollout_path.to_path_buf()), + ) + .await +} + +async fn copy_rollout_to_immutable_path( + source: &Path, + destination: &Path, +) -> ThreadStoreResult<()> { + let parent = destination + .parent() + .ok_or_else(|| ThreadStoreError::Internal { + message: format!( + "rotated rollout segment path {} does not have a parent", + destination.display() + ), + })?; + fs::create_dir_all(parent) + .await + .map_err(thread_store_io_error)?; + let staged_path = tempfile::NamedTempFile::new_in(parent) + .map_err(thread_store_io_error)? + .into_temp_path(); + copy_file_contents(source, staged_path.as_ref()) + .await + .map_err(thread_store_io_error)?; + staged_path + .persist_noclobber(destination) + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to publish immutable rollout segment {}: {}", + destination.display(), + err.error + ), + }) +} + +async fn remove_unreferenced_immutable_collision( + store: &LocalThreadStore, + immutable_path: &Path, +) -> ThreadStoreResult<()> { + if !fs::try_exists(immutable_path) + .await + .map_err(thread_store_io_error)? + { + return Ok(()); + } + super::segment_gc::collect_unreferenced_segments(store.config.codex_home.as_path()) + .await + .map_err(thread_store_io_error)?; + if fs::try_exists(immutable_path) + .await + .map_err(thread_store_io_error)? + { + return Err(ThreadStoreError::Conflict { + message: format!( + "immutable rollout segment {} is already referenced", + immutable_path.display() + ), + }); + } + Ok(()) +} + +async fn copy_file_contents(source: &Path, destination: &Path) -> io::Result<()> { + fs::copy(source, destination).await?; + fs::OpenOptions::new() + .write(true) + .open(destination) + .await? + .sync_all() + .await +} + +fn create_staged_rollout_path(live_rollout_path: &Path) -> io::Result { + let parent = live_rollout_path.parent().ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "live rollout path {} does not have a parent", + live_rollout_path.display() + ), + ) + })?; + Ok(tempfile::NamedTempFile::new_in(parent)?.into_temp_path()) +} + +async fn remove_immutable_after_failed_rotation(path: &Path) { + if let Err(err) = fs::remove_file(path).await + && err.kind() != io::ErrorKind::NotFound + { + warn!( + "failed to remove immutable rollout after failed rotation {}: {err}", + path.display() + ); + } +} + +fn rotated_segment_path( + codex_home: &Path, + thread_id: ThreadId, + segment_id: Option, + old_rollout_path: &Path, +) -> ThreadStoreResult { + let old_file_name = old_rollout_path + .file_name() + .ok_or_else(|| ThreadStoreError::Internal { + message: format!( + "previous rollout segment path {} does not have a file name", + old_rollout_path.display() + ), + })?; + let segment_key = segment_id + .map(|segment_id| segment_id.to_string()) + .unwrap_or_else(|| "initial".to_string()); + Ok(codex_home + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(thread_id.to_string()) + .join(segment_key) + .join(old_file_name)) +} + +fn rollout_timestamp_from_path(path: &Path) -> Option { + let file_name = path.file_name()?.to_str()?; + let core = file_name.strip_prefix("rollout-")?.strip_suffix(".jsonl")?; + core.match_indices('-').rev().find_map(|(index, _)| { + ThreadId::from_string(&core[index + 1..]) + .ok() + .map(|_| core[..index].to_string()) + }) +} + +#[cfg(test)] +#[path = "segment_writer_tests.rs"] +mod tests; diff --git a/codex-rs/thread-store/src/local/segment_writer_tests.rs b/codex-rs/thread-store/src/local/segment_writer_tests.rs new file mode 100644 index 0000000000..3bce82bd53 --- /dev/null +++ b/codex-rs/thread-store/src/local/segment_writer_tests.rs @@ -0,0 +1,459 @@ +use std::sync::Arc; + +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::UserMessageEvent; +use codex_rollout::RolloutRecorder; +use tempfile::TempDir; + +use super::*; +use crate::CreateThreadParams; +use crate::DeleteThreadParams; +use crate::LiveThread; +use crate::ResumeThreadParams; +use crate::ThreadPersistenceMetadata; +use crate::ThreadStore; +use crate::local::LocalThreadStore; +use crate::local::RotationTestHook; +use crate::local::test_support::test_config; +use crate::local::test_support::write_session_file; + +#[tokio::test] +async fn rotation_keeps_the_canonical_thread_active() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone()))); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("before rotation")]) + .await + .expect("append before rotation"); + let canonical_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + + let reference = store + .rotate_thread_segment(thread_id, rotation_params()) + .await + .expect("rotate thread segment"); + + assert!(canonical_path.exists()); + assert!(reference.rollout_path.exists()); + assert!( + reference.rollout_path.starts_with( + home.path() + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(thread_id.to_string()) + ) + ); + assert!( + !home + .path() + .join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR) + .join(canonical_path.file_name().expect("live rollout file name")) + .exists() + ); + let metadata = runtime + .get_thread(thread_id) + .await + .expect("sqlite metadata read") + .expect("sqlite metadata"); + assert_eq!(metadata.rollout_path, canonical_path); + assert_eq!(metadata.archived_at, None); +} + +#[tokio::test] +async fn deletion_collects_segments_after_the_last_reference() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let parent_thread_id = ThreadId::default(); + let parent = create_live_thread(store.clone(), parent_thread_id).await; + parent + .append_items(&[user_message_item("parent history")]) + .await + .expect("append parent history"); + let reference = store + .rotate_thread_segment(parent_thread_id, rotation_params()) + .await + .expect("rotate parent segment"); + let immutable_path = reference.rollout_path.clone(); + + let child_thread_id = ThreadId::default(); + let child = create_live_thread(store.clone(), child_thread_id).await; + child + .append_items(&[RolloutItem::RolloutReference(reference)]) + .await + .expect("persist child reference"); + + store + .delete_thread(DeleteThreadParams { + thread_id: parent_thread_id, + }) + .await + .expect("delete parent thread"); + assert!(immutable_path.exists()); + + store + .delete_thread(DeleteThreadParams { + thread_id: child_thread_id, + }) + .await + .expect("delete child thread"); + assert!(!immutable_path.exists()); +} + +#[tokio::test] +async fn failed_rotation_before_staging_keeps_the_writer_open() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("before failed rotation")]) + .await + .expect("append before failed rotation"); + tokio::fs::write( + home.path() + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR), + "block directory creation", + ) + .await + .expect("create rotation blocker"); + + store + .rotate_thread_segment(thread_id, rotation_params()) + .await + .expect_err("rotation should fail before replacing the live recorder"); + live_thread + .append_items(&[user_message_item("after failed rotation")]) + .await + .expect("live recorder should remain writable"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + assert_rollout_contains_message(rollout_path.as_path(), "after failed rotation").await; +} + +#[tokio::test] +async fn legacy_segment_without_id_is_sealed_at_an_immutable_path() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let uuid = uuid::Uuid::from_u128(401); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let canonical_path = + write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("legacy rollout"); + let live_thread = LiveThread::resume( + store.clone(), + ResumeThreadParams { + thread_id, + rollout_path: Some(canonical_path.clone()), + history: None, + include_archived: true, + metadata: thread_metadata(), + }, + ) + .await + .expect("resume legacy rollout"); + + let reference = store + .rotate_thread_segment(thread_id, rotation_params()) + .await + .expect("seal legacy segment"); + assert_eq!(reference.segment_id, None); + assert_ne!(reference.rollout_path, canonical_path); + assert!( + reference.rollout_path.starts_with( + home.path() + .join(codex_rollout::ROTATED_ROLLOUT_SEGMENTS_SUBDIR) + .join(thread_id.to_string()) + .join("initial") + ) + ); + + live_thread + .append_items(&[user_message_item("new live segment")]) + .await + .expect("append to successor segment"); + let immutable_json = rollout_json(reference.rollout_path.as_path()).await; + assert!(!immutable_json.contains("new live segment")); +} + +#[tokio::test] +async fn append_waits_for_atomic_rotation() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("before rotation")]) + .await + .expect("append before rotation"); + let canonical_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + + let reached = Arc::new(tokio::sync::Barrier::new(2)); + let release = Arc::new(tokio::sync::Barrier::new(2)); + *store.rotation_test_hook.lock().await = Some(RotationTestHook { + reached: Arc::clone(&reached), + release: Arc::clone(&release), + fail_before_install: false, + }); + let rotation_store = Arc::clone(&store); + let rotation = tokio::spawn(async move { + rotation_store + .rotate_thread_segment(thread_id, rotation_params()) + .await + }); + reached.wait().await; + assert!(canonical_path.exists()); + assert_rollout_contains_message(canonical_path.as_path(), "before rotation").await; + + let append_thread = live_thread.clone(); + let mut append = tokio::spawn(async move { + append_thread + .append_items(&[user_message_item("after rotation")]) + .await + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), &mut append) + .await + .is_err() + ); + release.wait().await; + let reference = rotation + .await + .expect("rotation task") + .expect("rotate segment"); + append + .await + .expect("append task") + .expect("append after rotation"); + + assert_rollout_contains_message(canonical_path.as_path(), "after rotation").await; + let immutable_json = rollout_json(reference.rollout_path.as_path()).await; + assert!(immutable_json.contains("before rotation")); + assert!(!immutable_json.contains("after rotation")); +} + +#[tokio::test] +async fn failed_atomic_install_restores_the_previous_writer() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("before failed install")]) + .await + .expect("append before failed install"); + + let reached = Arc::new(tokio::sync::Barrier::new(2)); + let release = Arc::new(tokio::sync::Barrier::new(2)); + *store.rotation_test_hook.lock().await = Some(RotationTestHook { + reached: Arc::clone(&reached), + release: Arc::clone(&release), + fail_before_install: true, + }); + let rotation_store = Arc::clone(&store); + let rotation = tokio::spawn(async move { + rotation_store + .rotate_thread_segment(thread_id, rotation_params()) + .await + }); + reached.wait().await; + release.wait().await; + rotation + .await + .expect("rotation task") + .expect_err("injected install failure"); + + live_thread + .append_items(&[user_message_item("after failed install")]) + .await + .expect("restored writer should accept appends"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + assert_rollout_contains_message(rollout_path.as_path(), "before failed install").await; + assert_rollout_contains_message(rollout_path.as_path(), "after failed install").await; +} + +#[tokio::test] +async fn rotation_recovers_unreferenced_immutable_left_by_crash() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("before interrupted rotation")]) + .await + .expect("append before interrupted rotation"); + let canonical_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + let meta = codex_rollout::read_session_meta_line(canonical_path.as_path()) + .await + .expect("read live rollout metadata"); + let immutable_path = rotated_segment_path( + home.path(), + thread_id, + meta.meta.segment_id, + canonical_path.as_path(), + ) + .expect("immutable path"); + tokio::fs::create_dir_all(immutable_path.parent().expect("immutable parent")) + .await + .expect("create immutable parent"); + tokio::fs::write(immutable_path.as_path(), "orphan from interrupted rotation") + .await + .expect("write orphaned immutable segment"); + + let reference = store + .rotate_thread_segment(thread_id, rotation_params()) + .await + .expect("retry rotation after interrupted publication"); + + assert_eq!(reference.rollout_path, immutable_path); + assert_rollout_contains_message( + reference.rollout_path.as_path(), + "before interrupted rotation", + ) + .await; +} + +#[tokio::test] +async fn snapshot_seals_mark_non_consuming_reference_boundaries() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let live_thread = create_live_thread(store.clone(), thread_id).await; + live_thread + .append_items(&[user_message_item("history retained across snapshots")]) + .await + .expect("append before snapshots"); + let canonical_path = store + .live_rollout_path(thread_id) + .await + .expect("live rollout path"); + + for _ in 0..4 { + let reference = store + .seal_thread_segment(thread_id, canonical_path.clone(), rotation_params()) + .await + .expect("seal snapshot segment"); + assert_eq!(reference.nth_user_message, Some(usize::MAX)); + let (successor_items, _, _) = RolloutRecorder::load_rollout_items(&canonical_path) + .await + .expect("load snapshot successor"); + assert!(successor_items.iter().any(|item| { + matches!( + item, + RolloutItem::RolloutReference(successor_reference) + if successor_reference.nth_user_message == Some(usize::MAX) + ) + })); + } +} + +async fn create_live_thread(store: Arc, thread_id: ThreadId) -> LiveThread { + LiveThread::create(store, create_thread_params(thread_id)) + .await + .expect("create live thread") +} + +fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams { + CreateThreadParams { + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + multi_agent_version: None, + metadata: thread_metadata(), + } +} + +fn thread_metadata() -> ThreadPersistenceMetadata { + ThreadPersistenceMetadata { + cwd: Some(std::env::current_dir().expect("cwd")), + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + } +} + +fn rotation_params() -> RotateThreadSegmentParams { + RotateThreadSegmentParams { + initial_items: Vec::new(), + previous_segment_reference_depth: 1, + } +} + +fn user_message_item(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + client_id: None, + message: message.to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + ..Default::default() + })) +} + +async fn assert_rollout_contains_message(path: &std::path::Path, expected: &str) { + let (items, _, _) = RolloutRecorder::load_rollout_items(path) + .await + .expect("load rollout items"); + assert!(items.iter().any(|item| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == expected + ) + })); +} + +async fn rollout_json(path: &std::path::Path) -> String { + let (items, _, _) = RolloutRecorder::load_rollout_items(path) + .await + .expect("load rollout items"); + serde_json::to_string(&items).expect("serialize rollout items") +} diff --git a/codex-rs/thread-store/src/local/unarchive_thread.rs b/codex-rs/thread-store/src/local/unarchive_thread.rs index ad41db69ac..e575bb017a 100644 --- a/codex-rs/thread-store/src/local/unarchive_thread.rs +++ b/codex-rs/thread-store/src/local/unarchive_thread.rs @@ -17,6 +17,7 @@ pub(super) async fn unarchive_thread( params: ArchiveThreadParams, ) -> ThreadStoreResult { let thread_id = params.thread_id; + let _segment_storage_guard = store.segment_storage_lock.clone().lock_owned().await; let state_db_ctx = store.state_db().await; let archived_path = find_archived_thread_path_by_id_str( store.config.codex_home.as_path(), diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 741ca12bc2..576efc632c 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -112,6 +112,15 @@ 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 { + /// 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 {