diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index e3f1072d1d..c054f23d1d 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -76,6 +76,7 @@ mod archive; mod catalog; mod layout; mod page_loading; +mod sources; #[cfg(test)] #[path = "resume_picker_color_tests.rs"] @@ -197,16 +198,20 @@ struct PageLoadRequest { enum PickerLoadRequest { Page(PageLoadRequest), Preview { + backend_id: crate::backend_id::BackendId, thread_id: ThreadId, }, Transcript { + backend_id: crate::backend_id::BackendId, thread_id: ThreadId, cancellation: oneshot::Receiver<()>, }, Archive { + backend_id: crate::backend_id::BackendId, thread_id: ThreadId, }, Unarchive { + backend_id: crate::backend_id::BackendId, thread_id: ThreadId, }, } @@ -322,10 +327,24 @@ impl From for SessionPickerViewMode { } type PickerLoader = Arc; + +struct PickerBackground { + loader: PickerLoader, + task: tokio::task::JoinHandle<()>, +} + enum BackgroundEvent { + PartialPage { + request_token: usize, + page: PickerPage, + }, + SourceError { + request_token: usize, + }, Page { request_token: usize, search_token: Option, + mode: PageLoadMode, page: std::io::Result, }, Preview { @@ -349,6 +368,7 @@ enum BackgroundEvent { #[derive(Clone)] enum PageCursor { AppServer(String), + Catalog, } struct PickerPage { @@ -582,6 +602,19 @@ pub async fn run_fork_picker_with_app_server( } async fn run_session_picker_with_loader( + tui: &mut Tui, + options: SessionPickerRunOptions, + background: PickerBackground, + bg_rx: mpsc::UnboundedReceiver, +) -> Result { + let PickerBackground { loader, task } = background; + let result = run_session_picker_event_loop(tui, options, loader, bg_rx).await; + + drop(task); + result +} + +async fn run_session_picker_event_loop( tui: &mut Tui, options: SessionPickerRunOptions, picker_loader: PickerLoader, @@ -752,33 +785,31 @@ fn spawn_app_server_page_loader( include_non_interactive: bool, raw_reasoning_visibility: RawReasoningVisibility, config: Option, - bg_tx: mpsc::UnboundedSender, -) -> PickerLoader { + output: mpsc::UnboundedSender, +) -> PickerBackground { + let backend_id = app_server.backend_id(); + let (bg_tx, events) = mpsc::unbounded_channel(); let (request_tx, mut request_rx) = mpsc::unbounded_channel::(); - tokio::spawn(async move { + let task = tokio::spawn(async move { let mut app_server = app_server; let mut page_cwd_filter = PageCwdFilter::default(); - let mut catalog = catalog::ThreadCatalog::default(); while let Some(request) = request_rx.recv().await { match request { PickerLoadRequest::Page(request) => { - if request.cursor.is_none() { - catalog.reset(); - } - let Some(catalog_request) = catalog.begin_page(app_server.backend_id()) else { - continue; - }; + let mode = request.mode; let cwd_filter = page_cwd_filter.for_request( request.cursor.as_ref(), request.cwd_filter.as_deref(), request.filter_mode, uses_remote_filesystem, ); - let cursor = catalog_request - .cursor - .clone() - .map(|PageCursor::AppServer(cursor)| cursor); + let cursor = request.cursor.map(|cursor| match cursor { + PageCursor::AppServer(cursor) => cursor, + PageCursor::Catalog => { + unreachable!("catalog cursors never reach a source") + } + }); let params = thread_list_params( cursor, cwd_filter, @@ -788,24 +819,28 @@ fn spawn_app_server_page_loader( include_non_interactive, matches!(request.mode, PageLoadMode::StateDbOnly), ); - let page = load_app_server_page(&mut app_server, params).await; - let Some(page) = catalog.finish_page(catalog_request, page) else { - continue; + let page = tokio::select! { + page = load_app_server_page(&mut app_server, params) => page, + _ = bg_tx.closed() => continue, }; let _ = bg_tx.send(BackgroundEvent::Page { request_token: request.request_token, search_token: request.search_token, + mode, page, }); } - PickerLoadRequest::Preview { thread_id } => { - let preview = - load_transcript_preview(&mut app_server, thread_id, config.as_ref()).await; + PickerLoadRequest::Preview { thread_id, .. } => { + let preview = tokio::select! { + preview = load_transcript_preview(&mut app_server, thread_id, config.as_ref()) => preview, + _ = bg_tx.closed() => continue, + }; let _ = bg_tx.send(BackgroundEvent::Preview { thread_id, preview }); } PickerLoadRequest::Transcript { thread_id, cancellation, + .. } => { tokio::select! { transcript = load_session_transcript( @@ -820,9 +855,10 @@ fn spawn_app_server_page_loader( }); } _ = cancellation => {} + _ = bg_tx.closed() => {} } } - PickerLoadRequest::Archive { thread_id } => { + PickerLoadRequest::Archive { thread_id, .. } => { let result = archive_request_handle .request_typed::(ClientRequest::ThreadArchive { request_id: RequestId::String(format!( @@ -838,7 +874,7 @@ fn spawn_app_server_page_loader( .map_err(std::io::Error::other); let _ = bg_tx.send(BackgroundEvent::Archive { thread_id, result }); } - PickerLoadRequest::Unarchive { thread_id } => { + PickerLoadRequest::Unarchive { thread_id, .. } => { let result = archive_request_handle .request_typed::(ClientRequest::ThreadUnarchive { request_id: RequestId::String(format!( @@ -867,9 +903,15 @@ fn spawn_app_server_page_loader( } }); - Arc::new(move |request: PickerLoadRequest| { - let _ = request_tx.send(request); - }) + let source = sources::Source { + backend_id, + loader: Arc::new(move |request: PickerLoadRequest| { + let _ = request_tx.send(request); + }), + events, + task, + }; + sources::spawn(vec![source], output) } /// Returns the human-readable column header for the given sort key. @@ -1262,6 +1304,7 @@ impl PickerState { self.request_frame(); return; }; + let backend_id = row.backend_id.unwrap_or(self.backend_id); match self.transcript_cells.get(&thread_id) { Some(SessionTranscriptState::Loaded(_)) => { self.begin_transcript_loading(thread_id); @@ -1276,6 +1319,7 @@ impl PickerState { let (cancellation_tx, cancellation) = oneshot::channel(); self.pending_transcript_cancellation = Some(cancellation_tx); (self.picker_loader)(PickerLoadRequest::Transcript { + backend_id, thread_id, cancellation, }); @@ -1385,7 +1429,7 @@ impl PickerState { let path = row.path.clone(); if let Some(thread_id) = row.thread_id { if self.status == SessionStatus::Archived { - self.request_unarchive(thread_id); + self.request_unarchive(thread_id, row.backend_id.unwrap_or(self.backend_id)); return Ok(None); } return Ok(Some(self.action.selection(SessionTarget { @@ -1569,20 +1613,36 @@ impl PickerState { event: BackgroundEvent, ) -> Result> { match event { + BackgroundEvent::PartialPage { + request_token, + page, + } => { + if self.pagination.is_pending(request_token) { + self.ingest_page(page); + } + } + BackgroundEvent::SourceError { request_token } => { + if self.pagination.is_pending(request_token) { + self.inline_error = + Some("Could not load one backend. Reopen the picker to retry.".into()); + self.request_frame(); + } + } BackgroundEvent::Page { request_token, search_token, + mode, page, } => { let Some(pending) = self.pagination.finish_load(request_token) else { return Ok(None); }; + self.pagination.next_page_mode = mode; 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; + let should_restart_from_store = + 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); @@ -2048,6 +2108,7 @@ impl PickerState { let Some(thread_id) = row.thread_id else { return; }; + let backend_id = row.backend_id.unwrap_or(self.backend_id); if self.expanded_thread_id == Some(thread_id) { self.expanded_thread_id = None; self.request_frame(); @@ -2058,7 +2119,10 @@ impl PickerState { self.transcript_previews.entry(thread_id) { e.insert(TranscriptPreviewState::Loading); - (self.picker_loader)(PickerLoadRequest::Preview { thread_id }); + (self.picker_loader)(PickerLoadRequest::Preview { + backend_id, + thread_id, + }); } self.request_frame(); } @@ -3754,6 +3818,7 @@ mod tests { .handle_background_event(BackgroundEvent::Page { request_token: request.request_token, search_token: request.search_token, + mode: request.mode, page, }) .await @@ -3974,7 +4039,7 @@ mod tests { let mut live_picker = PickerState::new( crate::backend_id::BackendId::new(), FrameRequester::test_dummy(), - loader, + loader.loader.clone(), ProviderFilter::MatchDefault(config.model_provider_id.clone()), /*show_all*/ false, Some(primary.clone()), @@ -3982,14 +4047,16 @@ mod tests { ); live_picker.local_filter_cwd = Some(primary.clone()); live_picker.start_initial_load(); - let event = tokio::time::timeout(std::time::Duration::from_secs(10), bg_rx.recv()) - .await - .expect("app-server page response") - .expect("page event"); - live_picker - .handle_background_event(event) - .await - .expect("ingest app-server page"); + while live_picker.pagination.is_loading() { + let event = tokio::time::timeout(std::time::Duration::from_secs(10), bg_rx.recv()) + .await + .expect("app-server page response") + .expect("page event"); + live_picker + .handle_background_event(event) + .await + .expect("ingest app-server page"); + } assert!( live_picker .filtered_rows @@ -3998,14 +4065,16 @@ mod tests { ); live_picker.filter_mode = SessionFilterMode::Repo; live_picker.start_initial_load(); - let event = tokio::time::timeout(std::time::Duration::from_secs(10), bg_rx.recv()) - .await - .expect("app-server repo page response") - .expect("repo page event"); - live_picker - .handle_background_event(event) - .await - .expect("ingest app-server repo page"); + while live_picker.pagination.is_loading() { + let event = tokio::time::timeout(std::time::Duration::from_secs(10), bg_rx.recv()) + .await + .expect("app-server repo page response") + .expect("repo page event"); + live_picker + .handle_background_event(event) + .await + .expect("ingest app-server repo page"); + } let linked_index = live_picker .filtered_rows .iter() @@ -4025,6 +4094,9 @@ mod tests { live_picker.handle_key(accept).await.expect("fork selection"), Some(SessionSelection::Fork(SessionTarget { thread_id, .. })) if thread_id == linked_id )); + drop(live_picker); + drop(bg_rx); + loader.task.await.unwrap(); let local = crate::latest_session_lookup_params( /*model_provider*/ None, Some(&primary), @@ -4205,6 +4277,15 @@ mod tests { "---- loading fallback ----\n{loading_snapshot}\n---- repaired results ----\n{repaired_snapshot}" ) ); + state.pagination.complete_page( + Some(PageCursor::AppServer("fallback-next".to_string())), + /*num_scanned_files*/ 0, + /*reached_scan_cap*/ false, + ); + assert!(matches!( + state.pagination.next_page(), + Some((_, PageLoadMode::StoreDefault)) + )); } #[tokio::test] @@ -5688,7 +5769,7 @@ session_picker_view = "dense" let recorded_requests: Arc>> = Arc::new(Mutex::new(Vec::new())); let request_sink = recorded_requests.clone(); let loader: PickerLoader = Arc::new(move |request| { - if let PickerLoadRequest::Preview { thread_id } = request { + if let PickerLoadRequest::Preview { thread_id, .. } = request { request_sink.lock().unwrap().push(thread_id); } }); @@ -7408,6 +7489,7 @@ session_picker_view = "dense" .handle_background_event(BackgroundEvent::Page { request_token: first_request.request_token, search_token: first_request.search_token, + mode: first_request.mode, page: Ok(page( vec![make_row("/tmp/beta.jsonl", "2025-01-02T00:00:00Z", "beta")], Some("2025-01-03T00:00:00Z"), @@ -7430,6 +7512,7 @@ session_picker_view = "dense" .handle_background_event(BackgroundEvent::Page { request_token: second_request.request_token, search_token: second_request.search_token, + mode: second_request.mode, page: Ok(page( vec![make_row( "/tmp/match.jsonl", @@ -7459,6 +7542,7 @@ session_picker_view = "dense" .handle_background_event(BackgroundEvent::Page { request_token: second_request.request_token, search_token: second_request.search_token, + mode: second_request.mode, page: Ok(page( Vec::new(), /*next_cursor*/ None, @@ -7474,6 +7558,7 @@ session_picker_view = "dense" .handle_background_event(BackgroundEvent::Page { request_token: active_request.request_token, search_token: active_request.search_token, + mode: active_request.mode, page: Ok(page( Vec::new(), /*next_cursor*/ None, diff --git a/codex-rs/tui/src/resume_picker/archive.rs b/codex-rs/tui/src/resume_picker/archive.rs index a3ccf54312..cdc81addb5 100644 --- a/codex-rs/tui/src/resume_picker/archive.rs +++ b/codex-rs/tui/src/resume_picker/archive.rs @@ -13,6 +13,7 @@ use super::SessionPickerLaunchContext; use super::SessionSelection; use super::SessionStatus; use super::SessionTarget; +use crate::backend_id::BackendId; use crate::keymap::KeymapContext; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -70,9 +71,13 @@ impl PickerState { return; } + let backend_id = row.backend_id.unwrap_or(self.backend_id); self.archive_state = ArchiveState::Pending { thread_id }; self.request_frame(); - (self.picker_loader)(PickerLoadRequest::Archive { thread_id }); + (self.picker_loader)(PickerLoadRequest::Archive { + backend_id, + thread_id, + }); } pub(super) fn handle_archive_result( @@ -137,14 +142,17 @@ impl PickerState { } } - pub(super) fn request_unarchive(&mut self, thread_id: ThreadId) { + pub(super) fn request_unarchive(&mut self, thread_id: ThreadId, backend_id: BackendId) { if !matches!(self.archive_state, ArchiveState::Idle) { return; } self.archive_state = ArchiveState::Restoring { thread_id }; self.request_frame(); - (self.picker_loader)(PickerLoadRequest::Unarchive { thread_id }); + (self.picker_loader)(PickerLoadRequest::Unarchive { + backend_id, + thread_id, + }); } pub(super) fn handle_unarchive_result( diff --git a/codex-rs/tui/src/resume_picker/archive_tests.rs b/codex-rs/tui/src/resume_picker/archive_tests.rs index fa07ff7e00..98c349c11d 100644 --- a/codex-rs/tui/src/resume_picker/archive_tests.rs +++ b/codex-rs/tui/src/resume_picker/archive_tests.rs @@ -27,8 +27,8 @@ fn archive_picker_state() -> (PickerState, Arc>>) { let requests = Arc::new(Mutex::new(Vec::new())); let request_sink = Arc::clone(&requests); let picker_loader: PickerLoader = Arc::new(move |request| { - if let PickerLoadRequest::Archive { thread_id } - | PickerLoadRequest::Unarchive { thread_id } = request + if let PickerLoadRequest::Archive { thread_id, .. } + | PickerLoadRequest::Unarchive { thread_id, .. } = request { request_sink .lock() @@ -66,15 +66,21 @@ fn set_selected_session(state: &mut PickerState, thread_id: ThreadId) { #[tokio::test] async fn archive_shortcut_archives_selected_session_once() { - let (mut state, requests) = archive_picker_state(); + let (mut state, _requests) = archive_picker_state(); + let (requests, mut actions) = tokio::sync::mpsc::unbounded_channel(); + state.picker_loader = Arc::new(move |request| requests.send(request).unwrap()); + let backend_id = crate::backend_id::BackendId::new(); let thread_id = ThreadId::new(); set_selected_session(&mut state, thread_id); + state.filtered_rows[0].backend_id = Some(backend_id); let shortcut = KeyEvent::new(KeyCode::Char('\u{0001}'), KeyModifiers::NONE); assert!(state.handle_key(shortcut).await.unwrap().is_none()); assert!(state.handle_key(shortcut).await.unwrap().is_none()); assert_eq!(state.archive_state, ArchiveState::Pending { thread_id }); - assert_eq!(*requests.lock().unwrap(), vec![thread_id]); + assert!( + matches!(actions.recv().await, Some(PickerLoadRequest::Archive { backend_id: owner, thread_id: archived }) if (owner, archived) == (backend_id, thread_id)) + ); assert!( state .handle_key(KeyEvent::from(KeyCode::Enter)) @@ -87,6 +93,17 @@ async fn archive_shortcut_archives_selected_session_once() { assert_eq!(state.archive_state, ArchiveState::Idle); assert!(state.filtered_rows.is_empty()); + + state.status = SessionStatus::Archived; + set_selected_session(&mut state, thread_id); + state.filtered_rows[0].backend_id = Some(backend_id); + state + .handle_key(KeyEvent::from(KeyCode::Enter)) + .await + .unwrap(); + assert!( + matches!(actions.recv().await, Some(PickerLoadRequest::Unarchive { backend_id: owner, thread_id: restored }) if (owner, restored) == (backend_id, thread_id)) + ); } #[test] @@ -287,7 +304,7 @@ fn archived_session_restore_failure_preserves_session_and_allows_retry() { let thread_id = ThreadId::new(); state.status = SessionStatus::Archived; set_selected_session(&mut state, thread_id); - state.request_unarchive(thread_id); + state.request_unarchive(thread_id, state.backend_id); let selection = state.handle_unarchive_result(thread_id, Err(std::io::Error::other("restore denied"))); @@ -300,7 +317,7 @@ fn archived_session_restore_failure_preserves_session_and_allows_retry() { ); assert_eq!(state.filtered_rows.len(), 1); - state.request_unarchive(thread_id); + state.request_unarchive(thread_id, state.backend_id); assert_eq!(*requests.lock().unwrap(), vec![thread_id, thread_id]); } diff --git a/codex-rs/tui/src/resume_picker/catalog.rs b/codex-rs/tui/src/resume_picker/catalog.rs index ccbe25d442..cfb5939c65 100644 --- a/codex-rs/tui/src/resume_picker/catalog.rs +++ b/codex-rs/tui/src/resume_picker/catalog.rs @@ -2,9 +2,11 @@ //! Resetting a listing invalidates pending pages without mixing source cursors. use super::PageCursor; +use super::PageLoadMode; use super::PickerPage; use crate::backend_id::BackendId; use std::collections::HashMap; +use std::collections::HashSet; #[derive(Default)] pub(super) struct ThreadCatalog { @@ -15,6 +17,7 @@ pub(super) struct ThreadCatalog { #[derive(Default)] struct Source { cursor: Option, + mode: Option, pending: Option, exhausted: bool, } @@ -23,19 +26,31 @@ pub(super) struct PageRequest { backend: BackendId, token: u64, pub(super) cursor: Option, + pub(super) mode: PageLoadMode, } impl ThreadCatalog { + pub(super) fn has_more(&self, excluded: &HashSet) -> bool { + self.sources + .iter() + .any(|(id, source)| !source.exhausted && !excluded.contains(id)) + } + pub(super) fn reset(&mut self) { self.sources.clear(); } /// A slow or exhausted source does not prevent another source from loading. - pub(super) fn begin_page(&mut self, backend: BackendId) -> Option { + pub(super) fn begin_page( + &mut self, + backend: BackendId, + initial_mode: PageLoadMode, + ) -> Option { let source = self.sources.entry(backend).or_default(); if source.pending.is_some() || source.exhausted { return None; } + let mode = *source.mode.get_or_insert(initial_mode); let token = self.next_request; self.next_request += 1; source.pending = Some(token); @@ -43,6 +58,7 @@ impl ThreadCatalog { backend, token, cursor: source.cursor.clone(), + mode, }) } @@ -57,6 +73,7 @@ impl ThreadCatalog { return None; } source.pending = None; + source.mode = Some(request.mode); Some(page.map(|mut page| { source.cursor = page.next_cursor.clone(); source.exhausted = source.cursor.is_none(); diff --git a/codex-rs/tui/src/resume_picker/catalog_tests.rs b/codex-rs/tui/src/resume_picker/catalog_tests.rs index 68a72d58ce..8c6543f410 100644 --- a/codex-rs/tui/src/resume_picker/catalog_tests.rs +++ b/codex-rs/tui/src/resume_picker/catalog_tests.rs @@ -1,3 +1,4 @@ +#![cfg(test)] //! Source pagination, stale completions, and row-specific source routing. use super::*; @@ -22,10 +23,14 @@ fn listed(thread: ThreadId, cursor: Option<&str>) -> PickerPage { } fn cursor(request: &PageRequest) -> Option<&str> { - request - .cursor - .as_ref() - .map(|PageCursor::AppServer(value)| value.as_str()) + request.cursor.as_ref().map(|cursor| match cursor { + PageCursor::AppServer(value) => value.as_str(), + PageCursor::Catalog => panic!("unexpected catalog cursor"), + }) +} + +fn begin_page(catalog: &mut ThreadCatalog, backend: BackendId) -> Option { + catalog.begin_page(backend, PageLoadMode::StoreDefault) } #[test] @@ -34,13 +39,13 @@ fn sources_progress_independently_and_retry_their_own_cursor() { let slow = BackendId::new(); let failed = BackendId::new(); let healthy = BackendId::new(); - let slow_request = catalog.begin_page(slow).unwrap(); - let first = catalog.begin_page(failed).unwrap(); + let slow_request = begin_page(&mut catalog, slow).unwrap(); + let first = begin_page(&mut catalog, failed).unwrap(); catalog .finish_page(first, Ok(listed(ThreadId::new(), Some("failed-next")))) .unwrap() .unwrap(); - let failing_request = catalog.begin_page(failed).unwrap(); + let failing_request = begin_page(&mut catalog, failed).unwrap(); assert_eq!(cursor(&failing_request), Some("failed-next")); let error = catalog .finish_page(failing_request, Err(std::io::Error::other("unavailable"))) @@ -48,16 +53,16 @@ fn sources_progress_independently_and_retry_their_own_cursor() { .err() .unwrap(); assert_eq!(error.to_string(), "unavailable"); - let healthy_request = catalog.begin_page(healthy).unwrap(); + let healthy_request = begin_page(&mut catalog, healthy).unwrap(); let thread = ThreadId::new(); let result = catalog .finish_page(healthy_request, Ok(listed(thread, /*cursor*/ None))) .unwrap() .unwrap(); assert_eq!(result.rows[0].backend_id, Some(healthy)); - assert!(catalog.begin_page(healthy).is_none()); - assert!(catalog.begin_page(slow).is_none()); - let retry = catalog.begin_page(failed).unwrap(); + assert!(begin_page(&mut catalog, healthy).is_none()); + assert!(begin_page(&mut catalog, slow).is_none()); + let retry = begin_page(&mut catalog, failed).unwrap(); assert_eq!(cursor(&retry), Some("failed-next")); catalog .finish_page(retry, Ok(listed(ThreadId::new(), /*cursor*/ None))) @@ -67,23 +72,23 @@ fn sources_progress_independently_and_retry_their_own_cursor() { .finish_page(slow_request, Ok(listed(ThreadId::new(), /*cursor*/ None))) .unwrap() .unwrap(); - assert!(catalog.begin_page(failed).is_none()); - assert!(catalog.begin_page(slow).is_none()); + assert!(begin_page(&mut catalog, failed).is_none()); + assert!(begin_page(&mut catalog, slow).is_none()); } #[test] fn reset_rejects_old_completions_without_finishing_new_requests() { let mut catalog = ThreadCatalog::default(); let backend = BackendId::new(); - let old = catalog.begin_page(backend).unwrap(); + let old = begin_page(&mut catalog, backend).unwrap(); catalog.reset(); - let current = catalog.begin_page(backend).unwrap(); + let current = begin_page(&mut catalog, backend).unwrap(); assert!( catalog .finish_page(old, Ok(listed(ThreadId::new(), /*cursor*/ None))) .is_none() ); - assert!(catalog.begin_page(backend).is_none()); + assert!(begin_page(&mut catalog, backend).is_none()); let thread = ThreadId::new(); let result = catalog .finish_page(current, Ok(listed(thread, /*cursor*/ None))) @@ -104,7 +109,7 @@ async fn picker_selections_retain_each_rows_source() { let owner = BackendId::new(); let thread = ThreadId::new(); let mut catalog = ThreadCatalog::default(); - let request = catalog.begin_page(owner).unwrap(); + let request = begin_page(&mut catalog, owner).unwrap(); let page = catalog .finish_page(request, Ok(listed(thread, /*cursor*/ None))) .unwrap() @@ -112,7 +117,7 @@ async fn picker_selections_retain_each_rows_source() { picker.ingest_page(page); let other = BackendId::new(); let other_thread = ThreadId::new(); - let other_request = catalog.begin_page(other).unwrap(); + let other_request = begin_page(&mut catalog, other).unwrap(); let other_page = catalog .finish_page(other_request, Ok(listed(other_thread, /*cursor*/ None))) .unwrap() diff --git a/codex-rs/tui/src/resume_picker/page_loading.rs b/codex-rs/tui/src/resume_picker/page_loading.rs index 2c074bb8a6..5795c61d06 100644 --- a/codex-rs/tui/src/resume_picker/page_loading.rs +++ b/codex-rs/tui/src/resume_picker/page_loading.rs @@ -52,7 +52,7 @@ pub(super) enum PageLoadMode { /// 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) next_page_mode: PageLoadMode, pub(super) num_scanned_files: usize, pub(super) reached_scan_cap: bool, loading: LoadingState, @@ -134,6 +134,10 @@ impl PaginationState { .map(|cursor| (cursor, self.next_page_mode)) } + pub(super) fn is_pending(&self, request_token: usize) -> bool { + matches!(self.loading, LoadingState::Pending(pending) if pending.request_token == request_token) + } + pub(super) fn is_loading(&self) -> bool { matches!(self.loading, LoadingState::Pending(_)) } diff --git a/codex-rs/tui/src/resume_picker/snapshots/codex_tui__resume_picker__sources__tests__partial_backend_failure.snap b/codex-rs/tui/src/resume_picker/snapshots/codex_tui__resume_picker__sources__tests__partial_backend_failure.snap new file mode 100644 index 0000000000..f9dfa63c89 --- /dev/null +++ b/codex-rs/tui/src/resume_picker/snapshots/codex_tui__resume_picker__sources__tests__partial_backend_failure.snap @@ -0,0 +1,6 @@ +--- +source: tui/src/resume_picker/sources_tests.rs +assertion_line: 206 +expression: terminal.backend().to_string() +--- +Could not load one backend. Reopen the picker to retry. diff --git a/codex-rs/tui/src/resume_picker/sources.rs b/codex-rs/tui/src/resume_picker/sources.rs new file mode 100644 index 0000000000..5a9efb7901 --- /dev/null +++ b/codex-rs/tui/src/resume_picker/sources.rs @@ -0,0 +1,194 @@ +//! Fans listing requests out to independent workers and routes actions to row sources. +//! Closing the picker cancels reads while accepted mutations and source shutdown continue +//! independently; stale listing replies cannot replace current pages. + +use super::*; +use crate::backend_id::BackendId; +use futures::StreamExt; +use futures::stream::SelectAll; +use std::collections::HashSet; + +pub(super) struct Source { + pub(super) backend_id: BackendId, + pub(super) loader: PickerLoader, + pub(super) events: mpsc::UnboundedReceiver, + pub(super) task: tokio::task::JoinHandle<()>, +} + +pub(super) fn spawn( + sources: Vec, + output: mpsc::UnboundedSender, +) -> PickerBackground { + assert!(!sources.is_empty()); + let (input, mut requests) = mpsc::unbounded_channel(); + let mut workers = HashMap::new(); + let mut events = SelectAll::new(); + let mut tasks = Vec::new(); + for Source { + backend_id, + loader, + events: receiver, + task, + } in sources + { + assert!(workers.insert(backend_id, loader).is_none()); + tasks.push(task); + events.push( + UnboundedReceiverStream::new(receiver) + .map(move |event| (backend_id, Some(event))) + .chain(futures::stream::once(async move { (backend_id, None) })) + .boxed(), + ); + } + let task = tokio::spawn(async move { + let mut catalog = catalog::ThreadCatalog::default(); + let mut pending = HashMap::new(); + let mut failed = HashSet::new(); + let mut current: Option = None; + let mut succeeded = false; + let mut error = None; + loop { + tokio::select! { + biased; + _ = output.closed(), if !requests.is_closed() => requests.close(), + request = requests.recv() => { + let Some(request) = request else { break }; + match request { + PickerLoadRequest::Page(request) => { + succeeded = false; + if request.cursor.is_none() { + catalog.reset(); + pending.clear(); + failed.retain(|id| !workers.contains_key(id)); + } + error = workers.is_empty().then(|| std::io::Error::other("No session backends are available.")); + for (id, worker) in &workers { + if failed.contains(id) { continue; } + if let Some(page) = catalog.begin_page(*id, request.mode) { + let mut source_request = request.clone(); + source_request.cursor = page.cursor.clone(); + source_request.mode = page.mode; + pending.insert(*id, page); + worker(PickerLoadRequest::Page(source_request)); + } + } + current = Some(request); + } + request => { + let backend = match &request { + PickerLoadRequest::Preview { backend_id, .. } + | PickerLoadRequest::Transcript { backend_id, .. } + | PickerLoadRequest::Archive { backend_id, .. } + | PickerLoadRequest::Unarchive { backend_id, .. } => *backend_id, + PickerLoadRequest::Page(_) => unreachable!(), + }; + if let Some(worker) = workers.get(&backend) { + worker(request); + } else { + let error = std::io::Error::other("The session's backend is unavailable."); + let event = match request { + PickerLoadRequest::Preview { thread_id, .. } => BackgroundEvent::Preview { thread_id, preview: Err(error) }, + PickerLoadRequest::Transcript { thread_id, .. } => BackgroundEvent::Transcript { thread_id, transcript: Err(error) }, + PickerLoadRequest::Archive { thread_id, .. } => BackgroundEvent::Archive { thread_id, result: Err(error) }, + PickerLoadRequest::Unarchive { thread_id, .. } => BackgroundEvent::Unarchive { thread_id, result: Err(error) }, + PickerLoadRequest::Page(_) => unreachable!(), + }; + let _ = output.send(event); + } + } + } + } + Some((id, event)) = events.next() => { + let event = match event { + Some(event) => event, + None => { + workers.remove(&id); + failed.insert(id); + let Some(request) = current.as_ref().filter(|_| pending.contains_key(&id)) else { continue }; + BackgroundEvent::Page { request_token: request.request_token, search_token: request.search_token, mode: request.mode, + page: Err(std::io::Error::other("The session backend disconnected.")) } + } + }; + match event { + BackgroundEvent::Page { request_token, page, .. } => { + if current.as_ref().is_none_or(|request| request.request_token != request_token) { continue; } + let Some(mut page_request) = pending.remove(&id) else { continue }; + let should_fallback = page_request.mode == PageLoadMode::StateDbOnly + && page_request.cursor.is_none() + && workers.contains_key(&id) + && (page.is_err() + || page.as_ref().is_ok_and(|page| page.rows.is_empty())); + if should_fallback { + let Some(request) = current.as_mut() else { + continue; + }; + request.mode = PageLoadMode::StoreDefault; + let mut fallback = request.clone(); + fallback.cursor = page_request.cursor.clone(); + page_request.mode = PageLoadMode::StoreDefault; + pending.insert(id, page_request); + workers[&id](PickerLoadRequest::Page(fallback)); + continue; + } + match page { + Ok(page) => { + let Some(Ok(mut page)) = catalog.finish_page(page_request, Ok(page)) else { continue }; + succeeded = true; + // Only the catalog owns source cursors. The UI receives a + // continuation marker once this round's workers have replied. + page.next_cursor = None; + let _ = output.send(BackgroundEvent::PartialPage { request_token, page }); + } + Err(err) => { + let Some(Err(err)) = catalog.finish_page(page_request, Err(err)) else { continue }; + failed.insert(id); + error = Some(err); + let _ = output.send(BackgroundEvent::SourceError { request_token }); + } + } + } + event => { let _ = output.send(event); } + } + } + } + if pending.is_empty() + && let Some(request) = current.take() + { + let page = if !succeeded && let Some(error) = error.take() { + Err(error) + } else { + Ok(PickerPage { + rows: Vec::new(), + history_modes: HashMap::new(), + next_cursor: catalog.has_more(&failed).then_some(PageCursor::Catalog), + num_scanned_files: 0, + reached_scan_cap: false, + }) + }; + let _ = output.send(BackgroundEvent::Page { + request_token: request.request_token, + search_token: request.search_token, + mode: request.mode, + page, + }); + } + } + drop(workers); + drop(events); + for task in tasks { + if let Err(err) = task.await { + tracing::warn!(%err, "Picker source task stopped unexpectedly"); + } + } + }); + PickerBackground { + loader: Arc::new(move |request| { + let _ = input.send(request); + }), + task, + } +} + +#[cfg(test)] +#[path = "sources_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/resume_picker/sources_tests.rs b/codex-rs/tui/src/resume_picker/sources_tests.rs new file mode 100644 index 0000000000..c37b87fd44 --- /dev/null +++ b/codex-rs/tui/src/resume_picker/sources_tests.rs @@ -0,0 +1,367 @@ +#![cfg(test)] +//! Exercises independent worker replies, stale pages, and action ownership. + +use super::*; +use crate::resume_picker::tests::make_row; +use crate::resume_picker::tests::page; +use pretty_assertions::assert_eq; +fn source() -> ( + Source, + mpsc::UnboundedReceiver, + mpsc::UnboundedSender, +) { + let (requests, receiver) = mpsc::unbounded_channel(); + let (events, responses) = mpsc::unbounded_channel(); + ( + Source { + backend_id: BackendId::new(), + loader: Arc::new(move |request| { + let _ = requests.send(request); + }), + events: responses, + task: tokio::spawn(async {}), + }, + receiver, + events, + ) +} +fn request(token: usize, cursor: Option) -> PickerLoadRequest { + PickerLoadRequest::Page(PageLoadRequest { + request_token: token, + cursor, + search_token: None, + mode: PageLoadMode::StoreDefault, + cwd_filter: None, + filter_mode: SessionFilterMode::All, + status: SessionStatus::Active, + provider_filter: ProviderFilter::Any, + sort_key: ThreadSortKey::UpdatedAt, + }) +} +async fn next(receiver: &mut mpsc::UnboundedReceiver) -> BackgroundEvent { + tokio::time::timeout(std::time::Duration::from_secs(2), receiver.recv()) + .await + .unwrap() + .unwrap() +} +fn reply(token: usize, thread: ThreadId, cursor: Option<&str>) -> BackgroundEvent { + let mut row = make_row("/tmp/worker.jsonl", "2025-01-01T00:00:00Z", "conversation"); + row.thread_id = Some(thread); + BackgroundEvent::Page { + request_token: token, + search_token: None, + mode: PageLoadMode::StoreDefault, + page: Ok(page( + vec![row], + cursor, + /*num_scanned_files*/ 1, + /*reached_scan_cap*/ false, + )), + } +} +#[tokio::test] +async fn sources_stream_rows_and_route_actions_to_each_source() { + let (first, mut first_requests, first_events) = source(); + let first_id = first.backend_id; + let (second, mut second_requests, second_events) = source(); + let second_id = second.backend_id; + let (third, mut third_requests, third_events) = source(); + let (output, mut events) = mpsc::unbounded_channel(); + let PickerBackground { loader, task } = spawn(vec![first, second, third], output); + let PickerLoadRequest::Page(mut initial_request) = + request(/*token*/ 1, /*cursor*/ None) + else { + unreachable!() + }; + initial_request.mode = PageLoadMode::StateDbOnly; + loader(PickerLoadRequest::Page(initial_request)); + for requests in [ + &mut first_requests, + &mut second_requests, + &mut third_requests, + ] { + assert!(matches!( + requests.recv().await, + Some(PickerLoadRequest::Page(_)) + )); + } + let thread = ThreadId::new(); + second_events + .send(BackgroundEvent::Page { + request_token: 1, + search_token: None, + mode: PageLoadMode::StateDbOnly, + page: Err(std::io::Error::other("indexed listing failed")), + }) + .unwrap(); + assert!(events.try_recv().is_err()); + let Some(PickerLoadRequest::Page(fallback_request)) = second_requests.recv().await else { + panic!("expected source-specific fallback") + }; + assert_eq!(fallback_request.mode, PageLoadMode::StoreDefault); + second_events + .send(reply(/*token*/ 1, thread, /*cursor*/ None)) + .unwrap(); + let BackgroundEvent::PartialPage { page, .. } = next(&mut events).await else { + panic!("expected responsive source page") + }; + assert_eq!(page.rows[0].backend_id, Some(second_id)); + first_events + .send(reply(/*token*/ 1, thread, Some("first-next"))) + .unwrap(); + let BackgroundEvent::PartialPage { page, .. } = next(&mut events).await else { + panic!("expected other source page") + }; + assert_eq!(page.rows[0].backend_id, Some(first_id)); + loader(PickerLoadRequest::Archive { + backend_id: first_id, + thread_id: thread, + }); + assert!( + matches!(first_requests.recv().await, Some(PickerLoadRequest::Archive { backend_id, thread_id }) if (backend_id, thread_id) == (first_id, thread)) + ); + let (_cancel, cancellation) = oneshot::channel(); + loader(PickerLoadRequest::Transcript { + backend_id: second_id, + thread_id: thread, + cancellation, + }); + assert!( + matches!(second_requests.recv().await, Some(PickerLoadRequest::Transcript { backend_id, thread_id, .. }) if (backend_id, thread_id) == (second_id, thread)) + ); + drop(third_events); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::SourceError { .. } + )); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::Page { + mode: PageLoadMode::StoreDefault, + page: Ok(_), + .. + } + )); + loader(request(/*token*/ 2, Some(PageCursor::Catalog))); + let Some(PickerLoadRequest::Page(next_page)) = first_requests.recv().await else { + panic!("expected next page") + }; + assert_eq!(next_page.mode, PageLoadMode::StateDbOnly); + assert!( + matches!(next_page.cursor, Some(PageCursor::AppServer(cursor)) if cursor == "first-next") + ); + assert!(second_requests.try_recv().is_err()); + first_events + .send(reply( + /*token*/ 2, + ThreadId::new(), + /*cursor*/ Some("last-page"), + )) + .unwrap(); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::PartialPage { .. } + )); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::Page { page: Ok(_), .. } + )); + loader(request(/*token*/ 3, Some(PageCursor::Catalog))); + let Some(PickerLoadRequest::Page(last_page)) = first_requests.recv().await else { + panic!("expected last page") + }; + assert!( + matches!(last_page.cursor, Some(PageCursor::AppServer(cursor)) if cursor == "last-page") + ); + first_events + .send(BackgroundEvent::Page { + request_token: 3, + search_token: None, + mode: PageLoadMode::StateDbOnly, + page: Err(std::io::Error::other("last page failed")), + }) + .unwrap(); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::SourceError { request_token: 3 } + )); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::Page { + request_token: 3, + page: Err(_), + .. + } + )); + drop(loader); + task.await.unwrap(); + assert!(first_requests.recv().await.is_none()); + assert!(second_requests.recv().await.is_none()); +} +#[tokio::test] +async fn reset_rejects_old_replies_and_missing_sources_never_fall_back() { + let (source, mut requests, replies) = source(); + let owner = source.backend_id; + let (output, mut events) = mpsc::unbounded_channel(); + let PickerBackground { loader, task } = spawn(vec![source], output); + loader(request(/*token*/ 1, /*cursor*/ None)); + requests.recv().await.unwrap(); + loader(request(/*token*/ 2, /*cursor*/ None)); + let stale = ThreadId::new(); + let current = ThreadId::new(); + replies + .send(reply(/*token*/ 1, stale, /*cursor*/ None)) + .unwrap(); + replies + .send(reply(/*token*/ 2, current, /*cursor*/ None)) + .unwrap(); + let BackgroundEvent::PartialPage { + request_token, + page, + } = next(&mut events).await + else { + panic!("expected current page") + }; + assert_eq!((request_token, page.rows[0].backend_id), (2, Some(owner))); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::Page { + request_token: 2, + .. + } + )); + loader(PickerLoadRequest::Archive { + backend_id: BackendId::new(), + thread_id: stale, + }); + assert!( + matches!(next(&mut events).await, BackgroundEvent::Archive { thread_id, result: Err(_), .. } if thread_id == stale) + ); + requests.recv().await.unwrap(); + loader(PickerLoadRequest::Archive { + backend_id: owner, + thread_id: current, + }); + drop(events); + assert!( + matches!(requests.recv().await, Some(PickerLoadRequest::Archive { thread_id, .. }) if thread_id == current) + ); + task.await.unwrap(); + assert!(requests.recv().await.is_none()); +} +#[tokio::test] +async fn reset_does_not_restore_a_disconnected_source() { + let (preferred, mut preferred_requests, preferred_events) = source(); + let (fallback, mut fallback_requests, fallback_events) = source(); + let owner = fallback.backend_id; + let (output, mut events) = mpsc::unbounded_channel(); + let PickerBackground { loader, task } = spawn(vec![preferred, fallback], output); + loader(request(/*token*/ 1, /*cursor*/ None)); + preferred_requests.recv().await.unwrap(); + fallback_requests.recv().await.unwrap(); + drop(preferred_events); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::SourceError { request_token: 1 } + )); + fallback_events + .send(reply( + /*token*/ 1, + ThreadId::new(), + /*cursor*/ None, + )) + .unwrap(); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::PartialPage { .. } + )); + next(&mut events).await; + loader(request(/*token*/ 2, /*cursor*/ None)); + let thread = ThreadId::new(); + fallback_requests.recv().await.unwrap(); + fallback_events + .send(reply(/*token*/ 2, thread, /*cursor*/ None)) + .unwrap(); + let BackgroundEvent::PartialPage { page, .. } = next(&mut events).await else { + panic!("expected fallback page after reset") + }; + assert_eq!(page.rows[0].backend_id, Some(owner)); + drop(loader); + task.await.unwrap(); +} +#[tokio::test] +async fn detached_picker_cleanup_drains_accepted_mutations() { + let backend_id = BackendId::new(); + let thread_id = ThreadId::new(); + let (requests, mut request_rx) = mpsc::unbounded_channel(); + let (source_events, source_responses) = mpsc::unbounded_channel(); + let (mutation_finished, mutation_completion) = oneshot::channel(); + let source_task = tokio::spawn(async move { + assert!(matches!( + request_rx.recv().await, + Some(PickerLoadRequest::Page(_)) + )); + source_events.closed().await; + assert!( + matches!(request_rx.recv().await, Some(PickerLoadRequest::Archive { backend_id: source, thread_id: thread }) if (source, thread) == (backend_id, thread_id)) + ); + mutation_finished.send(()).unwrap(); + assert!(request_rx.recv().await.is_none()); + }); + let source = Source { + backend_id, + loader: Arc::new(move |request| { + let _ = requests.send(request); + }), + events: source_responses, + task: source_task, + }; + let (output, events) = mpsc::unbounded_channel(); + let PickerBackground { loader, task } = spawn(vec![source], output); + loader(request(/*token*/ 1, /*cursor*/ None)); + loader(PickerLoadRequest::Archive { + backend_id, + thread_id, + }); + drop(events); + drop(task); + mutation_completion.await.unwrap(); +} +#[tokio::test] +async fn partial_failure_keeps_healthy_rows_visible() { + let (mut picker, _) = crate::resume_picker::tests::local_db_first_state(); + picker.pagination.start_load( + /*request_token*/ 42, + /*search_token*/ None, + PageLoadMode::StoreDefault, + ); + let thread = ThreadId::new(); + let BackgroundEvent::Page { page: Ok(page), .. } = + reply(/*token*/ 42, thread, /*cursor*/ None) + else { + unreachable!() + }; + picker + .handle_background_event(BackgroundEvent::PartialPage { + request_token: 42, + page, + }) + .await + .unwrap(); + picker + .handle_background_event(BackgroundEvent::SourceError { request_token: 42 }) + .await + .unwrap(); + assert_eq!(picker.all_rows.len(), 1); + assert!(picker.pagination.is_loading()); + let backend = crate::test_backend::VT100Backend::new(/*width*/ 80, /*height*/ 1); + let mut terminal = crate::custom_terminal::Terminal::with_options(backend).unwrap(); + terminal.set_viewport_area(Rect::new(0, 0, 80, 1)); + { + let mut frame = terminal.get_frame(); + let line = search_line(&picker, frame.area().width); + frame.render_widget_ref(&line, frame.area()); + } + terminal.flush().unwrap(); + insta::assert_snapshot!("partial_backend_failure", terminal.backend().to_string()); +} diff --git a/codex-rs/tui/src/resume_picker_pagination_error_tests.rs b/codex-rs/tui/src/resume_picker_pagination_error_tests.rs index 949d0c9e66..85848e1653 100644 --- a/codex-rs/tui/src/resume_picker_pagination_error_tests.rs +++ b/codex-rs/tui/src/resume_picker_pagination_error_tests.rs @@ -33,6 +33,7 @@ async fn later_page_error_preserves_loaded_sessions() { .handle_background_event(BackgroundEvent::Page { request_token: request.request_token, search_token: request.search_token, + mode: request.mode, page: Ok(PickerPage { rows: vec![Row { backend_id: None, @@ -67,6 +68,7 @@ async fn later_page_error_preserves_loaded_sessions() { .handle_background_event(BackgroundEvent::Page { request_token: request.request_token, search_token: request.search_token, + mode: request.mode, page: Err(std::io::Error::other( "failed to list threads from state database", )),