From 7bb13ab846a36341129bc431bbe277ded748a444 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Tue, 21 Jul 2026 15:06:51 +0000 Subject: [PATCH] Page through inherited thread history (#34563) ## What changed - Resolve a thread's rollout lineage when listing paginated turns and items. - Page across parent, child, and nested fork segments in either direction while respecting each segment's ordinal bounds. - Encode the physical segment in history cursors and reject cursors that belong to another requested thread or fall outside its lineage. - Load turn summaries and filtered items from inherited segments. ## Testing - Added coverage for forward and backward pagination across parent and child history, inherited summary and filtered items, cursor isolation, and nested forks. GitOrigin-RevId: b803fff781fcb6e21d0e151189290148fc8d8b18 --- .../thread-store/src/local/rollout_lineage.rs | 20 + .../thread-store/src/local/thread_history.rs | 1 + .../src/local/thread_history/read.rs | 360 +++++++----------- .../src/local/thread_history/read_tests.rs | 333 +++++++++++++++- .../src/local/thread_history/search.rs | 8 +- .../local/thread_history/segment_paging.rs | 302 +++++++++++++++ 6 files changed, 791 insertions(+), 233 deletions(-) create mode 100644 codex-rs/thread-store/src/local/thread_history/segment_paging.rs diff --git a/codex-rs/thread-store/src/local/rollout_lineage.rs b/codex-rs/thread-store/src/local/rollout_lineage.rs index aed97f55f1..312923afa8 100644 --- a/codex-rs/thread-store/src/local/rollout_lineage.rs +++ b/codex-rs/thread-store/src/local/rollout_lineage.rs @@ -102,6 +102,26 @@ impl LocalThreadStore { } } +impl RolloutLineage { + pub(super) fn segments(&self) -> &[RolloutLineageSegment] { + self.segments.as_slice() + } +} + +impl RolloutLineageSegment { + pub(super) fn thread_id(&self) -> ThreadId { + self.thread_id + } + + pub(super) fn start_ordinal(&self) -> u64 { + 1 + } + + pub(super) fn end_ordinal(&self) -> Option { + self.end.map(|end| end.end_ordinal_exclusive) + } +} + async fn validate_cutoff_bounds( requested_thread_id: ThreadId, rollout_path: &Path, diff --git a/codex-rs/thread-store/src/local/thread_history.rs b/codex-rs/thread-store/src/local/thread_history.rs index 11bca368ce..12a07c679b 100644 --- a/codex-rs/thread-store/src/local/thread_history.rs +++ b/codex-rs/thread-store/src/local/thread_history.rs @@ -10,6 +10,7 @@ use crate::ThreadStoreResult; mod read; mod search; +mod segment_paging; pub(super) use read::list_items; pub(super) use read::list_turns; diff --git a/codex-rs/thread-store/src/local/thread_history/read.rs b/codex-rs/thread-store/src/local/thread_history/read.rs index a119a13bfd..94e28eba3f 100644 --- a/codex-rs/thread-store/src/local/thread_history/read.rs +++ b/codex-rs/thread-store/src/local/thread_history/read.rs @@ -2,16 +2,15 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadHistoryMode; use serde::Deserialize; use serde::Serialize; -use sqlx::QueryBuilder; use sqlx::Row; -use sqlx::Sqlite; use super::super::LocalThreadStore; -use super::thread_history_error; +use super::segment_paging::page_item_rows; +use super::segment_paging::page_turn_rows; +use super::segment_paging::validate_page_size; use crate::ItemPage; use crate::ListItemsParams; use crate::ListTurnsParams; -use crate::SortDirection; use crate::StoredThreadItem; use crate::StoredTurn; use crate::StoredTurnError; @@ -25,15 +24,6 @@ use crate::TurnPage; #[path = "read_tests.rs"] mod tests; -#[derive(Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -struct HistoryCursor { - thread_id: ThreadId, - scope: CursorScope, - rollout_ordinal: i64, - include_anchor: bool, -} - #[derive(Clone, Deserialize, PartialEq, Eq, Serialize)] #[serde(tag = "kind", rename_all = "camelCase")] pub(super) enum CursorScope { @@ -41,21 +31,37 @@ pub(super) enum CursorScope { Items, } -struct StoredTurnRow { - turn_id: String, - rollout_ordinal: i64, - status: StoredTurnStatus, - error: Option, - started_at: Option, - completed_at: Option, - duration_ms: Option, - first_user_item_id: Option, - final_agent_item_id: Option, +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub(super) struct HistoryCursor { + pub requested_thread_id: ThreadId, + pub physical_thread_id: ThreadId, + pub rollout_ordinal: u64, + pub include_anchor: bool, + pub scope: CursorScope, } -struct StoredThreadItemRow { - item: StoredThreadItem, - rollout_ordinal: i64, +#[derive(Clone, Copy)] +pub(super) struct PhysicalHistoryPosition { + pub physical_thread_id: ThreadId, + pub rollout_ordinal: i64, +} + +pub(super) struct StoredTurnRow { + pub position: PhysicalHistoryPosition, + pub turn_id: String, + pub status: StoredTurnStatus, + pub error: Option, + pub started_at: Option, + pub completed_at: Option, + pub duration_ms: Option, + pub first_user_item_id: Option, + pub final_agent_item_id: Option, +} + +pub(super) struct StoredThreadItemRow { + pub position: PhysicalHistoryPosition, + pub item: StoredThreadItem, } pub(in crate::local) async fn list_turns( @@ -69,56 +75,25 @@ pub(in crate::local) async fn list_turns( "list_turns", ) .await?; - let scope = CursorScope::Turns; - let cursor = parse_cursor(params.cursor.as_deref(), params.thread_id, &scope)?; + validate_page_size(params.page_size)?; + let lineage = store.resolve_rollout_lineage(params.thread_id).await?; let pool = store.thread_history_db().await?; - let limit = page_limit(params.page_size)?; - let mut query = QueryBuilder::::new( - r#" -SELECT - turn_id, - rollout_ordinal, - status, - error_json, - started_at, - completed_at, - duration_ms, - first_user_item_id, - final_agent_item_id -FROM thread_turns -WHERE thread_id = - "#, - ); - query.push_bind(params.thread_id.to_string()); - push_pagination_clause(&mut query, params.sort_direction, cursor.as_ref(), limit); - let rows = query - .build() - .fetch_all(pool) - .await - .map_err(thread_history_error)?; - let mut turns = rows - .into_iter() - .map(stored_turn_row) - .collect::>>()?; - let has_more = turns.len() > params.page_size; - turns.truncate(params.page_size); - - let (next_cursor, backwards_cursor) = page_cursors( + let page = page_turn_rows( + pool, params.thread_id, - &scope, - turns.first().map(|turn| turn.rollout_ordinal), - turns.last().map(|turn| turn.rollout_ordinal), - has_more, - )?; - let mut stored_turns = Vec::with_capacity(turns.len()); - for turn in turns { + &lineage, + params.cursor.as_deref(), + params.page_size, + params.sort_direction, + ) + .await?; + let mut turns = Vec::with_capacity(page.rows.len()); + for turn in page.rows { let items = match params.items_view { StoredTurnItemsView::NotLoaded => Vec::new(), - StoredTurnItemsView::Summary => { - load_summary_items(pool, params.thread_id, &turn).await? - } + StoredTurnItemsView::Summary => load_summary_items(pool, &turn).await?, }; - stored_turns.push(StoredTurn { + turns.push(StoredTurn { turn_id: turn.turn_id, items, items_view: params.items_view, @@ -131,9 +106,9 @@ WHERE thread_id = } Ok(TurnPage { - turns: stored_turns, - next_cursor, - backwards_cursor, + turns, + next_cursor: page.next_cursor, + backwards_cursor: page.backwards_cursor, }) } @@ -148,46 +123,24 @@ pub(in crate::local) async fn list_items( "list_items", ) .await?; - let scope = CursorScope::Items; - let cursor = parse_cursor(params.cursor.as_deref(), params.thread_id, &scope)?; + validate_page_size(params.page_size)?; + let lineage = store.resolve_rollout_lineage(params.thread_id).await?; let pool = store.thread_history_db().await?; - let limit = page_limit(params.page_size)?; - let mut query = QueryBuilder::::new( - r#" -SELECT turn_id, item_id, rollout_ordinal, created_at_ms, item_json -FROM thread_items -WHERE thread_id = - "#, - ); - query.push_bind(params.thread_id.to_string()); - if let Some(turn_id) = params.turn_id.as_deref() { - query.push(" AND turn_id = ").push_bind(turn_id); - } - push_pagination_clause(&mut query, params.sort_direction, cursor.as_ref(), limit); - let rows = query - .build() - .fetch_all(pool) - .await - .map_err(thread_history_error)?; - let mut item_rows = rows - .into_iter() - .map(stored_thread_item_row) - .collect::>>()?; - let has_more = item_rows.len() > params.page_size; - item_rows.truncate(params.page_size); - let (next_cursor, backwards_cursor) = page_cursors( + let page = page_item_rows( + pool, params.thread_id, - &scope, - item_rows.first().map(|row| row.rollout_ordinal), - item_rows.last().map(|row| row.rollout_ordinal), - has_more, - )?; - let items = item_rows.into_iter().map(|row| row.item).collect(); + &lineage, + params.turn_id.as_deref(), + params.cursor.as_deref(), + params.page_size, + params.sort_direction, + ) + .await?; Ok(ItemPage { - items, - next_cursor, - backwards_cursor, + items: page.rows.into_iter().map(|row| row.item).collect(), + next_cursor: page.next_cursor, + backwards_cursor: page.backwards_cursor, }) } @@ -221,113 +174,68 @@ pub(super) async fn validate_thread_for_paginated_reads( } } -fn page_limit(page_size: usize) -> ThreadStoreResult { - if page_size == 0 { - return Err(ThreadStoreError::InvalidRequest { - message: "page size must be positive".to_string(), - }); - } - let limit = page_size - .checked_add(1) - .ok_or_else(|| ThreadStoreError::InvalidRequest { - message: "page size is too large".to_string(), - })?; - i64::try_from(limit).map_err(|_| ThreadStoreError::InvalidRequest { - message: "page size is too large".to_string(), - }) +async fn load_summary_items( + pool: &sqlx::SqlitePool, + turn: &StoredTurnRow, +) -> ThreadStoreResult> { + let rows = sqlx::query( + r#" +SELECT turn_id, item_id, created_at_ms, item_json +FROM thread_items +WHERE thread_id = ? + AND turn_id = ? + AND (item_id = ? OR item_id = ?) +ORDER BY rollout_ordinal ASC + "#, + ) + .bind(turn.position.physical_thread_id.to_string()) + .bind(turn.turn_id.as_str()) + .bind(turn.first_user_item_id.as_deref()) + .bind(turn.final_agent_item_id.as_deref()) + .fetch_all(pool) + .await + .map_err(super::thread_history_error)?; + rows.into_iter().map(stored_thread_item).collect() } -fn parse_cursor( +pub(super) fn parse_cursor( cursor: Option<&str>, - thread_id: ThreadId, - scope: &CursorScope, + requested_thread_id: ThreadId, + scope: CursorScope, ) -> ThreadStoreResult> { let Some(cursor) = cursor else { return Ok(None); }; let cursor_value: HistoryCursor = serde_json::from_str(cursor).map_err(|_| invalid_cursor(cursor))?; - if cursor_value.thread_id != thread_id || &cursor_value.scope != scope { + if cursor_value.requested_thread_id != requested_thread_id || cursor_value.scope != scope { return Err(invalid_cursor(cursor)); } Ok(Some(cursor_value)) } -fn push_pagination_clause( - query: &mut QueryBuilder, - direction: SortDirection, - cursor: Option<&HistoryCursor>, - limit: i64, -) { - if let Some(cursor) = cursor { - let comparator = match (direction, cursor.include_anchor) { - (SortDirection::Asc, true) => ">=", - (SortDirection::Asc, false) => ">", - (SortDirection::Desc, true) => "<=", - (SortDirection::Desc, false) => "<", - }; - query - .push(" AND rollout_ordinal ") - .push(comparator) - .push(" ") - .push_bind(cursor.rollout_ordinal); - } - let order = match direction { - SortDirection::Asc => "ASC", - SortDirection::Desc => "DESC", - }; - query - .push(" ORDER BY rollout_ordinal ") - .push(order) - .push(" LIMIT ") - .push_bind(limit); -} - -fn page_cursors( - thread_id: ThreadId, - scope: &CursorScope, - first_ordinal: Option, - last_ordinal: Option, - has_more: bool, -) -> ThreadStoreResult<(Option, Option)> { - let cursor = |rollout_ordinal, include_anchor| { - serialize_cursor(thread_id, scope, rollout_ordinal, include_anchor) - }; - let backwards_cursor = first_ordinal - .map(|rollout_ordinal| cursor(rollout_ordinal, /*include_anchor*/ true)) - .transpose()?; - let next_cursor = if has_more { - last_ordinal - .map(|rollout_ordinal| cursor(rollout_ordinal, /*include_anchor*/ false)) - .transpose()? - } else { - None - }; - Ok((next_cursor, backwards_cursor)) -} - pub(super) fn serialize_cursor( - thread_id: ThreadId, - scope: &CursorScope, - rollout_ordinal: i64, + requested_thread_id: ThreadId, + scope: CursorScope, + position: PhysicalHistoryPosition, include_anchor: bool, ) -> ThreadStoreResult { + let rollout_ordinal = u64::try_from(position.rollout_ordinal) + .map_err(|_| invalid_cursor("negative rollout ordinal"))?; serde_json::to_string(&HistoryCursor { - thread_id, - scope: scope.clone(), + requested_thread_id, + physical_thread_id: position.physical_thread_id, rollout_ordinal, include_anchor, + scope, }) - .map_err(thread_history_error) + .map_err(super::thread_history_error) } -fn invalid_cursor(cursor: &str) -> ThreadStoreError { - ThreadStoreError::InvalidRequest { - message: format!("invalid cursor: {cursor}"), - } -} - -fn stored_turn_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { +pub(super) fn stored_turn_row( + physical_thread_id: ThreadId, + row: sqlx::sqlite::SqliteRow, +) -> ThreadStoreResult { let status = match row.try_get::("status")?.as_str() { "completed" => StoredTurnStatus::Completed, "interrupted" => StoredTurnStatus::Interrupted, @@ -344,10 +252,13 @@ fn stored_turn_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult ThreadStoreResult ThreadStoreResult> { - let rows = sqlx::query( - r#" -SELECT turn_id, item_id, rollout_ordinal, created_at_ms, item_json -FROM thread_items -WHERE thread_id = ? - AND turn_id = ? - AND (item_id = ? OR item_id = ?) -ORDER BY rollout_ordinal ASC - "#, - ) - .bind(thread_id.to_string()) - .bind(turn.turn_id.as_str()) - .bind(turn.first_user_item_id.as_deref()) - .bind(turn.final_agent_item_id.as_deref()) - .fetch_all(pool) - .await - .map_err(thread_history_error)?; - rows.into_iter() - .map(|row| stored_thread_item_row(row).map(|row| row.item)) - .collect() -} - -fn stored_thread_item_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { +pub(super) fn stored_thread_item_row_for_thread( + physical_thread_id: ThreadId, + row: sqlx::sqlite::SqliteRow, +) -> ThreadStoreResult { let rollout_ordinal = row.try_get::("rollout_ordinal")?; if rollout_ordinal < 0 { return Err(ThreadStoreError::Internal { @@ -393,12 +280,25 @@ fn stored_thread_item_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult("item_json")?.into_bytes(), + position: PhysicalHistoryPosition { + physical_thread_id, + rollout_ordinal, }, - rollout_ordinal, + item: stored_thread_item(row)?, }) } + +fn stored_thread_item(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { + Ok(StoredThreadItem { + turn_id: row.try_get("turn_id")?, + item_id: row.try_get("item_id")?, + created_at_ms: row.try_get("created_at_ms")?, + item_json: row.try_get::("item_json")?.into_bytes(), + }) +} + +pub(super) fn invalid_cursor(cursor: &str) -> ThreadStoreError { + ThreadStoreError::InvalidRequest { + message: format!("invalid cursor: {cursor}"), + } +} diff --git a/codex-rs/thread-store/src/local/thread_history/read_tests.rs b/codex-rs/thread-store/src/local/thread_history/read_tests.rs index 70ee07d519..abb702b537 100644 --- a/codex-rs/thread-store/src/local/thread_history/read_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history/read_tests.rs @@ -1,12 +1,23 @@ +use std::fs; + use chrono::Utc; use codex_app_server_protocol::CodexErrorInfo; use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::HistoryPosition; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use tempfile::TempDir; use super::*; +use crate::SortDirection; +use crate::StoredTurnError; +use crate::StoredTurnStatus; use crate::local::test_support::test_config; #[tokio::test] @@ -250,10 +261,251 @@ async fn list_history_keeps_legacy_threads_unsupported() { )); } +#[tokio::test] +async fn lineage_reads_page_across_parent_and_child_segments() { + let (home, store, child_id) = store_with_mode(ThreadHistoryMode::Paginated).await; + let root_id = ThreadId::default(); + let root_path = rollout_path(home.path(), root_id); + write_rollout_with_end( + root_path.as_path(), + root_id, + /*history_base*/ None, + /*next_ordinal*/ 8, + ); + write_rollout_with_end( + rollout_path(home.path(), child_id).as_path(), + child_id, + Some(history_position( + root_path.as_path(), + root_id, + /*end_ordinal_exclusive*/ 6, + )), + /*next_ordinal*/ 3, + ); + let db = history_db(&store).await; + for (thread_id, turn_id, ordinal, first_user, final_agent) in [ + (root_id, "root-1", 1, Some("root-user"), Some("root-agent")), + (root_id, "root-2", 4, None, None), + (root_id, "excluded-root", 6, None, None), + (child_id, "child-1", 1, None, None), + ] { + insert_turn( + db, + thread_id, + turn_id, + ordinal, + "completed", + /*error_json*/ None, + first_user, + final_agent, + ) + .await; + } + for (thread_id, turn_id, item_id, ordinal) in [ + (root_id, "root-1", "root-user", 2), + (root_id, "root-1", "root-agent", 3), + (root_id, "root-2", "root-2-item", 5), + (root_id, "excluded-root", "excluded-item", 7), + (child_id, "child-1", "child-item", 2), + ] { + insert_item(db, thread_id, turn_id, item_id, ordinal).await; + } + + let first_turns = store + .list_turns(turn_params( + child_id, + /*cursor*/ None, + /*page_size*/ 2, + SortDirection::Asc, + StoredTurnItemsView::Summary, + )) + .await + .expect("first lineage turns page"); + assert_eq!(turn_ids(&first_turns), vec!["root-1", "root-2"]); + assert_eq!( + first_turns.turns[0].items, + vec![ + expected_item("root-1", "root-user", /*rollout_ordinal*/ 2), + expected_item("root-1", "root-agent", /*rollout_ordinal*/ 3), + ] + ); + let second_turns = store + .list_turns(turn_params( + child_id, + first_turns.next_cursor, + /*page_size*/ 2, + SortDirection::Asc, + StoredTurnItemsView::NotLoaded, + )) + .await + .expect("second lineage turns page"); + assert_eq!(turn_ids(&second_turns), vec!["child-1"]); + let backwards_turns = store + .list_turns(turn_params( + child_id, + second_turns.backwards_cursor, + /*page_size*/ 2, + SortDirection::Desc, + StoredTurnItemsView::NotLoaded, + )) + .await + .expect("backwards lineage turns page"); + assert_eq!(turn_ids(&backwards_turns), vec!["child-1", "root-2"]); + + let first_items = store + .list_items(item_params( + child_id, + /*turn_id*/ None, + /*cursor*/ None, + /*page_size*/ 2, + SortDirection::Asc, + )) + .await + .expect("first lineage items page"); + assert_eq!(item_ids(&first_items), vec!["root-user", "root-agent"]); + let second_items = store + .list_items(item_params( + child_id, + /*turn_id*/ None, + first_items.next_cursor, + /*page_size*/ 2, + SortDirection::Asc, + )) + .await + .expect("second lineage items page"); + assert_eq!(item_ids(&second_items), vec!["root-2-item", "child-item"]); + let descending_items = store + .list_items(item_params( + child_id, + /*turn_id*/ None, + /*cursor*/ None, + /*page_size*/ 2, + SortDirection::Desc, + )) + .await + .expect("descending lineage items page"); + assert_eq!( + item_ids(&descending_items), + vec!["child-item", "root-2-item"] + ); + let inherited_turn_items = store + .list_items(item_params( + child_id, + Some("root-1"), + /*cursor*/ None, + /*page_size*/ 2, + SortDirection::Asc, + )) + .await + .expect("inherited turn item page"); + assert_eq!( + item_ids(&inherited_turn_items), + vec!["root-user", "root-agent"] + ); + + let (_other_home, other_store, other_thread_id) = + store_with_mode(ThreadHistoryMode::Paginated).await; + let error = other_store + .list_items(item_params( + other_thread_id, + /*turn_id*/ None, + second_items.backwards_cursor, + /*page_size*/ 2, + SortDirection::Asc, + )) + .await + .expect_err("lineage cursor belongs to requested thread"); + assert!(matches!(error, ThreadStoreError::InvalidRequest { .. })); +} + +#[tokio::test] +async fn lineage_reads_nested_forks() { + let (home, store, child_id) = store_with_mode(ThreadHistoryMode::Paginated).await; + let root_id = ThreadId::default(); + let middle_id = ThreadId::default(); + let root_path = rollout_path(home.path(), root_id); + write_rollout_with_end( + root_path.as_path(), + root_id, + /*history_base*/ None, + /*next_ordinal*/ 3, + ); + let middle_path = rollout_path(home.path(), middle_id); + write_rollout_with_end( + middle_path.as_path(), + middle_id, + Some(history_position( + root_path.as_path(), + root_id, + /*end_ordinal_exclusive*/ 3, + )), + /*next_ordinal*/ 2, + ); + write_rollout_with_end( + rollout_path(home.path(), child_id).as_path(), + child_id, + Some(history_position( + middle_path.as_path(), + middle_id, + /*end_ordinal_exclusive*/ 2, + )), + /*next_ordinal*/ 2, + ); + let db = history_db(&store).await; + for (thread_id, turn_id) in [ + (root_id, "root"), + (middle_id, "middle"), + (child_id, "child"), + ] { + insert_turn( + db, + thread_id, + turn_id, + /*rollout_ordinal*/ 1, + "completed", + /*error_json*/ None, + /*first_user_item_id*/ None, + /*final_agent_item_id*/ None, + ) + .await; + } + + let first_descending_page = store + .list_turns(turn_params( + child_id, + /*cursor*/ None, + /*page_size*/ 2, + SortDirection::Desc, + StoredTurnItemsView::NotLoaded, + )) + .await + .expect("first nested descending page"); + assert_eq!(turn_ids(&first_descending_page), vec!["child", "middle"]); + let second_descending_page = store + .list_turns(turn_params( + child_id, + first_descending_page.next_cursor, + /*page_size*/ 2, + SortDirection::Desc, + StoredTurnItemsView::NotLoaded, + )) + .await + .expect("second nested descending page"); + assert_eq!(turn_ids(&second_descending_page), vec!["root"]); +} + async fn store_with_mode(history_mode: ThreadHistoryMode) -> (TempDir, LocalThreadStore, ThreadId) { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); let thread_id = ThreadId::default(); + let rollout_path = rollout_path(home.path(), thread_id); + if history_mode == ThreadHistoryMode::Paginated { + write_rollout( + rollout_path.as_path(), + thread_id, + /*history_base*/ None, + ); + } let runtime = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.default_model_provider_id.clone(), @@ -262,7 +514,7 @@ async fn store_with_mode(history_mode: ThreadHistoryMode) -> (TempDir, LocalThre .expect("state runtime"); let mut builder = codex_state::ThreadMetadataBuilder::new( thread_id, - home.path().join("missing-rollout.jsonl"), + rollout_path, Utc::now(), SessionSource::Cli, ); @@ -275,6 +527,85 @@ async fn store_with_mode(history_mode: ThreadHistoryMode) -> (TempDir, LocalThre (home, store, thread_id) } +fn write_rollout( + path: &std::path::Path, + thread_id: ThreadId, + history_base: Option, +) { + write_rollout_with_end(path, thread_id, history_base, /*next_ordinal*/ 1); +} + +fn write_rollout_with_end( + path: &std::path::Path, + thread_id: ThreadId, + history_base: Option, + next_ordinal: u64, +) { + fs::create_dir_all(path.parent().expect("rollout parent")).expect("create rollout parent"); + let mut lines = vec![RolloutLine { + timestamp: "2026-07-16T00:00:00.000Z".to_string(), + ordinal: Some(0), + item: RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + history_mode: ThreadHistoryMode::Paginated, + history_base, + ..SessionMeta::default() + }, + git: None, + }), + }]; + for ordinal in 1..next_ordinal { + lines.push(RolloutLine { + timestamp: "2026-07-16T00:00:00.000Z".to_string(), + ordinal: Some(ordinal), + item: RolloutItem::EventMsg(EventMsg::ShutdownComplete), + }); + } + fs::write( + path, + format!( + "{}\n", + lines + .iter() + .map(|line| serde_json::to_string(line).expect("serialize rollout")) + .collect::>() + .join("\n") + ), + ) + .expect("write rollout"); +} + +fn rollout_path(home: &std::path::Path, thread_id: ThreadId) -> std::path::PathBuf { + home.join(format!( + "sessions/2026/07/16/rollout-2026-07-16T00-00-00-{thread_id}.jsonl" + )) +} + +fn history_position( + path: &std::path::Path, + thread_id: ThreadId, + end_ordinal_exclusive: u64, +) -> HistoryPosition { + HistoryPosition { + thread_id, + end_ordinal_exclusive, + end_byte_offset: rollout_end_byte_offset(path, end_ordinal_exclusive), + } +} + +fn rollout_end_byte_offset(path: &std::path::Path, end_ordinal_exclusive: u64) -> u64 { + let line_count = usize::try_from(end_ordinal_exclusive).expect("ordinal fits usize"); + let bytes = fs::read(path).expect("read rollout"); + let end_byte_offset = bytes + .split_inclusive(|byte| *byte == b'\n') + .take(line_count) + .map(<[u8]>::len) + .sum::(); + u64::try_from(end_byte_offset).expect("rollout byte offset fits u64") +} + async fn history_db(store: &LocalThreadStore) -> &sqlx::SqlitePool { store .thread_history_db() diff --git a/codex-rs/thread-store/src/local/thread_history/search.rs b/codex-rs/thread-store/src/local/thread_history/search.rs index 89b2a81d38..46e61dd192 100644 --- a/codex-rs/thread-store/src/local/thread_history/search.rs +++ b/codex-rs/thread-store/src/local/thread_history/search.rs @@ -14,6 +14,7 @@ use sqlx::Row; use super::super::LocalThreadStore; use super::read::CursorScope; +use super::read::PhysicalHistoryPosition; use super::read::serialize_cursor; use super::read::validate_thread_for_paginated_reads; use super::thread_history_error; @@ -140,8 +141,11 @@ ORDER BY rollout_ordinal ASC .saturating_sub(items.len()); let turn_cursor = serialize_cursor( params.thread_id, - &CursorScope::Turns, - row.turn_rollout_ordinal, + CursorScope::Turns, + PhysicalHistoryPosition { + physical_thread_id: params.thread_id, + rollout_ordinal: row.turn_rollout_ordinal, + }, /*include_anchor*/ true, )?; for (occurrence_index, matched) in matcher diff --git a/codex-rs/thread-store/src/local/thread_history/segment_paging.rs b/codex-rs/thread-store/src/local/thread_history/segment_paging.rs new file mode 100644 index 0000000000..b85d796222 --- /dev/null +++ b/codex-rs/thread-store/src/local/thread_history/segment_paging.rs @@ -0,0 +1,302 @@ +use codex_protocol::ThreadId; +use sqlx::QueryBuilder; +use sqlx::Sqlite; + +use super::super::rollout_lineage::RolloutLineage; +use super::super::rollout_lineage::RolloutLineageSegment; +use super::read::CursorScope; +use super::read::HistoryCursor; +use super::read::PhysicalHistoryPosition; +use super::read::StoredThreadItemRow; +use super::read::StoredTurnRow; +use super::read::invalid_cursor; +use super::read::parse_cursor; +use super::read::serialize_cursor; +use super::read::stored_thread_item_row_for_thread; +use super::read::stored_turn_row; +use super::thread_history_error; +use crate::SortDirection; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +pub(super) struct SegmentPage { + pub rows: Vec, + pub next_cursor: Option, + pub backwards_cursor: Option, +} + +pub(super) fn validate_page_size(page_size: usize) -> ThreadStoreResult<()> { + if page_size == 0 { + return Err(ThreadStoreError::InvalidRequest { + message: "page size must be positive".to_string(), + }); + } + let limit = page_size.checked_add(1).ok_or_else(page_size_too_large)?; + i64::try_from(limit).map_err(|_| page_size_too_large())?; + Ok(()) +} + +pub(super) async fn page_turn_rows( + pool: &sqlx::SqlitePool, + requested_thread_id: ThreadId, + lineage: &RolloutLineage, + cursor: Option<&str>, + page_size: usize, + direction: SortDirection, +) -> ThreadStoreResult> { + let cursor = parse_cursor(cursor, requested_thread_id, CursorScope::Turns)?; + let mut rows = Vec::new(); + for (segment, segment_cursor) in segments_from_cursor(lineage, direction, cursor.as_ref())? { + let remaining = remaining_limit(page_size, rows.len())?; + if remaining == 0 { + break; + } + let mut query = QueryBuilder::::new( + r#" +SELECT + turn_id, + rollout_ordinal, + status, + error_json, + started_at, + completed_at, + duration_ms, + first_user_item_id, + final_agent_item_id +FROM thread_turns +WHERE thread_id = + "#, + ); + query.push_bind(segment.thread_id().to_string()); + push_segment_range(&mut query, segment)?; + push_cursor_clause(&mut query, direction, segment_cursor)?; + push_order_and_limit(&mut query, direction, remaining); + rows.extend( + query + .build() + .fetch_all(pool) + .await + .map_err(thread_history_error)? + .into_iter() + .map(|row| stored_turn_row(segment.thread_id(), row)) + .collect::>>()?, + ); + } + finish_page(requested_thread_id, CursorScope::Turns, rows, page_size) +} + +pub(super) async fn page_item_rows( + pool: &sqlx::SqlitePool, + requested_thread_id: ThreadId, + lineage: &RolloutLineage, + turn_id: Option<&str>, + cursor: Option<&str>, + page_size: usize, + direction: SortDirection, +) -> ThreadStoreResult> { + let cursor = parse_cursor(cursor, requested_thread_id, CursorScope::Items)?; + let mut rows = Vec::new(); + for (segment, segment_cursor) in segments_from_cursor(lineage, direction, cursor.as_ref())? { + let remaining = remaining_limit(page_size, rows.len())?; + if remaining == 0 { + break; + } + let mut query = QueryBuilder::::new( + r#" +SELECT turn_id, item_id, rollout_ordinal, created_at_ms, item_json +FROM thread_items +WHERE thread_id = + "#, + ); + query.push_bind(segment.thread_id().to_string()); + push_segment_range(&mut query, segment)?; + if let Some(turn_id) = turn_id { + query.push(" AND turn_id = ").push_bind(turn_id); + } + push_cursor_clause(&mut query, direction, segment_cursor)?; + push_order_and_limit(&mut query, direction, remaining); + rows.extend( + query + .build() + .fetch_all(pool) + .await + .map_err(thread_history_error)? + .into_iter() + .map(|row| stored_thread_item_row_for_thread(segment.thread_id(), row)) + .collect::>>()?, + ); + } + finish_page(requested_thread_id, CursorScope::Items, rows, page_size) +} + +fn segments_from_cursor<'a>( + lineage: &'a RolloutLineage, + direction: SortDirection, + cursor: Option<&'a HistoryCursor>, +) -> ThreadStoreResult)>> { + let segments = lineage.segments(); + let cursor_index = cursor + .map(|cursor| { + segments + .iter() + .position(|segment| segment.thread_id() == cursor.physical_thread_id) + .ok_or_else(|| invalid_cursor("unknown physical segment")) + }) + .transpose()?; + if let Some(cursor) = cursor + && let Some(index) = cursor_index + && !cursor_in_segment(cursor, &segments[index]) + { + return Err(invalid_cursor("position outside physical segment")); + } + let indexes: Vec = match direction { + SortDirection::Asc => (cursor_index.unwrap_or(0)..segments.len()).collect(), + SortDirection::Desc => { + let end = cursor_index.unwrap_or_else(|| segments.len().saturating_sub(1)); + (0..=end).rev().collect() + } + }; + Ok(indexes + .into_iter() + .map(|index| { + let segment_cursor = if Some(index) == cursor_index { + cursor + } else { + None + }; + (&segments[index], segment_cursor) + }) + .collect()) +} + +fn cursor_in_segment(cursor: &HistoryCursor, segment: &RolloutLineageSegment) -> bool { + let ordinal = cursor.rollout_ordinal; + ordinal >= segment.start_ordinal() + && segment + .end_ordinal() + .is_none_or(|end_ordinal| ordinal < end_ordinal) +} + +fn push_segment_range( + query: &mut QueryBuilder, + segment: &RolloutLineageSegment, +) -> ThreadStoreResult<()> { + query + .push(" AND rollout_ordinal >= ") + .push_bind(sqlite_integer(segment.start_ordinal())?); + if let Some(end_ordinal) = segment.end_ordinal() { + query + .push(" AND rollout_ordinal < ") + .push_bind(sqlite_integer(end_ordinal)?); + } + Ok(()) +} + +fn push_cursor_clause( + query: &mut QueryBuilder, + direction: SortDirection, + cursor: Option<&HistoryCursor>, +) -> ThreadStoreResult<()> { + if let Some(cursor) = cursor { + let comparator = match (direction, cursor.include_anchor) { + (SortDirection::Asc, true) => ">=", + (SortDirection::Asc, false) => ">", + (SortDirection::Desc, true) => "<=", + (SortDirection::Desc, false) => "<", + }; + query + .push(" AND rollout_ordinal ") + .push(comparator) + .push(" ") + .push_bind(sqlite_integer(cursor.rollout_ordinal)?); + } + Ok(()) +} + +fn push_order_and_limit(query: &mut QueryBuilder, direction: SortDirection, limit: i64) { + let order = match direction { + SortDirection::Asc => "ASC", + SortDirection::Desc => "DESC", + }; + query + .push(" ORDER BY rollout_ordinal ") + .push(order) + .push(" LIMIT ") + .push_bind(limit); +} + +fn remaining_limit(page_size: usize, row_count: usize) -> ThreadStoreResult { + let limit = page_size + .checked_add(1) + .and_then(|limit| limit.checked_sub(row_count)) + .ok_or_else(page_size_too_large)?; + i64::try_from(limit).map_err(|_| page_size_too_large()) +} + +fn finish_page( + requested_thread_id: ThreadId, + scope: CursorScope, + mut rows: Vec, + page_size: usize, +) -> ThreadStoreResult> { + let has_more = rows.len() > page_size; + rows.truncate(page_size); + let backwards_cursor = rows + .first() + .map(|row| { + serialize_cursor( + requested_thread_id, + scope.clone(), + row.position(), + /*include_anchor*/ true, + ) + }) + .transpose()?; + let next_cursor = if has_more { + rows.last() + .map(|row| { + serialize_cursor( + requested_thread_id, + scope, + row.position(), + /*include_anchor*/ false, + ) + }) + .transpose()? + } else { + None + }; + Ok(SegmentPage { + rows, + next_cursor, + backwards_cursor, + }) +} + +trait HasPosition { + fn position(&self) -> PhysicalHistoryPosition; +} + +impl HasPosition for StoredTurnRow { + fn position(&self) -> PhysicalHistoryPosition { + self.position + } +} + +impl HasPosition for StoredThreadItemRow { + fn position(&self) -> PhysicalHistoryPosition { + self.position + } +} + +fn sqlite_integer(value: u64) -> ThreadStoreResult { + i64::try_from(value).map_err(|_| ThreadStoreError::InvalidRequest { + message: "rollout ordinal exceeds SQLite integer range".to_string(), + }) +} + +fn page_size_too_large() -> ThreadStoreError { + ThreadStoreError::InvalidRequest { + message: "page size is too large".to_string(), + } +}