diff --git a/codex-rs/app-server-client/src/cloud_backend/connection_tests.rs b/codex-rs/app-server-client/src/cloud_backend/connection_tests.rs index dd59622278..fa69361e43 100644 --- a/codex-rs/app-server-client/src/cloud_backend/connection_tests.rs +++ b/codex-rs/app-server-client/src/cloud_backend/connection_tests.rs @@ -227,6 +227,42 @@ async fn failed_archive_loaded_list_tracks_attachment_not_cached_metadata() -> R #[tokio::test] async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> { + assert_abandoned_request_unblocks_worker( + "ListThreads", + CloudOperation::Request(Box::new(protocol::ClientRequest::ThreadList { + request_id: protocol::RequestId::Integer(1), + params: serde_json::from_value(serde_json::json!({"limit": 10}))?, + })), + ) + .await +} + +#[tokio::test] +async fn dropping_an_in_flight_start_unblocks_the_native_worker() -> Result<()> { + for operation in [ + CloudOperation::Request(Box::new(protocol::ClientRequest::ThreadStart { + request_id: protocol::RequestId::Integer(1), + params: Default::default(), + })), + CloudOperation::StartProjectThread { + request_id: protocol::RequestId::Integer(1), + params: Box::default(), + project: cloud::ProjectEnvironment { + name: "project".into(), + origin_urls: vec!["https://github.com/example/project.git".into()], + cwd: None, + }, + }, + ] { + assert_abandoned_request_unblocks_worker("Prewarm", operation).await?; + } + Ok(()) +} + +async fn assert_abandoned_request_unblocks_worker( + rpc: &str, + operation: CloudOperation, +) -> Result<()> { let server = MockServer::start().await; let home = tempfile::tempdir()?; let config = crate::legacy_core::config::ConfigBuilder::default() @@ -247,6 +283,7 @@ async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> { workspace_routing: None, }; let (client, commands, output) = CloudClient::channel(/*capacity*/ 32); + let handle = client.request_handle(); let client = AppServerClient::Cloud(client); let mut requests = requests::Requests { client: Arc::new(native), @@ -257,16 +294,16 @@ async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> { tokio::spawn(async move { run_worker(commands, &mut requests, &mut Replies::default(), &output).await; }); - let (list_arrived, mut arrival) = mpsc::unbounded_channel(); - let (release_list, held) = std::sync::mpsc::channel::<()>(); + let (request_arrived, mut arrival) = mpsc::unbounded_channel(); + let (release_request, held) = std::sync::mpsc::channel::<()>(); let held = Mutex::new(held); Mock::given(method("POST")) - .and(path( - "/oaiproto.codex_app_server_backend.ThreadService/ListThreads", - )) + .and(path(format!( + "/oaiproto.codex_app_server_backend.ThreadService/{rpc}" + ))) .and(header("authorization", "Bearer synthetic-bearer")) .respond_with(move |_: &wiremock::Request| { - let _ = list_arrived.send(()); + let _ = request_arrived.send(()); // Wiremock serves on a dedicated thread. Hold the RPC until the sender drops. let _ = held.lock().unwrap().recv(); ResponseTemplate::new(503) @@ -275,17 +312,28 @@ async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> { .mount(&server) .await; - let mut list = Box::pin(client.request(protocol::ClientRequest::ThreadList { - request_id: protocol::RequestId::Integer(1), - params: serde_json::from_value(serde_json::json!({"limit": 10}))?, - })); + let mut pending = Box::pin(async { + match operation { + CloudOperation::Request(request) => client.request(*request).await, + CloudOperation::StartProjectThread { + request_id, + params, + project, + } => { + handle + .start_project_thread(request_id, *params, project) + .await + } + _ => unreachable!("fixture supports Cloud lists and starts"), + } + }); tokio::select! { seen = tokio::time::timeout(crate::SHUTDOWN_TIMEOUT, arrival.recv()) => { assert_eq!(seen?, Some(())); } - _ = &mut list => panic!("list finished before the fixture held it"), + _ = &mut pending => panic!("request finished before the fixture held it"), } - drop(list); + drop(pending); let account_response = tokio::time::timeout( crate::SHUTDOWN_TIMEOUT, client.request_typed::(protocol::ClientRequest::GetAccount { @@ -296,7 +344,7 @@ async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> { }), ) .await; - drop(release_list); + drop(release_request); assert_eq!(account_response??, account); client.shutdown().await?; Ok(()) diff --git a/codex-rs/app-server-client/src/cloud_backend/mod.rs b/codex-rs/app-server-client/src/cloud_backend/mod.rs index 197c14f73a..6f128f0e54 100644 --- a/codex-rs/app-server-client/src/cloud_backend/mod.rs +++ b/codex-rs/app-server-client/src/cloud_backend/mod.rs @@ -282,21 +282,24 @@ async fn run_worker( () = replies.forward_next(output) => continue, command = commands.recv() => match command { Some(command) => command, None => return }, }; + let (is_thread_start, selects_attachment) = match &command.operation { + CloudOperation::Request(request) => match request.as_ref() { + protocol::ClientRequest::ThreadStart { .. } => (true, true), + protocol::ClientRequest::ThreadResume { .. } => (false, true), + _ => (false, false), + }, + CloudOperation::StartProjectThread { .. } => (true, true), + _ => (false, false), + }; + // A queued start that nobody owns must leave the existing recovery alone. + if is_thread_start && command.response.is_closed() { + continue; + } // A deliberate selection owns the next attachment, even if its handshake fails. - let canceled = match &command.operation { - CloudOperation::Request(request) - if matches!( - request.as_ref(), - protocol::ClientRequest::ThreadStart { .. } - | protocol::ClientRequest::ThreadResume { .. } - ) => - { - replies.idle.as_mut().and_then(idle::Recovery::cancel) - } - CloudOperation::StartProjectThread { .. } => { - replies.idle.as_mut().and_then(idle::Recovery::cancel) - } - _ => None, + let canceled = if selects_attachment { + replies.idle.as_mut().and_then(idle::Recovery::cancel) + } else { + None }; let request = match command.operation { CloudOperation::Shutdown => { @@ -336,14 +339,18 @@ async fn run_worker( }; let result = match request { Ok((request, mode, project)) => { - // Inventory is a read with no conversation state to commit. let is_thread_list = matches!(&request, protocol::ClientRequest::ThreadList { .. }); - if is_thread_list && command.response.is_closed() { - continue; - } + // When closure and the handshake are both ready, cancel before + // the handshake can replace the active attachment. tokio::select! { + biased; _ = output.closed() => return, - _ = command.response.closed(), if is_thread_list => continue, + _ = command.response.closed(), if is_thread_list || is_thread_start => { + if let Some(root) = canceled { + replies.report_loss(&root, output); + } + continue; + } result = handle_request(request, mode, project, requests, replies, output) => result, } } diff --git a/codex-rs/tui/src/app/backend_selection.rs b/codex-rs/tui/src/app/backend_selection.rs index 11325b0e83..5d1cb47be0 100644 --- a/codex-rs/tui/src/app/backend_selection.rs +++ b/codex-rs/tui/src/app/backend_selection.rs @@ -217,12 +217,17 @@ impl App { let bootstrap = destination .bootstrap_for_new_thread(&settings.config) .map_err(|error| format!("Failed to prepare the new conversation: {error}")); - let bootstrap = if responsive_cloud_switch { + let Some(bootstrap) = if responsive_cloud_switch { let events = tui.event_stream(); wait_for_cloud_switch(self, tui, bootstrap, events, clipboard_after_bootstrap).await } else { - StartupDraftPump::run_during_navigation(fresh_draft.as_mut(), tui, bootstrap).await - }?; + StartupDraftPump::run_during_navigation(fresh_draft.as_mut(), tui, bootstrap) + .await + .map(Some) + }? + else { + return Ok(()); + }; let destination_cwd = settings.config.cwd.to_path_buf(); if let Err(control) = self .confirm_directory_trust_for_target( @@ -243,7 +248,12 @@ impl App { { if responsive_cloud_switch { let events = tui.event_stream(); - finish_failed_cloud_switch(self, tui, events).await; + if finish_failed_cloud_switch(self, tui, events) + .await + .is_break() + { + return Ok(()); + } } if matches!(control, AppRunControl::Exit(_)) { self.app_event_tx.send(AppEvent::Exit(ExitMode::Immediate)); @@ -271,8 +281,11 @@ impl App { ) .await }; - let (config, local_settings) = - loaded.map_err(|error| format!("Failed to read new session defaults: {error}"))?; + let Some((config, local_settings)) = + loaded.map_err(|error| format!("Failed to read new session defaults: {error}"))? + else { + return Ok(()); + }; if let Some(draft) = fresh_draft.as_mut() { draft.apply_settings(&local_settings, config.cwd.as_path()); } @@ -287,7 +300,7 @@ impl App { host_project_cwd.as_deref(), ) .map_err(|error| format!("Failed to start the new conversation: {error}")); - let started = if responsive_cloud_switch { + let Some(started) = if responsive_cloud_switch { let events = tui.event_stream(); wait_for_cloud_switch( self, @@ -298,8 +311,13 @@ impl App { ) .await } else { - StartupDraftPump::run_during_navigation(fresh_draft.as_mut(), tui, start).await - }?; + StartupDraftPump::run_during_navigation(fresh_draft.as_mut(), tui, start) + .await + .map(Some) + }? + else { + return Ok(()); + }; if responsive_cloud_switch { draft = Some(self.chat_widget.composer_draft_snapshot()); } @@ -307,7 +325,12 @@ impl App { if let Err(error) = tui.clear_for_thread_switch() { if responsive_cloud_switch { let events = tui.event_stream(); - finish_failed_cloud_switch(self, tui, events).await; + if finish_failed_cloud_switch(self, tui, events) + .await + .is_break() + { + return Ok(()); + } } return Err(format!("Failed to open the new conversation: {error}")); } diff --git a/codex-rs/tui/src/app/cloud_switch_wait.rs b/codex-rs/tui/src/app/cloud_switch_wait.rs index 9e332a7093..5365da8ce5 100644 --- a/codex-rs/tui/src/app/cloud_switch_wait.rs +++ b/codex-rs/tui/src/app/cloud_switch_wait.rs @@ -4,14 +4,19 @@ //! session preparation and adoption, so no competing navigation can install a Cloud stream. use super::App; +use crate::app_event::AppEvent; +use crate::app_event::ExitMode; +use crate::bottom_pane::CancellationEvent; use crate::bottom_pane::RestrictedInputMode; use crate::chatwidget::ClipboardPasteContext; use crate::tui; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; +use crossterm::event::KeyEventKind; use crossterm::event::KeyModifiers; use ratatui::layout::Size; use std::future::Future; +use std::ops::ControlFlow; use tokio_stream::Stream; use tokio_stream::StreamExt; @@ -33,12 +38,12 @@ fn handle_cloud_switch_event( tui: &mut tui::Tui, event: tui::TuiEvent, progress: CloudSwitchProgress, -) { +) -> ControlFlow<()> { let screen_size = match tui.screen_size_for_event(&event) { Ok(size) => size, Err(error) => { tracing::warn!(%error, "failed to inspect terminal size during cloud switch"); - return; + return ControlFlow::Continue(()); } }; // Clipboard text is delivered on Draw after ordinary input invalidation and target checks. @@ -57,6 +62,21 @@ fn handle_cloud_switch_event( } match event { tui::TuiEvent::Key(mut key) => { + if key.kind == KeyEventKind::Press + && crate::key_hint::ctrl(KeyCode::Char('c')).is_press(key) + { + app.cancel_pending_key_chord(); + if app.chat_widget.cancel_cloud_handoff_input() == CancellationEvent::NotHandled + && app.chat_widget.no_modal_or_popup_active() + { + app.app_event_tx + .send(AppEvent::Exit(ExitMode::ShutdownFirst)); + clear_cloud_switch_hint(app); + return ControlFlow::Break(()); + } + show_cloud_switch_hint(app, progress); + return ControlFlow::Continue(()); + } let escape = if app.should_recover_vim_insert_escape(tui, key) && !(tui.is_owned_screen() && crate::transcript_view::JumpTarget::from_key(key).is_some()) @@ -91,7 +111,7 @@ fn handle_cloud_switch_event( .handle_handoff_paste_tick(tui.frame_requester()) { tui.defer_screen_size(screen_size); - return; + return ControlFlow::Continue(()); } app.chat_widget.pre_draw_tick(); draw_cloud_switch(app, tui, screen_size); @@ -104,6 +124,7 @@ fn handle_cloud_switch_event( app.start_right_click_paste(tui, mouse, ClipboardPasteContext::DraftHandoff); } } + ControlFlow::Continue(()) } fn show_cloud_switch_hint(app: &mut App, progress: CloudSwitchProgress) { @@ -115,6 +136,11 @@ fn show_cloud_switch_hint(app: &mut App, progress: CloudSwitchProgress) { .set_footer_hint_override(Some(vec![(label.into(), action.into())])); } +fn clear_cloud_switch_hint(app: &mut App) { + app.cancel_pending_key_chord(); + app.chat_widget.set_footer_hint_override(/*items*/ None); +} + fn draw_cloud_switch(app: &mut App, tui: &mut tui::Tui, screen_size: Size) { if let Err(error) = app.render_chat_widget_frame(tui, screen_size) { tracing::warn!(%error, "failed to render during cloud switch"); @@ -126,11 +152,11 @@ pub(super) async fn finish_failed_cloud_switch( app: &mut App, tui: &mut tui::Tui, mut events: impl Stream + Unpin, -) { +) -> ControlFlow<()> { let progress = CloudSwitchProgress::Failed; show_cloud_switch_hint(app, progress); draw_cloud_switch(app, tui, tui.terminal.last_known_screen_size); - drain_cloud_switch_events(app, tui, &mut events, progress).await; + drain_cloud_switch_events(app, tui, &mut events, progress).await } async fn drain_cloud_switch_events( @@ -138,7 +164,7 @@ async fn drain_cloud_switch_events( tui: &mut tui::Tui, events: &mut (impl Stream + Unpin), progress: CloudSwitchProgress, -) { +) -> ControlFlow<()> { // The current draft owns paste keys and any clipboard text until adoption. let settle_clipboard = matches!( progress, @@ -164,26 +190,27 @@ async fn drain_cloud_switch_events( let Ok(Some(event)) = tokio::time::timeout_at(deadline, events.next()).await else { if clipboard_deadline.is_some() { // The terminal stream can end or the deadline can beat its scheduled Draw. - handle_cloud_switch_event(app, tui, tui::TuiEvent::Draw, progress); + handle_cloud_switch_event(app, tui, tui::TuiEvent::Draw, progress)?; app.pending_right_click_paste = None; } break; }; - handle_cloud_switch_event(app, tui, event, progress); + handle_cloud_switch_event(app, tui, event, progress)?; } - app.cancel_pending_key_chord(); - app.chat_widget.set_footer_hint_override(/*items*/ None); + clear_cloud_switch_hint(app); + ControlFlow::Continue(()) } /// Keep the draft editable while Cloud prepares its destination. Other app actions wait /// until the handoff completes, so a second navigation cannot replace the pending stream. +/// Returns `Ok(None)` when the user requests exit, dropping the pending operation. pub(crate) async fn wait_for_cloud_switch( app: &mut App, tui: &mut tui::Tui, operation: impl Future>, events: impl Stream + Unpin, clipboard_handoff: ClipboardHandoff, -) -> Result { +) -> Result, E> { let progress = CloudSwitchProgress::Connecting(clipboard_handoff); app.cancel_pending_key_chord(); show_cloud_switch_hint(app, progress); @@ -198,13 +225,20 @@ pub(crate) async fn wait_for_cloud_switch( tokio::select! { biased; result = &mut operation => break result, - Some(event) = events.next() => handle_cloud_switch_event(app, tui, event, progress), + Some(event) = events.next() => { + if handle_cloud_switch_event(app, tui, event, progress).is_break() { + return Ok(None); + } + } } }; - if result.is_err() { - finish_failed_cloud_switch(app, tui, events).await; + let input = if result.is_err() { + finish_failed_cloud_switch(app, tui, events).await } else { - drain_cloud_switch_events(app, tui, &mut events, progress).await; + drain_cloud_switch_events(app, tui, &mut events, progress).await + }; + if input.is_break() { + return Ok(None); } - result + result.map(Some) } diff --git a/codex-rs/tui/src/app/new_session.rs b/codex-rs/tui/src/app/new_session.rs index a0721815fd..aa8e920d9b 100644 --- a/codex-rs/tui/src/app/new_session.rs +++ b/codex-rs/tui/src/app/new_session.rs @@ -68,7 +68,8 @@ impl App { let settings = self.backend_settings(); let cwd = self.chat_widget.config_ref().cwd.to_path_buf(); self.load_new_session_config_from(app_server, &settings, cwd, /*cloud_switch*/ None) - .await + .await? + .ok_or_else(|| color_eyre::eyre::eyre!("Only a Cloud switch can be canceled")) } pub(super) async fn load_new_session_config_from( @@ -77,7 +78,7 @@ impl App { settings: &BackendSettings, cwd: PathBuf, mut cloud_switch: Option<&mut tui::Tui>, - ) -> Result<(Config, crate::local_settings::LocalSettings)> { + ) -> Result> { let defaults_cwd = match app_server.thread_params_mode() { crate::app_server_session::ThreadParamsMode::Embedded => cwd.as_path(), crate::app_server_session::ThreadParamsMode::Remote @@ -99,7 +100,7 @@ impl App { .await?; Ok::<_, color_eyre::eyre::Report>((defaults, refresh.await)) }; - let (defaults, refreshed_config) = if let Some(tui) = cloud_switch.as_deref_mut() { + let Some((defaults, refreshed_config)) = if let Some(tui) = cloud_switch.as_deref_mut() { let events = tui.event_stream(); super::cloud_switch_wait::wait_for_cloud_switch( self, @@ -110,8 +111,11 @@ impl App { ) .await } else { - load.await - }?; + load.await.map(Some) + }? + else { + return Ok(None); + }; // Stage local preferences and permission carryover without changing the active task. let (mut config, local_settings) = match refreshed_config { Ok(config) => { @@ -139,7 +143,12 @@ impl App { ) { if let Some(tui) = cloud_switch { let events = tui.event_stream(); - super::cloud_switch_wait::finish_failed_cloud_switch(self, tui, events).await; + if super::cloud_switch_wait::finish_failed_cloud_switch(self, tui, events) + .await + .is_break() + { + return Ok(None); + } } return Err(error); } @@ -157,6 +166,6 @@ impl App { &self.cli_kv_overrides, &self.harness_overrides, ); - Ok((config, local_settings)) + Ok(Some((config, local_settings))) } } diff --git a/codex-rs/tui/src/app/tests/new_session_tests.rs b/codex-rs/tui/src/app/tests/new_session_tests.rs index f32fda29ff..0a6078c627 100644 --- a/codex-rs/tui/src/app/tests/new_session_tests.rs +++ b/codex-rs/tui/src/app/tests/new_session_tests.rs @@ -1708,6 +1708,84 @@ async fn cloud_switch_wait_keeps_draft_editable_without_submitting_to_local() -> Ok(()) } +#[tokio::test] +async fn cloud_switch_wait_quits_on_ctrl_c_before_adopting_the_thread() -> Result<()> { + use crate::app::cloud_switch_wait::ClipboardHandoff; + use crate::app::cloud_switch_wait::wait_for_cloud_switch; + use crate::app_event::ExitMode; + + for completed in [false, true] { + let (mut app, mut app_events, _operations) = make_test_app_with_channels().await; + let mut tui = crate::tui::test_support::make_test_tui()?; + let operation = async { + if !completed { + std::future::pending::<()>().await; + } + Ok::<_, std::convert::Infallible>(()) + }; + let ctrl_c = TuiEvent::Key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL)); + let wait = wait_for_cloud_switch( + &mut app, + &mut tui, + operation, + tokio_stream::iter([ctrl_c]).chain(tokio_stream::pending()), + ClipboardHandoff::SettleBeforeNextView, + ); + tokio::pin!(wait); + assert_eq!(futures::poll!(&mut wait), std::task::Poll::Ready(Ok(None))); + let exits: Vec<_> = std::iter::from_fn(|| app_events.try_recv().ok()) + .filter_map(|event| match event { + AppEvent::Exit(mode) => Some(mode), + _ => None, + }) + .collect(); + assert_eq!(exits, [ExitMode::ShutdownFirst]); + } + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn cloud_switch_wait_ctrl_c_clears_buffered_draft_before_quitting() -> Result<()> { + use crate::app::cloud_switch_wait::ClipboardHandoff; + use crate::app::cloud_switch_wait::wait_for_cloud_switch; + use crate::app_event::ExitMode; + + let (mut app, mut app_events, _operations) = make_test_app_with_channels().await; + let thread_id = ThreadId::new(); + let session = test_thread_session(thread_id, app.config.cwd.to_path_buf()); + app.chat_widget.handle_thread_session(session); + let mut tui = crate::tui::test_support::make_test_tui()?; + let (send, events) = tokio::sync::mpsc::unbounded_channel(); + let ctrl_c = TuiEvent::Key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL)); + send.send(TuiEvent::Key(KeyCode::Char('a').into()))?; + send.send(ctrl_c.clone())?; + let wait = wait_for_cloud_switch( + &mut app, + &mut tui, + std::future::pending::>(), + tokio_stream::wrappers::UnboundedReceiverStream::new(events), + ClipboardHandoff::SettleBeforeNextView, + ); + tokio::pin!(wait); + assert!(futures::poll!(&mut wait).is_pending()); + let history: Vec<_> = std::iter::from_fn(|| app_events.try_recv().ok()) + .filter_map(|event| match event { + AppEvent::AppendMessageHistoryEntry { thread_id, text } => Some((thread_id, text)), + AppEvent::Exit(_) => panic!("the first Ctrl+C should clear the draft"), + _ => None, + }) + .collect(); + assert_eq!(history, [(thread_id, "a".into())]); + + send.send(ctrl_c)?; + assert_eq!(futures::poll!(&mut wait), std::task::Poll::Ready(Ok(None))); + assert!( + std::iter::from_fn(|| app_events.try_recv().ok()) + .any(|event| matches!(event, AppEvent::Exit(ExitMode::ShutdownFirst))) + ); + Ok(()) +} + #[tokio::test(start_paused = true)] async fn cloud_switch_wait_routes_ctrl_d_to_the_draft_or_queues_shutdown() -> Result<()> { use crate::app::cloud_switch_wait::ClipboardHandoff; diff --git a/codex-rs/tui/src/chatwidget/reconnect.rs b/codex-rs/tui/src/chatwidget/reconnect.rs index a9b612bdd4..d4e867c33f 100644 --- a/codex-rs/tui/src/chatwidget/reconnect.rs +++ b/codex-rs/tui/src/chatwidget/reconnect.rs @@ -198,6 +198,16 @@ impl ChatWidget { self.bottom_pane.composer_paste_deadline() } + pub(crate) fn cancel_cloud_handoff_input(&mut self) -> CancellationEvent { + // Count buffered terminal input as part of the draft before deciding to quit. + self.bottom_pane.flush_composer_paste_burst(); + let cancellation = self.bottom_pane.on_ctrl_c(); + if cancellation == CancellationEvent::Handled { + self.refresh_startup_recovery(); + } + cancellation + } + pub(crate) fn handle_restricted_key( &mut self, key: KeyEvent,