diff --git a/codex-rs/state/thread_history_migrations/0003_turn_rollout_positions.sql b/codex-rs/state/thread_history_migrations/0003_turn_rollout_positions.sql new file mode 100644 index 0000000000..3cab073486 --- /dev/null +++ b/codex-rs/state/thread_history_migrations/0003_turn_rollout_positions.sql @@ -0,0 +1,3 @@ +ALTER TABLE thread_turns ADD COLUMN rollout_byte_offset INTEGER; +ALTER TABLE thread_turns ADD COLUMN rollout_end_ordinal INTEGER; +ALTER TABLE thread_turns ADD COLUMN rollout_end_byte_offset INTEGER; diff --git a/codex-rs/thread-store/src/local/thread_history.rs b/codex-rs/thread-store/src/local/thread_history.rs index b09568691b..11bca368ce 100644 --- a/codex-rs/thread-store/src/local/thread_history.rs +++ b/codex-rs/thread-store/src/local/thread_history.rs @@ -15,6 +15,18 @@ pub(super) use read::list_items; pub(super) use read::list_turns; pub(super) use search::search_thread_occurrences; +/// A valid complete rollout line with its absolute byte span in durable JSONL. +/// +/// `start_byte_offset..end_byte_offset` includes the terminating newline. Blank and rejected +/// lines do not produce a value here, but still advance later spans. +pub(super) struct ProjectedRolloutLine { + pub ordinal: u64, + pub start_byte_offset: u64, + pub end_byte_offset: u64, + pub created_at_ms: i64, + pub changes: ThreadHistoryChangeSet, +} + pub(super) async fn next_rollout_byte_offset( store: &LocalThreadStore, thread_id: ThreadId, @@ -46,7 +58,7 @@ pub(super) async fn apply_projection( thread_id: ThreadId, start_offset: u64, next_offset: u64, - projections: Vec<(Option, i64, ThreadHistoryChangeSet)>, + projections: Vec, ) -> ThreadStoreResult<()> { let pool = store.thread_history_db().await?; // Write the projected rows and advance the JSONL offset and ordinal in one transaction. If @@ -76,12 +88,8 @@ WHERE thread_id = ? }); } - for (ordinal, created_at_ms, changes) in projections { - let ordinal = ordinal - .ok_or_else(|| ThreadStoreError::Internal { - message: format!("paginated rollout line for {thread_id} is missing an ordinal"), - }) - .and_then(|ordinal| sqlite_integer(ordinal, "rollout ordinal"))?; + for projection in projections { + let ordinal = sqlite_integer(projection.ordinal, "rollout ordinal")?; if ordinal != next_ordinal { return Err(ThreadStoreError::Internal { message: format!( @@ -93,8 +101,10 @@ WHERE thread_id = ? &mut transaction, thread_id.as_str(), ordinal, - created_at_ms, - changes, + sqlite_integer(projection.start_byte_offset, "rollout byte offset")?, + sqlite_integer(projection.end_byte_offset, "rollout byte offset")?, + projection.created_at_ms, + projection.changes, ) .await?; next_ordinal = next_ordinal @@ -168,6 +178,8 @@ async fn apply_change_set( transaction: &mut sqlx::Transaction<'_, sqlx::Sqlite>, thread_id: &str, rollout_ordinal: i64, + rollout_byte_offset: i64, + rollout_end_byte_offset: i64, created_at_ms: i64, changes: ThreadHistoryChangeSet, ) -> ThreadStoreResult<()> { @@ -179,6 +191,12 @@ async fn apply_change_set( .map(serde_json::to_string) .transpose() .map_err(thread_history_error)?; + let (terminal_ordinal, terminal_byte_offset) = match &turn.status { + TurnStatus::Completed | TurnStatus::Interrupted | TurnStatus::Failed => { + (Some(rollout_ordinal), Some(rollout_end_byte_offset)) + } + TurnStatus::InProgress => (None, None), + }; // The same turn can appear again as it moves from started to completed. Update its latest // status, error, and timestamps, but keep the rollout ordinal from the first record that // created it. @@ -188,23 +206,33 @@ INSERT INTO thread_turns ( thread_id, turn_id, rollout_ordinal, + rollout_byte_offset, + rollout_end_ordinal, + rollout_end_byte_offset, status, error_json, started_at, completed_at, duration_ms -) VALUES (?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(thread_id, turn_id) DO UPDATE SET + rollout_end_ordinal = excluded.rollout_end_ordinal, + rollout_end_byte_offset = excluded.rollout_end_byte_offset, status = excluded.status, error_json = excluded.error_json, started_at = excluded.started_at, completed_at = excluded.completed_at, duration_ms = excluded.duration_ms +WHERE thread_turns.rollout_end_ordinal IS NULL + AND thread_turns.status = 'inProgress' "#, ) .bind(thread_id) .bind(turn_id.as_str()) .bind(rollout_ordinal) + .bind(rollout_byte_offset) + .bind(terminal_ordinal) + .bind(terminal_byte_offset) .bind(turn_status(&turn.status)) .bind(error_json) .bind(turn.started_at) @@ -257,7 +285,12 @@ SET END, final_agent_item_id ) -WHERE thread_id = ? AND turn_id = ? +WHERE thread_id = ? + AND turn_id = ? + AND ( + rollout_end_ordinal = ? + OR status = 'inProgress' + ) "#, ) .bind(thread_id) @@ -268,6 +301,7 @@ WHERE thread_id = ? AND turn_id = ? .bind(turn_id.as_str()) .bind(thread_id) .bind(turn_id.as_str()) + .bind(rollout_ordinal) .execute(&mut **transaction) .await .map_err(thread_history_error)?; @@ -314,7 +348,10 @@ ON CONFLICT(thread_id, turn_id, item_id) DO UPDATE SET r#" UPDATE thread_turns SET first_user_item_id = COALESCE(first_user_item_id, ?) -WHERE thread_id = ? AND turn_id = ? +WHERE thread_id = ? + AND turn_id = ? + AND rollout_end_ordinal IS NULL + AND status = 'inProgress' "#, ) .bind(item_id.as_str()) @@ -332,7 +369,10 @@ WHERE thread_id = ? AND turn_id = ? r#" UPDATE thread_turns SET final_agent_item_id = ? -WHERE thread_id = ? AND turn_id = ? +WHERE thread_id = ? + AND turn_id = ? + AND rollout_end_ordinal IS NULL + AND status = 'inProgress' "#, ) .bind(item_id.as_str()) diff --git a/codex-rs/thread-store/src/local/thread_history_materialization.rs b/codex-rs/thread-store/src/local/thread_history_materialization.rs index c3367cbf01..268d42a9db 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization.rs @@ -11,9 +11,16 @@ use tokio::io::AsyncSeekExt; use tracing::warn; use super::LocalThreadStore; +use super::thread_history::ProjectedRolloutLine; use crate::ThreadStoreError; use crate::ThreadStoreResult; +struct CompleteRolloutLine { + line: RolloutLine, + start_byte_offset: u64, + end_byte_offset: u64, +} + pub(super) async fn materialize_to_sqlite( store: &LocalThreadStore, thread_id: ThreadId, @@ -21,6 +28,7 @@ pub(super) async fn materialize_to_sqlite( ) -> ThreadStoreResult<()> { let start_offset = super::thread_history::next_rollout_byte_offset(store, thread_id).await?; let (lines, next_offset) = read_complete_rollout_lines(rollout_path, start_offset).await?; + // Empty valid records can still consume bytes through blank or rejected complete lines. if lines.is_empty() && start_offset == next_offset { return Ok(()); } @@ -32,18 +40,26 @@ pub(super) async fn materialize_to_sqlite( let projections = lines .iter() - .map(|line| { + .map(|record| { + let line = &record.line; + let ordinal = line.ordinal.ok_or_else(|| ThreadStoreError::Internal { + message: format!("paginated rollout line for {thread_id} is missing an ordinal"), + })?; let created_at_ms = DateTime::parse_from_rfc3339(line.timestamp.as_str()) .map(|timestamp| timestamp.timestamp_millis()) .map_err(thread_history_error)?; - let changes = if line.ordinal.is_some_and(|ordinal| { - subagent_history_start_ordinal.is_some_and(|start| ordinal < start) - }) { + let changes = if subagent_history_start_ordinal.is_some_and(|start| ordinal < start) { ThreadHistoryChangeSet::default() } else { project_rollout_line(line) }; - Ok((line.ordinal, created_at_ms, changes)) + Ok(ProjectedRolloutLine { + ordinal, + start_byte_offset: record.start_byte_offset, + end_byte_offset: record.end_byte_offset, + created_at_ms, + changes, + }) }) .collect::>>()?; super::thread_history::apply_projection( @@ -59,7 +75,7 @@ pub(super) async fn materialize_to_sqlite( async fn read_complete_rollout_lines( rollout_path: &Path, start_offset: u64, -) -> ThreadStoreResult<(Vec, u64)> { +) -> ThreadStoreResult<(Vec, u64)> { let next_offset = match tokio::fs::metadata(rollout_path).await { Ok(metadata) => metadata.len(), Err(err) if err.kind() == std::io::ErrorKind::NotFound && start_offset == 0 => { @@ -86,6 +102,8 @@ async fn read_complete_rollout_lines( file.read_exact(bytes.as_mut_slice()) .await .map_err(thread_store_io_error)?; + // Only project the newline-terminated prefix; leave a trailing partial record for the next + // pass. let complete_byte_count = bytes .iter() .rposition(|byte| *byte == b'\n') @@ -99,18 +117,39 @@ async fn read_complete_rollout_lines( .ok_or_else(|| ThreadStoreError::Internal { message: "durable rollout byte offset overflow".to_string(), })?; - let text = std::str::from_utf8(&bytes[..complete_byte_count]).map_err(thread_history_error)?; let mut lines = Vec::new(); - for line in text.lines().filter(|line| !line.is_empty()) { - match serde_json::from_str(line) { - Ok(line) => lines.push(line), - Err(err) => { - // A failed append can leave a partial record behind. The rollout writer repairs - // its newline before retrying, so skip rejected lines just like the canonical - // rollout loader and keep projecting the valid retry that follows. - warn!("skipping rejected rollout line while projecting {rollout_path:?}: {err}"); + let mut line_start_offset = start_offset; + // Preserve each complete physical line's trailing newline so byte offsets advance through + // every durable byte, including blank or rejected lines that do not project a row. + for line_bytes in bytes[..complete_byte_count].split_inclusive(|byte| *byte == b'\n') { + let line_end_offset = line_start_offset + .checked_add(u64::try_from(line_bytes.len()).map_err(|_| { + ThreadStoreError::Internal { + message: "durable rollout byte offset overflow".to_string(), + } + })?) + .ok_or_else(|| ThreadStoreError::Internal { + message: "durable rollout byte offset overflow".to_string(), + })?; + // Blank physical lines consume bytes but are not rollout records. + if !line_bytes.iter().all(u8::is_ascii_whitespace) { + match serde_json::from_slice(line_bytes) { + Ok(line) => lines.push(CompleteRolloutLine { + line, + start_byte_offset: line_start_offset, + end_byte_offset: line_end_offset, + }), + Err(err) => { + // A failed append can leave a partial record behind. The rollout writer + // repairs its newline before retrying, so skip rejected lines just like the + // canonical rollout loader and keep projecting the valid retry that follows. + warn!( + "skipping rejected rollout line while projecting {rollout_path:?}: {err}" + ); + } } } + line_start_offset = line_end_offset; } Ok((lines, next_offset)) } diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index 28208f8f0f..abeb1149a4 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -78,6 +78,14 @@ async fn paginated_live_append_materializes_turn_items_and_state() { .await .expect("append paginated items"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + let (turn_start_byte_offset, _) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 1); + let (_, turn_end_byte_offset) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 4); let pool = codex_state::open_thread_history_db(home.path()) .await .expect("open thread history db"); @@ -85,6 +93,9 @@ async fn paginated_live_append_materializes_turn_items_and_state() { _, ( i64, + Option, + Option, + Option, String, Option, Option, @@ -96,6 +107,9 @@ async fn paginated_live_append_materializes_turn_items_and_state() { r#" SELECT rollout_ordinal, + rollout_byte_offset, + rollout_end_ordinal, + rollout_end_byte_offset, status, started_at, completed_at, @@ -115,6 +129,9 @@ WHERE thread_id = ? AND turn_id = ? turn, ( 1, + Some(turn_start_byte_offset), + Some(4), + Some(turn_end_byte_offset), "completed".to_string(), Some(10), Some(20), @@ -141,10 +158,6 @@ ORDER BY rollout_ordinal vec![("user-1".to_string(), 2), ("agent-1".to_string(), 3)] ); - let rollout_path = store - .live_rollout_path(thread_id) - .await - .expect("rollout path"); let rollout_len = i64::try_from(fs::metadata(rollout_path).expect("rollout metadata").len()) .expect("rollout length"); let projection_state = sqlx::query_as::<_, (i64, i64)>( @@ -161,6 +174,41 @@ WHERE thread_id = ? assert_eq!(projection_state, (rollout_len, 5)); } +#[tokio::test] +async fn active_turn_stores_only_its_start_position() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let thread_id = ThreadId::default(); + create_paginated_thread(&store, thread_id).await; + + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![turn_started("turn-1")], + }) + .await + .expect("append active turn"); + + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + let (turn_start_byte_offset, _) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 1); + let pool = codex_state::open_thread_history_db(home.path()) + .await + .expect("open thread history db"); + let turn_position = sqlx::query_as::<_, (Option, Option, Option)>( + "SELECT rollout_byte_offset, rollout_end_ordinal, rollout_end_byte_offset FROM thread_turns WHERE thread_id = ? AND turn_id = ?", + ) + .bind(thread_id.to_string()) + .bind("turn-1") + .fetch_one(&pool) + .await + .expect("read active turn position"); + assert_eq!(turn_position, (Some(turn_start_byte_offset), None, None)); +} + #[tokio::test] async fn subagent_prefix_advances_projection_without_materializing_history() { let home = TempDir::new().expect("temp dir"); @@ -208,17 +256,26 @@ async fn subagent_prefix_advances_projection_without_materializing_history() { .await .expect("append inherited prefix and child history"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + let (child_start_byte_offset, _) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 4); let pool = codex_state::open_thread_history_db(home.path()) .await .expect("open thread history db"); - let turns = sqlx::query_as::<_, (String, i64)>( - "SELECT turn_id, rollout_ordinal FROM thread_turns WHERE thread_id = ?", + let turns = sqlx::query_as::<_, (String, i64, Option)>( + "SELECT turn_id, rollout_ordinal, rollout_byte_offset FROM thread_turns WHERE thread_id = ?", ) .bind(thread_id.to_string()) .fetch_all(&pool) .await .expect("read projected turns"); - assert_eq!(turns, vec![("child-turn".to_string(), 4)]); + assert_eq!( + turns, + vec![("child-turn".to_string(), 4, Some(child_start_byte_offset))] + ); let items = sqlx::query_as::<_, (String, i64)>( "SELECT item_id, rollout_ordinal FROM thread_items WHERE thread_id = ?", ) @@ -308,6 +365,91 @@ WHERE thread_id = ? AND turn_id = ? AND item_id = ? ); } +#[tokio::test] +async fn terminal_turn_does_not_change_after_later_records() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let thread_id = ThreadId::default(); + create_paginated_thread(&store, thread_id).await; + + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![turn_started("turn-1"), turn_completed("turn-1")], + }) + .await + .expect("append terminal turn"); + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![ + turn_started("turn-1"), + completed_item( + thread_id, + "turn-1", + TurnItem::UserMessage(UserMessageItem { + id: "late-user".to_string(), + client_id: None, + content: Vec::new(), + }), + ), + ], + }) + .await + .expect("append later records"); + + let pool = codex_state::open_thread_history_db(home.path()) + .await + .expect("open thread history db"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + let (turn_start_byte_offset, _) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 1); + let (_, turn_end_byte_offset) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 2); + let turn = sqlx::query_as::< + _, + ( + i64, + Option, + Option, + Option, + String, + Option, + ), + >( + r#" +SELECT + rollout_ordinal, + rollout_byte_offset, + rollout_end_ordinal, + rollout_end_byte_offset, + status, + first_user_item_id +FROM thread_turns +WHERE thread_id = ? AND turn_id = ? + "#, + ) + .bind(thread_id.to_string()) + .bind("turn-1") + .fetch_one(&pool) + .await + .expect("read projected turn"); + assert_eq!( + turn, + ( + 1, + Some(turn_start_byte_offset), + Some(2), + Some(turn_end_byte_offset), + "completed".to_string(), + None, + ) + ); +} + #[tokio::test] async fn summary_items_use_final_answers_and_ignore_commentary() { let home = TempDir::new().expect("temp dir"); @@ -552,7 +694,7 @@ async fn synchronized_catch_up_does_not_replay_old_rows() { } #[tokio::test] -async fn catch_up_leaves_trailing_partial_line_unprojected() { +async fn catch_up_preserves_trailing_partial_line_boundaries() { let home = TempDir::new().expect("temp dir"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let thread_id = ThreadId::default(); @@ -567,18 +709,7 @@ async fn catch_up_leaves_trailing_partial_line_unprojected() { .expect("open thread history db"); let before = projection_state(&pool, thread_id).await; let complete_line = rollout_line(Some(1), turn_started("turn-1")); - let partial_line = rollout_line( - Some(2), - completed_item( - thread_id, - "turn-1", - TurnItem::UserMessage(UserMessageItem { - id: "user-1".to_string(), - client_id: None, - content: Vec::new(), - }), - ), - ); + let partial_line = rollout_line(Some(2), turn_completed("turn-1")); let complete_suffix = format!("{complete_line}\n"); let rollout_path = store .live_rollout_path(thread_id) @@ -599,19 +730,22 @@ async fn catch_up_leaves_trailing_partial_line_unprojected() { projection_state(&pool, thread_id).await, (expected_offset, 2) ); - let counts = sqlx::query_as::<_, (i64, i64)>( - r#" -SELECT - (SELECT COUNT(*) FROM thread_turns WHERE thread_id = ?), - (SELECT COUNT(*) FROM thread_items WHERE thread_id = ?) - "#, + append_suffix(rollout_path.as_path(), "\n"); + super::materialize_to_sqlite(&store, thread_id, rollout_path.as_path()) + .await + .expect("catch up completed partial suffix"); + + let rollout_len = i64::try_from(fs::metadata(rollout_path).expect("rollout metadata").len()) + .expect("rollout length"); + let turn_position = sqlx::query_as::<_, (Option, Option, Option)>( + "SELECT rollout_byte_offset, rollout_end_ordinal, rollout_end_byte_offset FROM thread_turns WHERE thread_id = ? AND turn_id = ?", ) .bind(thread_id.to_string()) - .bind(thread_id.to_string()) + .bind("turn-1") .fetch_one(&pool) .await - .expect("read projected row counts"); - assert_eq!(counts, (1, 0)); + .expect("read completed turn position"); + assert_eq!(turn_position, (Some(before.0), Some(2), Some(rollout_len))); } #[tokio::test] @@ -762,7 +896,7 @@ async fn sqlite_failure_does_not_fail_durable_jsonl_write() { } #[tokio::test] -async fn rejected_rollout_line_does_not_poison_projection() { +async fn blank_and_rejected_rollout_lines_do_not_poison_projection() { let home = TempDir::new().expect("temp dir"); let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); let thread_id = ThreadId::default(); @@ -780,8 +914,8 @@ async fn rejected_rollout_line_does_not_poison_projection() { .append(true) .open(rollout_path.as_path()) .expect("open rollout for rejected line"); - file.write_all(b"{not json}\n") - .expect("append rejected line"); + file.write_all(b"\n \t\r\n{not json}\n\xff\n") + .expect("append blank and rejected lines"); file.flush().expect("flush rejected line"); let recorder = store .live_recorders @@ -804,15 +938,17 @@ async fn rejected_rollout_line_does_not_poison_projection() { let pool = codex_state::open_thread_history_db(home.path()) .await .expect("open thread history db"); - let projected_turns = sqlx::query_scalar::<_, i64>( - "SELECT COUNT(*) FROM thread_turns WHERE thread_id = ? AND turn_id = ?", + let (expected_start_byte_offset, _) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 1); + let start_byte_offset = sqlx::query_scalar::<_, Option>( + "SELECT rollout_byte_offset FROM thread_turns WHERE thread_id = ? AND turn_id = ?", ) .bind(thread_id.to_string()) .bind("turn-1") .fetch_one(&pool) .await - .expect("read projected turns"); - assert_eq!(projected_turns, 1); + .expect("read projected turn byte offset"); + assert_eq!(start_byte_offset, Some(expected_start_byte_offset)); } #[tokio::test] @@ -992,6 +1128,26 @@ fn agent_message(id: &str, phase: MessagePhase) -> TurnItem { }) } +fn rollout_line_byte_offsets(path: &std::path::Path, ordinal: u64) -> (i64, i64) { + let bytes = fs::read(path).expect("read rollout"); + let mut start_byte_offset = 0; + for line in bytes.split_inclusive(|byte| *byte == b'\n') { + let end_byte_offset = start_byte_offset + line.len(); + if serde_json::from_slice::(line) + .ok() + .and_then(|line| line.ordinal) + == Some(ordinal) + { + return ( + i64::try_from(start_byte_offset).expect("start byte offset fits i64"), + i64::try_from(end_byte_offset).expect("end byte offset fits i64"), + ); + } + start_byte_offset = end_byte_offset; + } + panic!("missing rollout ordinal {ordinal}"); +} + async fn projection_state(pool: &sqlx::SqlitePool, thread_id: ThreadId) -> (i64, i64) { sqlx::query_as::<_, (i64, i64)>( r#"