diff --git a/codex-rs/app-server-client/src/cloud.rs b/codex-rs/app-server-client/src/cloud.rs index 1a356ece4c..c0438499da 100644 --- a/codex-rs/app-server-client/src/cloud.rs +++ b/codex-rs/app-server-client/src/cloud.rs @@ -11,7 +11,12 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::Result as JsonRpcResult; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnStartParams; +use std::collections::HashMap; use std::io; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::PoisonError; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -59,7 +64,54 @@ pub struct CloudRequestHandle { commands: mpsc::Sender, } +/// Tracks only losses queued for delivery. The worker remains the sole owner +/// of the active attachment. A successful replacement supersedes earlier losses +/// for that conversation, even if the replacement itself later ends. +#[derive(Clone, Default)] +pub(super) struct CloudStreamLosses(Arc>>); + +#[derive(Default)] +struct PendingLosses { + remaining: usize, + superseded: usize, +} + +impl CloudStreamLosses { + fn pending(&self) -> MutexGuard<'_, HashMap> { + self.0.lock().unwrap_or_else(PoisonError::into_inner) + } + + pub(super) fn record(&self, thread_id: &str) { + self.pending() + .entry(thread_id.to_owned()) + .or_default() + .remaining += 1; + } + + pub(super) fn supersede(&self, thread_id: &str) { + if let Some(losses) = self.pending().get_mut(thread_id) { + losses.superseded = losses.remaining; + } + } + + fn should_deliver(&self, thread_id: &str) -> bool { + let mut pending = self.pending(); + let Some(losses) = pending.get_mut(thread_id) else { + return true; + }; + // A single worker sends losses in the same order as it records them. + let deliver = losses.superseded == 0; + losses.superseded = losses.superseded.saturating_sub(1); + losses.remaining -= 1; + if losses.remaining == 0 { + pending.remove(thread_id); + } + deliver + } +} + pub struct CloudClient { + pub(super) stream_losses: CloudStreamLosses, handle: CloudRequestHandle, events: mpsc::UnboundedReceiver, } @@ -79,6 +131,7 @@ impl CloudClient { let (events, event_receiver) = mpsc::unbounded_channel(); ( Self { + stream_losses: CloudStreamLosses::default(), handle: CloudRequestHandle { commands }, events: event_receiver, }, @@ -92,7 +145,15 @@ impl CloudClient { } pub(crate) async fn next_event(&mut self) -> Option { - self.events.recv().await + while let Some(event) = self.events.recv().await { + if let AppServerEvent::ThreadStreamDisconnected(error) = &event + && !self.stream_losses.should_deliver(&error.thread_id) + { + continue; + } + return Some(event); + } + None } pub(crate) async fn request(&self, request: ClientRequest) -> io::Result { 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 d67380b25c..928d967598 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 @@ -44,19 +44,62 @@ async fn failed_replacement_keeps_old_thread_until_new_attachment_succeeds() -> let _ = close.send(()); assert!(replies.admit_turn("new", &output).is_err()); - let AppServerEvent::ServerNotification(notification) = notifications.try_recv()? else { + let AppServerEvent::ThreadStreamDisconnected(notification) = notifications.try_recv()? else { panic!("expected a disconnect notification"); }; assert_eq!( - serde_json::to_value(notification)?, - serde_json::to_value(cloud_error( + serde_json::to_value(protocol::ServerNotification::Error(*notification))?, + serde_json::to_value(protocol::ServerNotification::Error(cloud_error( "new", "Cloud stream ended. Reopen the conversation to reconnect." - ))? + )))? ); Ok(()) } +#[tokio::test] +async fn successful_reopen_supersedes_only_earlier_losses_for_that_thread() -> Result<()> { + let (mut client, _commands, output) = CloudClient::channel(/*capacity*/ 1); + let mut replies = Replies { + stream_losses: client.stream_losses.clone(), + ..Default::default() + }; + + for id in ["reopened", "different", "reopened"] { + let (sender, stream) = futures::channel::mpsc::unbounded(); + replies + .replace(id, async { Ok(stream.boxed()) }, &output) + .await?; + drop(sender); + replies.forward_next(&output).await; + } + assert!( + replies + .replace("reopened", async { Ok(stream::empty().boxed()) }, &output) + .await + .is_err() + ); + + // The first EOF was followed by a successful reopen. A different thread's + // EOF and the reopened thread's newer EOF must still be delivered, in order. + drop(output); + for expected in ["different", "reopened"] { + let Some(AppServerEvent::ThreadStreamDisconnected(error)) = client.next_event().await + else { + panic!("expected a thread stream disconnect"); + }; + assert_eq!( + *error, + cloud_error( + expected, + "Cloud stream ended. Reopen the conversation to reconnect." + ) + ); + } + assert!(client.next_event().await.is_none()); + Ok(()) +} + #[tokio::test] async fn attachment_status_follows_its_server_notifications() -> Result<()> { let (output, _notifications) = mpsc::unbounded_channel(); 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 7454fdb302..2b8188553c 100644 --- a/codex-rs/app-server-client/src/cloud_backend/mod.rs +++ b/codex-rs/app-server-client/src/cloud_backend/mod.rs @@ -9,6 +9,7 @@ use crate::AppServerClient; use crate::AppServerEvent; use crate::CloudClient; use crate::CloudOperation; +use crate::cloud::CloudStreamLosses; use codex_app_server_protocol as protocol; use codex_cloud_client as cloud; use color_eyre::eyre::Result; @@ -41,6 +42,7 @@ struct Attachment { #[derive(Default)] struct Replies { active: Option, + stream_losses: CloudStreamLosses, // Keep the last observed status through EOF until this conversation is unloaded. latest_status: Option<(String, cloud::ThreadStatus)>, } @@ -53,14 +55,13 @@ impl Replies { match active.stream.next().await { Some(Ok(event)) => forward(&active.thread_id, event, output, &mut self.latest_status), Some(Err(_)) | None => { - send_notification( - output, - cloud_error( - &active.thread_id, - "Cloud stream ended. Reopen the conversation to reconnect.", - ), + let error = cloud_error( + &active.thread_id, + "Cloud stream ended. Reopen the conversation to reconnect.", ); self.active = None; + self.stream_losses.record(&error.thread_id); + let _ = output.send(AppServerEvent::ThreadStreamDisconnected(Box::new(error))); } } } @@ -96,14 +97,15 @@ impl Replies { Some(Some(Ok(event))) => Some(event), None => None, }; - self.latest_status = None; - if let Some(event) = first_event { - forward(id, event, output, &mut self.latest_status); - } self.active = Some(Attachment { thread_id: id.to_owned(), stream, }); + self.stream_losses.supersede(id); + self.latest_status = None; + if let Some(event) = first_event { + forward(id, event, output, &mut self.latest_status); + } Ok(()) } @@ -155,8 +157,12 @@ fn start_worker(client: cloud::Client, account: protocol::GetAccountResponse) -> loaded: None, warned_partial_history: None, }; + let mut replies = Replies { + stream_losses: app_client.stream_losses.clone(), + ..Default::default() + }; tokio::spawn(async move { - run_worker(commands, &mut requests, &mut Replies::default(), &output).await; + run_worker(commands, &mut requests, &mut replies, &output).await; }); AppServerClient::Cloud(app_client) } @@ -298,7 +304,10 @@ fn forward( cloud::Event::Transcript(event) => match events::transcript(event.thread_id, event.update) { Ok(Some(notification)) => notification, Ok(None) => return, - Err(_) => cloud_error(id, "The cloud sent an unsupported transcript event."), + Err(_) => protocol::ServerNotification::Error(cloud_error( + id, + "The cloud sent an unsupported transcript event.", + )), }, cloud::Event::ServerRequest(_) => protocol::ServerNotification::Warning(protocol::WarningNotification { thread_id: Some(id.to_owned()), @@ -316,8 +325,8 @@ fn send_notification( let _ = output.send(AppServerEvent::ServerNotification(Box::new(notification))); } -fn cloud_error(id: &str, message: &str) -> protocol::ServerNotification { - protocol::ServerNotification::Error(protocol::ErrorNotification { +fn cloud_error(id: &str, message: &str) -> protocol::ErrorNotification { + protocol::ErrorNotification { thread_id: id.to_owned(), turn_id: String::new(), will_retry: false, @@ -327,7 +336,7 @@ fn cloud_error(id: &str, message: &str) -> protocol::ServerNotification { codex_error_info: None, misalignment: None, }, - }) + } } #[cfg(test)] diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 5b011863f1..3d25037697 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -112,10 +112,16 @@ pub type RequestResult = std::result::Result; #[derive(Debug, Clone)] pub enum AppServerEvent { - Lagged { skipped: usize }, + Lagged { + skipped: usize, + }, ServerNotification(Box), ServerRequest(Box), - Disconnected { message: String }, + /// A native Cloud reply stream ended. The underlying thread can still be running. + ThreadStreamDisconnected(Box), + Disconnected { + message: String, + }, } impl From for AppServerEvent { diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index bdd7a651f8..3743d0a15e 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -156,6 +156,13 @@ impl App { .await; self.repaint_agents_overview(); } + AppServerEvent::ThreadStreamDisconnected(error) => { + self.handle_server_notification_event( + app_server_client, + ServerNotification::Error(*error), + ) + .await; + } AppServerEvent::Disconnected { message } => { // The request result distinguishes a refused start from a lost reply. if self.pending_startup_thread_start diff --git a/codex-rs/tui/src/app/tests/backend_event_tests.rs b/codex-rs/tui/src/app/tests/backend_event_tests.rs index e2c3fb69b8..6f74cb1c6f 100644 --- a/codex-rs/tui/src/app/tests/backend_event_tests.rs +++ b/codex-rs/tui/src/app/tests/backend_event_tests.rs @@ -11,26 +11,32 @@ use pretty_assertions::assert_eq; #[tokio::test] async fn foreign_backend_events_leave_thread_cache_and_connection_untouched() -> Result<()> { let (mut app, mut events, _) = make_test_app_with_channels().await; - let (server, requests, proxy) = + let (mut server, requests, proxy) = session_lifecycle_requests::start_recording_remote_app_server(&app.config).await?; let id = ThreadId::new(); app.primary_thread_id = Some(id); - app.active_thread_id = Some(id); set_test_initial_prompt(&mut app, String::new()); app.chat_widget .handle_thread_session(test_thread_session(id, app.config.cwd.to_path_buf())); app.chat_widget .restore_user_message_to_composer("keep draft".into()); app.ensure_thread_channel(id); + app.activate_thread_channel(id).await; while events.try_recv().is_ok() {} requests.lock().unwrap().clear(); let input = app.chat_widget.capture_thread_input_state(); let hard_stop_generation = app.rate_limit_hard_stop_generation; let foreign = BackendId::new(); let request = exec_approval_request(id, "foreign-turn", "command", /*approval_id*/ None); + let loss: codex_app_server_protocol::ErrorNotification = + serde_json::from_value(serde_json::json!({ + "threadId": id.to_string(), "turnId": "", "willRetry": false, + "error": { "message": "Cloud stream ended. Reopen the conversation to reconnect." } + }))?; for event in [ AppServerEvent::ServerNotification(Box::new(turn_started_notification(id, "foreign-turn"))), AppServerEvent::ServerRequest(Box::new(request.clone())), + AppServerEvent::ThreadStreamDisconnected(Box::new(loss.clone())), AppServerEvent::Lagged { skipped: 1 }, AppServerEvent::Disconnected { message: "foreign disconnected".into(), @@ -107,6 +113,21 @@ async fn foreign_backend_events_leave_thread_cache_and_connection_untouched() -> .active_turn_id(), Some("owned-turn") ); + let mut tui = crate::tui::test_support::make_test_tui()?; + app.drain_active_thread_events(&mut tui).await?; + app.handle_backend_event( + &server, + AppServerSessionEvent { + backend_id: server.backend_id(), + event: AppServerEvent::ThreadStreamDisconnected(Box::new(loss)), + }, + ) + .await; + app.drain_active_thread_events(&mut tui).await?; + let history = + super::active_reconnect::drain_history(&mut app, &mut tui, &mut server, &mut events) + .await?; + insta::assert_snapshot!(history, @"■ Cloud stream ended. Reopen the conversation to reconnect."); server.shutdown().await?; proxy.await??; Ok(()) diff --git a/codex-rs/tui/src/dynamic_tools_tests.rs b/codex-rs/tui/src/dynamic_tools_tests.rs index 86049ce6cd..5190b56410 100644 --- a/codex-rs/tui/src/dynamic_tools_tests.rs +++ b/codex-rs/tui/src/dynamic_tools_tests.rs @@ -256,6 +256,7 @@ stream_max_retries = 0 return Ok::<_, color_eyre::eyre::Report>(()); } } + codex_app_server_client::AppServerEvent::ThreadStreamDisconnected(_) => {} codex_app_server_client::AppServerEvent::Lagged { .. } | codex_app_server_client::AppServerEvent::Disconnected { .. } => { panic!("lost app-server events during the tool round trip"); diff --git a/codex-rs/tui/src/onboarding/onboarding_screen.rs b/codex-rs/tui/src/onboarding/onboarding_screen.rs index 83c426a11f..88adbefa7a 100644 --- a/codex-rs/tui/src/onboarding/onboarding_screen.rs +++ b/codex-rs/tui/src/onboarding/onboarding_screen.rs @@ -794,7 +794,8 @@ async fn run_onboarding_screen_inner( AppServerEvent::Disconnected { message } => { return Err(color_eyre::eyre::eyre!(message)); } - AppServerEvent::Lagged { .. } + AppServerEvent::ThreadStreamDisconnected(_) + | AppServerEvent::Lagged { .. } | AppServerEvent::ServerRequest(_) => {} } }