Load resume picker data concurrently

codex-rs/tui/src/resume_picker.rs

@@ -24,7 +24,6 @@ use crate::terminal_palette::default_bg;
 use crate::text_formatting::truncate_text;
 use crate::thread_transcript::RawReasoningVisibility;
 use crate::thread_transcript::TranscriptCells;
-use crate::thread_transcript::load_session_transcript;
 use crate::tui::FrameRequester;
 use crate::tui::Tui;
 use crate::tui::TuiEvent;
@@ -633,16 +632,11 @@ impl LoadingState {
     }
 }
 
-async fn load_transcript_preview(
-    app_server: &mut AppServerSession,
-    thread_id: ThreadId,
-) -> std::io::Result<Vec<TranscriptPreviewLine>> {
+fn transcript_preview_lines(
+    thread: &codex_app_server_protocol::Thread,
+) -> Vec<TranscriptPreviewLine> {
     const MAX_PREVIEW_LINES: usize = 6;
 
-    let thread = app_server
-        .thread_read(thread_id, /*include_turns*/ true)
-        .await
-        .map_err(std::io::Error::other)?;
     let cwd = thread.cwd.as_path();
     let mut lines = thread
         .turns
@@ -682,7 +676,7 @@ async fn load_transcript_preview(
     if lines.len() > MAX_PREVIEW_LINES {
         lines.drain(..lines.len() - MAX_PREVIEW_LINES);
     }
-    Ok(lines)
+    lines
 }
 
 impl SearchState {

codex-rs/tui/src/resume_picker/loading.rs

@@ -1,18 +1,28 @@
 //! Session picker loading and background event handling.
 
+use std::future::Future;
 use std::io;
 use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use chrono::Utc;
+use codex_app_server_client::AppServerRequestHandle;
+use codex_app_server_protocol::ClientRequest;
+use codex_app_server_protocol::RequestId;
 use codex_app_server_protocol::Thread;
 use codex_app_server_protocol::ThreadListCwdFilter;
 use codex_app_server_protocol::ThreadListParams;
+use codex_app_server_protocol::ThreadListResponse;
+use codex_app_server_protocol::ThreadReadParams;
+use codex_app_server_protocol::ThreadReadResponse;
 use codex_app_server_protocol::ThreadSortKey;
 use codex_protocol::ThreadId;
 use tokio::sync::mpsc;
+use tokio::task::JoinError;
+use tokio::task::JoinSet;
 use tracing::warn;
+use uuid::Uuid;
 
 use super::AppServerSession;
 use super::LoadingState;
@@ -24,10 +34,12 @@ use super::SessionTranscriptState;
 use super::TranscriptCells;
 use super::TranscriptPreviewLine;
 use super::TranscriptPreviewState;
-use super::load_session_transcript;
-use super::load_transcript_preview;
+use super::transcript_preview_lines;
+use crate::thread_transcript::thread_to_transcript_cells;
 
 const PAGE_SIZE: usize = 25;
+// Expanded rows read full transcripts, so keep preview I/O narrowly bounded.
+const MAX_CONCURRENT_PREVIEW_READS: usize = 2;
 
 #[derive(Clone)]
 pub(super) struct PageLoadRequest {
@@ -81,47 +93,20 @@ pub(super) fn spawn_app_server_page_loader(
     raw_reasoning_visibility: RawReasoningVisibility,
     bg_tx: mpsc::UnboundedSender<BackgroundEvent>,
 ) -> PickerLoader {
-    let (request_tx, mut request_rx) = mpsc::unbounded_channel::<PickerLoadRequest>();
+    let (request_tx, request_rx) = mpsc::unbounded_channel::<PickerLoadRequest>();
+    let request_handle = app_server.request_handle();
 
     tokio::spawn(async move {
-        let mut app_server = app_server;
-        while let Some(request) = request_rx.recv().await {
-            match request {
-                PickerLoadRequest::Page(request) => {
-                    let cursor = request.cursor.map(|PageCursor::AppServer(cursor)| cursor);
-                    let page = load_app_server_page(
-                        &mut app_server,
-                        cursor,
-                        request.cwd_filter.as_deref(),
-                        request.provider_filter,
-                        request.sort_key,
-                        include_non_interactive,
-                    )
-                    .await;
-                    let _ = bg_tx.send(BackgroundEvent::Page {
-                        request_token: request.request_token,
-                        search_token: request.search_token,
-                        page,
-                    });
-                }
-                PickerLoadRequest::Preview { thread_id } => {
-                    let preview = load_transcript_preview(&mut app_server, thread_id).await;
-                    let _ = bg_tx.send(BackgroundEvent::Preview { thread_id, preview });
-                }
-                PickerLoadRequest::Transcript { thread_id } => {
-                    let transcript = load_session_transcript(
-                        &mut app_server,
-                        thread_id,
-                        raw_reasoning_visibility,
-                    )
-                    .await;
-                    let _ = bg_tx.send(BackgroundEvent::Transcript {
-                        thread_id,
-                        transcript,
-                    });
-                }
-            }
-        }
+        run_picker_loader(request_rx, move |request| {
+            handle_picker_load_request(
+                request,
+                request_handle.clone(),
+                include_non_interactive,
+                raw_reasoning_visibility,
+                bg_tx.clone(),
+            )
+        })
+        .await;
         if let Err(err) = app_server.shutdown().await {
             warn!(%err, "Failed to shut down app-server picker session");
         }
@@ -132,22 +117,194 @@ pub(super) fn spawn_app_server_page_loader(
     })
 }
 
+async fn run_picker_loader<F, Fut>(
+    mut request_rx: mpsc::UnboundedReceiver<PickerLoadRequest>,
+    load_request: F,
+) where
+    F: Fn(PickerLoadRequest) -> Fut + Clone + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    let (page_tx, page_rx) = mpsc::unbounded_channel();
+    let page_load_request = load_request.clone();
+    let page_task = tokio::spawn(run_page_loader(page_rx, move |request| {
+        page_load_request(PickerLoadRequest::Page(request))
+    }));
+    let (preview_tx, preview_rx) = mpsc::unbounded_channel();
+    let preview_load_request = load_request.clone();
+    let preview_task = tokio::spawn(run_preview_loader(preview_rx, move |thread_id| {
+        preview_load_request(PickerLoadRequest::Preview { thread_id })
+    }));
+    let mut tasks = JoinSet::new();
+    loop {
+        tokio::select! {
+            request = request_rx.recv() => {
+                let Some(request) = request else {
+                    break;
+                };
+                match request {
+                    PickerLoadRequest::Page(request) => {
+                        let _ = page_tx.send(request);
+                    }
+                    PickerLoadRequest::Preview { thread_id } => {
+                        let _ = preview_tx.send(thread_id);
+                    }
+                    request @ PickerLoadRequest::Transcript { .. } => {
+                        tasks.spawn(load_request(request));
+                    }
+                }
+            }
+            result = tasks.join_next(), if !tasks.is_empty() => {
+                if let Some(result) = result {
+                    log_loader_task_result(result);
+                }
+            }
+        }
+    }
+
+    drop(page_tx);
+    drop(preview_tx);
+    page_task.abort();
+    log_loader_task_result(page_task.await);
+    preview_task.abort();
+    log_loader_task_result(preview_task.await);
+    tasks.abort_all();
+    while let Some(result) = tasks.join_next().await {
+        log_loader_task_result(result);
+    }
+}
+
+async fn run_preview_loader<F, Fut>(
+    mut request_rx: mpsc::UnboundedReceiver<ThreadId>,
+    mut load_preview: F,
+) where
+    F: FnMut(ThreadId) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    let mut tasks = JoinSet::new();
+    let mut request_channel_open = true;
+    while request_channel_open || !tasks.is_empty() {
+        tokio::select! {
+            request = request_rx.recv(), if request_channel_open && tasks.len() < MAX_CONCURRENT_PREVIEW_READS => {
+                match request {
+                    Some(thread_id) => {
+                        tasks.spawn(load_preview(thread_id));
+                    }
+                    None => request_channel_open = false,
+                }
+            }
+            result = tasks.join_next(), if !tasks.is_empty() => {
+                if let Some(result) = result {
+                    log_loader_task_result(result);
+                }
+            }
+        }
+    }
+}
+
+async fn run_page_loader<F, Fut>(
+    mut request_rx: mpsc::UnboundedReceiver<PageLoadRequest>,
+    load_page: F,
+) where
+    F: Fn(PageLoadRequest) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    let Some(mut request) = request_rx.recv().await else {
+        return;
+    };
+    loop {
+        let load = load_page(request);
+        tokio::pin!(load);
+        tokio::select! {
+            () = &mut load => {
+                let Some(next_request) = request_rx.recv().await else {
+                    break;
+                };
+                request = next_request;
+            }
+            next_request = request_rx.recv() => {
+                let Some(next_request) = next_request else {
+                    load.await;
+                    break;
+                };
+                request = next_request;
+            }
+        }
+        while let Ok(next_request) = request_rx.try_recv() {
+            request = next_request;
+        }
+    }
+}
+
+fn log_loader_task_result(result: Result<(), JoinError>) {
+    if let Err(err) = result
+        && !err.is_cancelled()
+    {
+        warn!(%err, "Session picker loader task failed");
+    }
+}
+
+async fn handle_picker_load_request(
+    request: PickerLoadRequest,
+    request_handle: AppServerRequestHandle,
+    include_non_interactive: bool,
+    raw_reasoning_visibility: RawReasoningVisibility,
+    bg_tx: mpsc::UnboundedSender<BackgroundEvent>,
+) {
+    match request {
+        PickerLoadRequest::Page(request) => {
+            let cursor = request.cursor.map(|PageCursor::AppServer(cursor)| cursor);
+            let page = load_app_server_page(
+                &request_handle,
+                cursor,
+                request.cwd_filter.as_deref(),
+                request.provider_filter,
+                request.sort_key,
+                include_non_interactive,
+            )
+            .await;
+            let _ = bg_tx.send(BackgroundEvent::Page {
+                request_token: request.request_token,
+                search_token: request.search_token,
+                page,
+            });
+        }
+        PickerLoadRequest::Preview { thread_id } => {
+            let preview = read_app_server_thread(&request_handle, thread_id)
+                .await
+                .map(|thread| transcript_preview_lines(&thread));
+            let _ = bg_tx.send(BackgroundEvent::Preview { thread_id, preview });
+        }
+        PickerLoadRequest::Transcript { thread_id } => {
+            let transcript = read_app_server_thread(&request_handle, thread_id)
+                .await
+                .map(|thread| thread_to_transcript_cells(&thread, raw_reasoning_visibility));
+            let _ = bg_tx.send(BackgroundEvent::Transcript {
+                thread_id,
+                transcript,
+            });
+        }
+    }
+}
+
 async fn load_app_server_page(
-    app_server: &mut AppServerSession,
+    request_handle: &AppServerRequestHandle,
     cursor: Option<String>,
     cwd_filter: Option<&Path>,
     provider_filter: ProviderFilter,
     sort_key: ThreadSortKey,
     include_non_interactive: bool,
 ) -> io::Result<PickerPage> {
-    let response = app_server
-        .thread_list(thread_list_params(
-            cursor,
-            cwd_filter,
-            provider_filter,
-            sort_key,
-            include_non_interactive,
-        ))
+    let response: ThreadListResponse = request_handle
+        .request_typed(ClientRequest::ThreadList {
+            request_id: RequestId::String(format!("resume-picker-thread-list-{}", Uuid::new_v4())),
+            params: thread_list_params(
+                cursor,
+                cwd_filter,
+                provider_filter,
+                sort_key,
+                include_non_interactive,
+            ),
+        })
         .await
         .map_err(io::Error::other)?;
     let num_scanned_files = response.data.len();
@@ -164,6 +321,23 @@ async fn load_app_server_page(
     })
 }
 
+async fn read_app_server_thread(
+    request_handle: &AppServerRequestHandle,
+    thread_id: ThreadId,
+) -> io::Result<Thread> {
+    let response: ThreadReadResponse = request_handle
+        .request_typed(ClientRequest::ThreadRead {
+            request_id: RequestId::String(format!("resume-picker-thread-read-{}", Uuid::new_v4())),
+            params: ThreadReadParams {
+                thread_id: thread_id.to_string(),
+                include_turns: true,
+            },
+        })
+        .await
+        .map_err(io::Error::other)?;
+    Ok(response.thread)
+}
+
 fn row_from_app_server_thread(thread: Thread) -> Option<Row> {
     let thread_id = match ThreadId::from_string(&thread.id) {
         Ok(thread_id) => thread_id,

codex-rs/tui/src/resume_picker/loading_tests.rs

@@ -1,4 +1,10 @@
+use std::future::pending;
 use std::path::Path;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
 
 use codex_app_server_protocol::Thread;
 use codex_app_server_protocol::ThreadListCwdFilter;
@@ -12,6 +18,206 @@ use pretty_assertions::assert_eq;
 use super::*;
 use crate::resume_picker::picker_cwd_filter;
 
+fn page_request(request_token: usize) -> PageLoadRequest {
+    PageLoadRequest {
+        cursor: None,
+        request_token,
+        search_token: None,
+        cwd_filter: None,
+        provider_filter: ProviderFilter::Any,
+        sort_key: ThreadSortKey::UpdatedAt,
+    }
+}
+
+#[tokio::test]
+async fn loader_does_not_block_followup_requests_and_cancels_on_close() {
+    struct SetOnDrop(Arc<AtomicBool>);
+
+    impl Drop for SetOnDrop {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    let (request_tx, request_rx) = mpsc::unbounded_channel();
+    let (signal_tx, mut signal_rx) = mpsc::unbounded_channel();
+    let page_dropped = Arc::new(AtomicBool::new(false));
+    let worker_page_dropped = page_dropped.clone();
+    let worker = tokio::spawn(run_picker_loader(request_rx, move |request| {
+        let signal_tx = signal_tx.clone();
+        let page_dropped = worker_page_dropped.clone();
+        async move {
+            match request {
+                PickerLoadRequest::Page(_) => {
+                    let _drop_guard = SetOnDrop(page_dropped);
+                    let _ = signal_tx.send("page");
+                    pending::<()>().await;
+                }
+                PickerLoadRequest::Transcript { .. } => {
+                    let _ = signal_tx.send("transcript");
+                }
+                PickerLoadRequest::Preview { .. } => {}
+            }
+        }
+    }));
+
+    request_tx
+        .send(PickerLoadRequest::Page(page_request(
+            /*request_token*/ 1,
+        )))
+        .expect("send page request");
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(1), signal_rx.recv())
+            .await
+            .expect("page request should start"),
+        Some("page")
+    );
+
+    request_tx
+        .send(PickerLoadRequest::Transcript {
+            thread_id: ThreadId::new(),
+        })
+        .expect("send transcript request");
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(1), signal_rx.recv())
+            .await
+            .expect("transcript request should not wait for page load"),
+        Some("transcript")
+    );
+
+    drop(request_tx);
+    tokio::time::timeout(Duration::from_secs(1), worker)
+        .await
+        .expect("loader should stop promptly")
+        .expect("loader task should not panic");
+    assert!(page_dropped.load(Ordering::SeqCst));
+}
+
+#[tokio::test]
+async fn page_loader_cancels_and_coalesces_obsolete_requests() {
+    struct ActiveLoad(Arc<AtomicUsize>);
+
+    impl ActiveLoad {
+        fn start(active_loads: Arc<AtomicUsize>) -> Self {
+            assert_eq!(active_loads.fetch_add(1, Ordering::SeqCst), 0);
+            Self(active_loads)
+        }
+    }
+
+    impl Drop for ActiveLoad {
+        fn drop(&mut self) {
+            self.0.fetch_sub(1, Ordering::SeqCst);
+        }
+    }
+
+    let (request_tx, request_rx) = mpsc::unbounded_channel();
+    let (started_tx, mut started_rx) = mpsc::unbounded_channel();
+    let release = Arc::new(tokio::sync::Semaphore::new(0));
+    let worker_release = release.clone();
+    let active_loads = Arc::new(AtomicUsize::new(0));
+    let worker_active_loads = active_loads.clone();
+    let worker = tokio::spawn(run_page_loader(request_rx, move |request| {
+        let started_tx = started_tx.clone();
+        let release = worker_release.clone();
+        let active_loads = worker_active_loads.clone();
+        async move {
+            let _active_load = ActiveLoad::start(active_loads);
+            let _ = started_tx.send(request.request_token);
+            let _permit = release
+                .acquire_owned()
+                .await
+                .expect("release semaphore should remain open");
+        }
+    }));
+
+    request_tx
+        .send(page_request(/*request_token*/ 1))
+        .expect("send first page");
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
+            .await
+            .expect("first page should start"),
+        Some(1)
+    );
+    request_tx
+        .send(page_request(/*request_token*/ 2))
+        .expect("send second page");
+    request_tx
+        .send(page_request(/*request_token*/ 3))
+        .expect("send third page");
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
+            .await
+            .expect("latest page should supersede the active load"),
+        Some(3)
+    );
+    assert!(
+        tokio::time::timeout(Duration::from_millis(50), started_rx.recv())
+            .await
+            .is_err()
+    );
+
+    drop(request_tx);
+    release.add_permits(1);
+    tokio::time::timeout(Duration::from_secs(1), worker)
+        .await
+        .expect("page loader should stop")
+        .expect("page loader should not panic");
+    assert_eq!(active_loads.load(Ordering::SeqCst), 0);
+}
+
+#[tokio::test]
+async fn loader_bounds_concurrent_preview_reads() {
+    let (request_tx, request_rx) = mpsc::unbounded_channel();
+    let (started_tx, mut started_rx) = mpsc::unbounded_channel();
+    let release = Arc::new(tokio::sync::Semaphore::new(0));
+    let worker_release = release.clone();
+    let worker = tokio::spawn(run_picker_loader(request_rx, move |request| {
+        let started_tx = started_tx.clone();
+        let release = worker_release.clone();
+        async move {
+            if let PickerLoadRequest::Preview { thread_id } = request {
+                let _ = started_tx.send(thread_id);
+                let _permit = release
+                    .acquire_owned()
+                    .await
+                    .expect("release semaphore should remain open");
+            }
+        }
+    }));
+
+    for _ in 0..=MAX_CONCURRENT_PREVIEW_READS {
+        request_tx
+            .send(PickerLoadRequest::Preview {
+                thread_id: ThreadId::new(),
+            })
+            .expect("send preview request");
+    }
+    for _ in 0..MAX_CONCURRENT_PREVIEW_READS {
+        tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
+            .await
+            .expect("preview should start")
+            .expect("preview signal channel should remain open");
+    }
+    assert!(
+        tokio::time::timeout(Duration::from_millis(50), started_rx.recv())
+            .await
+            .is_err()
+    );
+
+    release.add_permits(MAX_CONCURRENT_PREVIEW_READS + 1);
+    tokio::time::timeout(Duration::from_secs(1), started_rx.recv())
+        .await
+        .expect("queued preview should start")
+        .expect("preview signal channel should remain open");
+
+    drop(request_tx);
+    tokio::time::timeout(Duration::from_secs(1), worker)
+        .await
+        .expect("loader should stop")
+        .expect("loader task should not panic");
+}
+
 #[test]
 fn local_picker_thread_list_params_include_cwd_filter() {
     let cwd_filter = picker_cwd_filter(

codex-rs/tui/src/thread_transcript.rs

@@ -2,7 +2,6 @@
 
 use std::sync::Arc;
 
-use crate::app_server_session::AppServerSession;
 use crate::git_action_directives::parse_assistant_markdown;
 use crate::history_cell::AgentMarkdownCell;
 use crate::history_cell::HistoryCell;
@@ -12,7 +11,6 @@ use crate::history_cell::UserHistoryCell;
 use crate::multi_agents::sub_agent_activity_summary;
 use codex_app_server_protocol::Thread;
 use codex_app_server_protocol::ThreadItem;
-use codex_protocol::ThreadId;
 use codex_protocol::items::UserMessageItem;
 use ratatui::style::Stylize as _;
 use ratatui::text::Line;
@@ -25,21 +23,6 @@ pub(crate) enum RawReasoningVisibility {
     Visible,
 }
 
-pub(crate) async fn load_session_transcript(
-    app_server: &mut AppServerSession,
-    thread_id: ThreadId,
-    raw_reasoning_visibility: RawReasoningVisibility,
-) -> std::io::Result<TranscriptCells> {
-    let thread = app_server
-        .thread_read(thread_id, /*include_turns*/ true)
-        .await
-        .map_err(std::io::Error::other)?;
-    Ok(thread_to_transcript_cells(
-        &thread,
-        raw_reasoning_visibility,
-    ))
-}
-
 pub(crate) fn thread_to_transcript_cells(
     thread: &Thread,
     raw_reasoning_visibility: RawReasoningVisibility,