From 436a2a8dd367400c1d43a0cc1597bb173a90fa71 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 2 Jul 2026 12:05:26 -0400 Subject: [PATCH] Speed up reverse history search --- codex-rs/tui/src/app/event_dispatch.rs | 8 + codex-rs/tui/src/app/tests.rs | 73 ++++- codex-rs/tui/src/app/thread_routing.rs | 77 ++++- codex-rs/tui/src/app_event.rs | 40 ++- codex-rs/tui/src/bottom_pane/chat_composer.rs | 42 +++ .../chat_composer/history_search.rs | 42 ++- .../src/bottom_pane/chat_composer_history.rs | 209 ++++++++---- .../chat_composer_history/search_batch.rs | 204 ++++++++++++ .../search_batch_tests.rs | 308 ++++++++++++++++++ codex-rs/tui/src/bottom_pane/mod.rs | 33 +- codex-rs/tui/src/chatwidget.rs | 8 +- .../src/chatwidget/tests/history_replay.rs | 2 +- 12 files changed, 938 insertions(+), 108 deletions(-) create mode 100644 codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch.rs create mode 100644 codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch_tests.rs diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 8713537fbc..96e6b2cf48 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -432,6 +432,14 @@ impl App { self.lookup_message_history_entry(thread_id, offset, log_id) .await?; } + AppEvent::LookupMessageHistoryBatch { + thread_id, + cursor, + log_id, + } => { + self.lookup_message_history_batch(thread_id, cursor, log_id) + .await?; + } AppEvent::ApproveRecentAutoReviewDenial { thread_id, id } => { self.chat_widget .approve_recent_auto_review_denial(thread_id, id); diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index d8d63575bb..5e4a3b5035 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -9,6 +9,7 @@ use super::*; use crate::app_backtrack::BacktrackSelection; use crate::app_backtrack::BacktrackState; use crate::app_backtrack::user_count; +use crate::app_event::HistoryBatchEntryResponse; use crate::chatwidget::ChatWidgetInit; use crate::chatwidget::create_initial_user_message; @@ -450,9 +451,34 @@ async fn history_lookup_response_is_routed_to_requesting_thread() -> Result<()> panic!("expected thread-routed history response"); }; assert_eq!(routed_thread_id, thread_id); - assert_eq!(event.offset, 0); - assert_eq!(event.log_id, 1); - assert!(event.entry.is_none()); + assert_eq!( + event, + HistoryLookupResponse::Entry { + offset: 0, + log_id: 1, + entry: None, + } + ); + + let cursor = codex_message_history::HistoryBatchCursor::new(/*end_offset*/ 10); + app.lookup_message_history_batch(thread_id, cursor, /*log_id*/ 1) + .await?; + let app_event = tokio::time::timeout(Duration::from_secs(1), app_event_rx.recv()) + .await + .expect("history batch lookup should emit an app event") + .expect("app event channel should stay open"); + let AppEvent::ThreadHistoryEntryResponse { + thread_id: routed_thread_id, + event, + } = app_event + else { + panic!("expected thread-routed history batch response"); + }; + assert_eq!(routed_thread_id, thread_id); + assert_eq!( + event, + HistoryLookupResponse::BatchError { cursor, log_id: 1 } + ); Ok(()) } @@ -496,6 +522,47 @@ async fn enqueue_thread_event_does_not_block_when_channel_full() -> Result<()> { Ok(()) } +#[tokio::test] +async fn active_history_batch_is_delivered_without_replay_buffering() -> Result<()> { + let mut app = make_test_app().await; + let thread_id = ThreadId::new(); + app.thread_event_channels + .insert(thread_id, ThreadEventChannel::new(/*capacity*/ 4)); + app.set_thread_active(thread_id, /*active*/ true).await; + + let cursor = codex_message_history::HistoryBatchCursor::new(/*end_offset*/ 1); + let event = HistoryLookupResponse::Batch { + cursor, + log_id: 1, + entries: vec![HistoryBatchEntryResponse { + offset: 1, + entry: Some("history entry".to_string()), + }], + next_older_cursor: Some(codex_message_history::HistoryBatchCursor::new( + /*end_offset*/ 0, + )), + }; + app.enqueue_thread_history_entry_response(thread_id, event.clone()) + .await?; + + let channel = app + .thread_event_channels + .get_mut(&thread_id) + .expect("missing thread channel"); + assert!(channel.store.lock().await.buffer.is_empty()); + let mut receiver = channel.receiver.take().expect("missing receiver"); + let delivered = time::timeout(Duration::from_millis(50), receiver.recv()) + .await + .expect("timed out waiting for history batch") + .expect("missing history batch"); + let ThreadBufferedEvent::HistoryEntryResponse(delivered) = delivered else { + panic!("expected history batch response"); + }; + assert_eq!(delivered, event); + + Ok(()) +} + #[tokio::test] async fn replay_thread_snapshot_restores_draft_and_queued_input() { let mut app = make_test_app().await; diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index b149fd53f3..d984ccea95 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -504,7 +504,7 @@ impl App { app_event_tx.send(AppEvent::ThreadHistoryEntryResponse { thread_id, - event: HistoryLookupResponse { + event: HistoryLookupResponse::Entry { offset, log_id, entry: entry_opt.map(|entry| entry.text), @@ -514,6 +514,55 @@ impl App { Ok(()) } + /// Fetch one bounded local cross-session history batch for the requesting thread. + pub(super) async fn lookup_message_history_batch( + &mut self, + thread_id: ThreadId, + cursor: codex_message_history::HistoryBatchCursor, + log_id: u64, + ) -> Result<()> { + let history_config = codex_message_history::HistoryConfig::new( + self.chat_widget.config_ref().codex_home.clone(), + &self.chat_widget.config_ref().history, + ); + let app_event_tx = self.app_event_tx.clone(); + tokio::spawn(async move { + let event = match tokio::task::spawn_blocking(move || { + codex_message_history::lookup_batch(log_id, cursor, &history_config) + }) + .await + { + Ok(Ok(batch)) => { + let entries = batch + .entries + .into_iter() + .map(|entry| crate::app_event::HistoryBatchEntryResponse { + offset: entry.offset, + entry: entry.entry.map(|entry| entry.text), + }) + .collect(); + HistoryLookupResponse::Batch { + cursor, + log_id, + entries, + next_older_cursor: batch.next_older_cursor, + } + } + Ok(Err(err)) => { + tracing::warn!(error = %err, "history batch lookup failed"); + HistoryLookupResponse::BatchError { cursor, log_id } + } + Err(err) => { + tracing::warn!(error = %err, "history batch lookup task failed"); + HistoryLookupResponse::BatchError { cursor, log_id } + } + }; + + app_event_tx.send(AppEvent::ThreadHistoryEntryResponse { thread_id, event }); + }); + Ok(()) + } + pub(super) async fn try_submit_active_thread_op_via_app_server( &mut self, app_server: &mut AppServerSession, @@ -1074,18 +1123,24 @@ impl App { let should_send = { let mut guard = store.lock().await; - guard - .buffer - .push_back(ThreadBufferedEvent::HistoryEntryResponse(event.clone())); - if guard.buffer.len() > guard.capacity - && let Some(removed) = guard.buffer.pop_front() - && let ThreadBufferedEvent::Request(request) = &removed - { + let should_send = guard.active; + // Active batch responses remain queued in the receiver across a concurrent detach, so + // retaining another deep copy for thread replay only accumulates already-delivered + // history data. Inactive responses still need the buffer because they are not sent. + if !should_send || !matches!(&event, HistoryLookupResponse::Batch { .. }) { guard - .pending_interactive_replay - .note_evicted_server_request(request); + .buffer + .push_back(ThreadBufferedEvent::HistoryEntryResponse(event.clone())); + if guard.buffer.len() > guard.capacity + && let Some(removed) = guard.buffer.pop_front() + && let ThreadBufferedEvent::Request(request) = &removed + { + guard + .pending_interactive_replay + .note_evicted_server_request(request); + } } - guard.active + should_send }; if should_send { diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index fa2c9f29fc..f68f6fd1e9 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -30,6 +30,7 @@ use codex_app_server_protocol::SkillsListResponse; use codex_app_server_protocol::ThreadGoalStatus; use codex_connectors::AppInfo; use codex_file_search::FileMatch; +use codex_message_history::HistoryBatchCursor; use codex_protocol::ThreadId; use codex_protocol::openai_models::ModelPreset; use codex_utils_absolute_path::AbsolutePathBuf; @@ -63,13 +64,39 @@ pub(crate) enum ThreadGoalSetMode { }, } -#[derive(Debug, Clone)] -pub(crate) struct HistoryLookupResponse { +/// One absolute history offset returned by a batch lookup. +/// +/// Malformed rows retain their offset with `entry` set to `None` so the composer can cache the gap +/// without shifting every older record. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct HistoryBatchEntryResponse { pub(crate) offset: usize, - pub(crate) log_id: u64, pub(crate) entry: Option, } +/// Persistent-history data routed back to the thread that requested it. +/// +/// Batch responses preserve absolute offsets and malformed-row gaps so the composer can cache the +/// data independently of whichever search query is active when the response arrives. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HistoryLookupResponse { + Entry { + offset: usize, + log_id: u64, + entry: Option, + }, + Batch { + cursor: HistoryBatchCursor, + log_id: u64, + entries: Vec, + next_older_cursor: Option, + }, + BatchError { + cursor: HistoryBatchCursor, + log_id: u64, + }, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ConsolidationScrollbackReflow { IfResizeReflowRan, @@ -197,6 +224,13 @@ pub(crate) enum AppEvent { log_id: u64, }, + /// Fetch a bounded batch of persistent history entries for reverse search. + LookupMessageHistoryBatch { + thread_id: ThreadId, + cursor: HistoryBatchCursor, + log_id: u64, + }, + /// Start a new session. NewSession, diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 76a5760658..f29bdfa9a7 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -151,6 +151,7 @@ use crate::key_hint::KeyBinding; use crate::key_hint::has_ctrl_or_alt; use crate::line_truncation::truncate_line_with_ellipsis_if_overflow; use crate::ui_consts::FOOTER_INDENT_COLS; +use codex_message_history::HistoryBatchCursor; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -173,6 +174,7 @@ use ratatui::widgets::WidgetRef; use super::chat_composer_history::ChatComposerHistory; use super::chat_composer_history::HistoryEntry; use super::chat_composer_history::HistoryEntryResponse; +use super::chat_composer_history::HistorySearchResult; use super::command_popup::CommandItem; use super::file_search_popup::FileSearchPopup; use super::footer::CollaborationModeIndicator; @@ -863,6 +865,46 @@ impl ChatComposer { } } + pub(crate) fn on_history_batch_response( + &mut self, + log_id: u64, + cursor: HistoryBatchCursor, + entries: Vec, + next_older_cursor: Option, + ) -> bool { + let result = self.history.on_batch_response( + log_id, + cursor, + entries, + next_older_cursor, + &self.app_event_tx, + ); + self.apply_history_batch_result(result) + } + + /// Applies a failed batch lookup without conflating it with history exhaustion. + /// + /// The history state machine either schedules a bounded retry, preserves an existing match, or + /// returns the search UI to its idle draft when no match has been selected yet. + pub(crate) fn on_history_batch_error( + &mut self, + log_id: u64, + cursor: HistoryBatchCursor, + ) -> bool { + let result = self + .history + .on_batch_error(log_id, cursor, &self.app_event_tx); + self.apply_history_batch_result(result) + } + + fn apply_history_batch_result(&mut self, result: Option) -> bool { + let Some(result) = result else { + return false; + }; + self.apply_history_search_result(result); + true + } + pub(crate) fn record_replayed_user_message_history(&mut self, entry: HistoryEntry) { self.history.record_replayed_submission(entry); } diff --git a/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs b/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs index 3c2663a659..5217f5aac9 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs @@ -305,9 +305,10 @@ impl ChatComposer { /// /// `Found` previews the matching entry, `Pending` keeps the footer in a waiting state while an /// async persistent entry lookup is outstanding, `AtBoundary` preserves the current match, and - /// `NotFound` restores the original draft while keeping the query available for further edits. - /// Treating `AtBoundary` like `NotFound` would produce the visible "no match" flicker at the - /// end of a one-result search and desynchronize Up/Down counts. + /// `NotFound` restores the original draft while keeping the query available for further edits, + /// and `Unavailable` does the same without claiming there was no match. Treating `AtBoundary` + /// like `NotFound` would produce the visible "no match" flicker at the end of a one-result + /// search and desynchronize Up/Down counts. pub(super) fn apply_history_search_result(&mut self, result: HistorySearchResult) { match result { HistorySearchResult::Found(entry) => { @@ -326,13 +327,17 @@ impl ChatComposer { search.status = HistorySearchStatus::Match; } } - HistorySearchResult::NotFound => { + result @ (HistorySearchResult::NotFound | HistorySearchResult::Unavailable) => { let original_draft = self .history_search .as_ref() .map(|search| search.original_draft.clone()); if let Some(search) = self.history_search.as_mut() { - search.status = HistorySearchStatus::NoMatch; + search.status = if matches!(result, HistorySearchResult::NotFound) { + HistorySearchStatus::NoMatch + } else { + HistorySearchStatus::Idle + }; } if let Some(original_draft) = original_draft { self.restore_draft(original_draft); @@ -494,6 +499,7 @@ mod tests { use tokio::sync::mpsc::unbounded_channel; use super::super::super::chat_composer_history::HistoryEntry; + use super::super::super::chat_composer_history::HistorySearchResult; use super::super::super::footer::FooterMode; use super::super::ChatComposer; use super::HistorySearchStatus; @@ -524,6 +530,32 @@ mod tests { assert_eq!(composer.footer_mode(), FooterMode::HistorySearch); } + #[test] + fn unavailable_history_search_restores_idle_draft() { + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + /*has_input_focus*/ true, + sender, + /*enhanced_keys_supported*/ false, + "Ask Codex to do anything".to_string(), + /*disable_paste_burst*/ false, + ); + composer.set_text_content("draft".to_string(), Vec::new(), Vec::new()); + composer.begin_history_search(); + composer.apply_history_search_result(HistorySearchResult::Pending); + + composer.apply_history_search_result(HistorySearchResult::Unavailable); + + assert_eq!(composer.draft.textarea.text(), "draft"); + assert!( + composer + .history_search + .as_ref() + .is_some_and(|search| matches!(search.status, HistorySearchStatus::Idle)) + ); + } + #[test] fn history_search_match_ranges_are_case_insensitive() { assert_eq!( diff --git a/codex-rs/tui/src/bottom_pane/chat_composer_history.rs b/codex-rs/tui/src/bottom_pane/chat_composer_history.rs index ff6dd9983c..a0af332528 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer_history.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer_history.rs @@ -1,9 +1,13 @@ //! The chat composer history module owns shell-style recall and incremental search traversal. //! //! It combines persistent cross-session entries with local in-session entries into one offset -//! space. Persistent entries are fetched lazily and re-enter this state machine through -//! [`ChatComposerHistory::on_entry_response`], while local entries are already available with full -//! draft metadata. +//! space. Normal navigation fetches persistent entries individually through +//! [`ChatComposerHistory::on_entry_response`]. Reverse search switches to bounded, +//! query-independent batches through [`ChatComposerHistory::on_batch_response`] after probing the +//! newest entry. Batch responses populate the shared cache even when the active search has moved +//! on, but only the awaited cursor resumes a search; stale log IDs are ignored, and batch read +//! failures follow a bounded retry path. Local entries are already available with full draft +//! metadata. //! //! Ctrl+R search is modeled separately from normal Up/Down navigation because it has different //! guarantees: query edits restart from the newest match, repeated Older/Newer keys move through @@ -19,9 +23,18 @@ use crate::app_event::AppEvent; use crate::app_event_sender::AppEventSender; use crate::bottom_pane::MentionBinding; use crate::mention_codec::decode_history_mentions_with_at_mentions; +use codex_message_history::HistoryBatchCursor; use codex_protocol::ThreadId; use codex_protocol::user_input::TextElement; +#[path = "chat_composer_history/search_batch.rs"] +mod search_batch; +#[cfg(test)] +#[path = "chat_composer_history/search_batch_tests.rs"] +mod search_batch_tests; + +const MAX_BATCH_READ_RETRIES: u8 = 2; + /// A composer history entry that can rehydrate draft state. #[derive(Debug, Clone, PartialEq)] pub(crate) struct HistoryEntry { @@ -124,8 +137,8 @@ pub(crate) struct ChatComposerHistory { /// Local entries seeded from resumed transcript replay. replay_seeded_history: Vec, - /// Cache of persistent history entries fetched on-demand (text-only). - fetched_history: HashMap, + /// Persistent offsets fetched on demand, with `None` for malformed batch rows. + fetched_history: HashMap>, /// Current cursor within the combined (persistent + local) history. `None` /// indicates the user is *not* currently browsing history. @@ -157,13 +170,15 @@ pub(crate) enum HistorySearchDirection { /// `Pending` means a persistent entry lookup has been requested and the caller should keep the /// visible search session open until [`ChatComposerHistory::on_entry_response`] supplies the next /// result. `AtBoundary` means the current selected match is still valid but the requested direction -/// has no further unique match; callers should avoid treating it like a query miss. +/// has no further unique match; callers should avoid treating it like a query miss. `Unavailable` +/// ends a failed lookup without claiming the query has no matching history. #[derive(Clone, Debug, PartialEq)] pub(crate) enum HistorySearchResult { Found(HistoryEntry), Pending, AtBoundary, NotFound, + Unavailable, } /// Result of integrating an asynchronous persistent history response. @@ -183,7 +198,8 @@ pub(crate) enum HistoryEntryResponse { /// continue scanning, while `selected_match_index` points into `unique_matches` so already /// discovered unique results can be revisited without rescanning duplicate offsets. `seen_texts` /// intentionally keys on exact prompt text because the UI previews and accepts text, not the -/// storage identity of each historical record. +/// storage identity of each historical record. `next_older_cursor` retains the query-independent +/// batch boundary so a match does not force the next Older search back onto the prefix-scan path. #[derive(Clone, Debug)] struct HistorySearchState { query: String, @@ -193,6 +209,7 @@ struct HistorySearchState { selected_match_index: Option, seen_texts: HashSet, awaiting: Option, + next_older_cursor: Option, exhausted_older: bool, exhausted_newer: bool, } @@ -210,14 +227,22 @@ struct UniqueHistoryMatch { /// Persistent-history lookup currently blocking an incremental search scan. /// -/// The pending request records the direction and boundary behavior that were active when the fetch -/// was issued so the response can either return a unique match or continue scanning as if no async -/// gap had occurred. -#[derive(Clone, Copy, Debug)] -struct PendingHistorySearch { - offset: usize, - direction: HistorySearchDirection, - boundary_if_exhausted: bool, +/// The pending request records the boundary behavior that was active when the fetch was issued so +/// the response can either return a unique match or continue scanning as if no async gap had +/// occurred. Single-entry requests retain their direction; batch requests are older-only by +/// construction. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PendingHistorySearch { + Entry { + offset: usize, + direction: HistorySearchDirection, + boundary_if_exhausted: bool, + }, + Batch { + cursor: HistoryBatchCursor, + boundary_if_exhausted: bool, + read_failures: u8, + }, } impl ChatComposerHistory { @@ -445,36 +470,54 @@ impl ChatComposerHistory { HistoryEntry::new_with_at_mentions(entry, self.at_mention_restore_enabled) }); if let Some(entry) = entry.clone() { - self.fetched_history.insert(offset, entry); + self.fetched_history.insert(offset, Some(entry)); } if self .search .as_ref() .and_then(|search| search.awaiting) - .is_some_and(|pending| pending.offset == offset) + .is_some_and(|pending| { + matches!(pending, PendingHistorySearch::Entry { offset: awaited, .. } if awaited == offset) + }) { let pending = self .search .as_ref() .and_then(|search| search.awaiting) - .unwrap_or(PendingHistorySearch { + .unwrap_or(PendingHistorySearch::Entry { offset, direction: HistorySearchDirection::Older, boundary_if_exhausted: false, }); + let PendingHistorySearch::Entry { + direction, + boundary_if_exhausted, + .. + } = pending + else { + return HistoryEntryResponse::Ignored; + }; if let Some(entry) = entry && self.search_matches(&entry) && self.search_result_is_unique(&entry) { return HistoryEntryResponse::Search(self.search_match(offset, entry)); } - return HistoryEntryResponse::Search(self.advance_search_after( - offset, - pending.direction, - pending.boundary_if_exhausted, - app_event_tx, - )); + let result = match direction { + HistorySearchDirection::Older => self.advance_older_search_after_entry_miss( + offset, + boundary_if_exhausted, + app_event_tx, + ), + HistorySearchDirection::Newer => self.advance_search_after( + offset, + direction, + boundary_if_exhausted, + app_event_tx, + ), + }; + return HistoryEntryResponse::Search(result); } if self.history_cursor == Some(offset as isize) { @@ -655,22 +698,37 @@ impl ChatComposerHistory { if self.search_matches(&entry) && self.search_result_is_unique(&entry) { return self.search_match(offset, entry); } - } else if offset < self.persistent_entry_count - && let (Some(thread_id), Some(log_id)) = (self.thread_id, self.persistent_log_id) + } else if !self.fetched_history.contains_key(&offset) + && offset < self.persistent_entry_count { - if let Some(search) = self.search.as_mut() { - search.awaiting = Some(PendingHistorySearch { - offset, - direction, + if direction == HistorySearchDirection::Older + && let Some(cursor) = self + .search + .as_ref() + .and_then(|search| search.next_older_cursor) + && cursor.end_offset() == offset + { + return self.request_older_search_batch( + cursor, boundary_if_exhausted, - }); + app_event_tx, + ); + } + if let (Some(thread_id), Some(log_id)) = (self.thread_id, self.persistent_log_id) { + if let Some(search) = self.search.as_mut() { + search.awaiting = Some(PendingHistorySearch::Entry { + offset, + direction, + boundary_if_exhausted, + }); + } + app_event_tx.send(AppEvent::LookupMessageHistoryEntry { + thread_id, + offset, + log_id, + }); + return HistorySearchResult::Pending; } - app_event_tx.send(AppEvent::LookupMessageHistoryEntry { - thread_id, - offset, - log_id, - }); - return HistorySearchResult::Pending; } let next_offset = match direction { @@ -694,7 +752,7 @@ impl ChatComposerHistory { .get(offset - self.persistent_entry_count) .cloned() } else { - self.fetched_history.get(&offset).cloned() + self.fetched_history.get(&offset).cloned().flatten() } } @@ -839,6 +897,7 @@ impl HistorySearchState { selected_match_index: None, seen_texts: HashSet::new(), awaiting: None, + next_older_cursor: None, exhausted_older: false, exhausted_newer: false, } @@ -898,6 +957,7 @@ impl HistorySearchState { mod tests { use super::*; use crate::app_event::AppEvent; + use crate::app_event::HistoryBatchEntryResponse; use pretty_assertions::assert_eq; use tokio::sync::mpsc::unbounded_channel; @@ -906,6 +966,13 @@ mod tests { .expect("thread id should parse") } + fn batch_entry(offset: usize, entry: &str) -> HistoryBatchEntryResponse { + HistoryBatchEntryResponse { + offset, + entry: Some(entry.to_string()), + } + } + #[test] fn duplicate_submissions_are_not_recorded() { let mut history = ChatComposerHistory::new(); @@ -1251,13 +1318,19 @@ mod tests { &tx, ) ); - let _ = rx.try_recv().expect("expected oldest lookup"); + let AppEvent::LookupMessageHistoryBatch { cursor, .. } = + rx.try_recv().expect("expected oldest batch") + else { + panic!("unexpected event variant"); + }; + assert_eq!(cursor.end_offset(), 0); assert_eq!( - HistoryEntryResponse::Search(HistorySearchResult::AtBoundary), - history.on_entry_response( + Some(HistorySearchResult::AtBoundary), + history.on_batch_response( /*log_id*/ 1, - /*offset*/ 0, - Some("also not a match".into()), + cursor, + vec![batch_entry(/*offset*/ 0, "also not a match")], + /*next_older_cursor*/ None, &tx, ) ); @@ -1314,26 +1387,27 @@ mod tests { &tx ) ); - let AppEvent::LookupMessageHistoryEntry { + let AppEvent::LookupMessageHistoryBatch { thread_id: response_thread_id, - offset, + cursor, log_id, } = rx.try_recv().expect("expected next lookup") else { panic!("unexpected event variant"); }; assert_eq!(response_thread_id, thread_id); - assert_eq!(offset, 1); + assert_eq!(cursor.end_offset(), 1); assert_eq!(log_id, 1); assert_eq!( - HistoryEntryResponse::Search(HistorySearchResult::Found(HistoryEntry::new( - "older command".to_string() + Some(HistorySearchResult::Found(HistoryEntry::new( + "OLDER command".to_string() ))), - history.on_entry_response( + history.on_batch_response( /*log_id*/ 1, - /*offset*/ 1, - Some("older command".into()), + cursor, + vec![batch_entry(/*offset*/ 1, "OLDER command")], + Some(HistoryBatchCursor::new(/*end_offset*/ 0)), &tx ) ); @@ -1388,25 +1462,24 @@ mod tests { &tx, ) ); - let _ = rx.try_recv().expect("expected next lookup after duplicate"); + let AppEvent::LookupMessageHistoryBatch { cursor, .. } = + rx.try_recv().expect("expected next batch after duplicate") + else { + panic!("unexpected event variant"); + }; + assert_eq!(cursor.end_offset(), 1); assert_eq!( - HistoryEntryResponse::Search(HistorySearchResult::Pending), - history.on_entry_response( - /*log_id*/ 1, - /*offset*/ 1, - Some("not a match".into()), - &tx, - ) - ); - let _ = rx.try_recv().expect("expected oldest lookup"); - assert_eq!( - HistoryEntryResponse::Search(HistorySearchResult::Found(HistoryEntry::new( + Some(HistorySearchResult::Found(HistoryEntry::new( "needle older".to_string() ))), - history.on_entry_response( + history.on_batch_response( /*log_id*/ 1, - /*offset*/ 0, - Some("needle older".into()), + cursor, + vec![ + batch_entry(/*offset*/ 1, "not a match"), + batch_entry(/*offset*/ 0, "needle older"), + ], + /*next_older_cursor*/ None, &tx, ) ); @@ -1467,10 +1540,10 @@ mod tests { history.set_metadata(test_thread_id(), /*log_id*/ 1, /*entry_count*/ 3); history .fetched_history - .insert(1, HistoryEntry::new("command2".to_string())); + .insert(1, Some(HistoryEntry::new("command2".to_string()))); history .fetched_history - .insert(2, HistoryEntry::new("command3".to_string())); + .insert(2, Some(HistoryEntry::new("command3".to_string()))); assert_eq!( Some(HistoryEntry::new("command3".to_string())), diff --git a/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch.rs b/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch.rs new file mode 100644 index 0000000000..95db7b80cf --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch.rs @@ -0,0 +1,204 @@ +use super::ChatComposerHistory; +use super::HistoryEntry; +use super::HistorySearchDirection; +use super::HistorySearchResult; +use super::MAX_BATCH_READ_RETRIES; +use super::PendingHistorySearch; +use crate::app_event::AppEvent; +use crate::app_event::HistoryBatchEntryResponse; +use crate::app_event_sender::AppEventSender; +use codex_message_history::HistoryBatchCursor; + +impl ChatComposerHistory { + /// Applies a query-independent batch to the persistent cache and, when still applicable, + /// resumes the active reverse search. + pub(crate) fn on_batch_response( + &mut self, + log_id: u64, + cursor: HistoryBatchCursor, + entries: Vec, + next_older_cursor: Option, + app_event_tx: &AppEventSender, + ) -> Option { + if self.persistent_log_id != Some(log_id) { + return None; + } + + let entries: Vec<_> = entries + .into_iter() + .map(|response| { + let entry = response.entry.map(|text| { + HistoryEntry::new_with_at_mentions(text, self.at_mention_restore_enabled) + }); + if entry.is_some() { + self.fetched_history.insert(response.offset, entry.clone()); + } else { + self.fetched_history.entry(response.offset).or_insert(None); + } + (response.offset, entry) + }) + .collect(); + + let (boundary_if_exhausted, _) = self.pending_batch(cursor)?; + if let Some(search) = self.search.as_mut() { + search.next_older_cursor = next_older_cursor; + } + + for (offset, entry) in entries { + if let Some(entry) = entry + && self.search_matches(&entry) + && self.search_result_is_unique(&entry) + { + return Some(self.search_match(offset, entry)); + } + } + + let result = if let Some(next_cursor) = next_older_cursor { + self.advance_older_search_with_batches_from( + next_cursor, + boundary_if_exhausted, + app_event_tx, + ) + } else { + self.exhausted_search_result(HistorySearchDirection::Older, boundary_if_exhausted) + }; + Some(result) + } + + /// Retries a failed batch lookup up to a fixed limit without treating the failure as history + /// exhaustion. + pub(crate) fn on_batch_error( + &mut self, + log_id: u64, + cursor: HistoryBatchCursor, + app_event_tx: &AppEventSender, + ) -> Option { + if self.persistent_log_id != Some(log_id) { + return None; + } + + let (boundary_if_exhausted, read_failures) = self.pending_batch(cursor)?; + + if read_failures < MAX_BATCH_READ_RETRIES + && let Some(thread_id) = self.thread_id + { + if let Some(search) = self.search.as_mut() { + search.awaiting = Some(PendingHistorySearch::Batch { + cursor, + boundary_if_exhausted, + read_failures: read_failures + 1, + }); + } + app_event_tx.send(AppEvent::LookupMessageHistoryBatch { + thread_id, + cursor, + log_id, + }); + return Some(HistorySearchResult::Pending); + } + + Some(if boundary_if_exhausted { + if let Some(search) = self.search.as_mut() { + search.awaiting = None; + } + HistorySearchResult::AtBoundary + } else { + self.search = None; + HistorySearchResult::Unavailable + }) + } + + fn pending_batch(&self, cursor: HistoryBatchCursor) -> Option<(bool, u8)> { + let Some(PendingHistorySearch::Batch { + cursor: awaited_cursor, + boundary_if_exhausted, + read_failures, + }) = self.search.as_ref().and_then(|search| search.awaiting) + else { + return None; + }; + (awaited_cursor == cursor).then_some((boundary_if_exhausted, read_failures)) + } + + /// Switches an older search from the single newest-entry probe to bounded batch lookups. + /// + /// Keeping the first probe on the single-entry path avoids fetching a batch when the newest + /// persistent entry matches, while every later miss is amortized across a bounded batch. + pub(super) fn advance_older_search_after_entry_miss( + &mut self, + offset: usize, + boundary_if_exhausted: bool, + app_event_tx: &AppEventSender, + ) -> HistorySearchResult { + let Some(next_offset) = offset.checked_sub(1) else { + return self + .exhausted_search_result(HistorySearchDirection::Older, boundary_if_exhausted); + }; + self.advance_older_search_with_batches_from( + HistoryBatchCursor::new(next_offset), + boundary_if_exhausted, + app_event_tx, + ) + } + + /// Scans cached offsets from a batch cursor and requests the first uncached older range. + /// + /// The byte anchor remains usable only while scanning begins at its exact `end_offset`; moving + /// past cached entries creates an offset-only cursor because no validated byte boundary exists + /// for those intermediate positions. + fn advance_older_search_with_batches_from( + &mut self, + mut cursor: HistoryBatchCursor, + boundary_if_exhausted: bool, + app_event_tx: &AppEventSender, + ) -> HistorySearchResult { + let mut offset = cursor.end_offset(); + loop { + if let Some(entry) = self.entry_at_cached_offset(offset) { + if self.search_matches(&entry) && self.search_result_is_unique(&entry) { + return self.search_match(offset, entry); + } + } else if !self.fetched_history.contains_key(&offset) + && offset < self.persistent_entry_count + { + return self.request_older_search_batch( + cursor, + boundary_if_exhausted, + app_event_tx, + ); + } + + let Some(next_offset) = offset.checked_sub(1) else { + return self + .exhausted_search_result(HistorySearchDirection::Older, boundary_if_exhausted); + }; + offset = next_offset; + cursor = HistoryBatchCursor::new(offset); + } + } + + pub(super) fn request_older_search_batch( + &mut self, + cursor: HistoryBatchCursor, + boundary_if_exhausted: bool, + app_event_tx: &AppEventSender, + ) -> HistorySearchResult { + let (Some(thread_id), Some(log_id)) = (self.thread_id, self.persistent_log_id) else { + return self + .exhausted_search_result(HistorySearchDirection::Older, boundary_if_exhausted); + }; + if let Some(search) = self.search.as_mut() { + search.awaiting = Some(PendingHistorySearch::Batch { + cursor, + boundary_if_exhausted, + read_failures: 0, + }); + } + app_event_tx.send(AppEvent::LookupMessageHistoryBatch { + thread_id, + cursor, + log_id, + }); + HistorySearchResult::Pending + } +} diff --git a/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch_tests.rs b/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch_tests.rs new file mode 100644 index 0000000000..883b5df90c --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/chat_composer_history/search_batch_tests.rs @@ -0,0 +1,308 @@ +use pretty_assertions::assert_eq; +use tokio::sync::mpsc::UnboundedReceiver; +use tokio::sync::mpsc::unbounded_channel; + +use super::*; +use crate::app_event::AppEvent; +use crate::app_event::HistoryBatchEntryResponse; + +fn thread_id(value: u8) -> ThreadId { + ThreadId::from_string(&format!("00000000-0000-0000-0000-{value:012}")) + .expect("thread id should parse") +} + +fn batch_entry(offset: usize, entry: &str) -> HistoryBatchEntryResponse { + HistoryBatchEntryResponse { + offset, + entry: Some(entry.to_string()), + } +} + +fn found(entry: &str) -> HistorySearchResult { + HistorySearchResult::Found(HistoryEntry::new(entry.to_string())) +} + +fn history( + entry_count: usize, +) -> ( + ChatComposerHistory, + AppEventSender, + UnboundedReceiver, +) { + let (tx, rx) = unbounded_channel(); + let tx = AppEventSender::new(tx); + let mut history = ChatComposerHistory::new(); + history.set_metadata(thread_id(/*value*/ 1), /*log_id*/ 42, entry_count); + (history, tx, rx) +} + +fn recv_batch_request(rx: &mut UnboundedReceiver) -> (HistoryBatchCursor, u64) { + let AppEvent::LookupMessageHistoryBatch { cursor, log_id, .. } = + rx.try_recv().expect("batch request") + else { + panic!("expected batch request"); + }; + (cursor, log_id) +} + +fn recv_entry_request(rx: &mut UnboundedReceiver) -> (usize, u64) { + let AppEvent::LookupMessageHistoryEntry { offset, log_id, .. } = + rx.try_recv().expect("entry request") + else { + panic!("expected entry request"); + }; + (offset, log_id) +} + +fn start_older_search( + history: &mut ChatComposerHistory, + tx: &AppEventSender, + rx: &mut UnboundedReceiver, + query: &str, + newest_offset: usize, +) -> HistoryBatchCursor { + assert_eq!( + history.search( + query, + HistorySearchDirection::Older, + /*restart*/ true, + tx + ), + HistorySearchResult::Pending + ); + let (offset, log_id) = recv_entry_request(rx); + assert_eq!((offset, log_id), (newest_offset, 42)); + assert_eq!( + history.on_entry_response(log_id, offset, Some("unrelated entry".to_string()), tx), + HistoryEntryResponse::Search(HistorySearchResult::Pending) + ); + let (cursor, _) = recv_batch_request(rx); + assert_eq!(cursor.end_offset(), newest_offset - 1); + cursor +} + +#[test] +fn search_batch_late_data_is_cache_only_after_cancel_or_query_edit() { + let (mut cancelled, tx, mut rx) = history(/*entry_count*/ 5); + start_older_search( + &mut cancelled, + &tx, + &mut rx, + "cached", + /*newest_offset*/ 4, + ); + cancelled.reset_search(); + assert_eq!( + cancelled.on_batch_response( + /*log_id*/ 42, + HistoryBatchCursor::new(/*end_offset*/ 3), + vec![batch_entry(/*offset*/ 3, "cached match")], + Some(HistoryBatchCursor::new(/*end_offset*/ 2)), + &tx, + ), + None + ); + assert_eq!( + cancelled.search( + "cached", + HistorySearchDirection::Older, + /*restart*/ true, + &tx + ), + found("cached match") + ); + assert!(rx.try_recv().is_err()); + + let (mut edited, tx, mut rx) = history(/*entry_count*/ 5); + start_older_search(&mut edited, &tx, &mut rx, "old", /*newest_offset*/ 4); + assert_eq!( + edited.search( + "new", + HistorySearchDirection::Older, + /*restart*/ true, + &tx + ), + HistorySearchResult::Pending + ); + let (offset, _) = recv_entry_request(&mut rx); + assert_eq!(offset, 3); + assert_eq!( + edited.on_batch_response( + /*log_id*/ 42, + HistoryBatchCursor::new(/*end_offset*/ 3), + vec![batch_entry(/*offset*/ 3, "old data")], + Some(HistoryBatchCursor::new(/*end_offset*/ 2)), + &tx, + ), + None + ); + assert_eq!( + edited.on_entry_response( + /*log_id*/ 42, + /*offset*/ 3, + Some("new current match".to_string()), + &tx, + ), + HistoryEntryResponse::Search(found("new current match")) + ); +} + +#[test] +fn search_batch_rejects_stale_thread_and_log_metadata() { + let (mut history, tx, mut rx) = history(/*entry_count*/ 5); + start_older_search( + &mut history, + &tx, + &mut rx, + "stale", + /*newest_offset*/ 4, + ); + history.set_metadata( + thread_id(/*value*/ 2), + /*log_id*/ 43, + /*entry_count*/ 5, + ); + assert_eq!( + history.on_batch_response( + /*log_id*/ 42, + HistoryBatchCursor::new(/*end_offset*/ 3), + vec![batch_entry(/*offset*/ 3, "stale match")], + Some(HistoryBatchCursor::new(/*end_offset*/ 2)), + &tx, + ), + None + ); + assert!(history.fetched_history.is_empty()); +} + +#[test] +fn search_batch_read_failure_stops_after_bounded_retries() { + let (mut history, tx, mut rx) = history(/*entry_count*/ 5); + let cursor = start_older_search( + &mut history, + &tx, + &mut rx, + "retry", + /*newest_offset*/ 4, + ); + + for _ in 0..MAX_BATCH_READ_RETRIES { + assert_eq!( + history.on_batch_error(/*log_id*/ 42, cursor, &tx), + Some(HistorySearchResult::Pending) + ); + let (retry_cursor, log_id) = recv_batch_request(&mut rx); + assert_eq!((retry_cursor, log_id), (cursor, 42)); + } + + assert_eq!( + history.on_batch_error(/*log_id*/ 42, cursor, &tx), + Some(HistorySearchResult::Unavailable) + ); + assert!(rx.try_recv().is_err()); + assert!(history.search.is_none()); + assert_eq!( + history.search( + "retry", + HistorySearchDirection::Older, + /*restart*/ false, + &tx, + ), + HistorySearchResult::Pending + ); + assert_eq!(recv_entry_request(&mut rx), (3, 42)); +} + +#[test] +fn search_batch_match_preserves_cursor_for_next_older_search() { + let (mut history, tx, mut rx) = history(/*entry_count*/ 7); + let cursor = start_older_search( + &mut history, + &tx, + &mut rx, + "needle", + /*newest_offset*/ 6, + ); + let continuation = HistoryBatchCursor::new(/*end_offset*/ 2); + + assert_eq!( + history.on_batch_response( + /*log_id*/ 42, + cursor, + vec![ + batch_entry(/*offset*/ 5, "unrelated newer"), + batch_entry(/*offset*/ 4, "needle first"), + batch_entry(/*offset*/ 3, "unrelated older"), + ], + Some(continuation), + &tx, + ), + Some(found("needle first")) + ); + + assert_eq!( + history.search( + "needle", + HistorySearchDirection::Older, + /*restart*/ false, + &tx, + ), + HistorySearchResult::Pending + ); + let (requested_cursor, _) = recv_batch_request(&mut rx); + assert_eq!(requested_cursor, continuation); + assert!(rx.try_recv().is_err()); + + assert_eq!( + history.on_batch_response( + /*log_id*/ 42, + continuation, + vec![batch_entry(/*offset*/ 2, "needle second")], + /*next_older_cursor*/ None, + &tx, + ), + Some(found("needle second")) + ); +} + +#[test] +fn search_batch_absent_1024_uses_one_single_and_eight_batches() { + let (mut history, tx, mut rx) = history(/*entry_count*/ 1_024); + let mut cursor = start_older_search( + &mut history, + &tx, + &mut rx, + "absent", + /*newest_offset*/ 1_023, + ); + let mut batches = 0; + + loop { + let end_offset = cursor.end_offset(); + let start_offset = end_offset.saturating_sub(127); + let entries = (start_offset..=end_offset) + .rev() + .map(|offset| batch_entry(offset, "unrelated entry")) + .collect(); + let next_older_cursor = start_offset.checked_sub(1).map(HistoryBatchCursor::new); + let expected = if next_older_cursor.is_some() { + HistorySearchResult::Pending + } else { + HistorySearchResult::NotFound + }; + assert_eq!( + history.on_batch_response(/*log_id*/ 42, cursor, entries, next_older_cursor, &tx,), + Some(expected) + ); + batches += 1; + let Some(expected_cursor) = next_older_cursor else { + break; + }; + let (next, _) = recv_batch_request(&mut rx); + assert_eq!(next, expected_cursor); + cursor = next; + } + + assert_eq!((1 + batches, batches), (9, 8)); + assert!(rx.try_recv().is_err()); +} diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index ad30dedf5a..85db3b79e1 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -19,6 +19,7 @@ use std::path::PathBuf; use crate::app::app_server_requests::ResolvedAppServerRequest; use crate::app_event::AppEvent; use crate::app_event::ConnectorsSnapshot; +use crate::app_event::HistoryLookupResponse; use crate::app_event_sender::AppEventSender; use crate::bottom_pane::pending_input_preview::PendingInputPreview; use crate::bottom_pane::pending_thread_approvals::PendingThreadApprovals; @@ -1612,16 +1613,28 @@ impl BottomPane { || self.composer.is_in_paste_burst() } - pub(crate) fn on_history_entry_response( - &mut self, - log_id: u64, - offset: usize, - entry: Option, - ) { - let updated = self - .composer - .on_history_entry_response(log_id, offset, entry); - + pub(crate) fn on_history_lookup_response(&mut self, response: HistoryLookupResponse) { + let updated = match response { + HistoryLookupResponse::Entry { + offset, + log_id, + entry, + } => self + .composer + .on_history_entry_response(log_id, offset, entry), + HistoryLookupResponse::Batch { + cursor, + log_id, + entries, + next_older_cursor, + } => { + self.composer + .on_history_batch_response(log_id, cursor, entries, next_older_cursor) + } + HistoryLookupResponse::BatchError { cursor, log_id } => { + self.composer.on_history_batch_error(log_id, cursor) + } + }; if updated { self.composer.sync_popups(); self.request_redraw(); diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index fa745c652e..5b7cb88692 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1176,13 +1176,7 @@ impl ChatWidget { } pub(crate) fn handle_history_entry_response(&mut self, event: HistoryLookupResponse) { - let HistoryLookupResponse { - offset, - log_id, - entry, - } = event; - self.bottom_pane - .on_history_entry_response(log_id, offset, entry); + self.bottom_pane.on_history_lookup_response(event); } pub(crate) fn pre_draw_tick(&mut self) { diff --git a/codex-rs/tui/src/chatwidget/tests/history_replay.rs b/codex-rs/tui/src/chatwidget/tests/history_replay.rs index f8626dfcc0..e610cc60b4 100644 --- a/codex-rs/tui/src/chatwidget/tests/history_replay.rs +++ b/codex-rs/tui/src/chatwidget/tests/history_replay.rs @@ -145,7 +145,7 @@ async fn replayed_user_messages_seed_composer_history() { }; offset }; - let response = |offset, entry: &str| HistoryLookupResponse { + let response = |offset, entry: &str| HistoryLookupResponse::Entry { offset, log_id: 1, entry: Some(entry.to_string()),