diff --git a/codex-rs/app-server-client/src/cloud.rs b/codex-rs/app-server-client/src/cloud.rs index d99562207d..db028fa8f8 100644 --- a/codex-rs/app-server-client/src/cloud.rs +++ b/codex-rs/app-server-client/src/cloud.rs @@ -25,6 +25,14 @@ use tokio::sync::oneshot; /// Unsupported operations must return an explicit error; they must not claim success. pub enum CloudOperation { Request(Box), + /// Opens a provisional attachment and serves its initial history on a scoped channel. + PrepareThreadResume { + request_id: RequestId, + params: Box, + history: mpsc::Receiver, + }, + /// Sent on the scoped channel after the caller has prepared the replacement view. + CommitPreparedResume, /// Host project metadata stays on the private transport, outside the app-server wire. StartProjectThread { request_id: RequestId, @@ -309,6 +317,38 @@ impl CloudRequestHandle { crate::decode_typed_response("cloud/submitUserMessage", response) } + /// Returns the activation snapshot and a handle restricted to this resume's history. + /// Dropping the scoped handle without committing discards the new attachment. + pub async fn prepare_thread_resume( + &self, + request_id: RequestId, + params: codex_app_server_protocol::ThreadResumeParams, + ) -> Result<(codex_app_server_protocol::ThreadResumeResponse, Self), TypedRequestError> { + let (commands, history) = mpsc::channel(/*capacity*/ 1); + let response = self + .send(CloudOperation::PrepareThreadResume { + request_id, + params: Box::new(params), + history, + }) + .await; + let snapshot = crate::decode_typed_response("thread/resume", response)?; + Ok(( + snapshot, + Self { + commands, + reconnect: self.reconnect.clone(), + catalog: None, + }, + )) + } + + /// Installs the prepared attachment on the worker before its next ordinary command. + pub async fn commit_prepared_resume(&self) -> Result<(), TypedRequestError> { + let response = self.send(CloudOperation::CommitPreparedResume).await; + crate::decode_typed_response("thread/resume", response) + } + async fn send(&self, operation: CloudOperation) -> io::Result { let (response, result) = oneshot::channel(); self.commands 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 fa69361e43..4fea016f4b 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 @@ -384,6 +384,13 @@ async fn failed_replacement_keeps_old_thread_until_new_attachment_succeeds() -> .await?; replies.detach("old"); assert!(replies.admit_turn("old", &output).is_err()); + replies + .replace( + "new", + async { panic!("an already subscribed thread must reuse its reply stream") }, + &output, + ) + .await?; replies.admit_turn("new", &output)?; let _ = close.send(()); @@ -401,6 +408,163 @@ async fn failed_replacement_keeps_old_thread_until_new_attachment_succeeds() -> Ok(()) } +#[tokio::test] +async fn prepared_resume_keeps_the_previous_stream_until_handoff() -> Result<()> { + let server = MockServer::start().await; + let home = tempfile::tempdir()?; + let settings = |id: &str| { + serde_json::from_value(serde_json::json!({ + "thread": { + "id": id, "sessionId": id, "preview": "", "ephemeral": false, + "modelProvider": "openai", "createdAt": 0, "updatedAt": 0, + "status": {"type": "idle"}, "cwd": home.path(), "cliVersion": "", + "source": "cli", "turns": [], + }, + "model": "test", "modelProvider": "openai", "cwd": home.path(), + "approvalPolicy": "never", "approvalsReviewer": "user", "sandbox": {"type": "readOnly"}, + })) + }; + let native = Arc::new(cloud::Client::new( + &codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + &server.uri(), + cloud::Credentials::new("synthetic-bearer", "synthetic-account")?, + )?); + // One opaque history item with a continuation. A's mutation deliberately loses its receipt. + let partial_page: &[u8] = b"\0\0\0\0\x0f\x0a\x0a\x0a\x08\x0a\x01u\x12\x03\xa2\x01\0\x12\x01x"; + let list = ResponseTemplate::new(200) + .insert_header("grpc-status", "0") + .set_body_raw(partial_page, "application/grpc"); + let add = ResponseTemplate::new(503); + for (rpc, response, count) in [("ListItems", list, 8), ("AddUserMessage", add, 4)] { + let path_name = format!("/oaiproto.codex_app_server_backend.ThreadService/{rpc}"); + Mock::given(method("POST")) + .and(path(path_name)) + .respond_with(response) + .expect(count) + .mount(&server) + .await; + } + let waiting = b"\x1a\x0d\x12\x0b\x0a\x06\x22\x04\x0a\x02\x08\x01\x12\x01b"; + let idle = b"\x1a\x09\x12\x07\x0a\x02\x12\x00\x12\x01b"; + for outcome in ["eof", "burst", "shutdown", "abandon", "reused", "stale"] { + let (mut client, commands, output) = CloudClient::channel(/*capacity*/ 1); + let (scoped, history, _) = CloudClient::channel(/*capacity*/ 1); + let mut requests = requests::Requests { + client: Arc::clone(&native), + account: serde_json::from_value(serde_json::json!({"requiresOpenaiAuth": true}))?, + loaded: Some(settings("a")?), + warned_partial_history: Some("a".into()), + }; + let mut replies = Replies::default(); + replies + .replace("a", async { Ok(stream::pending().boxed()) }, &output) + .await?; + let (incoming, stream) = mpsc::unbounded_channel(); + incoming.send(Ok(cloud::Event::Notification(idle.as_slice().into())))?; + if outcome == "burst" { + incoming.send(Ok(cloud::Event::Notification(waiting.as_slice().into())))?; + incoming.send(Ok(cloud::Event::Notification(idle.as_slice().into())))?; + for _ in 0..PREPARED_EVENT_LIMIT { + incoming.send(Ok(cloud::Event::ServerRequest(Default::default())))?; + } + } + let stream = stream::unfold(stream, |mut rx| async { + rx.recv().await.map(|event| (event, rx)) + }); + let reused = matches!(outcome, "reused" | "stale"); + let first = cloud::Event::ServerRequest(Default::default()); + let candidate = PreparedResume { + settings: settings(if reused { "a" } else { "b" })?, + opened: (!reused).then(|| (stream.boxed(), Some(first))), + history, + generation: replies.generation - u64::from(outcome == "stale"), + canceled_recovery: None, + warning: None, + }; + // The EOF case closes after a second, distinguishable buffered event. + let incoming = (outcome != "eof").then_some(incoming); + let drive = async { + let handle = scoped.request_handle(); + let page = |cursor: Option<&str>| { + handle.request(protocol::ClientRequest::ThreadItemsList { + request_id: protocol::RequestId::Integer(1), + params: serde_json::from_value( + serde_json::json!({"threadId": "b", "cursor": cursor}), + ) + .unwrap(), + }) + }; + if !reused { + assert_eq!(page(None).await?.unwrap()["nextCursor"], "x"); + let params = serde_json::from_value(serde_json::json!({ + "threadId": "a", "clientUserMessageId": "send", + "input": [{"type": "text", "text": "hello"}], + }))?; + let sent = client.request_handle().submit_user_message(params).await; + assert!( + matches!(sent, Err(crate::TypedRequestError::Server { source, .. }) + if source.code == -32000 && source.message.contains("mutation outcome unknown")) + ); + assert!(page(Some("x")).await?.is_ok()); + assert!(client.next_event().now_or_never().is_none()); + } + if outcome == "abandon" { + drop((scoped, handle)); + incoming.as_ref().unwrap().closed().await; + } else if outcome != "shutdown" { + let commit = handle.commit_prepared_resume().await; + assert_eq!(commit.is_ok(), matches!(outcome, "burst" | "reused")); + if outcome == "eof" { + assert!(client.next_event().now_or_never().is_none()); + } + } + if outcome == "burst" { + // One partial-history warning plus the entire live burst, including its unread tail. + for _ in 0..PREPARED_EVENT_LIMIT + 2 { + assert!( + matches!(client.next_event().await, Some(AppServerEvent::ServerNotification(notification)) + if matches!(*notification, protocol::ServerNotification::Warning(_))) + ); + } + } + client.shutdown().await?; + Ok::<(), color_eyre::Report>(()) + }; + let worker = run_worker_with_prepared( + commands, + &mut requests, + &mut replies, + &output, + Some(candidate), + ); + let ((), observed) = tokio::time::timeout(crate::SHUTDOWN_TIMEOUT, async { + tokio::join!(worker, drive) + }) + .await?; + observed?; + let committed = outcome == "burst"; + assert_eq!( + requests.loaded.unwrap().thread.id, + if committed { "b" } else { "a" } + ); + if committed { + assert_eq!(replies.thread_status("b"), Some(&cloud::ThreadStatus::Idle)); + } + assert_eq!(replies.generation, if committed { 2 } else { 1 }); + assert_eq!( + requests.warned_partial_history.as_deref(), + match outcome { + "burst" => Some("b"), + "reused" => None, + _ => Some("a"), + } + ); + } + Ok(()) +} + #[tokio::test] async fn successful_reopen_supersedes_only_earlier_losses_for_that_thread() -> Result<()> { let (mut client, _commands, output) = CloudClient::channel(/*capacity*/ 1); 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 6f128f0e54..0b2741b36a 100644 --- a/codex-rs/app-server-client/src/cloud_backend/mod.rs +++ b/codex-rs/app-server-client/src/cloud_backend/mod.rs @@ -23,10 +23,12 @@ use futures::stream::BoxStream; use serde_json::Value; use std::sync::Arc; use tokio::sync::mpsc; +use tokio::sync::oneshot; pub use auth::connect; type ReplyStream = BoxStream<'static, Result>; +type OpenReplyStream = (ReplyStream, Option); /// Rejected locally before an admission request can be sent. #[derive(Debug, thiserror::Error)] @@ -46,6 +48,7 @@ struct Attachment { struct Replies { idle: Option, active: Option, + generation: u64, stream_losses: CloudStreamLosses, // Keep the last observed status through EOF until this conversation is unloaded. latest_status: Option<(String, cloud::ThreadStatus)>, @@ -143,25 +146,38 @@ impl Replies { open: impl Future>, output: &mpsc::UnboundedSender, ) -> Result<()> { + if let Some(opened) = self.open_attachment(id, open, output).await? { + self.install_replacement(id, opened, output); + } + Ok(()) + } + + async fn open_attachment( + &mut self, + id: &str, + open: impl Future>, + output: &mpsc::UnboundedSender, + ) -> Result> { + if self.admit_turn(id, output).is_ok() { + return Ok(None); + } let mut stream = self.while_listening(open, output).await?; - // A server can close Attach immediately after returning successful headers. - // A healthy idle stream needs no event to prove that it is ready. - let first_event = match tokio::task::unconstrained(stream.next()).now_or_never() { - Some(None | Some(Err(_))) => { - bail!( - "Cloud conversation activated but its stream did not connect. Reopen it before sending." - ); - } - Some(Some(Ok(event))) => Some(event), - None => None, - }; + let first = probe_reply_stream(&mut stream)?; + Ok(Some((stream, first))) + } + + fn install_replacement( + &mut self, + id: &str, + (stream, first): OpenReplyStream, + output: &mpsc::UnboundedSender, + ) { if let Some(idle) = &mut self.idle { idle.reset(); } self.stream_losses.supersede(id); self.latest_status = None; - self.install(id, stream, first_event, output); - Ok(()) + self.install(id, stream, first, output); } fn install( @@ -171,6 +187,7 @@ impl Replies { first_event: Option, output: &mpsc::UnboundedSender, ) { + self.generation = self.generation.wrapping_add(1); self.active = Some(Attachment { thread_id: id.to_owned(), stream, @@ -237,6 +254,19 @@ impl Replies { } } +fn probe_reply_stream(stream: &mut ReplyStream) -> Result> { + // Attach can close after successful headers. A healthy idle stream need not emit an event. + match tokio::task::unconstrained(stream.next()).now_or_never() { + Some(None | Some(Err(_))) => { + bail!( + "Cloud conversation activated but its stream did not connect. Reopen it before sending." + ); + } + Some(Some(Ok(event))) => Ok(Some(event)), + None => Ok(None), + } +} + /// Opens an app-server client backed by the authenticated native Cloud client. fn start_worker( client: cloud::Client, @@ -271,16 +301,38 @@ fn start_worker( } async fn run_worker( + commands: mpsc::Receiver, + requests: &mut requests::Requests, + replies: &mut Replies, + output: &mpsc::UnboundedSender, +) { + run_worker_with_prepared(commands, requests, replies, output, /*prepared*/ None).await; +} + +async fn run_worker_with_prepared( mut commands: mpsc::Receiver, requests: &mut requests::Requests, replies: &mut Replies, output: &mpsc::UnboundedSender, + mut prepared: Option, ) { loop { let mut command = tokio::select! { _ = output.closed() => return, () = replies.forward_next(output) => continue, command = commands.recv() => match command { Some(command) => command, None => return }, + command = async { + match &mut prepared { + Some(candidate) => candidate.history.recv().await, + None => std::future::pending().await, + } + } => { + let Some(candidate) = prepared.take() else { + continue; + }; + prepared = advance_prepared_resume(candidate, command, requests, replies, output).await; + continue; + } }; let (is_thread_start, selects_attachment) = match &command.operation { CloudOperation::Request(request) => match request.as_ref() { @@ -289,6 +341,8 @@ async fn run_worker( _ => (false, false), }, CloudOperation::StartProjectThread { .. } => (true, true), + CloudOperation::PrepareThreadResume { .. } if command.response.is_closed() => continue, + CloudOperation::PrepareThreadResume { .. } => (false, true), _ => (false, false), }; // A queued start that nobody owns must leave the existing recovery alone. @@ -297,7 +351,14 @@ async fn run_worker( } // A deliberate selection owns the next attachment, even if its handshake fails. let canceled = if selects_attachment { - replies.idle.as_mut().and_then(idle::Recovery::cancel) + let previous = prepared + .take() + .and_then(|candidate| candidate.canceled_recovery); + replies + .idle + .as_mut() + .and_then(idle::Recovery::cancel) + .or(previous) } else { None }; @@ -309,6 +370,28 @@ async fn run_worker( CloudOperation::Request(request) => { Ok((*request, cloud::SubmissionMode::StartIfIdle, None)) } + CloudOperation::PrepareThreadResume { + request_id, + params, + history, + } => { + prepared = tokio::select! { + _ = output.closed() => return, + candidate = prepare_thread_resume( + protocol::ClientRequest::ThreadResume { request_id, params: *params }, + history, command.response, requests, replies, output, + ) => candidate, + }; + if let Some(candidate) = &mut prepared { + candidate.canceled_recovery = canceled; + } else if let Some(root) = canceled { + replies.report_loss(&root, output); + } + continue; + } + CloudOperation::CommitPreparedResume => Err(eyre!( + "Cloud resume commit requires its scoped history handle." + )), CloudOperation::StartProjectThread { request_id, params, @@ -339,10 +422,22 @@ async fn run_worker( }; let result = match request { Ok((request, mode, project)) => { + let unloads = prepared.as_ref().is_some_and(|candidate| { + let id = match &request { + protocol::ClientRequest::ThreadArchive { params, .. } => ¶ms.thread_id, + protocol::ClientRequest::ThreadUnsubscribe { params, .. } => { + ¶ms.thread_id + } + _ => return false, + }; + let loaded = requests.loaded.as_deref(); + candidate.settings.thread.id == *id + || loaded.is_some_and(|loaded| loaded.thread.id == *id) + }); let is_thread_list = matches!(&request, protocol::ClientRequest::ThreadList { .. }); // When closure and the handshake are both ready, cancel before // the handshake can replace the active attachment. - tokio::select! { + let result = tokio::select! { biased; _ = output.closed() => return, _ = command.response.closed(), if is_thread_list || is_thread_start => { @@ -352,7 +447,17 @@ async fn run_worker( continue; } result = handle_request(request, mode, project, requests, replies, output) => result, + }; + if unloads + && result.is_ok() + && let Some(root) = prepared + .take() + .and_then(|candidate| candidate.canceled_recovery) + && requests.loaded.is_some() + { + replies.report_loss(&root, output); } + result } Err(error) => Err(error), }; @@ -361,24 +466,192 @@ async fn run_worker( { replies.report_loss(&root, output); } - let result = result.map_err(|error| protocol::JSONRPCErrorError { - code: if error - .downcast_ref::() - .is_some() - { - -32601 - } else if error.is::() || error.is::() { - -32602 - } else { - -32000 - }, - message: error.to_string(), - data: None, - }); - let _ = command.response.send(Ok(result)); + let _ = command.response.send(Ok(request_result(result))); } } +fn request_result(result: Result) -> crate::RequestResult { + result.map_err(|error| protocol::JSONRPCErrorError { + code: if error + .downcast_ref::() + .is_some() + { + -32601 + } else if error.is::() || error.is::() { + -32602 + } else { + -32000 + }, + message: error.to_string(), + data: None, + }) +} + +struct PreparedResume { + settings: Box, + opened: Option, + history: mpsc::Receiver, + generation: u64, + canceled_recovery: Option>, + warning: Option, +} + +const PREPARED_EVENT_LIMIT: usize = 8; + +fn check_prepared_stream((mut stream, first): OpenReplyStream) -> Result { + let mut prefix = Vec::new(); + // The regular stream path forwards the prefix and unread tail in order. + for _ in usize::from(first.is_some())..PREPARED_EVENT_LIMIT { + let Some(event) = probe_reply_stream(&mut stream)? else { + break; + }; + prefix.push(Ok(event)); + } + Ok((futures::stream::iter(prefix).chain(stream).boxed(), first)) +} + +async fn prepare_thread_resume( + request: protocol::ClientRequest, + history: mpsc::Receiver, + mut ready: oneshot::Sender>, + requests: &mut requests::Requests, + replies: &mut Replies, + output: &mpsc::UnboundedSender, +) -> Option { + let prepared = tokio::select! { + biased; + _ = ready.closed() => return None, + prepared = async { + let response = replies.while_listening( + requests.handle(request, cloud::SubmissionMode::StartIfIdle, /*project*/ None, output), + output, + ).await?; + let requests::Response::Attach(settings) = response else { + bail!("Cloud resume did not return attachment settings."); + }; + let value = serde_json::to_value(&settings)?; + let id = &settings.thread.id; + let opened = replies.open_attachment(id, requests.client.attach(id), output).await?; + Ok((settings, value, opened)) + } => prepared, + }; + let (settings, value, opened) = match prepared { + Ok(prepared) => prepared, + Err(error) => { + let _ = ready.send(Ok(request_result(Err(error)))); + return None; + } + }; + ready.send(Ok(Ok(value))).ok()?; + Some(PreparedResume { + settings, + opened, + history, + generation: replies.generation, + canceled_recovery: None, + warning: None, + }) +} + +async fn advance_prepared_resume( + mut candidate: PreparedResume, + command: Option, + requests: &mut requests::Requests, + replies: &mut Replies, + output: &mpsc::UnboundedSender, +) -> Option { + if let Some(crate::CloudCommand { + operation, + mut response, + }) = command.filter(|command| !command.response.is_closed()) + { + let id = &candidate.settings.thread.id; + match operation { + CloudOperation::CommitPreparedResume => { + let ready = match candidate.opened { + Some(opened) => check_prepared_stream(opened).map(Some), + None if replies.admit_turn(id, output).is_ok() + && replies.generation == candidate.generation => + { + Ok(None) + } + None => Err(MissingReplyAttachment.into()), + }; + match ready { + Ok(opened) => { + // No await between accepting the handoff and replacing the state. + if response.send(Ok(Ok(Value::Null))).is_ok() { + if let Some(opened) = opened { + replies.install_replacement(id, opened, output); + let status = &candidate.settings.thread.status; + if let Some(recovery) = &mut replies.idle { + let initial_idle = *status == protocol::ThreadStatus::Idle; + recovery.observed_idle.get_or_insert(initial_idle); + } + } + requests.warned_partial_history = + candidate.warning.as_ref().map(|_| id.clone()); + requests.loaded = Some(candidate.settings); + if let Some(warning) = candidate.warning { + let _ = output.send(warning); + } + return None; + } + } + Err(error) => { + let _ = response.send(Ok(request_result(Err(error)))); + } + } + } + CloudOperation::Request(request) + if match request.as_ref() { + protocol::ClientRequest::ThreadTurnsList { params, .. } => { + params.thread_id == *id + } + protocol::ClientRequest::ThreadItemsList { params, .. } => { + params.thread_id == *id + } + _ => false, + } => + { + // Ordinary requests keep their own warning marker between scoped pages. + let previous = requests.warned_partial_history.take(); + let (warning_output, mut warnings) = mpsc::unbounded_channel(); + let page = tokio::select! { + biased; + _ = response.closed() => None, + _ = output.closed() => None, + page = replies.while_listening( + requests.handle(*request, cloud::SubmissionMode::StartIfIdle, /*project*/ None, &warning_output), + output, + ) => Some(page), + }; + requests.warned_partial_history = previous; + if let Some(page) = page { + let page = page.and_then(|response| match response { + requests::Response::Value(value) => Ok(value), + _ => Err(eyre!("Expected a Cloud history page.")), + }); + let ok = page.is_ok(); + if response.send(Ok(request_result(page))).is_ok() && ok { + candidate.warning = candidate.warning.or_else(|| warnings.try_recv().ok()); + return Some(candidate); + } + } + } + _ => { + let _ = response.send(Ok(request_result(Err(eyre!( + "Only the prepared Cloud conversation's history may be read here." + ))))); + } + } + } + if let Some(root) = candidate.canceled_recovery { + replies.report_loss(&root, output); + } + None +} + async fn handle_request( request: protocol::ClientRequest, mode: cloud::SubmissionMode, diff --git a/codex-rs/cloud-client/src/client_tests.rs b/codex-rs/cloud-client/src/client_tests.rs index c8247b1b74..fe521c78a9 100644 --- a/codex-rs/cloud-client/src/client_tests.rs +++ b/codex-rs/cloud-client/src/client_tests.rs @@ -478,10 +478,21 @@ async fn attach_preserves_opaque_events_and_native_error_details() { secret.clone(), ); let data = [frame(&[10, 3, 10, 1, b'n']), frame(&[18, 3, 10, 1, b'r'])].concat(); - let mut fixture = fixture(response( - data.iter().map(|b| Bytes::from(vec![*b])).collect(), - status, - )) + let large = vec![b'x'; 256 * 1024]; + let burst = frame( + &wire::AttachResponse { + event: Some(wire::Payload::ServerRequest(large.clone().into())), + } + .encode_to_vec(), + ); + let fragmented = data.iter().map(|b| Bytes::from(vec![*b])); + let mut fixture = fixture_responses(vec![ + response( + fragmented.chain(std::iter::repeat_n(burst, 16)).collect(), + status, + ), + response(vec![frame(&[0x0a, 0x00])], tonic::Status::ok("")), + ]) .await; let mut events = fixture.client.attach("t").await.unwrap(); assert_request(&mut fixture, "Attach", &[10, 1, b't']).await; @@ -495,6 +506,27 @@ async fn attach_preserves_opaque_events_and_native_error_details() { panic!("server request expected") }; assert_eq!(bytes.as_ref(), &[10, 1, b'r']); + // History completes with several megabytes still unread on the Attach body. + let request = ListItemsRequest { + thread_id: "t".into(), + turn_id: None, + page: HistoryQuery::default(), + }; + let history = fixture.client.list_items(&request); + assert!( + tokio::time::timeout(Duration::from_secs(5), history) + .await + .unwrap() + .unwrap() + .data + .is_empty() + ); + for _ in 0..16 { + let Event::ServerRequest(bytes) = events.next().await.unwrap().unwrap() else { + panic!("burst event expected") + }; + assert_eq!(bytes.as_ref(), large.as_slice()); + } let error = events.next().await.unwrap().unwrap_err(); assert!(!format!("{error:?} {error}").contains("private")); let Error::Rpc(status) = error else { diff --git a/codex-rs/tui/src/app/agent_center_tests.rs b/codex-rs/tui/src/app/agent_center_tests.rs index 1f72f2540d..a2fe2e9c1c 100644 --- a/codex-rs/tui/src/app/agent_center_tests.rs +++ b/codex-rs/tui/src/app/agent_center_tests.rs @@ -914,13 +914,26 @@ async fn live_center_show_more_navigation_search_and_loading() { view.handle_paste("Older task".into()); let id = assert_matches!(rx.try_recv(), Ok(AppEvent::SearchAgentsOverview(id, _)) => id); insta::assert_snapshot!(screen(&view, /*width*/ 80, /*height*/ 12).lines().nth(6).unwrap(), @" › Searching history…"); - app.apply_agents_overview_search_page(id, Vec::new(), Default::default(), /*ok*/ false); + let backend_id = crate::backend_id::BackendId::new(); + app.apply_agents_overview_search_page( + id, + backend_id, + Vec::new(), + Default::default(), + /*ok*/ false, + ); insta::assert_snapshot!( "live_center_show_more_search_retry", screen(&view, /*width*/ 80, /*height*/ 12) ); view.handle_key_event(KeyCode::Enter.into()); let id = assert_matches!(rx.try_recv(), Ok(AppEvent::SearchAgentsOverview(id, _)) => id); - app.apply_agents_overview_search_page(id, Vec::new(), Default::default(), /*ok*/ true); + app.apply_agents_overview_search_page( + id, + backend_id, + Vec::new(), + Default::default(), + /*ok*/ true, + ); insta::assert_snapshot!(screen(&view, /*width*/ 80, /*height*/ 12).lines().nth(6).unwrap(), @" › Show more matches"); } diff --git a/codex-rs/tui/src/app/agents_overview.rs b/codex-rs/tui/src/app/agents_overview.rs index 974f1a2ad9..f1f9e8b90c 100644 --- a/codex-rs/tui/src/app/agents_overview.rs +++ b/codex-rs/tui/src/app/agents_overview.rs @@ -524,10 +524,23 @@ impl App { app_server: &mut AppServerSession, thread_id: ThreadId, ) -> Result { + // Selecting an already visible chat can just dismiss the overview. + let viewing_task = self.current_displayed_thread_id() == Some(thread_id) + && (!self.thread_unavailable(thread_id) || self.chat_widget.is_external_writer_view()); if let Some(owner) = self.backend_registry.owner(thread_id) - && owner != app_server.backend_id() + && (owner != app_server.backend_id() + || (app_server.thread_params_mode() + == crate::app_server_session::ThreadParamsMode::Cloud + && !self.chat_widget.fork_in_progress + && !viewing_task)) { - let Some(Some(thread)) = self.agents_overview.threads.get(&thread_id) else { + let Some(thread) = self + .agents_overview + .threads + .get(&thread_id) + .and_then(Option::as_ref) + .or_else(|| self.agents_overview.search_matches.get(&thread_id)) + else { self.add_agents_overview_error(format!( "Agent session {thread_id} is unavailable. Refresh Command Center and try again." )); diff --git a/codex-rs/tui/src/app/agents_overview_discovery.rs b/codex-rs/tui/src/app/agents_overview_discovery.rs index 37404300c9..98e8838bbc 100644 --- a/codex-rs/tui/src/app/agents_overview_discovery.rs +++ b/codex-rs/tui/src/app/agents_overview_discovery.rs @@ -3,6 +3,8 @@ use super::App; use crate::app_event::AppEvent; +use crate::app_server_session::AppServerSession; +use crate::backend_id::BackendId; use codex_app_server_client::AppServerRequestHandle; use codex_app_server_client::TypedRequestError; use codex_app_server_protocol::ClientRequest; @@ -277,10 +279,12 @@ pub(super) fn matches_query(thread: &Thread, query: &str) -> bool { impl App { pub(super) fn search_agents_overview( &mut self, - handle: AppServerRequestHandle, + app_server: &AppServerSession, id: Uuid, limit: usize, ) { + let backend_id = app_server.backend_id(); + let handle = app_server.request_handle(); let shared = &self.agents_overview.view_state; let mut state = shared.lock().unwrap_or_else(PoisonError::into_inner); let query = state.search.to_lowercase(); @@ -312,6 +316,7 @@ impl App { } tx.send(AppEvent::AgentsOverviewSearchPage { id, + backend_id, rows, cursor, ok, @@ -324,6 +329,7 @@ impl App { pub(super) fn apply_agents_overview_search_page( &mut self, id: Uuid, + backend_id: BackendId, rows: Vec<(Thread, bool)>, cursor: AgentsOverviewDiscovery, ok: bool, @@ -336,15 +342,14 @@ impl App { let Some(search) = state.history.as_mut().filter(|s| s.id == id) else { return; }; - let backend_id = self.app_server_backend_id; for (thread, unread) in rows { if let Ok(id) = codex_protocol::ThreadId::from_string(&thread.id) - && self.backend_registry.owner(id).or(backend_id) == backend_id && !overview.threads.get(&id).is_some_and(Option::is_some) && !overview.removed_threads.contains(&id) && !overview.hidden_threads.contains(&id) && let std::collections::hash_map::Entry::Vacant(entry) = overview.search_matches.entry(id) + && self.backend_registry.observe(backend_id, id) { entry.insert(thread); if unread { diff --git a/codex-rs/tui/src/app/agents_overview_tests.rs b/codex-rs/tui/src/app/agents_overview_tests.rs index c7f2ff1b39..1a7ac065a1 100644 --- a/codex-rs/tui/src/app/agents_overview_tests.rs +++ b/codex-rs/tui/src/app/agents_overview_tests.rs @@ -219,6 +219,80 @@ async fn refreshing_one_backends_pins_preserves_cross_backend_rank() { } } +#[tokio::test] +async fn overview_opens_a_search_only_task_on_the_search_backend() -> Result<()> { + let mut app = make_test_app().await; + let (client, _commands, _events) = CloudClient::channel(/*capacity*/ 1); + let mut displayed = AppServerSession::new( + BackendId::new(), + AppServerClient::Cloud(client), + ThreadParamsMode::Cloud, + ); + let source = BackendId::new(); + let [hit, collision] = std::array::from_fn(|_| ThreadId::new()); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + app.app_event_tx = AppEventSender::new(tx); + app.app_server_backend_id = Some(displayed.backend_id()); + assert!( + app.backend_registry + .observe(displayed.backend_id(), collision) + ); + + let mut view = app.agents_overview_view(Vec::new(), /*selected_thread_id*/ None); + view.handle_key_event(KeyCode::Char('/').into()); + view.handle_paste("Older task".into()); + let AppEvent::SearchAgentsOverview(search_id, _) = rx.try_recv()? else { + panic!("expected an older-history search"); + }; + app.agents_overview.search_id = search_id; + let row = |id| { + overview_thread( + id, + /*parent_thread_id*/ None, + "Older task", + ThreadStatus::NotLoaded, + ) + }; + let found = row(hit); + app.apply_agents_overview_search_page( + search_id, + source, + vec![(found.clone(), false), (row(collision), false)], + Default::default(), + /*ok*/ true, + ); + assert_eq!( + app.backend_registry.owner(collision), + Some(displayed.backend_id()) + ); + assert!(!app.agents_overview.search_matches.contains_key(&collision)); + assert!(!app.agents_overview.threads.contains_key(&hit)); + + let mut tui = crate::tui::test_support::make_test_tui()?; + app.select_agents_overview_thread(&mut tui, &mut displayed, hit) + .await?; + let AppEvent::ResumeBackendSession(target) = rx.try_recv()? else { + panic!("expected the search-only task to resume"); + }; + assert_eq!( + ( + target.backend_id, + target.thread_id, + target.path, + target.cwd, + target.history_mode + ), + ( + source, + hit, + found.path, + Some(found.cwd.to_path_buf()), + Some(found.history_mode) + ) + ); + Ok(()) +} + #[tokio::test] async fn overview_keeps_rows_and_routes_actions_to_the_owner() -> Result<()> { let mut app = make_test_app().await; @@ -227,7 +301,7 @@ async fn overview_keeps_rows_and_routes_actions_to_the_owner() -> Result<()> { let session = AppServerSession::new(owner, AppServerClient::Cloud(client), mode); (session, (commands, events)) }; - let (mut displayed, _displayed_transport) = stub(BackendId::new(), ThreadParamsMode::Cloud); + let (mut displayed, (mut commands, _events)) = stub(BackendId::new(), ThreadParamsMode::Cloud); let remote_owner = BackendId::new(); let [local_id, remote_id, colliding_id] = std::array::from_fn(|_| ThreadId::new()); let [local, remote, colliding] = [ @@ -313,17 +387,24 @@ async fn overview_keeps_rows_and_routes_actions_to_the_owner() -> Result<()> { Some(&remote) ); + app.active_thread_id = Some(ThreadId::new()); let mut tui = crate::tui::test_support::make_test_tui()?; - assert!(matches!( - app.select_agents_overview_thread(&mut tui, &mut displayed, remote_id) - .await?, - AppRunControl::Continue - )); - assert!(matches!( - rx.try_recv()?, - AppEvent::ResumeBackendSession(target) - if target.backend_id == remote_owner && target.thread_id == remote_id - )); + for (owner, thread_id) in [ + (remote_owner, remote_id), + (displayed.backend_id(), local_id), + ] { + let result = tokio::select! { + result = app.select_agents_overview_thread(&mut tui, &mut displayed, thread_id) => result?, + _ = commands.recv() => panic!("selecting a listed task must resume before reading it"), + }; + assert!(matches!(result, AppRunControl::Continue)); + assert!(matches!( + rx.try_recv()?, + AppEvent::ResumeBackendSession(target) + if target.backend_id == owner && target.thread_id == thread_id + )); + } + assert!(commands.try_recv().is_err()); assert!(app.agents_overview.threads.contains_key(&local_id)); let mut view = app.agents_overview_view(vec![remote], Some(remote_id)); assert!(!view.pin_action_supported); diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 6b21f3168b..61de38b6f4 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -3282,9 +3282,11 @@ impl App { } } AppEvent::OpenAgentsOverview => self.open_agents_overview(app_server), - AppEvent::SearchAgentsOverview(id, limit) => self.search_agents_overview(app_server.request_handle(), id, limit), - AppEvent::AgentsOverviewSearchPage { id, rows, cursor, ok } => { - self.apply_agents_overview_search_page(id, rows, cursor, ok); + AppEvent::SearchAgentsOverview(id, limit) => { + self.search_agents_overview(app_server, id, limit); + } + AppEvent::AgentsOverviewSearchPage { id, backend_id, rows, cursor, ok } => { + self.apply_agents_overview_search_page(id, backend_id, rows, cursor, ok); } AppEvent::ShowMoreAgentsOverview => self.show_more_agents_overview(app_server), AppEvent::RefreshViewedThread { backend_id, thread_id, revision } => { diff --git a/codex-rs/tui/src/app/resume_preparation.rs b/codex-rs/tui/src/app/resume_preparation.rs index 2735654b35..d54d57fc0a 100644 --- a/codex-rs/tui/src/app/resume_preparation.rs +++ b/codex-rs/tui/src/app/resume_preparation.rs @@ -151,7 +151,7 @@ impl App { Some(thread) => Ok(thread), None => { app_server - .resume_thread_with_permission_overrides( + .resume_thread_for_navigation( &local_settings, resume_config.clone(), thread_id, diff --git a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs index 90a4f36ddb..461c1821f0 100644 --- a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs +++ b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs @@ -327,7 +327,7 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { use codex_protocol::models::MessagePhase; let (mut app, mut events, mut ops) = make_test_app_with_channels().await; - let [original, other] = [ThreadId::new(), ThreadId::new()]; + let [original, other, failed] = std::array::from_fn(|_| ThreadId::new()); app.app_server_target = AppServerTarget::Cloud; app.enqueue_primary_thread_session( test_thread_session(original, app.config.cwd.to_path_buf()), @@ -421,10 +421,30 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { crate::ThreadParamsMode::Cloud, ); let backend_id = server.backend_id(); + let mut failed_row = other_thread.clone(); + failed_row.id = failed.to_string(); + failed_row.session_id = failed_row.id.clone(); + failed_row.history_mode = ThreadHistoryMode::Paginated; + let unavailable = failed_row.clone(); let serve = tokio::spawn(async move { let mut resumed_threads = Vec::new(); while let Some(command) = commands.recv().await { - let reply = match command.operation { + let (operation, scoped_history) = match command.operation { + CloudOperation::PrepareThreadResume { + request_id, + params, + history, + } => ( + CloudOperation::Request(Box::new(ClientRequest::ThreadResume { + request_id, + params: *params, + })), + Some(history), + ), + operation => (operation, None), + }; + let mut history_fails = false; + let reply = match operation { CloudOperation::Request(request) => match *request { ClientRequest::ThreadRead { params, .. } => { let mut thread = if params.thread_id == original.to_string() { @@ -438,7 +458,10 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { Ok(serde_json::json!({"thread": thread})) } ClientRequest::ThreadResume { params, .. } => { - let mut thread = if params.thread_id == original.to_string() { + history_fails = params.thread_id == failed.to_string(); + let mut thread = if history_fails { + unavailable.clone() + } else if params.thread_id == original.to_string() { original_thread.clone() } else { other_thread.clone() @@ -452,7 +475,9 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { vec![agent_item("missed-item", "Missed response", None)], )); } - resumed_threads.push(params.thread_id); + if !history_fails { + resumed_threads.push(params.thread_id); + } Ok(serde_json::json!({ "thread": thread, "model": "gpt-5.5", @@ -490,6 +515,34 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { _ => panic!("unexpected Cloud operation while resuming"), }; command.response.send(Ok(reply)).unwrap(); + if let Some(mut scoped_history) = scoped_history { + let commit = scoped_history.recv().await.unwrap(); + if history_fails { + assert!(matches!(commit.operation, CloudOperation::Request(request) + if matches!(request.as_ref(), ClientRequest::ThreadTurnsList { params, .. } if params.thread_id == failed.to_string()))); + commit + .response + .send(Ok(Err(JSONRPCErrorError { + code: -32000, + message: "history unavailable".into(), + data: None, + }))) + .unwrap(); + assert!( + scoped_history.recv().await.is_none(), + "failed B must not commit" + ); + continue; + } + assert!(matches!( + commit.operation, + CloudOperation::CommitPreparedResume + )); + commit + .response + .send(Ok(Ok(serde_json::Value::Null))) + .unwrap(); + } } resumed_threads }); @@ -633,7 +686,34 @@ async fn cloud_resume_picker_reconciles_admitted_steers() -> Result<()> { assert!( !std::iter::from_fn(|| ops.try_recv().ok()).any(|op| matches!(op, Op::UserTurn { .. })) ); - server.shutdown().await?; + // Use the regular Command Center Enter action to open a task whose history fails. + app.app_server_backend_id = Some(backend_id); + app.chat_widget + .apply_external_edit("Keep this draft".into()); + app.backend_registry.observe(backend_id, failed); + app.agents_overview + .threads + .insert(failed, Some(failed_row.clone())); + let view = app.agents_overview_view(vec![failed_row], Some(failed)); + app.agents_overview.visible_thread_ids = view.thread_ids(); + app.chat_widget.show_bottom_pane_view(Box::new(view)); + while events.try_recv().is_ok() {} + app.chat_widget.handle_key_event(KeyCode::Enter.into()); + Box::pin(app.handle_event(&mut tui, &mut server, events.try_recv()?)).await?; + let AppEvent::ResumeBackendSession(target) = events.try_recv()? else { + color_eyre::eyre::bail!("Enter must use the owner-routed resume path"); + }; + assert_eq!((target.backend_id, target.thread_id), (backend_id, failed)); + let mut backends = crate::app::backend_sessions::BackendSessions::from_initial(server, vec![]); + app.resume_backend_session(&mut tui, &mut backends, target) + .await?; + assert_eq!(backends.displayed_owner(), backend_id); + assert_eq!(app.chat_widget.thread_id(), Some(original)); + assert_eq!( + app.chat_widget.composer_text_with_pending(), + "Keep this draft" + ); + backends.shutdown().await?; assert_eq!( serve.await?, [ diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index fc69731656..a2e020b9a9 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -376,6 +376,7 @@ pub(crate) enum AppEvent { SearchAgentsOverview(Uuid, usize), AgentsOverviewSearchPage { id: Uuid, + backend_id: crate::backend_id::BackendId, rows: Vec<(Thread, bool)>, cursor: crate::app::agents_overview_discovery::AgentsOverviewDiscovery, ok: bool, diff --git a/codex-rs/tui/src/app_server_session/rollout_history.rs b/codex-rs/tui/src/app_server_session/rollout_history.rs index 381be916c8..6921853093 100644 --- a/codex-rs/tui/src/app_server_session/rollout_history.rs +++ b/codex-rs/tui/src/app_server_session/rollout_history.rs @@ -5,6 +5,7 @@ use super::AppServerStartedThread; use super::HistoryHydrationScope; use super::ResumeModelSettings; use super::ThreadHistorySupport; +use super::ThreadParamsMode; use super::ThreadSettingsOwner; use super::bootstrap_request_error; use super::is_history_pagination_unsupported; @@ -12,11 +13,13 @@ use super::started_thread_from_resume_response; use super::thread_resume_params_from_config; use super::thread_session_state_from_thread_response; use crate::legacy_core::config::Config; +use codex_app_server_client::AppServerRequestHandle; use codex_app_server_client::TypedRequestError; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::SortDirection; use codex_app_server_protocol::ThreadHistoryMode; +use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadTurnsListParams; use codex_app_server_protocol::ThreadTurnsListResponse; @@ -150,6 +153,67 @@ impl AppServerSession { .await } + /// Prepare a saved Cloud task without displacing the displayed attachment on failure. + pub(crate) async fn resume_thread_for_navigation( + &mut self, + local_settings: &crate::local_settings::LocalSettings, + config: Config, + thread_id: ThreadId, + model_settings: ResumeModelSettings, + permission_overrides: crate::resume_permissions::ResumePermissions, + ) -> Result { + let (ThreadParamsMode::Cloud, AppServerRequestHandle::Cloud(handle)) = + (self.thread_params_mode(), self.request_handle()) + else { + return Box::pin(self.resume_thread_with_permission_overrides( + local_settings, + config, + thread_id, + model_settings, + permission_overrides, + )) + .await; + }; + let (mut response, scoped) = handle + .prepare_thread_resume( + self.next_request_id(), + ThreadResumeParams { + thread_id: thread_id.to_string(), + exclude_turns: true, + ..ThreadResumeParams::default() + }, + ) + .await + .map_err(|err| { + bootstrap_request_error("thread/resume failed during TUI bootstrap", err) + })?; + let mut history = self.request_view(); + history.request_handle = AppServerRequestHandle::Cloud(scoped.clone()); + history + .hydrate_initial_thread_history( + &mut response.thread, + response.turns_backwards_cursor.clone(), + response.items_backwards_cursor.clone(), + Some(&config), + Some(local_settings), + HistoryHydrationScope::Initial, + ) + .await?; + let mut started = + started_thread_from_resume_response(response, local_settings, ThreadParamsMode::Cloud) + .await?; + scoped.commit_prepared_resume().await.map_err(|err| { + bootstrap_request_error("thread/resume failed during TUI bootstrap", err) + })?; + self.unread_revisions.remove(&thread_id); + self.history_pagination.extend(history.history_pagination); + if self.task_tools_available(thread_id) { + self.remember_task_tool_thread(thread_id); + started.task_tools_available = true; + } + Ok(started) + } + pub(crate) async fn resume_thread_with_permission_overrides( &mut self, local_settings: &crate::local_settings::LocalSettings,