openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add feature-gated session segmentation writers

Browse Source
This commit is contained in:
Friel
2026-06-13 00:51:11 +00:00
parent 4f4e4b8e76
commit 848575dc0d
27 changed files with 2320 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/config.schema.json
+7 -1
View File
@@ -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"
}
}
codex-rs/core/src/agent/control/spawn.rs
+52 -17
View File
@@ -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!(
codex-rs/core/src/agent/control_tests.rs
+191 -2
View File
@@ -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<Item = RolloutItem>) {
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;
codex-rs/core/src/codex_thread.rs
+10
View File
@@ -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<Option<RolloutReferenceItem>> {
self.codex
.session
.seal_rollout_segment_for_reference()
.await
}
pub async fn submit_with_trace(
&self,
op: Op,
codex-rs/core/src/session/mod.rs
+78 -12
View File
@@ -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<Vec<RolloutItem>>,
) {
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<TurnContextItem>,
) {
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<Option<RolloutReferenceItem>> {
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<RolloutItem>,
) -> ThreadStoreResult<Option<RolloutReferenceItem>> {
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<RolloutItem> = 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;
codex-rs/core/src/session/session.rs
+16 -1
View File
@@ -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);
codex-rs/core/src/session/tests.rs
+194
View File
@@ -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",
codex-rs/core/src/thread_manager.rs
+117 -13
View File
@@ -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<InitialHistory> {
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::<LocalThreadStore>()?;
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))
}
codex-rs/core/src/thread_manager_tests.rs
+42 -4
View File
@@ -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();
codex-rs/core/tests/suite/fork_thread.rs
+52 -9
View File
@@ -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<RolloutItem> = 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<RolloutItem> {
}
items
}
fn without_session_meta(items: Vec<RolloutItem>) -> Vec<RolloutItem> {
items
.into_iter()
.filter(|item| !matches!(item, RolloutItem::SessionMeta(_)))
.collect()
}
codex-rs/core/tests/suite/resume.rs
+15 -2
View File
@@ -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::<RolloutLine>(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(
codex-rs/features/src/lib.rs
+12
View File
@@ -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",
codex-rs/features/src/tests.rs
+17
View File
@@ -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));
codex-rs/rollout/src/recorder.rs
+122 -35
View File
@@ -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<DynamicToolSpec>,
multi_agent_version: Option<MultiAgentVersion>,
},
CreateAtPath {
path: PathBuf,
conversation_id: ThreadId,
forked_from_id: Option<ThreadId>,
parent_thread_id: Option<ThreadId>,
source: SessionSource,
thread_source: Option<ThreadSource>,
base_instructions: BaseInstructions,
dynamic_tools: Vec<DynamicToolSpec>,
multi_agent_version: Option<MultiAgentVersion>,
session_timestamp: Option<String>,
},
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<ThreadId>,
parent_thread_id: Option<ThreadId>,
source: SessionSource,
thread_source: Option<ThreadSource>,
base_instructions: BaseInstructions,
dynamic_tools: Vec<DynamicToolSpec>,
multi_agent_version: Option<MultiAgentVersion>,
timestamp_override: Option<String>,
}
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<File> {
.open(path)
}
fn create_session_meta(
config: &impl RolloutConfigView,
log_file_info: &LogFileInfo,
params: SessionMetaParams,
) -> std::io::Result<SessionMeta> {
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
codex-rs/rollout/src/recorder_tests.rs
+14
View File
@@ -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");
codex-rs/thread-store/Cargo.toml
+1
View File
@@ -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 }
codex-rs/thread-store/src/lib.rs
+1
View File
@@ -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;
codex-rs/thread-store/src/live_thread.rs
+38
View File
@@ -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<Option<RolloutReferenceItem>> {
let Some(local_store) = self
.thread_store
.as_any()
.downcast_ref::<LocalThreadStore>()
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<Option<RolloutReferenceItem>> {
let Some(local_store) = self
.thread_store
.as_any()
.downcast_ref::<LocalThreadStore>()
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
codex-rs/thread-store/src/local/archive_thread.rs
+1
View File
@@ -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(),
codex-rs/thread-store/src/local/delete_thread.rs
+19 -2
View File
@@ -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(())
}
codex-rs/thread-store/src/local/live_writer.rs
+93 -42
View File
@@ -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<PathBuf> {
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<RolloutRecorder>,
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(),
}
}
codex-rs/thread-store/src/local/mod.rs
+60 -4
View File
@@ -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<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
live_recorders: Arc<Mutex<HashMap<ThreadId, LiveRecorderHandle>>>,
segment_storage_lock: Arc<Mutex<()>>,
#[cfg(test)]
rotation_test_hook: Arc<Mutex<Option<RotationTestHook>>>,
state_db: Option<StateDbHandle>,
}
pub(super) type LiveRecorderHandle = Arc<Mutex<LiveRecorderState>>;
pub(super) struct LiveRecorderState {
pub(super) recorder: Option<RolloutRecorder>,
}
#[cfg(test)]
#[derive(Clone)]
pub(super) struct RotationTestHook {
pub(super) reached: Arc<tokio::sync::Barrier>,
pub(super) release: Arc<tokio::sync::Barrier>,
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<RolloutRecorder> {
params: RotateThreadSegmentParams,
) -> ThreadStoreResult<codex_protocol::protocol::RolloutReferenceItem> {
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<codex_protocol::protocol::RolloutReferenceItem> {
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<codex_protocol::protocol::RolloutReferenceItem> {
segment_writer::seal_thread_segment(self, thread_id, rollout_path, params).await
}
pub(super) async fn live_recorder_handle(
&self,
thread_id: ThreadId,
) -> ThreadStoreResult<LiveRecorderHandle> {
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(())
}
}
codex-rs/thread-store/src/local/segment_gc.rs
+136
View File
@@ -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<Vec<PathBuf>> {
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()
);
}
}
}
codex-rs/thread-store/src/local/segment_writer.rs
+563
View File
@@ -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<RolloutReferenceItem> {
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<RolloutReferenceItem> {
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<RolloutReferenceItem> {
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<RolloutReferenceItem> {
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<RolloutRecorder>,
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<RolloutRecorder> {
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<TempPath> {
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<SegmentId>,
old_rollout_path: &Path,
) -> ThreadStoreResult<PathBuf> {
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<String> {
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;
codex-rs/thread-store/src/local/segment_writer_tests.rs
+459
View File
@@ -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<LocalThreadStore>, 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")
}
codex-rs/thread-store/src/local/unarchive_thread.rs
+1
View File
@@ -17,6 +17,7 @@ pub(super) async fn unarchive_thread(
params: ArchiveThreadParams,
) -> ThreadStoreResult<StoredThread> {
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(),
codex-rs/thread-store/src/types.rs
+9
View File
@@ -112,6 +112,15 @@ pub struct AppendThreadItemsParams {
pub items: Vec<RolloutItem>,
}
/// 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<RolloutItem>,
/// 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 {