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

Load turn summaries with paginated queries (#36384)

Browse Source 
## 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
This commit is contained in:
Owen Lin
2026-07-31 20:12:28 +00:00
committed by copyberry
parent c42ea41ee0
commit 332eac4b85
2 changed files with 225 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/thread-store/src/local/thread_history/read.rs
+120 -41
View File
@@ -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<i64>,
pub first_user_item_id: Option<String>,
pub final_agent_item_id: Option<String>,
pub summary_items: Vec<StoredThreadItem>,
}
#[derive(sqlx::FromRow)]
pub(super) struct StoredSummaryColumns {
summary_first_user_turn_id: Option<String>,
summary_first_user_item_id: Option<String>,
summary_first_user_rollout_ordinal: Option<i64>,
summary_first_user_updated_at_ordinal: Option<i64>,
summary_first_user_created_at_ms: Option<i64>,
summary_first_user_item_json: Option<String>,
summary_final_agent_turn_id: Option<String>,
summary_final_agent_item_id: Option<String>,
summary_final_agent_rollout_ordinal: Option<i64>,
summary_final_agent_updated_at_ordinal: Option<i64>,
summary_final_agent_created_at_ms: Option<i64>,
summary_final_agent_item_json: Option<String>,
}
struct StoredSummaryItemColumns {
turn_id: Option<String>,
item_id: Option<String>,
rollout_ordinal: Option<i64>,
updated_at_ordinal: Option<i64>,
created_at_ms: Option<i64>,
item_json: Option<String>,
}
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<Vec<StoredThreadItem>> {
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<StoredTurnRow> {
pub(super) fn stored_turn_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult<StoredTurnRow> {
let status = match row.try_get::<String, _>("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<Vec<StoredThreadItem>> {
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::<Vec<_>>();
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<Option<(i64, StoredThreadItem)>> {
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<StoredThreadItemRow> {
let rollout_ordinal = row.try_get::<i64, _>("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<StoredThreadItem> {
let updated_at_ordinal = row.try_get::<i64, _>("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<StoredT
})
}
fn stored_updated_at_ordinal(updated_at_ordinal: i64) -> ThreadStoreResult<u64> {
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}"),
codex-rs/thread-store/src/local/thread_history/segment_paging.rs
+105 -8
View File
@@ -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<SegmentPage<StoredTurnRow>> {
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::<Sqlite>::new(
r#"
let mut query =
QueryBuilder::<Sqlite>::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::<ThreadStoreResult<Vec<_>>>()?,
);
}
finish_page(requested_thread_id, CursorScope::Turns, rows, page_size)
}
fn push_summary_item_join(
query: &mut QueryBuilder<Sqlite>,
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::<ThreadStoreResult<Vec<_>>>()?,
);
}
@@ -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::<ThreadStoreResult<Vec<_>>>()?;
finish_page(
params.thread_id,