From d62353a3127335172479112e8a775f9fbb36a44c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 31 Jul 2026 18:53:04 +0000 Subject: [PATCH] Load local session pickers from the state DB first (#36378) ## What changed - Start local resume and fork picker listings from indexed state DB metadata, while preserving the existing store-backed behavior for remote workspaces. - Keep the selected listing mode across pagination and fall back to the normal local store listing when the initial state DB request returns no rows or fails. - Add coverage for pagination, fallback rendering, and selecting an indexed thread without validating its rollout file. GitOrigin-RevId: 9005fbe73134e261bc6dd3f588c4fff494bdeea0 --- codex-rs/tui/src/resume_picker.rs | 362 ++++++++++++++---- .../tui/src/resume_picker/page_loading.rs | 100 +++++ ..._resume_picker_db_fallback_transition.snap | 15 + 3 files changed, 393 insertions(+), 84 deletions(-) create mode 100644 codex-rs/tui/src/resume_picker/page_loading.rs create mode 100644 codex-rs/tui/src/snapshots/codex_tui__resume_picker__tests__resume_picker_db_fallback_transition.snap diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index 1e06b1cd88..1131ce9509 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -64,6 +64,11 @@ use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::warn; use unicode_width::UnicodeWidthStr; +mod page_loading; + +use page_loading::PageLoadMode; +use page_loading::PaginationState; + const PAGE_SIZE: usize = 25; const LOAD_NEAR_THRESHOLD: usize = 5; const SESSION_META_INDENT_WIDTH: usize = 2; @@ -143,6 +148,7 @@ struct PageLoadRequest { cursor: Option, request_token: usize, search_token: Option, + mode: PageLoadMode, cwd_filter: Option, provider_filter: ProviderFilter, sort_key: ThreadSortKey, @@ -284,6 +290,7 @@ struct SessionPickerRunOptions { view_persistence: Option, pager_keymap: PagerKeymap, list_keymap: ListKeymap, + initial_page_mode: PageLoadMode, } /// Interactive session picker that lists app-server threads with simple search, @@ -370,6 +377,11 @@ async fn run_resume_picker_with_launch_context( }), pager_keymap: runtime_keymap.pager, list_keymap: runtime_keymap.list, + initial_page_mode: if uses_remote_workspace { + PageLoadMode::StoreDefault + } else { + PageLoadMode::StateDbOnly + }, }; run_session_picker_with_loader( tui, @@ -416,6 +428,11 @@ pub async fn run_fork_picker_with_app_server( }), pager_keymap: runtime_keymap.pager, list_keymap: runtime_keymap.list, + initial_page_mode: if uses_remote_workspace { + PageLoadMode::StoreDefault + } else { + PageLoadMode::StateDbOnly + }, }; run_session_picker_with_loader( tui, @@ -453,6 +470,7 @@ async fn run_session_picker_with_loader( state.pager_keymap = options.pager_keymap; state.list_keymap = options.list_keymap; state.launch_context = options.launch_context; + state.initial_page_mode = options.initial_page_mode; state.start_initial_load(); state.request_frame(); @@ -572,6 +590,7 @@ fn spawn_app_server_page_loader( request.provider_filter, request.sort_key, include_non_interactive, + matches!(request.mode, PageLoadMode::StateDbOnly), ) .await; let _ = bg_tx.send(BackgroundEvent::Page { @@ -676,25 +695,7 @@ struct PickerState { overlay: Option, pager_keymap: PagerKeymap, list_keymap: ListKeymap, -} - -struct PaginationState { - next_cursor: Option, - num_scanned_files: usize, - reached_scan_cap: bool, - loading: LoadingState, -} - -#[derive(Clone, Copy, Debug)] -enum LoadingState { - Idle, - Pending(PendingLoad), -} - -#[derive(Clone, Copy, Debug)] -struct PendingLoad { - request_token: usize, - search_token: Option, + initial_page_mode: PageLoadMode, } #[derive(Clone, Copy, Debug)] @@ -733,12 +734,6 @@ enum LoadTrigger { Search { token: usize }, } -impl LoadingState { - fn is_pending(&self) -> bool { - matches!(self, LoadingState::Pending(_)) - } -} - async fn load_app_server_page( app_server: &mut AppServerSession, cursor: Option, @@ -746,6 +741,7 @@ async fn load_app_server_page( provider_filter: ProviderFilter, sort_key: ThreadSortKey, include_non_interactive: bool, + use_state_db_only: bool, ) -> std::io::Result { let response = app_server .thread_list(thread_list_params( @@ -754,6 +750,7 @@ async fn load_app_server_page( provider_filter, sort_key, include_non_interactive, + use_state_db_only, )) .await .map_err(std::io::Error::other)?; @@ -934,12 +931,7 @@ impl PickerState { Self { requester, relative_time_reference: None, - pagination: PaginationState { - next_cursor: None, - num_scanned_files: 0, - reached_scan_cap: false, - loading: LoadingState::Idle, - }, + pagination: PaginationState::new(), all_rows: Vec::new(), filtered_rows: Vec::new(), seen_rows: HashSet::new(), @@ -973,6 +965,7 @@ impl PickerState { overlay: None, pager_keymap: RuntimeKeymap::defaults().pager, list_keymap: RuntimeKeymap::defaults().list, + initial_page_mode: PageLoadMode::StoreDefault, } } @@ -1295,16 +1288,16 @@ impl PickerState { }; let request_token = self.allocate_request_token(); - self.pagination.loading = LoadingState::Pending(PendingLoad { - request_token, - search_token, - }); + let mode = self.initial_page_mode; + self.pagination + .start_load(request_token, search_token, mode); self.request_frame(); (self.picker_loader)(PickerLoadRequest::Page(PageLoadRequest { cursor: None, request_token, search_token, + mode, cwd_filter: self.active_cwd_filter(), provider_filter: self.provider_filter.clone(), sort_key: self.sort_key, @@ -1318,14 +1311,35 @@ impl PickerState { search_token, page, } => { - let pending = match self.pagination.loading { - LoadingState::Pending(pending) => pending, - LoadingState::Idle => return Ok(()), + let Some(pending) = self.pagination.finish_load(request_token) else { + return Ok(()); }; - if pending.request_token != request_token { + let page_has_rows = matches!(&page, Ok(page) if !page.rows.is_empty()); + // Fall back only when the initial DB listing is unusable. Once SQLite returns + // rows, its pagination is authoritative and an empty later page ends the list. + let should_restart_from_store = pending.mode == PageLoadMode::StateDbOnly + && self.all_rows.is_empty() + && !page_has_rows; + if should_restart_from_store { + let request_token = self.allocate_request_token(); + let search_token = pending.search_token.or(search_token); + self.pagination.reset(); + self.pagination.start_load( + request_token, + search_token, + PageLoadMode::StoreDefault, + ); + (self.picker_loader)(PickerLoadRequest::Page(PageLoadRequest { + cursor: None, + request_token, + search_token, + mode: PageLoadMode::StoreDefault, + cwd_filter: self.active_cwd_filter(), + provider_filter: self.provider_filter.clone(), + sort_key: self.sort_key, + })); return Ok(()); } - self.pagination.loading = LoadingState::Idle; let page = page.map_err(color_eyre::Report::from)?; self.ingest_page(page); self.complete_pending_page_down(); @@ -1371,28 +1385,21 @@ impl PickerState { } fn reset_pagination(&mut self) { - self.pagination.next_cursor = None; - self.pagination.num_scanned_files = 0; - self.pagination.reached_scan_cap = false; - self.pagination.loading = LoadingState::Idle; + self.pagination.reset(); self.frozen_footer_percent = None; } fn ingest_page(&mut self, page: PickerPage) { - if let Some(cursor) = page.next_cursor.clone() { - self.pagination.next_cursor = Some(cursor); - } else { - self.pagination.next_cursor = None; - } - self.pagination.num_scanned_files = self - .pagination - .num_scanned_files - .saturating_add(page.num_scanned_files); - if page.reached_scan_cap { - self.pagination.reached_scan_cap = true; - } + let PickerPage { + rows, + next_cursor, + num_scanned_files, + reached_scan_cap, + } = page; + self.pagination + .complete_page(next_cursor, num_scanned_files, reached_scan_cap); - for row in page.rows { + for row in rows { if let Some(seen_key) = row.seen_key() { if self.seen_rows.insert(seen_key) { self.all_rows.push(row); @@ -1552,7 +1559,7 @@ impl PickerState { if minimum_rows == 0 { return; } - if self.pagination.loading.is_pending() || self.pagination.next_cursor.is_none() { + if self.pagination.is_loading() || self.pagination.next_cursor.is_none() { return; } let rendered_rows = if self.filtered_rows.is_empty() { @@ -1577,10 +1584,7 @@ impl PickerState { } fn maybe_load_more_for_scroll(&mut self) { - if self.pagination.loading.is_pending() { - return; - } - if self.pagination.next_cursor.is_none() { + if self.pagination.is_loading() || self.pagination.next_cursor.is_none() { return; } if self.filtered_rows.is_empty() { @@ -1593,10 +1597,7 @@ impl PickerState { } fn load_more_if_needed(&mut self, trigger: LoadTrigger) { - if self.pagination.loading.is_pending() { - return; - } - let Some(cursor) = self.pagination.next_cursor.clone() else { + let Some((cursor, mode)) = self.pagination.next_page() else { return; }; self.freeze_footer_percent(); @@ -1605,16 +1606,15 @@ impl PickerState { LoadTrigger::Scroll => None, LoadTrigger::Search { token } => Some(token), }; - self.pagination.loading = LoadingState::Pending(PendingLoad { - request_token, - search_token, - }); + self.pagination + .start_load(request_token, search_token, mode); self.request_frame(); (self.picker_loader)(PickerLoadRequest::Page(PageLoadRequest { cursor: Some(cursor), request_token, search_token, + mode, cwd_filter: self.active_cwd_filter(), provider_filter: self.provider_filter.clone(), sort_key: self.sort_key, @@ -1843,6 +1843,7 @@ fn thread_list_params( provider_filter: ProviderFilter, sort_key: ThreadSortKey, include_non_interactive: bool, + use_state_db_only: bool, ) -> ThreadListParams { ThreadListParams { cursor, @@ -1859,7 +1860,7 @@ fn thread_list_params( parent_thread_id: None, ancestor_thread_id: None, cwd: cwd_filter.map(|cwd| ThreadListCwdFilter::One(cwd.to_string_lossy().into_owned())), - use_state_db_only: false, + use_state_db_only, search_term: None, } } @@ -2114,7 +2115,7 @@ fn picker_footer_progress_label(state: &PickerState, list_height: u16, width: u1 } else { state.selected.saturating_add(1) }; - let total = if state.pagination.loading.is_pending() { + let total = if state.pagination.is_loading() { format!("{}…", state.filtered_rows.len()) } else { state.filtered_rows.len().to_string() @@ -2132,7 +2133,7 @@ fn picker_footer_progress_label(state: &PickerState, list_height: u16, width: u1 } fn picker_footer_percent(state: &PickerState, list_height: u16) -> u8 { - if state.pagination.loading.is_pending() { + if state.pagination.is_loading() { return state.frozen_footer_percent.unwrap_or_else(|| { if state.filtered_rows.is_empty() { 0 @@ -2508,15 +2509,13 @@ fn render_list(frame: &mut crate::custom_terminal::Frame, area: Rect, state: &Pi } } - if state.pagination.loading.is_pending() - && y < content_area.y.saturating_add(content_area.height) - { + if state.pagination.is_loading() && y < content_area.y.saturating_add(content_area.height) { let loading_line: Line = vec![" ".into(), "Loading older sessions…".italic().dim()].into(); let rect = Rect::new(area.x, y, area.width, 1); frame.render_widget_ref(&loading_line, rect); } if show_more_below { - let label = if state.pagination.loading.is_pending() { + let label = if state.pagination.is_loading() { "↓ loading more" } else { "↓ more" @@ -3209,7 +3208,7 @@ fn format_timestamp(ts: DateTime) -> String { fn render_empty_state_line(state: &PickerState) -> Line<'static> { if !state.query.is_empty() { if state.search_state.is_active() - || (state.pagination.loading.is_pending() && state.pagination.next_cursor.is_some()) + || (state.pagination.is_loading() && state.pagination.next_cursor.is_some()) { return vec!["Searching…".italic().dim()].into(); } @@ -3223,7 +3222,7 @@ fn render_empty_state_line(state: &PickerState) -> Line<'static> { return vec!["No results for your search".italic().dim()].into(); } - if state.pagination.loading.is_pending() { + if state.pagination.is_loading() { if state.all_rows.is_empty() && state.pagination.num_scanned_files == 0 { return vec!["Loading sessions…".italic().dim()].into(); } @@ -3276,6 +3275,26 @@ mod tests { }) } + async fn deliver_page( + state: &mut PickerState, + request: &PageLoadRequest, + page: std::io::Result, + ) { + state + .handle_background_event(BackgroundEvent::Page { + request_token: request.request_token, + search_token: request.search_token, + page, + }) + .await + .expect("page event"); + } + + fn ok_page(rows: Vec, next_cursor: Option<&str>) -> std::io::Result { + let n = rows.len(); + Ok(page(rows, next_cursor, n, /*reached_scan_cap*/ false)) + } + fn make_row(path: &str, ts: &str, preview: &str) -> Row { let timestamp = parse_timestamp_str(ts).expect("timestamp should parse"); Row { @@ -3290,6 +3309,24 @@ mod tests { } } + fn local_db_first_state() -> (PickerState, Arc>>) { + let requests = Arc::new(Mutex::new(Vec::new())); + let request_sink = Arc::clone(&requests); + let loader = page_only_loader(move |request| { + request_sink.lock().unwrap().push(request); + }); + let mut state = PickerState::new( + FrameRequester::test_dummy(), + loader, + ProviderFilter::MatchDefault(String::from("openai")), + /*show_all*/ true, + /*filter_cwd*/ None, + SessionPickerAction::Resume, + ); + state.initial_page_mode = PageLoadMode::StateDbOnly; + (state, requests) + } + fn footer_lines_text(state: &PickerState, width: u16) -> String { footer_hint_lines(state, width) .into_iter() @@ -3352,12 +3389,165 @@ mod tests { ProviderFilter::MatchDefault(String::from("openai")), ThreadSortKey::UpdatedAt, /*include_non_interactive*/ false, + /*use_state_db_only*/ true, ); assert_eq!( params.cwd, Some(ThreadListCwdFilter::One(String::from("/tmp/project"))) ); + assert!(params.use_state_db_only); + } + + #[tokio::test] + async fn local_db_first_treats_empty_later_page_as_end_of_db_listing() { + let (mut state, requests) = local_db_first_state(); + state.start_initial_load(); + let db_request = requests.lock().unwrap()[0].clone(); + deliver_page( + &mut state, + &db_request, + ok_page( + vec![make_row( + "/tmp/indexed.jsonl", + "2025-01-03T00:00:00Z", + "indexed metadata", + )], + Some("db-cursor"), + ), + ) + .await; + + state.maybe_load_more_for_scroll(); + let later_db_request = requests.lock().unwrap()[1].clone(); + assert_eq!(later_db_request.mode, PageLoadMode::StateDbOnly); + deliver_page( + &mut state, + &later_db_request, + ok_page(Vec::new(), /*next_cursor*/ None), + ) + .await; + + assert_eq!(requests.lock().unwrap().len(), 2); + assert_eq!(state.all_rows[0].preview, "indexed metadata"); + assert!(state.pagination.next_cursor.is_none()); + } + + #[tokio::test] + async fn local_db_first_falls_back_when_initial_db_page_is_empty() { + let (mut state, requests) = local_db_first_state(); + state.start_initial_load(); + let db_request = requests.lock().unwrap()[0].clone(); + deliver_page( + &mut state, + &db_request, + ok_page(Vec::new(), /*next_cursor*/ None), + ) + .await; + + let fallback_request = requests.lock().unwrap()[1].clone(); + assert_eq!(fallback_request.mode, PageLoadMode::StoreDefault); + assert!(fallback_request.cursor.is_none()); + state.relative_time_reference = + Some(parse_timestamp_str("2025-01-04T00:00:00Z").expect("timestamp")); + state.update_viewport(/*rows*/ 12, /*width*/ 80); + let render_state = |state: &PickerState| { + use crate::custom_terminal::Terminal; + use crate::test_backend::VT100Backend; + + let backend = VT100Backend::new(/*width*/ 80, /*height*/ 12); + let mut terminal = Terminal::with_options(backend).expect("terminal"); + terminal.set_viewport_area(Rect::new( + /*x*/ 0, /*y*/ 0, /*width*/ 80, /*height*/ 12, + )); + { + let mut frame = terminal.get_frame(); + let area = frame.area(); + render_list(&mut frame, area, state); + } + terminal.flush().expect("flush"); + terminal.backend().to_string() + }; + let loading_snapshot = render_state(&state); + + deliver_page( + &mut state, + &fallback_request, + ok_page( + vec![ + make_row( + "/tmp/first.jsonl", + "2025-01-03T00:00:00Z", + "first store result", + ), + make_row( + "/tmp/repaired.jsonl", + "2025-01-03T00:00:00Z", + "repaired store result", + ), + make_row( + "/tmp/third.jsonl", + "2025-01-03T00:00:00Z", + "third store result", + ), + ], + /*next_cursor*/ None, + ), + ) + .await; + + assert_eq!( + state + .all_rows + .iter() + .map(|row| row.preview.as_str()) + .collect::>(), + vec![ + "first store result", + "repaired store result", + "third store result" + ] + ); + assert_eq!(state.pagination.num_scanned_files, 3); + let repaired_snapshot = render_state(&state); + assert_snapshot!( + "resume_picker_db_fallback_transition", + format!( + "---- loading fallback ----\n{loading_snapshot}\n---- repaired results ----\n{repaired_snapshot}" + ) + ); + } + + #[tokio::test] + async fn local_db_first_accepts_db_thread_id_without_rollout_validation() { + let (mut state, requests) = local_db_first_state(); + state.start_initial_load(); + let db_request = requests.lock().unwrap()[0].clone(); + let thread_id = ThreadId::new(); + let mut row = make_row( + "/tmp/missing-rollout.jsonl", + "2025-01-01T00:00:00Z", + "indexed metadata", + ); + row.thread_id = Some(thread_id); + deliver_page( + &mut state, + &db_request, + ok_page(vec![row], /*next_cursor*/ None), + ) + .await; + + let selection = state + .handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)) + .await + .expect("enter should not abort the picker"); + assert!(matches!( + selection, + Some(SessionSelection::Resume(SessionTarget { + thread_id: selected_thread_id, + .. + })) if selected_thread_id == thread_id + )); } #[test] @@ -3577,6 +3767,7 @@ mod tests { ProviderFilter::Any, ThreadSortKey::UpdatedAt, /*include_non_interactive*/ false, + /*use_state_db_only*/ false, ); assert_eq!(params.cursor, Some(String::from("cursor-1"))); @@ -3599,6 +3790,7 @@ mod tests { ProviderFilter::Any, ThreadSortKey::UpdatedAt, /*include_non_interactive*/ true, + /*use_state_db_only*/ false, ); assert_eq!(params.cursor, Some(String::from("cursor-1"))); @@ -3632,6 +3824,7 @@ mod tests { let guard = recorded_requests.lock().unwrap(); assert_eq!(guard.len(), 1); assert_eq!(guard[0].cwd_filter, remote_cwd); + assert_eq!(guard[0].mode, PageLoadMode::StoreDefault); } let row = Row { @@ -4065,10 +4258,11 @@ mod tests { state.selected = 9; state.scroll_top = 9; state.pagination.next_cursor = Some(PageCursor::AppServer(String::from("cursor-1"))); - state.pagination.loading = LoadingState::Pending(PendingLoad { - request_token: 1, - search_token: None, - }); + state.pagination.start_load( + /*request_token*/ 1, + /*search_token*/ None, + PageLoadMode::StoreDefault, + ); state.frozen_footer_percent = Some(37); let label = picker_footer_progress_label(&state, /*list_height*/ 6, /*width*/ 80); @@ -5669,7 +5863,7 @@ session_picker_view = "dense" .unwrap(); assert_eq!(state.selected, 9); - assert!(state.pagination.loading.is_pending()); + assert!(state.pagination.is_loading()); assert_eq!(recorded_requests.lock().unwrap().len(), 1); assert_eq!( picker_footer_progress_label(&state, /*list_height*/ 5, /*width*/ 80), diff --git a/codex-rs/tui/src/resume_picker/page_loading.rs b/codex-rs/tui/src/resume_picker/page_loading.rs new file mode 100644 index 0000000000..c580a61505 --- /dev/null +++ b/codex-rs/tui/src/resume_picker/page_loading.rs @@ -0,0 +1,100 @@ +use super::PageCursor; + +/// Selects whether thread listing trusts indexed metadata or uses the store's normal behavior. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum PageLoadMode { + /// Return only the store's indexed metadata without repairing it from rollout files. + StateDbOnly, + /// Use the store's normal listing behavior, which may include repair for local stores. + StoreDefault, +} + +/// Tracks the current request and the cursor policy inherited by its next page. +pub(super) struct PaginationState { + pub(super) next_cursor: Option, + next_page_mode: PageLoadMode, + pub(super) num_scanned_files: usize, + pub(super) reached_scan_cap: bool, + loading: LoadingState, +} + +#[derive(Clone, Copy, Debug)] +enum LoadingState { + Idle, + Pending(PendingLoad), +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct PendingLoad { + request_token: usize, + pub(super) search_token: Option, + pub(super) mode: PageLoadMode, +} + +impl PaginationState { + pub(super) fn new() -> Self { + Self { + next_cursor: None, + next_page_mode: PageLoadMode::StoreDefault, + num_scanned_files: 0, + reached_scan_cap: false, + loading: LoadingState::Idle, + } + } + + pub(super) fn reset(&mut self) { + *self = Self::new(); + } + + pub(super) fn start_load( + &mut self, + request_token: usize, + search_token: Option, + mode: PageLoadMode, + ) { + self.loading = LoadingState::Pending(PendingLoad { + request_token, + search_token, + mode, + }); + } + + /// Completes only the matching in-flight request, leaving newer requests untouched. + pub(super) fn finish_load(&mut self, request_token: usize) -> Option { + let LoadingState::Pending(pending) = self.loading else { + return None; + }; + if pending.request_token != request_token { + return None; + } + self.loading = LoadingState::Idle; + self.next_page_mode = pending.mode; + Some(pending) + } + + /// Records a page cursor and accumulates scan progress across the current listing. + pub(super) fn complete_page( + &mut self, + next_cursor: Option, + num_scanned_files: usize, + reached_scan_cap: bool, + ) { + self.next_cursor = next_cursor; + self.num_scanned_files = self.num_scanned_files.saturating_add(num_scanned_files); + self.reached_scan_cap |= reached_scan_cap; + } + + /// Returns the next cursor together with the listing mode that produced it. + pub(super) fn next_page(&self) -> Option<(PageCursor, PageLoadMode)> { + if self.is_loading() { + return None; + } + self.next_cursor + .clone() + .map(|cursor| (cursor, self.next_page_mode)) + } + + pub(super) fn is_loading(&self) -> bool { + matches!(self.loading, LoadingState::Pending(_)) + } +} diff --git a/codex-rs/tui/src/snapshots/codex_tui__resume_picker__tests__resume_picker_db_fallback_transition.snap b/codex-rs/tui/src/snapshots/codex_tui__resume_picker__tests__resume_picker_db_fallback_transition.snap new file mode 100644 index 0000000000..61158efc1a --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__resume_picker__tests__resume_picker_db_fallback_transition.snap @@ -0,0 +1,15 @@ +--- +source: tui/src/resume_picker.rs +expression: "format!(\"---- loading fallback ----\\n{loading_snapshot}\\n---- repaired results ----\\n{repaired_snapshot}\")" +--- +---- loading fallback ---- +Loading sessions… +---- repaired results ---- +❯ first store result + 1d ago ⌁ no cwd  no branch + + repaired store result + 1d ago ⌁ no cwd  no branch + + third store result + 1d ago ⌁ no cwd  no branch