Fix compacted rollout segment scrollback

This commit is contained in:
Friel
2026-05-04 23:00:02 +00:00
parent de1c034f67
commit b053597281
10 changed files with 372 additions and 12 deletions
@@ -4945,7 +4945,11 @@ impl CodexMessageProcessor {
thread.id = thread_id.to_string();
thread.path = Some(rollout_path.to_path_buf());
if include_turns {
let history_items = thread_history.get_rollout_items();
let history_items = codex_core::materialize_rollout_items_for_replay(
self.config.codex_home.as_path(),
&thread_history.get_rollout_items(),
)
.await;
populate_thread_turns_from_history(
&mut thread,
&history_items,
@@ -8409,7 +8413,7 @@ async fn handle_thread_listener_command(
async fn handle_pending_thread_resume_request(
conversation_id: ThreadId,
conversation: &Arc<CodexThread>,
_codex_home: &Path,
codex_home: &Path,
thread_state_manager: &ThreadStateManager,
thread_state: &Arc<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
@@ -8438,12 +8442,14 @@ async fn handle_pending_thread_resume_request(
let request_id = pending.request_id;
let connection_id = request_id.connection_id;
let mut thread = pending.thread_summary;
let history_items = if pending.include_turns {
codex_core::materialize_rollout_items_for_replay(codex_home, &pending.history_items).await
} else {
Vec::new()
};
if pending.include_turns
&& let Err(message) = populate_thread_turns_from_history(
&mut thread,
&pending.history_items,
active_turn.as_ref(),
)
&& let Err(message) =
populate_thread_turns_from_history(&mut thread, &history_items, active_turn.as_ref())
{
outgoing
.send_error(request_id, internal_error(message))
@@ -8531,7 +8537,7 @@ async fn handle_pending_thread_resume_request(
// paying the cost of turn reconstruction for historical usage replay.
if let Some(token_usage_thread) = token_usage_thread {
let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items(
&pending.history_items,
&history_items,
token_usage_thread.turns.as_slice(),
);
// Rejoining a loaded thread has the same UI contract as a cold resume, but
@@ -469,6 +469,7 @@ async fn thread_read_unloaded_include_turns_materializes_rollout_reference() ->
thread_id: Some(codex_protocol::ThreadId::from_string(
&referenced_thread_id,
)?),
rollout_timestamp: None,
segment_id: None,
max_depth: 2,
}),
@@ -52,6 +52,7 @@ use codex_app_server_protocol::UserInput;
use codex_config::types::AuthCredentialsStoreMode;
use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
use codex_protocol::AgentPath;
use codex_protocol::SegmentId;
use codex_protocol::ThreadId;
use codex_protocol::config_types::Personality;
use codex_protocol::models::ContentItem;
@@ -61,10 +62,12 @@ use codex_protocol::protocol::AgentMessageEvent;
use codex_protocol::protocol::AgentStatus;
use codex_protocol::protocol::CollabAgentSpawnEndEvent;
use codex_protocol::protocol::CollabCloseEndEvent;
use codex_protocol::protocol::ContextCompactedEvent;
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::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource as RolloutSessionSource;
@@ -74,6 +77,7 @@ use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::protocol::UserMessageEvent;
use codex_protocol::user_input::ByteRange;
use codex_protocol::user_input::TextElement;
use codex_state::StateRuntime;
@@ -112,6 +116,49 @@ fn normalized_existing_path(path: impl AsRef<Path>) -> Result<PathBuf> {
Ok(AbsolutePathBuf::from_absolute_path(path.as_ref().canonicalize()?)?.into_path_buf())
}
fn session_meta_rollout_item(
thread_id: ThreadId,
segment_id: Option<SegmentId>,
timestamp: &str,
) -> RolloutItem {
RolloutItem::SessionMeta(SessionMetaLine {
meta: SessionMeta {
id: thread_id,
segment_id,
timestamp: timestamp.to_string(),
cwd: PathBuf::from("/"),
originator: "codex".to_string(),
cli_version: "0.0.0".to_string(),
source: RolloutSessionSource::Cli,
model_provider: Some("mock_provider".to_string()),
..SessionMeta::default()
},
git: None,
})
}
fn user_message_item(text: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: text.to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
}))
}
fn write_rollout_items(path: &Path, timestamp: &str, items: &[RolloutItem]) -> Result<()> {
let mut jsonl = String::new();
for item in items {
jsonl.push_str(&serde_json::to_string(&RolloutLine {
timestamp: timestamp.to_string(),
item: item.clone(),
})?);
jsonl.push('\n');
}
std::fs::write(path, jsonl)?;
Ok(())
}
async fn wait_for_responses_request_count(
server: &wiremock::MockServer,
expected_count: usize,
@@ -357,6 +404,108 @@ async fn thread_resume_returns_rollout_history() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn thread_resume_materializes_rollout_references_for_scrollback() -> Result<()> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri())?;
let thread_id = ThreadId::new();
let thread_id_str = thread_id.to_string();
let old_ts = "2025-01-05T12-00-00";
let current_ts = "2025-01-05T12-05-00";
let old_rfc3339 = "2025-01-05T12:00:00Z";
let current_rfc3339 = "2025-01-05T12:05:00Z";
let active_old_path = rollout_path(codex_home.path(), old_ts, &thread_id_str);
let archived_old_path = codex_home
.path()
.join("archived_sessions")
.join(active_old_path.file_name().expect("old rollout file name"));
std::fs::create_dir_all(archived_old_path.parent().expect("archived rollout parent"))?;
let current_path = rollout_path(codex_home.path(), current_ts, &thread_id_str);
std::fs::create_dir_all(current_path.parent().expect("current rollout parent"))?;
let old_items = vec![
session_meta_rollout_item(thread_id, Some(SegmentId::new()), old_rfc3339),
user_message_item("before compaction"),
];
write_rollout_items(archived_old_path.as_path(), old_rfc3339, &old_items)?;
let current_items = vec![
session_meta_rollout_item(thread_id, Some(SegmentId::new()), current_rfc3339),
RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: active_old_path,
thread_id: Some(thread_id),
rollout_timestamp: Some(old_ts.to_string()),
segment_id: None,
max_depth: 2,
}),
RolloutItem::EventMsg(EventMsg::ContextCompacted(ContextCompactedEvent {})),
user_message_item("after compaction"),
];
write_rollout_items(current_path.as_path(), current_rfc3339, &current_items)?;
let mut mcp = McpProcess::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let resume_id = mcp
.send_thread_resume_request(ThreadResumeParams {
thread_id: thread_id_str.clone(),
..Default::default()
})
.await?;
let resume_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(resume_id)),
)
.await??;
let ThreadResumeResponse { thread, .. } = to_response::<ThreadResumeResponse>(resume_resp)?;
let user_messages: Vec<&str> = thread
.turns
.iter()
.flat_map(|turn| turn.items.iter())
.filter_map(|item| match item {
ThreadItem::UserMessage { content, .. } => {
content.iter().find_map(|input| match input {
UserInput::Text { text, .. } => Some(text.as_str()),
UserInput::Image { .. }
| UserInput::LocalImage { .. }
| UserInput::Skill { .. }
| UserInput::Mention { .. } => None,
})
}
ThreadItem::AgentMessage { .. }
| ThreadItem::Reasoning { .. }
| ThreadItem::CommandExecution { .. }
| ThreadItem::FileChange { .. }
| ThreadItem::McpToolCall { .. }
| ThreadItem::WebSearch { .. }
| ThreadItem::ImageView { .. }
| ThreadItem::ImageGeneration { .. }
| ThreadItem::EnteredReviewMode { .. }
| ThreadItem::ExitedReviewMode { .. }
| ThreadItem::ContextCompaction { .. }
| ThreadItem::HookPrompt { .. }
| ThreadItem::CollabAgentToolCall { .. }
| ThreadItem::DynamicToolCall { .. }
| ThreadItem::Plan { .. }
| ThreadItem::RawResponseItem { .. } => None,
})
.collect();
assert_eq!(user_messages, vec!["before compaction", "after compaction"]);
assert!(
thread.turns.iter().any(|turn| {
turn.items
.iter()
.any(|item| matches!(item, ThreadItem::ContextCompaction { .. }))
}),
"resume scrollback should include the compaction marker from the current segment"
);
Ok(())
}
#[tokio::test]
async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
+18 -1
View File
@@ -2910,6 +2910,23 @@ async fn replace_compacted_history_rotates_local_rollout_segment() {
.expect("load current rollout path")
.expect("rollout path after compaction");
assert_ne!(new_rollout_path, old_rollout_path);
let config = sess.get_config().await;
let archived_old_rollout_path = config
.codex_home
.join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR)
.join(
old_rollout_path
.file_name()
.expect("old rollout path should have a file name"),
);
assert!(!old_rollout_path.exists());
assert!(archived_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.conversation_id)))
.expect("old rollout timestamp");
let (new_items, new_thread_id, _) =
RolloutRecorder::load_rollout_items(new_rollout_path.as_path())
.await
@@ -2921,6 +2938,7 @@ async fn replace_compacted_history_rotates_local_rollout_segment() {
RolloutItem::RolloutReference(reference)
if reference.rollout_path == old_rollout_path
&& reference.thread_id == Some(sess.conversation_id)
&& reference.rollout_timestamp.as_deref() == Some(old_rollout_timestamp)
)
}));
assert!(new_items.iter().any(|item| {
@@ -2933,7 +2951,6 @@ async fn replace_compacted_history_rotates_local_rollout_segment() {
)
}));
let config = sess.get_config().await;
let replay_items = crate::thread_rollout_truncation::materialize_rollout_items_for_replay(
config.codex_home.as_path(),
&new_items,
@@ -271,6 +271,7 @@ async fn materializes_fork_reference_before_truncating_rollout_references() {
RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: old_path,
thread_id: None,
rollout_timestamp: None,
segment_id: None,
max_depth: 2,
}),
@@ -327,6 +328,7 @@ async fn materializes_rollout_reference_with_bounded_depth() {
RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: oldest_path,
thread_id: None,
rollout_timestamp: None,
segment_id: None,
max_depth: 2,
}),
@@ -341,6 +343,7 @@ async fn materializes_rollout_reference_with_bounded_depth() {
RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: old_path,
thread_id: None,
rollout_timestamp: None,
segment_id: None,
max_depth: 2,
}),
@@ -353,6 +356,7 @@ async fn materializes_rollout_reference_with_bounded_depth() {
RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: middle_path,
thread_id: None,
rollout_timestamp: None,
segment_id: None,
max_depth: 2,
}),
+2
View File
@@ -2792,6 +2792,8 @@ pub struct RolloutReferenceItem {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thread_id: Option<ThreadId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rollout_timestamp: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub segment_id: Option<SegmentId>,
#[serde(default = "default_rollout_reference_depth")]
pub max_depth: usize,
+76 -1
View File
@@ -940,6 +940,19 @@ pub(crate) fn parse_timestamp_uuid_from_filename(name: &str) -> Option<(OffsetDa
Some((ts, uuid))
}
fn parse_timestamp_string_uuid_from_filename(name: &str) -> Option<(&str, Uuid)> {
// Expected: rollout-YYYY-MM-DDThh-mm-ss-<uuid>.jsonl
let core = name.strip_prefix("rollout-")?.strip_suffix(".jsonl")?;
core.match_indices('-')
.rev()
.find_map(|(index, _)| {
Uuid::parse_str(&core[index + 1..])
.ok()
.map(|uuid| (index, uuid))
})
.map(|(index, uuid)| (&core[..index], uuid))
}
struct ThreadCandidate {
path: PathBuf,
id: Uuid,
@@ -1445,6 +1458,13 @@ pub async fn resolve_fork_reference_rollout_path(
let Some((_, uuid)) = parse_timestamp_uuid_from_filename(file_name) else {
return Ok(rollout_path.to_path_buf());
};
let archived_path = codex_home.join(ARCHIVED_SESSIONS_SUBDIR).join(file_name);
if tokio::fs::try_exists(archived_path.as_path())
.await
.unwrap_or(false)
{
return Ok(archived_path);
}
let id = uuid.to_string();
if let Some(path) = find_thread_path_by_id_str(codex_home, id.as_str()).await? {
return Ok(path);
@@ -1459,6 +1479,27 @@ pub async fn resolve_rollout_reference_rollout_path(
codex_home: &Path,
reference: &codex_protocol::protocol::RolloutReferenceItem,
) -> io::Result<PathBuf> {
if let (Some(thread_id), Some(rollout_timestamp)) =
(reference.thread_id, reference.rollout_timestamp.as_deref())
{
let file_name = format!("rollout-{rollout_timestamp}-{thread_id}.jsonl");
if let Some(active_path) =
rollout_path_for_timestamp_file(codex_home, rollout_timestamp, &file_name)
&& tokio::fs::try_exists(active_path.as_path())
.await
.unwrap_or(false)
{
return Ok(active_path);
}
let archived_path = codex_home.join(ARCHIVED_SESSIONS_SUBDIR).join(&file_name);
if tokio::fs::try_exists(archived_path.as_path())
.await
.unwrap_or(false)
{
return Ok(archived_path);
}
}
if let (Some(thread_id), Some(segment_id)) = (reference.thread_id, reference.segment_id)
&& let Some(path) =
find_rollout_path_by_segment_id(codex_home, thread_id, segment_id).await?
@@ -1477,9 +1518,25 @@ pub async fn resolve_rollout_reference_rollout_path(
else {
return Ok(rollout_path.to_path_buf());
};
let Some((_, uuid)) = parse_timestamp_uuid_from_filename(file_name) else {
let Some((rollout_timestamp, uuid)) = parse_timestamp_string_uuid_from_filename(file_name)
else {
return Ok(rollout_path.to_path_buf());
};
let archived_path = codex_home.join(ARCHIVED_SESSIONS_SUBDIR).join(file_name);
if tokio::fs::try_exists(archived_path.as_path())
.await
.unwrap_or(false)
{
return Ok(archived_path);
}
if let Some(active_path) =
rollout_path_for_timestamp_file(codex_home, rollout_timestamp, file_name)
&& tokio::fs::try_exists(active_path.as_path())
.await
.unwrap_or(false)
{
return Ok(active_path);
}
let id = uuid.to_string();
if let Some(path) = find_thread_path_by_id_str(codex_home, id.as_str()).await? {
return Ok(path);
@@ -1490,6 +1547,24 @@ pub async fn resolve_rollout_reference_rollout_path(
Ok(rollout_path.to_path_buf())
}
fn rollout_path_for_timestamp_file(
codex_home: &Path,
rollout_timestamp: &str,
file_name: &str,
) -> Option<PathBuf> {
let year = rollout_timestamp.get(0..4)?;
let month = rollout_timestamp.get(5..7)?;
let day = rollout_timestamp.get(8..10)?;
Some(
codex_home
.join(SESSIONS_SUBDIR)
.join(year)
.join(month)
.join(day)
.join(file_name),
)
}
/// Extract the `YYYY/MM/DD` directory components from a rollout filename.
pub fn rollout_date_parts(file_name: &OsStr) -> Option<(String, String, String)> {
let name = file_name.to_string_lossy();
+29
View File
@@ -34,6 +34,7 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::RolloutLine;
use codex_protocol::protocol::RolloutReferenceItem;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
@@ -294,6 +295,34 @@ fn rollout_date_parts_extracts_directory_components() {
);
}
#[tokio::test]
async fn rollout_reference_resolves_archived_file_by_stable_thread_and_timestamp() {
let temp = TempDir::new().expect("tempdir");
let home = temp.path();
let uuid = Uuid::new_v4();
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id");
let ts = "2025-01-03T13-00-00";
let active_path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl"));
let archived_path = home.join(format!("archived_sessions/rollout-{ts}-{uuid}.jsonl"));
fs::create_dir_all(archived_path.parent().expect("archived parent")).unwrap();
fs::write(&archived_path, "").unwrap();
let resolved = crate::resolve_rollout_reference_rollout_path(
home,
&RolloutReferenceItem {
rollout_path: active_path,
thread_id: Some(thread_id),
rollout_timestamp: Some(ts.to_string()),
segment_id: None,
max_depth: 2,
},
)
.await
.expect("resolve rollout reference");
assert_eq!(resolved, archived_path);
}
async fn assert_state_db_rollout_path(
home: &Path,
thread_id: ThreadId,
@@ -9,6 +9,7 @@ use codex_rollout::RolloutRecorder;
use codex_rollout::RolloutRecorderParams;
use codex_rollout::builder_from_items;
use codex_rollout::read_session_meta_line;
use tokio::fs;
use tracing::warn;
use super::LocalThreadStore;
@@ -210,6 +211,7 @@ pub(super) async fn rotate_thread_segment(
initial_items.push(RolloutItem::RolloutReference(RolloutReferenceItem {
rollout_path: old_rollout_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,
}));
@@ -246,6 +248,46 @@ pub(super) async fn rotate_thread_segment(
);
}
let current_path = store
.live_recorders
.lock()
.await
.get(&thread_id)
.ok_or(ThreadStoreError::ThreadNotFound { thread_id })?
.rollout_path()
.to_path_buf();
if current_path != old_rollout_path {
return Err(ThreadStoreError::Conflict {
message: format!("live writer for thread {thread_id} changed during segment rotation"),
});
}
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 archived_root = store
.config
.codex_home
.join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR);
fs::create_dir_all(archived_root.as_path())
.await
.map_err(thread_store_io_error)?;
let archived_path = archived_root.join(old_file_name);
fs::rename(old_rollout_path.as_path(), archived_path.as_path())
.await
.map_err(|err| ThreadStoreError::Internal {
message: format!(
"failed to archive previous rollout segment {} to {}: {err}",
old_rollout_path.display(),
archived_path.display()
),
})?;
let mut live_recorders = store.live_recorders.lock().await;
let current_path = live_recorders
.get(&thread_id)
@@ -266,3 +308,13 @@ fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError {
message: err.to_string(),
}
}
fn rollout_timestamp_from_path(path: &std::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())
})
}
+27 -2
View File
@@ -428,6 +428,29 @@ mod tests {
.flush_thread(thread_id)
.await
.expect("flush post-rotation item");
let archived_old_rollout_path = home
.path()
.join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR)
.join(
old_rollout_path
.file_name()
.expect("old rollout path should have a file name"),
);
assert!(!old_rollout_path.exists());
assert!(archived_old_rollout_path.exists());
let old_segment_id = old_segment_id.expect("old rollout should have a segment id");
let resolved_old_rollout_path =
codex_rollout::find_rollout_path_by_segment_id(home.path(), thread_id, old_segment_id)
.await
.expect("resolve archived segment by id")
.expect("archived segment should resolve by id");
assert_eq!(resolved_old_rollout_path, archived_old_rollout_path);
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!("-{thread_id}.jsonl")))
.expect("old rollout timestamp");
let (new_items, new_thread_id, _) =
RolloutRecorder::load_rollout_items(new_rollout_path.as_path())
@@ -440,7 +463,8 @@ mod tests {
RolloutItem::RolloutReference(reference)
if reference.rollout_path == old_rollout_path
&& reference.thread_id == Some(thread_id)
&& reference.segment_id == old_segment_id
&& reference.rollout_timestamp.as_deref() == Some(old_rollout_timestamp)
&& reference.segment_id == Some(old_segment_id)
&& reference.max_depth == 2
)
}));
@@ -458,7 +482,8 @@ mod tests {
if event.message == "after rotation"
)
}));
assert_rollout_contains_message(old_rollout_path.as_path(), "before rotation").await;
assert_rollout_contains_message(archived_old_rollout_path.as_path(), "before rotation")
.await;
}
#[tokio::test]