openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Coordinate resume picker sources independently (#57215)

Browse Source 
## What changed

- Add a source coordinator that streams rows as backends reply, preserves each source's pagination cursor and listing mode, and rejects stale replies.
- Route previews, transcripts, archive, and restore requests to the selected row's backend.
- Keep healthy results visible when a source fails, and fall back from an empty or failed initial database listing per source.
- Cancel reads when the picker closes while allowing accepted mutations to finish.

The app-server picker now uses the coordinator with a single source.

## Testing

Add regression coverage for independent source replies, pagination and fallback, action ownership, stale replies, disconnected sources, and mutation completion after picker shutdown. Add a snapshot for the partial-backend failure warning.

GitOrigin-RevId: 1b565c680133d460f4a29c9946a8398e9c950773
This commit is contained in:
Benjamin Carlsson
2026-10-26 20:00:50 +00:00
committed by copyberry
parent dd5521406a
commit 9756b49379
10 changed files with 782 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/tui/src/resume_picker.rs
+133 -48
View File
@@ -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<SessionListDensity> for SessionPickerViewMode {
}
type PickerLoader = Arc<dyn Fn(PickerLoadRequest) + Send + Sync>;
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<usize>,
mode: PageLoadMode,
page: std::io::Result<PickerPage>,
},
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<BackgroundEvent>,
) -> Result<SessionSelection> {
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<Config>,
bg_tx: mpsc::UnboundedSender<BackgroundEvent>,
) -> PickerLoader {
output: mpsc::UnboundedSender<BackgroundEvent>,
) -> PickerBackground {
let backend_id = app_server.backend_id();
let (bg_tx, events) = mpsc::unbounded_channel();
let (request_tx, mut request_rx) = mpsc::unbounded_channel::<PickerLoadRequest>();
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::<ThreadArchiveResponse>(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::<ThreadUnarchiveResponse>(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<Option<SessionSelection>> {
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<Mutex<Vec<ThreadId>>> = 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,
codex-rs/tui/src/resume_picker/archive.rs
+11 -3
View File
@@ -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(
codex-rs/tui/src/resume_picker/archive_tests.rs
+23 -6
View File
@@ -27,8 +27,8 @@ fn archive_picker_state() -> (PickerState, Arc<Mutex<Vec<ThreadId>>>) {
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]);
}
codex-rs/tui/src/resume_picker/catalog.rs
+18 -1
View File
@@ -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<PageCursor>,
mode: Option<PageLoadMode>,
pending: Option<u64>,
exhausted: bool,
}
@@ -23,19 +26,31 @@ pub(super) struct PageRequest {
backend: BackendId,
token: u64,
pub(super) cursor: Option<PageCursor>,
pub(super) mode: PageLoadMode,
}
impl ThreadCatalog {
pub(super) fn has_more(&self, excluded: &HashSet<BackendId>) -> 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<PageRequest> {
pub(super) fn begin_page(
&mut self,
backend: BackendId,
initial_mode: PageLoadMode,
) -> Option<PageRequest> {
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();
codex-rs/tui/src/resume_picker/catalog_tests.rs
+23 -18
View File
@@ -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<PageRequest> {
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()
codex-rs/tui/src/resume_picker/page_loading.rs
+5 -1
View File
@@ -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<PageCursor>,
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(_))
}
codex-rs/tui/src/resume_picker/snapshots/codex_tui__resume_picker__sources__tests__partial_backend_failure.snap
+6
View File
@@ -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.
codex-rs/tui/src/resume_picker/sources.rs
+194
View File
@@ -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<BackgroundEvent>,
pub(super) task: tokio::task::JoinHandle<()>,
}
pub(super) fn spawn(
sources: Vec<Source>,
output: mpsc::UnboundedSender<BackgroundEvent>,
) -> 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<PageLoadRequest> = 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;
codex-rs/tui/src/resume_picker/sources_tests.rs
+367
View File
@@ -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<PickerLoadRequest>,
mpsc::UnboundedSender<BackgroundEvent>,
) {
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<PageCursor>) -> 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>) -> 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());
}
codex-rs/tui/src/resume_picker_pagination_error_tests.rs
+2
View File
@@ -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",
)),