diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 369ea4222f..6d2b1d8e09 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -398,6 +398,37 @@ impl ThreadHistoryBuilder { .map(|turn| turn.rollout_start_index) } + /// Keeps the latest assistant item in memory for live summaries. Full reconnects + /// load older text from history. Rollout replay uses the regular reducer. + pub fn handle_live_paginated_event(&mut self, event: &EventMsg) { + let previous_len = self + .current_turn + .as_ref() + .map_or(0, |turn| turn.items.len()); + self.handle_event_for_history_mode(event, ThreadHistoryMode::Paginated); + if let EventMsg::ItemCompleted(completed) = event + && let codex_protocol::items::TurnItem::AgentMessage(item) = &completed.item + && let Some(turn) = self + .current_turn + .as_mut() + .filter(|turn| turn.id == completed.turn_id) + { + if turn.items.len() > previous_len { + if let Some(index) = turn.evictable_agent_item.replace(previous_len) + && let ThreadItemPayload::AgentMessage { text, .. } = &mut turn.items[index] + { + *text = String::new(); + } + } else if turn + .evictable_agent_item + .is_some_and(|index| turn.items[index].id() == item.id) + { + // A read by ID could still find the prior version of a corrected item. + turn.evictable_agent_item = None; + } + } + } + /// Shared reducer for persisted rollout replay and in-memory current-turn /// tracking used by running thread resume/rejoin. /// @@ -1561,6 +1592,7 @@ impl ThreadHistoryBuilder { root_turn_id: None, items: Vec::new(), item_index: TurnItemIndex::default(), + evictable_agent_item: None, item_times: HashMap::new(), error: None, status: TurnStatus::Completed, @@ -1889,6 +1921,7 @@ struct PendingTurn { root_turn_id: Option, items: Vec, item_index: TurnItemIndex, + evictable_agent_item: Option, item_times: HashMap, error: Option, status: TurnStatus, @@ -4344,6 +4377,43 @@ mod tests { ); } + #[test] + fn live_paginated_turn_releases_superseded_agent_text() { + let mut builder = ThreadHistoryBuilder::new(); + let turn_id = builder.ensure_turn().id.clone(); + for (id, text) in [ + ("a", "old"), + ("b", "draft"), + ("b", "corrected"), + ("c", "latest"), + ] { + builder.handle_live_paginated_event(&EventMsg::ItemCompleted(ItemCompletedEvent { + thread_id: ThreadId::default(), + turn_id: turn_id.clone(), + item: CoreTurnItem::AgentMessage(codex_protocol::items::AgentMessageItem { + id: id.into(), + content: vec![codex_protocol::items::AgentMessageContent::Text { + text: text.into(), + }], + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }), + started_at_ms: None, + completed_at_ms: 1, + })); + } + let mut texts = builder.current_turn.unwrap().items.into_iter().map(|item| { + let ThreadItemPayload::AgentMessage { text, .. } = item else { + panic!("expected assistant item") + }; + text + }); + assert_eq!(texts.next().unwrap().capacity(), 0); + assert_eq!(texts.collect::>(), ["corrected", "latest"]); + } + #[test] fn active_turn_summary_preserves_selection_and_completed_metadata() { let mut builder = ThreadHistoryBuilder::new(); diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 555671fb78..b9062d6e7a 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -1,6 +1,9 @@ +use super::thread_processor::THREAD_ITEMS_MAX_LIMIT; +use super::thread_processor::thread_turns_page_size; use super::*; use crate::extensions::send_thread_warning; use crate::thread_state::ListenerStateCandidate; +use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadQueueChangedNotification; use codex_core::ThreadEvictionOutcome; use codex_core::ThreadRegistration; @@ -8,6 +11,7 @@ use codex_extension_api::ThreadIdleCause; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; +use std::borrow::Cow; #[cfg(test)] #[path = "thread_listener_shutdown_tests.rs"] @@ -1027,20 +1031,49 @@ pub(super) async fn handle_pending_thread_resume_request( pending_thread_unloads: &PendingThreadUnloads, mut pending: crate::thread_state::PendingThreadResumeRequest, ) { - let (active_turn_metadata, active_turn) = { + let (active_turn_metadata, mut active_turn, active_page_params) = { let state = thread_state.lock().await; + let active_page_params = pending.initial_turns_page.as_ref().filter(|params| { + pending + .paginated_initial_turns_page + .as_ref() + .is_none_or(|page| { + params.sort_direction != Some(SortDirection::Asc) + || (page.next_cursor.is_none() + && page.data.len() < thread_turns_page_size(params.limit)) + || state.active_turn_id().is_some_and(|active_turn_id| { + page.data.iter().any(|turn| turn.id == active_turn_id) + }) + }) + }); let items_view = if pending.include_turns { Some(TurnItemsView::Full) } else { - pending - .initial_turns_page - .as_ref() - .map(|page| page.items_view.unwrap_or(TurnItemsView::Summary)) + active_page_params.map(|params| params.items_view.unwrap_or(TurnItemsView::Summary)) }; let active_turn = items_view.and_then(|view| state.active_turn_snapshot_with_items_view(view)); - (state.active_turn_metadata_snapshot(), active_turn) + ( + state.active_turn_metadata_snapshot(), + active_turn, + active_page_params, + ) }; + if let Some(store) = pending.resume_cursor_store.as_deref() + && let Some(turn) = &mut active_turn + && turn.items_view == TurnItemsView::Full + { + let stored = pending.paginated_turns.iter_mut().flatten().chain( + pending + .paginated_initial_turns_page + .iter_mut() + .flat_map(|page| &mut page.data), + ); + if let Err(error) = restore_evicted_agent_text(store, conversation_id, turn, stored).await { + outgoing.send_error(pending.request_id, error).await; + return; + } + } tracing::debug!( thread_id = %conversation_id, request_id = ?pending.request_id, @@ -1082,16 +1115,9 @@ pub(super) async fn handle_pending_thread_resume_request( thread_status.clone(), has_live_in_progress_turn, ); - let active_turn = if pending.initial_turns_page.is_some() { - active_turn.or_else(|| active_turn_metadata.map(Turn::from)) - } else { - None - }; let mut initial_turns_page = if let Some(mut page) = pending.paginated_initial_turns_page.take() { - if let (Some(active_turn), Some(params)) = - (active_turn, pending.initial_turns_page.as_ref()) - { + if let (Some(active_turn), Some(params)) = (active_turn, active_page_params) { let sort_direction = params.sort_direction.unwrap_or(SortDirection::Desc); let active_turn_is_in_page = page.data.iter().any(|turn| turn.id == active_turn.id); if matches!(sort_direction, SortDirection::Desc) @@ -1328,6 +1354,70 @@ pub(super) async fn resolve_pending_server_request( outgoing.resolve_server_request(request_id).await; } +#[derive(Clone, Eq, Hash, PartialEq)] +struct HistoryItemId<'a>(Cow<'a, str>); + +async fn restore_evicted_agent_text<'a>( + store: &dyn ThreadStore, + thread_id: ThreadId, + turn: &mut Turn, + stored: impl Iterator, +) -> Result<(), JSONRPCErrorError> { + // Live history skips empty assistant completions, so every empty slot was evicted. + let mut missing: HashMap<_, _> = turn + .items + .iter_mut() + .filter_map(|item| match &mut item.payload { + ThreadItemPayload::AgentMessage { id, text, .. } if text.is_empty() => { + Some((HistoryItemId(Cow::Borrowed(id)), text)) + } + _ => None, + }) + .collect(); + if missing.is_empty() { + return Ok(()); + } + for stored in stored.filter(|stored| stored.id == turn.id) { + fill_agent_text(&mut missing, std::mem::take(&mut stored.items)); + if missing.is_empty() { + return Ok(()); + } + } + // The prepared page can predate some of the listener's items. Fetch only those gaps. + let ids: Vec<_> = missing.keys().cloned().collect(); + for chunk in ids.chunks(THREAD_ITEMS_MAX_LIMIT) { + let params = ThreadItemsReadParams { + thread_id: thread_id.to_string(), + turn_id: turn.id.clone(), + item_ids: chunk.iter().map(|id| id.0.to_string()).collect(), + }; + let saved = ThreadRequestProcessor::read_thread_items(store, params).await?; + fill_agent_text(&mut missing, saved.data.into_iter().map(|entry| entry.item)); + } + if missing.is_empty() { + Ok(()) + } else { + Err(internal_error(format!( + "saved history is missing completed assistant items for turn {}", + turn.id + ))) + } +} + +fn fill_agent_text( + missing: &mut HashMap, &mut String>, + items: impl IntoIterator, +) { + for item in items { + if let ThreadItemPayload::AgentMessage { id, text, .. } = item.payload + && !text.is_empty() + && let Some(slot) = missing.remove(&HistoryItemId(Cow::Owned(id))) + { + *slot = text; + } + } +} + pub(super) fn merge_turn_history_with_active_turn(turns: &mut Vec, active_turn: Turn) { turns.retain(|turn| turn.id != active_turn.id); turns.push(active_turn); @@ -1343,17 +1433,9 @@ fn merge_active_turn_into_page( params.items_view.unwrap_or(TurnItemsView::Summary), ); let sort_direction = params.sort_direction.unwrap_or(SortDirection::Desc); - let page_size = super::thread_processor::thread_turns_page_size(params.limit); - let active_turn_is_in_page = page.data.iter().any(|turn| turn.id == active_turn.id); page.data.retain(|turn| turn.id != active_turn.id); match sort_direction { - SortDirection::Asc - if active_turn_is_in_page - || (page.data.len() < page_size && page.next_cursor.is_none()) => - { - page.data.push(active_turn); - } - SortDirection::Asc => {} + SortDirection::Asc => page.data.push(active_turn), SortDirection::Desc => page.data.insert(0, active_turn), } } diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 45153612f6..36d9f435ae 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1007,7 +1007,7 @@ impl ThreadRequestProcessor { &self, params: ThreadItemsReadParams, ) -> Result, JSONRPCErrorError> { - self.thread_items_read_response_inner(params) + Self::read_thread_items(self.thread_store.as_ref(), params) .await .map(|response| Some(response.into())) } @@ -3762,11 +3762,11 @@ impl ThreadRequestProcessor { }) } - async fn thread_items_read_response_inner( - &self, + pub(super) async fn read_thread_items( + store: &dyn ThreadStore, params: ThreadItemsReadParams, ) -> Result { - if !(1..=100).contains(¶ms.item_ids.len()) { + if !(1..=THREAD_ITEMS_MAX_LIMIT).contains(¶ms.item_ids.len()) { return Err(invalid_request( "itemIds must contain 1 to 100 IDs".to_string(), )); @@ -3779,8 +3779,7 @@ impl ThreadRequestProcessor { let thread_id = ThreadId::from_string(&thread_id) .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; let page_size = item_ids.len(); - let page = self - .thread_store + let page = store .list_items(StoreListItemsParams { item_ids: Some(item_ids), thread_id, @@ -6117,7 +6116,7 @@ fn xcode_26_4_mcp_elicitations_auto_deny( const THREAD_TURNS_DEFAULT_LIMIT: usize = 25; const THREAD_TURNS_MAX_LIMIT: usize = 100; const THREAD_ITEMS_DEFAULT_LIMIT: usize = 25; -const THREAD_ITEMS_MAX_LIMIT: usize = 100; +pub(super) const THREAD_ITEMS_MAX_LIMIT: usize = 100; const THREAD_SEARCH_OCCURRENCES_DEFAULT_LIMIT: usize = 50; const THREAD_SEARCH_OCCURRENCES_MAX_LIMIT: usize = 250; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index e47c9ed81c..a712ee2f47 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -139,6 +139,7 @@ pub(crate) struct ThreadState { listener_command_tx: Option>, current_turn_history: ThreadHistoryBuilder, history_mode: ThreadHistoryMode, + evict_superseded_agent_text: bool, listener_thread: Option>, watch_registration: WatchRegistration, } @@ -181,6 +182,8 @@ impl ThreadState { self.listener_generation = self.listener_generation.wrapping_add(1); self.last_thread_settings = Some(thread_settings_baseline); self.history_mode = history_mode; + self.evict_superseded_agent_text = + history_mode == ThreadHistoryMode::Paginated && conversation.rollout_path().is_some(); let (listener_command_tx, listener_command_rx) = mpsc::unbounded_channel(); self.listener_command_tx = Some(listener_command_tx); self.listener_thread = Some(Arc::downgrade(conversation)); @@ -243,8 +246,12 @@ impl ThreadState { { self.turn_summary.last_agent_message = Some(ThreadItem::from(payload)); } - self.current_turn_history - .handle_event_for_history_mode(event, self.history_mode); + if self.evict_superseded_agent_text { + self.current_turn_history.handle_live_paginated_event(event); + } else { + self.current_turn_history + .handle_event_for_history_mode(event, self.history_mode); + } if matches!(event, EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_)) { self.last_terminal_turn_id = Some(event_turn_id.to_string()); if !self.current_turn_history.has_active_turn() { diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 8badd11648..03e498f2f0 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -29,6 +29,7 @@ use codex_app_server_protocol::DeprecationNoticeNotification; use codex_app_server_protocol::FileChangeApprovalDecision; use codex_app_server_protocol::FileChangeRequestApprovalResponse; use codex_app_server_protocol::GeneratedImageRef; +use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; @@ -133,7 +134,6 @@ use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; use pretty_assertions::assert_eq; use serde_json::json; -use std::collections::HashMap; use std::fs::FileTimes; use std::io::Write; use std::path::Path; @@ -5512,6 +5512,8 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R gate: None, body: responses::sse(vec![ responses::ev_response_created("resp-2"), + responses::ev_message_item_added("earlier", ""), + responses::ev_assistant_message("earlier", "First progress update"), responses::ev_message_item_added("progress", ""), responses::ev_assistant_message("progress", "Checking"), responses::ev_message_item_added("msg-2", ""), @@ -5588,20 +5590,18 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R primary.read_stream_until_notification_message("turn/started"), ) .await??; - // Starting the gated reply means the prompt and the progress message have completed. - let mut message_starts = HashMap::new(); - timeout(DEFAULT_READ_TIMEOUT, async { - loop { - let started: ItemStartedNotification = - primary.read_notification("item/started").await?; - assert_eq!(started.turn_id, running_turn.id); - if started.item.id() == "msg-2" { - return Ok::<(), anyhow::Error>(()); - } - message_starts.insert(started.item.id().to_owned(), started.started_at_ms); - } - }) - .await??; + // The final item is gated. Record exactly what subscribers saw before reconnecting. + let mut live_items = Vec::new(); + for _ in 0..3 { + let completed: ItemCompletedNotification = timeout( + DEFAULT_READ_TIMEOUT, + primary.read_notification("item/completed"), + ) + .await??; + assert_eq!(completed.turn_id, running_turn.id); + live_items.push(completed.item); + } + let summary_items = [live_items[0].clone(), live_items[2].clone()]; let resume_id = primary .send_thread_resume_request(ThreadResumeParams { @@ -5673,6 +5673,10 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R "hot paginated resume should wait for a real token usage update" ); + let history_db = codex_state::open_thread_history_db( + &codex_state::SqliteConfig::new_for_testing(codex_home.path().abs()), + ) + .await?; for (exclude_turns, items_view) in [ (true, None), (true, Some(TurnItemsView::Full)), @@ -5700,26 +5704,51 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R page.data[0].items_view, items_view.unwrap_or(TurnItemsView::Summary) ); - assert!(matches!( - page.data[0].items.as_slice(), - [codex_app_server_protocol::ThreadItem { payload: ThreadItemPayload::UserMessage { .. }, .. }, codex_app_server_protocol::ThreadItem { payload: ThreadItemPayload::AgentMessage { text, .. }, .. }] if text == "Checking" - )); - assert_eq!( - page.data[0] - .items - .iter() - .filter_map(|item| Some((item.id().to_owned(), item.started_at_ms?))) - .collect::>(), - message_starts - ); + let page_turn = &page.data[0]; + let expected = if page_turn.items_view == TurnItemsView::Full { + live_items.as_slice() + } else { + &summary_items + }; + assert_eq!(page_turn.items, expected); + if page_turn.items_view == TurnItemsView::Full { + // Simulate a prepared page that does not yet include the live turn. + let deleted = sqlx::query("DELETE FROM thread_turns WHERE turn_id = ?") + .bind(&running_turn.id) + .execute(&history_db) + .await?; + assert_eq!(deleted.rows_affected(), 1); + } if !exclude_turns { let full_turn = resumed.thread.turns.last().expect("full active turn"); assert_eq!(full_turn.id, running_turn.id); assert_eq!(full_turn.items_view, TurnItemsView::Full); - assert!(!full_turn.items.is_empty()); + assert_eq!(full_turn.items, live_items); } } + // An unreadable item only fails requests that need its text. + sqlx::query("UPDATE thread_items SET item_json = json_set(item_json, '$.text', '') WHERE turn_id = ? AND item_id = 'earlier'") + .bind(&running_turn.id).execute(&history_db).await?; + let missing_id = primary + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + }) + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + primary.read_stream_until_error_message(RequestId::Integer(missing_id)), + ) + .await??; + assert_eq!( + error.error.message, + format!( + "saved history is missing completed assistant items for turn {}", + running_turn.id + ) + ); + let asc_resume_id = primary .send_thread_resume_request(ThreadResumeParams { thread_id: thread.id.clone(), @@ -5727,7 +5756,7 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { limit: Some(1), sort_direction: Some(SortDirection::Asc), - items_view: Some(TurnItemsView::NotLoaded), + items_view: Some(TurnItemsView::Full), }), ..Default::default() }) 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 24a23e64c4..69cc63244c 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 @@ -294,10 +294,10 @@ pub(super) async fn page_item_rows( SELECT turn_id, item_id, rollout_ordinal, updated_at_ordinal, created_at_ms, started_at_ms, completed_at_ms, item_json FROM thread_items"#, ); - if params.turn_id.is_some() - && params.item_ids.is_none() - && params.after_updated_at_ordinal.is_none() - { + if params.item_ids.is_some() { + // Seek each requested ID instead of scanning the turn in rollout order. + query.push(" INDEXED BY sqlite_autoindex_thread_items_1"); + } else if params.turn_id.is_some() && params.after_updated_at_ordinal.is_none() { // A thread-wide index can scan earlier turns when a turn's final page is underfilled. query.push(" INDEXED BY idx_thread_items_by_turn_page"); }