From 332eac4b858a40575a55f21026ecf8cb662e539c Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 31 Jul 2026 20:05:42 +0000 Subject: [PATCH] Load turn summaries with paginated queries (#36384) ## Why Loading the summary view issued a separate item query for every returned turn. ## What changed Join each turn's first user and final agent items into the paginated turn query when `StoredTurnItemsView::Summary` is requested. Preserve the existing fallback that resolves inherited summary items for synthetic fork-boundary rows. GitOrigin-RevId: add7f2e03c0809bbee85dc9f07bb6f91af92dfe7 --- .../src/local/thread_history/read.rs | 161 +++++++++++++----- .../local/thread_history/segment_paging.rs | 113 +++++++++++- 2 files changed, 225 insertions(+), 49 deletions(-) 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 bc58e329c6..59d817ee94 100644 --- a/codex-rs/thread-store/src/local/thread_history/read.rs +++ b/codex-rs/thread-store/src/local/thread_history/read.rs @@ -46,7 +46,6 @@ pub(super) struct HistoryCursor { #[derive(Clone, Copy)] pub(super) struct PhysicalHistoryPosition { - pub physical_thread_id: ThreadId, pub rollout_ordinal: i64, } @@ -60,6 +59,32 @@ pub(super) struct StoredTurnRow { pub duration_ms: Option, pub first_user_item_id: Option, pub final_agent_item_id: Option, + pub summary_items: Vec, +} + +#[derive(sqlx::FromRow)] +pub(super) struct StoredSummaryColumns { + summary_first_user_turn_id: Option, + summary_first_user_item_id: Option, + summary_first_user_rollout_ordinal: Option, + summary_first_user_updated_at_ordinal: Option, + summary_first_user_created_at_ms: Option, + summary_first_user_item_json: Option, + summary_final_agent_turn_id: Option, + summary_final_agent_item_id: Option, + summary_final_agent_rollout_ordinal: Option, + summary_final_agent_updated_at_ordinal: Option, + summary_final_agent_created_at_ms: Option, + summary_final_agent_item_json: Option, +} + +struct StoredSummaryItemColumns { + turn_id: Option, + item_id: Option, + rollout_ordinal: Option, + updated_at_ordinal: Option, + created_at_ms: Option, + item_json: Option, } pub(super) struct StoredThreadItemRow { @@ -88,13 +113,23 @@ pub(in crate::local) async fn list_turns( params.cursor.as_deref(), params.page_size, params.sort_direction, + params.items_view, ) .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, &lineage, &turn).await?, + StoredTurnItemsView::Summary + if matches!(turn.status, StoredTurnStatus::Interrupted) + && turn.first_user_item_id.is_none() + && turn.final_agent_item_id.is_none() => + { + // Synthetic fork-boundary rows are interrupted without local summary IDs. + // Load their summary from the earliest visible source turn. + load_inherited_summary_items(pool, &lineage, &turn).await? + } + StoredTurnItemsView::Summary => turn.summary_items, }; turns.push(StoredTurn { turn_id: turn.turn_id, @@ -168,33 +203,16 @@ pub(super) async fn validate_thread_for_paginated_reads( } } -async fn load_summary_items( +async fn load_inherited_summary_items( pool: &sqlx::SqlitePool, lineage: &RolloutLineage, turn: &StoredTurnRow, ) -> ThreadStoreResult> { - let (item_owner, first_user_item_id, final_agent_item_id) = - if matches!(turn.status, StoredTurnStatus::Interrupted) - && turn.first_user_item_id.is_none() - && turn.final_agent_item_id.is_none() - { - let source = find_source_turn(pool, lineage, turn.turn_id.as_str()).await?; - ( - source.physical_thread_id, - source.first_user_item_id, - source.final_agent_item_id, - ) - } else { - ( - turn.position.physical_thread_id, - turn.first_user_item_id.clone(), - turn.final_agent_item_id.clone(), - ) - }; + let source = find_source_turn(pool, lineage, turn.turn_id.as_str()).await?; let Some(segment) = lineage .segments() .iter() - .find(|segment| segment.thread_id() == item_owner) + .find(|segment| segment.thread_id() == source.physical_thread_id) else { return Ok(Vec::new()); }; @@ -215,13 +233,13 @@ WHERE thread_id = ? ORDER BY rollout_ordinal ASC "#, ) - .bind(item_owner.to_string()) + .bind(source.physical_thread_id.to_string()) .bind(turn.turn_id.as_str()) .bind(start_ordinal) .bind(end_ordinal) .bind(end_ordinal) - .bind(first_user_item_id.as_deref()) - .bind(final_agent_item_id.as_deref()) + .bind(source.first_user_item_id) + .bind(source.final_agent_item_id) .fetch_all(pool) .await .map_err(super::thread_history_error)?; @@ -261,10 +279,7 @@ pub(super) fn serialize_cursor( .map_err(super::thread_history_error) } -pub(super) fn stored_turn_row( - physical_thread_id: ThreadId, - row: sqlx::sqlite::SqliteRow, -) -> ThreadStoreResult { +pub(super) fn stored_turn_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { let status = match row.try_get::("status")?.as_str() { "completed" => StoredTurnStatus::Completed, "interrupted" => StoredTurnStatus::Interrupted, @@ -284,7 +299,6 @@ pub(super) fn stored_turn_row( .map_err(super::thread_history_error)?; Ok(StoredTurnRow { position: PhysicalHistoryPosition { - physical_thread_id, rollout_ordinal: row.try_get("rollout_ordinal")?, }, turn_id: row.try_get("turn_id")?, @@ -295,11 +309,77 @@ pub(super) fn stored_turn_row( duration_ms: row.try_get("duration_ms")?, first_user_item_id: row.try_get("first_user_item_id")?, final_agent_item_id: row.try_get("final_agent_item_id")?, + summary_items: Vec::new(), }) } -pub(super) fn stored_thread_item_row_for_thread( - physical_thread_id: ThreadId, +impl StoredSummaryColumns { + pub(super) fn into_stored_items(self) -> ThreadStoreResult> { + let mut summary_items = [ + StoredSummaryItemColumns { + turn_id: self.summary_first_user_turn_id, + item_id: self.summary_first_user_item_id, + rollout_ordinal: self.summary_first_user_rollout_ordinal, + updated_at_ordinal: self.summary_first_user_updated_at_ordinal, + created_at_ms: self.summary_first_user_created_at_ms, + item_json: self.summary_first_user_item_json, + } + .into_stored_item()?, + StoredSummaryItemColumns { + turn_id: self.summary_final_agent_turn_id, + item_id: self.summary_final_agent_item_id, + rollout_ordinal: self.summary_final_agent_rollout_ordinal, + updated_at_ordinal: self.summary_final_agent_updated_at_ordinal, + created_at_ms: self.summary_final_agent_created_at_ms, + item_json: self.summary_final_agent_item_json, + } + .into_stored_item()?, + ] + .into_iter() + .flatten() + .collect::>(); + summary_items.sort_by_key(|(rollout_ordinal, _)| *rollout_ordinal); + Ok(summary_items.into_iter().map(|(_, item)| item).collect()) + } +} + +impl StoredSummaryItemColumns { + fn into_stored_item(self) -> ThreadStoreResult> { + let Some(item_id) = self.item_id else { + return Ok(None); + }; + let ( + Some(turn_id), + Some(rollout_ordinal), + Some(updated_at_ordinal), + Some(created_at_ms), + Some(item_json), + ) = ( + self.turn_id, + self.rollout_ordinal, + self.updated_at_ordinal, + self.created_at_ms, + self.item_json, + ) + else { + return Err(ThreadStoreError::Internal { + message: "stored summary item is missing joined columns".to_string(), + }); + }; + Ok(Some(( + rollout_ordinal, + StoredThreadItem { + turn_id, + item_id, + updated_at_ordinal: stored_updated_at_ordinal(updated_at_ordinal)?, + created_at_ms, + item_json: item_json.into_bytes(), + }, + ))) + } +} + +pub(super) fn stored_thread_item_row( row: sqlx::sqlite::SqliteRow, ) -> ThreadStoreResult { let rollout_ordinal = row.try_get::("rollout_ordinal")?; @@ -309,20 +389,13 @@ pub(super) fn stored_thread_item_row_for_thread( }); } Ok(StoredThreadItemRow { - position: PhysicalHistoryPosition { - physical_thread_id, - rollout_ordinal, - }, + position: PhysicalHistoryPosition { rollout_ordinal }, item: stored_thread_item(row)?, }) } fn stored_thread_item(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { - let updated_at_ordinal = row.try_get::("updated_at_ordinal")?; - let updated_at_ordinal = - u64::try_from(updated_at_ordinal).map_err(|_| ThreadStoreError::Internal { - message: format!("invalid stored item updated-at ordinal: {updated_at_ordinal}"), - })?; + let updated_at_ordinal = stored_updated_at_ordinal(row.try_get("updated_at_ordinal")?)?; Ok(StoredThreadItem { turn_id: row.try_get("turn_id")?, item_id: row.try_get("item_id")?, @@ -332,6 +405,12 @@ fn stored_thread_item(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult ThreadStoreResult { + u64::try_from(updated_at_ordinal).map_err(|_| ThreadStoreError::Internal { + message: format!("invalid stored item updated-at ordinal: {updated_at_ordinal}"), + }) +} + 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/segment_paging.rs b/codex-rs/thread-store/src/local/thread_history/segment_paging.rs index fb7ed411fe..9171c1c5a6 100644 --- a/codex-rs/thread-store/src/local/thread_history/segment_paging.rs +++ b/codex-rs/thread-store/src/local/thread_history/segment_paging.rs @@ -1,4 +1,5 @@ use codex_protocol::ThreadId; +use sqlx::FromRow; use sqlx::QueryBuilder; use sqlx::Sqlite; @@ -7,17 +8,19 @@ use super::super::rollout_lineage::RolloutLineageSegment; use super::read::CursorScope; use super::read::HistoryCursor; use super::read::PhysicalHistoryPosition; +use super::read::StoredSummaryColumns; 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_thread_item_row; use super::read::stored_turn_row; use super::thread_history_error; use crate::ItemSortKey; use crate::ListItemsParams; use crate::SortDirection; +use crate::StoredTurnItemsView; use crate::ThreadStoreError; use crate::ThreadStoreResult; @@ -45,6 +48,7 @@ pub(super) async fn page_turn_rows( cursor: Option<&str>, page_size: usize, direction: SortDirection, + items_view: StoredTurnItemsView, ) -> ThreadStoreResult> { let cursor = parse_cursor(cursor, requested_thread_id, CursorScope::Turns)?; let mut rows = Vec::new(); @@ -55,8 +59,10 @@ pub(super) async fn page_turn_rows( if remaining == 0 { break; } - let mut query = QueryBuilder::::new( - r#" + let mut query = + QueryBuilder::::new(if matches!(items_view, StoredTurnItemsView::Summary) { + r#" +WITH page_turns AS ( SELECT turn_id, rollout_ordinal, @@ -69,8 +75,23 @@ SELECT final_agent_item_id FROM thread_turns WHERE thread_id = - "#, - ); + "# + } else { + 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)?; for newer_segment in &lineage.segments()[segment_index + 1..] { @@ -88,6 +109,45 @@ WHERE thread_id = } push_cursor_clause(&mut query, direction, segment_cursor)?; push_order_and_limit(&mut query, direction, remaining); + if matches!(items_view, StoredTurnItemsView::Summary) { + query.push( + r#" +) +SELECT + page_turns.*, + first_user.turn_id AS summary_first_user_turn_id, + first_user.item_id AS summary_first_user_item_id, + first_user.rollout_ordinal AS summary_first_user_rollout_ordinal, + first_user.updated_at_ordinal AS summary_first_user_updated_at_ordinal, + first_user.created_at_ms AS summary_first_user_created_at_ms, + first_user.item_json AS summary_first_user_item_json, + final_agent.turn_id AS summary_final_agent_turn_id, + final_agent.item_id AS summary_final_agent_item_id, + final_agent.rollout_ordinal AS summary_final_agent_rollout_ordinal, + final_agent.updated_at_ordinal AS summary_final_agent_updated_at_ordinal, + final_agent.created_at_ms AS summary_final_agent_created_at_ms, + final_agent.item_json AS summary_final_agent_item_json +FROM page_turns +LEFT JOIN thread_items AS first_user + ON first_user.thread_id = + "#, + ); + push_summary_item_join(&mut query, segment, "first_user", "first_user_item_id")?; + query.push( + r#" +LEFT JOIN thread_items AS final_agent + ON final_agent.thread_id = + "#, + ); + push_summary_item_join(&mut query, segment, "final_agent", "final_agent_item_id")?; + let order = match direction { + SortDirection::Asc => "ASC", + SortDirection::Desc => "DESC", + }; + query + .push(" ORDER BY page_turns.rollout_ordinal ") + .push(order); + } rows.extend( query .build() @@ -95,13 +155,50 @@ WHERE thread_id = .await .map_err(thread_history_error)? .into_iter() - .map(|row| stored_turn_row(segment.thread_id(), row)) + .map(|row| { + let summary_items = if matches!(items_view, StoredTurnItemsView::Summary) { + StoredSummaryColumns::from_row(&row)?.into_stored_items()? + } else { + Vec::new() + }; + let mut turn = stored_turn_row(row)?; + turn.summary_items = summary_items; + Ok(turn) + }) .collect::>>()?, ); } finish_page(requested_thread_id, CursorScope::Turns, rows, page_size) } +fn push_summary_item_join( + query: &mut QueryBuilder, + segment: &RolloutLineageSegment, + alias: &str, + item_id_column: &str, +) -> ThreadStoreResult<()> { + query + .push_bind(segment.thread_id().to_string()) + .push(" AND ") + .push(alias) + .push(".turn_id = page_turns.turn_id AND ") + .push(alias) + .push(".item_id = page_turns.") + .push(item_id_column) + .push(" AND ") + .push(alias) + .push(".rollout_ordinal >= ") + .push_bind(sqlite_integer(segment.start_ordinal())?); + if let Some(end_ordinal) = segment.end_ordinal() { + query + .push(" AND ") + .push(alias) + .push(".rollout_ordinal < ") + .push_bind(sqlite_integer(end_ordinal)?); + } + Ok(()) +} + pub(super) async fn page_item_rows( pool: &sqlx::SqlitePool, lineage: &RolloutLineage, @@ -161,7 +258,7 @@ WHERE thread_id = .await .map_err(thread_history_error)? .into_iter() - .map(|row| stored_thread_item_row_for_thread(segment.thread_id(), row)) + .map(stored_thread_item_row) .collect::>>()?, ); } @@ -225,7 +322,7 @@ WHERE thread_id = .await .map_err(thread_history_error)? .into_iter() - .map(|row| stored_thread_item_row_for_thread(params.thread_id, row)) + .map(stored_thread_item_row) .collect::>>()?; finish_page( params.thread_id,