Add local thread turn item views

This commit is contained in:
Tom Wiltzius
2026-05-14 16:18:45 -07:00
parent 51b0e94105
commit 594eaa1904
3 changed files with 663 additions and 0 deletions
@@ -0,0 +1,457 @@
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::TurnContextItem;
use serde::Deserialize;
use serde::Serialize;
use super::LocalThreadStore;
use super::read_thread;
use crate::ListTurnsParams;
use crate::ReadThreadParams;
use crate::SortDirection;
use crate::StoredTurn;
use crate::StoredTurnError;
use crate::StoredTurnItemsView;
use crate::StoredTurnStatus;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
use crate::TurnPage;
const DEFAULT_TURN_PAGE_SIZE: usize = 50;
const MAX_TURN_PAGE_SIZE: usize = 200;
pub(super) async fn list_turns(
store: &LocalThreadStore,
params: ListTurnsParams,
) -> ThreadStoreResult<TurnPage> {
let thread = read_thread::read_thread(
store,
ReadThreadParams {
thread_id: params.thread_id,
include_archived: params.include_archived,
include_history: true,
},
)
.await?;
let history = thread.history.ok_or_else(|| ThreadStoreError::Internal {
message: format!("failed to load history for thread {}", params.thread_id),
})?;
let turns = build_stored_turns_from_rollout_items(&history.items, params.items_view);
paginate_turns(
turns,
params.cursor.as_deref(),
params.page_size,
params.sort_direction,
)
}
fn build_stored_turns_from_rollout_items(
items: &[RolloutItem],
items_view: StoredTurnItemsView,
) -> Vec<StoredTurn> {
let mut builder = StoredTurnBuilder::new(items_view);
for (rollout_index, item) in items.iter().enumerate() {
builder.handle_rollout_item(rollout_index, item);
}
builder.finish()
}
struct StoredTurnBuilder {
turns: Vec<PendingStoredTurn>,
current_turn: Option<PendingStoredTurn>,
items_view: StoredTurnItemsView,
}
impl StoredTurnBuilder {
fn new(items_view: StoredTurnItemsView) -> Self {
Self {
turns: Vec::new(),
current_turn: None,
items_view,
}
}
fn finish(mut self) -> Vec<StoredTurn> {
self.finish_current_turn();
self.turns
.into_iter()
.map(|turn| turn.into_stored_turn(self.items_view))
.collect()
}
fn handle_rollout_item(&mut self, rollout_index: usize, item: &RolloutItem) {
match item {
RolloutItem::SessionMeta(_) => {}
RolloutItem::TurnContext(payload) => {
self.handle_turn_context(rollout_index, payload, item);
}
RolloutItem::Compacted(_) => {
let turn = self.ensure_turn(rollout_index);
turn.saw_compaction = true;
turn.items.push(item.clone());
}
RolloutItem::ResponseItem(payload) => {
if response_item_starts_turn(payload) {
self.maybe_finish_implicit_turn();
}
self.ensure_turn(rollout_index).items.push(item.clone());
}
RolloutItem::EventMsg(payload) => self.handle_event(rollout_index, payload, item),
}
}
fn handle_turn_context(
&mut self,
rollout_index: usize,
payload: &TurnContextItem,
item: &RolloutItem,
) {
if let Some(turn_id) = payload.turn_id.as_deref()
&& let Some(turn) = self.find_turn_mut(turn_id)
{
turn.items.push(item.clone());
return;
}
let turn = self.ensure_turn(rollout_index);
if turn.items.is_empty()
&& let Some(turn_id) = payload.turn_id.as_deref()
{
turn.turn_id = turn_id.to_string();
}
turn.items.push(item.clone());
}
fn handle_event(&mut self, rollout_index: usize, event: &EventMsg, item: &RolloutItem) {
match event {
EventMsg::TurnStarted(payload) => {
self.finish_current_turn();
let mut turn = PendingStoredTurn::new(payload.turn_id.clone());
turn.status = StoredTurnStatus::InProgress;
turn.started_at = payload.started_at;
turn.opened_explicitly = true;
turn.items.push(item.clone());
self.current_turn = Some(turn);
}
EventMsg::TurnComplete(payload) => {
let should_finish_current = self
.current_turn
.as_ref()
.is_some_and(|turn| turn.turn_id == payload.turn_id);
let turn = self.turn_for_completion(rollout_index, payload.turn_id.as_str());
turn.items.push(item.clone());
if matches!(
turn.status,
StoredTurnStatus::Completed | StoredTurnStatus::InProgress
) {
turn.status = StoredTurnStatus::Completed;
}
turn.completed_at = payload.completed_at;
turn.duration_ms = payload.duration_ms;
if should_finish_current {
self.finish_current_turn();
}
}
EventMsg::TurnAborted(payload) => {
let should_finish_current = payload.turn_id.as_ref().is_none_or(|turn_id| {
self.current_turn
.as_ref()
.is_some_and(|turn| turn.turn_id == *turn_id)
});
let turn = match payload.turn_id.as_deref() {
Some(turn_id) => self.turn_for_completion(rollout_index, turn_id),
None => self.ensure_turn(rollout_index),
};
turn.items.push(item.clone());
turn.status = StoredTurnStatus::Interrupted;
turn.completed_at = payload.completed_at;
turn.duration_ms = payload.duration_ms;
if should_finish_current {
self.finish_current_turn();
}
}
EventMsg::ThreadRolledBack(payload) => {
self.finish_current_turn();
let n = usize::try_from(payload.num_turns).unwrap_or(usize::MAX);
if n >= self.turns.len() {
self.turns.clear();
} else {
self.turns.truncate(self.turns.len().saturating_sub(n));
}
}
EventMsg::UserMessage(_) => {
self.maybe_finish_implicit_turn();
self.ensure_turn(rollout_index).items.push(item.clone());
}
EventMsg::Error(payload) => {
let turn = self.ensure_turn(rollout_index);
turn.items.push(item.clone());
if payload.affects_turn_status() {
turn.status = StoredTurnStatus::Failed;
turn.error = Some(StoredTurnError {
message: payload.message.clone(),
additional_details: None,
});
}
}
_ => {
self.ensure_turn(rollout_index).items.push(item.clone());
}
}
}
fn ensure_turn(&mut self, rollout_index: usize) -> &mut PendingStoredTurn {
self.current_turn
.get_or_insert_with(|| PendingStoredTurn::new(format!("rollout-{rollout_index}")))
}
fn finish_current_turn(&mut self) {
if let Some(turn) = self.current_turn.take()
&& (!turn.items.is_empty() || turn.opened_explicitly || turn.saw_compaction)
{
self.turns.push(turn);
}
}
fn maybe_finish_implicit_turn(&mut self) {
if self
.current_turn
.as_ref()
.is_some_and(|turn| !turn.opened_explicitly && !turn.items.is_empty())
{
self.finish_current_turn();
}
}
fn turn_for_completion(
&mut self,
rollout_index: usize,
turn_id: &str,
) -> &mut PendingStoredTurn {
if self
.current_turn
.as_ref()
.is_some_and(|turn| turn.turn_id == turn_id)
{
let Some(turn) = self.current_turn.as_mut() else {
unreachable!("current turn exists after matching above");
};
return turn;
}
if let Some(index) = self.turns.iter().position(|turn| turn.turn_id == turn_id) {
return &mut self.turns[index];
}
self.ensure_turn(rollout_index)
}
fn find_turn_mut(&mut self, turn_id: &str) -> Option<&mut PendingStoredTurn> {
if self
.current_turn
.as_ref()
.is_some_and(|turn| turn.turn_id == turn_id)
{
return self.current_turn.as_mut();
}
self.turns.iter_mut().find(|turn| turn.turn_id == turn_id)
}
}
struct PendingStoredTurn {
turn_id: String,
items: Vec<RolloutItem>,
status: StoredTurnStatus,
error: Option<StoredTurnError>,
started_at: Option<i64>,
completed_at: Option<i64>,
duration_ms: Option<i64>,
opened_explicitly: bool,
saw_compaction: bool,
}
impl PendingStoredTurn {
fn new(turn_id: String) -> Self {
Self {
turn_id,
items: Vec::new(),
status: StoredTurnStatus::Completed,
error: None,
started_at: None,
completed_at: None,
duration_ms: None,
opened_explicitly: false,
saw_compaction: false,
}
}
fn into_stored_turn(self, items_view: StoredTurnItemsView) -> StoredTurn {
StoredTurn {
turn_id: self.turn_id,
items: filter_items_for_view(self.items, items_view),
items_view,
status: self.status,
error: self.error,
started_at: self.started_at,
completed_at: self.completed_at,
duration_ms: self.duration_ms,
}
}
}
fn filter_items_for_view(
items: Vec<RolloutItem>,
items_view: StoredTurnItemsView,
) -> Vec<RolloutItem> {
match items_view {
StoredTurnItemsView::NotLoaded => Vec::new(),
StoredTurnItemsView::Summary => summary_items(items),
StoredTurnItemsView::Full
| StoredTurnItemsView::ResponseItems
| StoredTurnItemsView::EventItems
| StoredTurnItemsView::ResponseAndEventItems => items
.into_iter()
.filter(|item| items_view.includes_persisted_item(item))
.collect(),
}
}
fn summary_items(items: Vec<RolloutItem>) -> Vec<RolloutItem> {
let first_user_index = items.iter().position(is_user_message_item);
let final_agent_index = items.iter().rposition(is_agent_message_item);
let mut summary = Vec::new();
if let Some(index) = first_user_index {
summary.push(items[index].clone());
}
if let Some(index) = final_agent_index
&& Some(index) != first_user_index
{
summary.push(items[index].clone());
}
summary
}
fn response_item_starts_turn(item: &ResponseItem) -> bool {
matches!(
item,
ResponseItem::Message { role, .. } if role == "user"
)
}
fn is_user_message_item(item: &RolloutItem) -> bool {
match item {
RolloutItem::EventMsg(EventMsg::UserMessage(_)) => true,
RolloutItem::ResponseItem(ResponseItem::Message { role, .. }) => role == "user",
_ => false,
}
}
fn is_agent_message_item(item: &RolloutItem) -> bool {
match item {
RolloutItem::EventMsg(EventMsg::AgentMessage(_)) => true,
RolloutItem::ResponseItem(ResponseItem::Message { role, .. }) => role == "assistant",
_ => false,
}
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct TurnCursor {
turn_id: String,
include_anchor: bool,
}
fn paginate_turns(
turns: Vec<StoredTurn>,
cursor: Option<&str>,
page_size: usize,
sort_direction: SortDirection,
) -> ThreadStoreResult<TurnPage> {
if turns.is_empty() {
return Ok(TurnPage {
turns: Vec::new(),
next_cursor: None,
backwards_cursor: None,
});
}
let anchor = cursor.map(parse_turn_cursor).transpose()?;
let page_size = if page_size == 0 {
DEFAULT_TURN_PAGE_SIZE
} else {
page_size
}
.clamp(1, MAX_TURN_PAGE_SIZE);
let anchor_index = anchor
.as_ref()
.and_then(|anchor| turns.iter().position(|turn| turn.turn_id == anchor.turn_id));
if anchor.is_some() && anchor_index.is_none() {
return Err(ThreadStoreError::InvalidRequest {
message: "invalid cursor: anchor turn is no longer present".to_string(),
});
}
let mut keyed_turns: Vec<_> = turns.into_iter().enumerate().collect();
match sort_direction {
SortDirection::Asc => {
if let (Some(anchor), Some(anchor_index)) = (anchor.as_ref(), anchor_index) {
keyed_turns.retain(|(index, _)| {
if anchor.include_anchor {
*index >= anchor_index
} else {
*index > anchor_index
}
});
}
}
SortDirection::Desc => {
keyed_turns.reverse();
if let (Some(anchor), Some(anchor_index)) = (anchor.as_ref(), anchor_index) {
keyed_turns.retain(|(index, _)| {
if anchor.include_anchor {
*index <= anchor_index
} else {
*index < anchor_index
}
});
}
}
}
let more_turns_available = keyed_turns.len() > page_size;
keyed_turns.truncate(page_size);
let backwards_cursor = keyed_turns
.first()
.map(|(_, turn)| serialize_turn_cursor(&turn.turn_id, /*include_anchor*/ true))
.transpose()?;
let next_cursor = if more_turns_available {
keyed_turns
.last()
.map(|(_, turn)| serialize_turn_cursor(&turn.turn_id, /*include_anchor*/ false))
.transpose()?
} else {
None
};
let turns = keyed_turns.into_iter().map(|(_, turn)| turn).collect();
Ok(TurnPage {
turns,
next_cursor,
backwards_cursor,
})
}
fn serialize_turn_cursor(turn_id: &str, include_anchor: bool) -> ThreadStoreResult<String> {
serde_json::to_string(&TurnCursor {
turn_id: turn_id.to_string(),
include_anchor,
})
.map_err(|err| ThreadStoreError::Internal {
message: format!("failed to serialize cursor: {err}"),
})
}
fn parse_turn_cursor(cursor: &str) -> ThreadStoreResult<TurnCursor> {
serde_json::from_str(cursor).map_err(|_| ThreadStoreError::InvalidRequest {
message: format!("invalid cursor: {cursor}"),
})
}
+139
View File
@@ -2,6 +2,7 @@ mod archive_thread;
mod create_thread;
mod helpers;
mod list_threads;
mod list_turns;
mod live_writer;
mod read_thread;
mod unarchive_thread;
@@ -24,6 +25,7 @@ use crate::AppendThreadItemsParams;
use crate::ArchiveThreadParams;
use crate::CreateThreadParams;
use crate::ListThreadsParams;
use crate::ListTurnsParams;
use crate::LoadThreadHistoryParams;
use crate::ReadThreadByRolloutPathParams;
use crate::ReadThreadParams;
@@ -34,6 +36,7 @@ use crate::ThreadPage;
use crate::ThreadStore;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
use crate::TurnPage;
use crate::UpdateThreadMetadataParams;
/// Local filesystem/SQLite-backed implementation of [`ThreadStore`].
@@ -261,6 +264,10 @@ impl ThreadStore for LocalThreadStore {
list_threads::list_threads(self, params).await
}
async fn list_turns(&self, params: ListTurnsParams) -> ThreadStoreResult<TurnPage> {
list_turns::list_turns(self, params).await
}
async fn update_thread_metadata(
&self,
params: UpdateThreadMetadataParams,
@@ -286,15 +293,23 @@ mod tests {
use codex_protocol::ThreadId;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::CompactedItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::protocol::UserMessageEvent;
use tempfile::TempDir;
use super::*;
use crate::ListTurnsParams;
use crate::LiveThread;
use crate::SortDirection;
use crate::StoredTurnItemsView;
use crate::ThreadEventPersistenceMode;
use crate::ThreadPersistenceMetadata;
use crate::local::test_support::test_config;
@@ -1009,6 +1024,110 @@ mod tests {
);
}
#[tokio::test]
async fn list_turns_filters_items_view_locally() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let thread_id = ThreadId::default();
store
.create_thread(create_thread_params(thread_id))
.await
.expect("create thread");
store
.append_items(AppendThreadItemsParams {
thread_id,
items: vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent {
turn_id: "turn-1".to_string(),
started_at: Some(1),
model_context_window: None,
collaboration_mode_kind: Default::default(),
})),
RolloutItem::ResponseItem(ResponseItem::Message {
id: None,
role: "assistant".to_string(),
content: vec![ContentItem::OutputText {
text: "response assistant".to_string(),
}],
phase: None,
}),
user_message_item("event user"),
RolloutItem::Compacted(CompactedItem {
message: "summary".to_string(),
replacement_history: None,
}),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-1".to_string(),
last_agent_message: None,
completed_at: Some(2),
duration_ms: Some(1000),
time_to_first_token_ms: None,
})),
],
})
.await
.expect("append items");
store.flush_thread(thread_id).await.expect("flush thread");
let response_items =
list_single_turn(&store, thread_id, StoredTurnItemsView::ResponseItems)
.await
.items;
assert_eq!(
response_items
.iter()
.filter(|item| matches!(item, RolloutItem::ResponseItem(_)))
.count(),
1
);
assert!(
response_items
.iter()
.all(|item| { matches!(item, RolloutItem::ResponseItem(_)) })
);
let event_items = list_single_turn(&store, thread_id, StoredTurnItemsView::EventItems)
.await
.items;
assert_eq!(
event_items
.iter()
.filter(|item| matches!(item, RolloutItem::EventMsg(_)))
.count(),
3
);
assert!(
event_items
.iter()
.all(|item| matches!(item, RolloutItem::EventMsg(_)))
);
let response_and_event_items = list_single_turn(
&store,
thread_id,
StoredTurnItemsView::ResponseAndEventItems,
)
.await
.items;
assert_eq!(response_and_event_items.len(), 4);
assert!(
!response_and_event_items
.iter()
.any(|item| matches!(item, RolloutItem::Compacted(_)))
);
let full_items = list_single_turn(&store, thread_id, StoredTurnItemsView::Full)
.await
.items;
assert_eq!(full_items.len(), 5);
assert!(
full_items
.iter()
.any(|item| matches!(item, RolloutItem::Compacted(_)))
);
}
fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams {
CreateThreadParams {
thread_id,
@@ -1039,6 +1158,26 @@ mod tests {
}))
}
async fn list_single_turn(
store: &LocalThreadStore,
thread_id: ThreadId,
items_view: StoredTurnItemsView,
) -> crate::StoredTurn {
let page = store
.list_turns(ListTurnsParams {
thread_id,
include_archived: true,
cursor: None,
page_size: 10,
sort_direction: SortDirection::Asc,
items_view,
})
.await
.expect("list turns");
assert_eq!(page.turns.len(), 1);
page.turns.into_iter().next().expect("turn")
}
async fn assert_rollout_contains_message(path: &std::path::Path, expected: &str) {
let (items, _, _) = RolloutRecorder::load_rollout_items(path)
.await
+67
View File
@@ -218,6 +218,31 @@ pub enum StoredTurnItemsView {
Summary,
/// Return every persisted item available for each turn.
Full,
/// Return only persisted Responses API items for each turn.
ResponseItems,
/// Return only persisted app-server event items for each turn.
EventItems,
/// Return persisted Responses API items and app-server event items for each turn.
ResponseAndEventItems,
}
impl StoredTurnItemsView {
/// Returns whether this view includes the persisted rollout item verbatim.
///
/// `Summary` is intentionally not handled here because it is a projected display view rather
/// than a simple rollout-item filter.
pub fn includes_persisted_item(self, item: &RolloutItem) -> bool {
match self {
StoredTurnItemsView::NotLoaded | StoredTurnItemsView::Summary => false,
StoredTurnItemsView::Full => true,
StoredTurnItemsView::ResponseItems => matches!(item, RolloutItem::ResponseItem(_)),
StoredTurnItemsView::EventItems => matches!(item, RolloutItem::EventMsg(_)),
StoredTurnItemsView::ResponseAndEventItems => matches!(
item,
RolloutItem::ResponseItem(_) | RolloutItem::EventMsg(_)
),
}
}
}
/// Store-owned status for a persisted turn.
@@ -626,6 +651,10 @@ pub struct ArchiveThreadParams {
#[cfg(test)]
mod tests {
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::CompactedItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::UserMessageEvent;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -735,4 +764,42 @@ mod tests {
})
);
}
#[test]
fn stored_turn_items_view_filters_persisted_rollout_item_kinds() {
let response_item = RolloutItem::ResponseItem(ResponseItem::Other);
let event_item = RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: "hello".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
}));
let compacted_item = RolloutItem::Compacted(CompactedItem {
message: "summary".to_string(),
replacement_history: None,
});
let cases = [
(StoredTurnItemsView::NotLoaded, [false, false, false]),
(StoredTurnItemsView::Summary, [false, false, false]),
(StoredTurnItemsView::Full, [true, true, true]),
(StoredTurnItemsView::ResponseItems, [true, false, false]),
(StoredTurnItemsView::EventItems, [false, true, false]),
(
StoredTurnItemsView::ResponseAndEventItems,
[true, true, false],
),
];
for (view, expected) in cases {
assert_eq!(
[
view.includes_persisted_item(&response_item),
view.includes_persisted_item(&event_item),
view.includes_persisted_item(&compacted_item),
],
expected
);
}
}
}