From ce58201cb091be72bb2d03b05742030f16f45e13 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 15 Jun 2026 20:52:05 -0400 Subject: [PATCH] Load resume picker data concurrently --- codex-rs/tui/src/resume_picker.rs | 14 +- codex-rs/tui/src/resume_picker/loading.rs | 274 ++++++++++++++---- .../tui/src/resume_picker/loading_tests.rs | 206 +++++++++++++ codex-rs/tui/src/thread_transcript.rs | 17 -- 4 files changed, 434 insertions(+), 77 deletions(-) diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index 8abcf8adb7..fba83df723 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/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> { +fn transcript_preview_lines( + thread: &codex_app_server_protocol::Thread, +) -> Vec { 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 { diff --git a/codex-rs/tui/src/resume_picker/loading.rs b/codex-rs/tui/src/resume_picker/loading.rs index 1af4f2c0ed..663d7020f1 100644 --- a/codex-rs/tui/src/resume_picker/loading.rs +++ b/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, ) -> PickerLoader { - let (request_tx, mut request_rx) = mpsc::unbounded_channel::(); + let (request_tx, request_rx) = mpsc::unbounded_channel::(); + 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( + mut request_rx: mpsc::UnboundedReceiver, + load_request: F, +) where + F: Fn(PickerLoadRequest) -> Fut + Clone + Send + 'static, + Fut: Future + 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( + mut request_rx: mpsc::UnboundedReceiver, + mut load_preview: F, +) where + F: FnMut(ThreadId) -> Fut + Send + 'static, + Fut: Future + 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( + mut request_rx: mpsc::UnboundedReceiver, + load_page: F, +) where + F: Fn(PageLoadRequest) -> Fut + Send + 'static, + Fut: Future + 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, +) { + 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, cwd_filter: Option<&Path>, provider_filter: ProviderFilter, sort_key: ThreadSortKey, include_non_interactive: bool, ) -> io::Result { - 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 { + 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 { let thread_id = match ThreadId::from_string(&thread.id) { Ok(thread_id) => thread_id, diff --git a/codex-rs/tui/src/resume_picker/loading_tests.rs b/codex-rs/tui/src/resume_picker/loading_tests.rs index 028eb3fe12..796e1c9acd 100644 --- a/codex-rs/tui/src/resume_picker/loading_tests.rs +++ b/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); + + 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); + + impl ActiveLoad { + fn start(active_loads: Arc) -> 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( diff --git a/codex-rs/tui/src/thread_transcript.rs b/codex-rs/tui/src/thread_transcript.rs index cdffaab27c..ba528c0312 100644 --- a/codex-rs/tui/src/thread_transcript.rs +++ b/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 { - 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,