From eab2603712f90ec08ac44207a04eed158be293a3 Mon Sep 17 00:00:00 2001 From: Benjamin Carlsson Date: Fri, 30 Oct 2026 07:26:33 +0000 Subject: [PATCH] Reuse connected backends in the in-app resume picker (#58280) ## What changed - List sessions from listening backends and the displayed session's backend using their existing connections. - Give picker views distinct request ID namespaces while keeping connection ownership and event handling with the live sessions. - Start multi-backend listings without a working-directory filter, since local and remote workspaces cannot share a single path filter. ## Testing Add coverage for distinct picker request IDs, preserving the live event stream after picker shutdown, listing all providers after a provider switch while retaining the remote working-directory filter, and clearing the pending cursor when a source fails late. GitOrigin-RevId: a0b8780e1f86dc187f088080d9213afe8e18f77f --- codex-rs/tui/src/app.rs | 4 + codex-rs/tui/src/app/backend_sessions.rs | 8 + codex-rs/tui/src/app/session_picker.rs | 59 ++--- codex-rs/tui/src/app/startup.rs | 8 +- .../app/tests/session_lifecycle_requests.rs | 83 +++---- codex-rs/tui/src/app_server_session.rs | 182 ++++++++++++---- .../external_agent_config.rs | 4 +- .../tui/src/app_server_session/history.rs | 6 +- .../tui/src/app_server_session/realtime.rs | 6 +- .../src/app_server_session/rollout_history.rs | 6 +- .../tui/src/app_server_session/thread_list.rs | 2 +- codex-rs/tui/src/lib.rs | 1 + codex-rs/tui/src/resume_picker.rs | 203 ++++++++++++++---- .../tui/src/resume_picker/sources_tests.rs | 45 ++++ 14 files changed, 418 insertions(+), 199 deletions(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 2d31721bd2..c442f48793 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -549,6 +549,10 @@ pub(crate) struct App { launch_cwd: PathBuf, /// Resume anchor selected by `/cd`; ordinary resumes retain the immutable launch cwd. runtime_working_directory_override: Option, + #[cfg_attr( + not(test), + expect(dead_code, reason = "retain the state runtime for the app lifetime") + )] pub(crate) state_db: Option, cli_kv_overrides: Vec<(String, TomlValue)>, harness_overrides: ConfigOverrides, diff --git a/codex-rs/tui/src/app/backend_sessions.rs b/codex-rs/tui/src/app/backend_sessions.rs index bc5ba758e0..4cfecb12c9 100644 --- a/codex-rs/tui/src/app/backend_sessions.rs +++ b/codex-rs/tui/src/app/backend_sessions.rs @@ -94,6 +94,14 @@ impl BackendSessions { &self.entries[slot.0].session } + pub(super) fn picker_views(&self, displayed_owner: BackendId) -> Vec { + self.entries + .iter() + .filter(|entry| entry.listening || entry.session.backend_id() == displayed_owner) + .map(|entry| entry.session.picker_view()) + .collect() + } + pub(super) fn session_at_mut(&mut self, slot: SessionSlot) -> &mut AppServerSession { &mut self.entries[slot.0].session } diff --git a/codex-rs/tui/src/app/session_picker.rs b/codex-rs/tui/src/app/session_picker.rs index 6230c5e831..be8edb7089 100644 --- a/codex-rs/tui/src/app/session_picker.rs +++ b/codex-rs/tui/src/app/session_picker.rs @@ -4,50 +4,15 @@ use super::*; use crate::chatwidget::ThreadInputStateRestoreMode; impl App { - pub(super) async fn open_resume_picker( + pub(super) async fn open_resume_picker_for_backends( &mut self, tui: &mut tui::Tui, - app_server: &mut AppServerSession, + backends: &mut super::backend_sessions::BackendSessions, + displayed_session: super::backend_sessions::SessionSlot, ) -> Result { - // Keep embedded-server initialization on its own runtime task as well as opening the - // picker on a fresh event-loop iteration; its auth manager needs more stack headroom. - let backend_id = app_server.backend_id(); - let picker_config = self.config.clone(); - let picker_target = self.app_server_target.clone(); - let picker_cli_kv_overrides = self.cli_kv_overrides.clone(); - let picker_loader_overrides = self.loader_overrides.clone(); - let picker_state_db = self.state_db.clone(); - let picker_environment_manager = Arc::clone(&self.environment_manager); - let picker_app_server = tokio::spawn(async move { - crate::start_app_server_for_picker( - backend_id, - &picker_config, - &picker_target, - picker_cli_kv_overrides, - picker_loader_overrides, - picker_state_db, - picker_environment_manager, - ) - .await - }) - .await; - let picker_app_server = match picker_app_server { - Ok(result) => result, - Err(err) => Err(err.into()), - }; - let mut picker_app_server = match picker_app_server { - Ok(app_server) => app_server, - Err(err) => { - self.add_session_picker_error(format!("Failed to start TUI session picker: {err}")); - self.chat_widget.maybe_send_next_queued_input(); - return Ok(AppRunControl::Continue); - } - }; - picker_app_server = picker_app_server - .with_remote_cwd_override(app_server.remote_cwd_override().map(Path::to_path_buf)); - picker_app_server.model_provider_override = self.harness_overrides.model_provider.clone(); + let displayed_owner = backends.session_at(displayed_session).backend_id(); let selection = - crate::resume_picker::run_resume_picker_from_existing_session_with_app_server( + crate::resume_picker::run_resume_picker_from_existing_session_with_app_servers( crate::uses_remote_workspace_or_environment( &self.app_server_target, &self.environment_manager, @@ -57,16 +22,22 @@ impl App { &self.local_settings, /*show_all*/ false, /*include_non_interactive*/ false, - picker_app_server, - app_server.request_handle(), + crate::resume_picker::ResumePickerAppServers::new( + backends.picker_views(displayed_owner), + displayed_owner, + ), self.primary_thread_id .or(self.current_displayed_thread_id()), ) .await; match selection { Ok(selection) => { - self.apply_resume_picker_selection(tui, app_server, selection) - .await + self.apply_resume_picker_selection( + tui, + backends.session_at_mut(displayed_session), + selection, + ) + .await } Err(err) => { self.add_session_picker_error(format!("Failed to open session picker: {err}")); diff --git a/codex-rs/tui/src/app/startup.rs b/codex-rs/tui/src/app/startup.rs index c47633c83c..7afd8e4791 100644 --- a/codex-rs/tui/src/app/startup.rs +++ b/codex-rs/tui/src/app/startup.rs @@ -1311,9 +1311,11 @@ See the Codex keymap documentation for supported actions and examples." } if app.pending_open_resume_picker { app.pending_open_resume_picker = false; - match Box::pin( - app.open_resume_picker(tui, backends.session_at_mut(displayed_session)), - ) + match Box::pin(app.open_resume_picker_for_backends( + tui, + &mut backends, + displayed_session, + )) .await { Ok(AppRunControl::Continue) => {} 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 ed82fc15a5..0d6cd632c8 100644 --- a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs +++ b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs @@ -946,85 +946,50 @@ async fn make_history_test_app() -> Result<(Box, tempfile::TempDir)> { } #[tokio::test] -async fn in_app_remote_resume_picker_preserves_remote_cwd_filter() -> Result<()> { +async fn in_app_remote_resume_picker_lists_all_providers_after_provider_switch() -> Result<()> { let (mut app, _codex_home) = make_history_test_app().await?; let remote_root = tempdir()?; let remote_cwd = remote_root.path().join("project"); std::fs::create_dir(&remote_cwd)?; - let (app_server, _requests, active_proxy) = - start_recording_remote_app_server(&app.config).await?; + let (app_server, requests, proxy) = start_recording_remote_app_server(&app.config).await?; + let backend_id = app_server.backend_id(); let mut app_server = app_server.with_remote_cwd_override(Some(remote_cwd.clone())); - app.harness_overrides.model_provider = Some(app.config.model_provider_id.clone()); - - let listener = TcpListener::bind("127.0.0.1:0").await?; - app.app_server_target = AppServerTarget::Remote { - endpoint: crate::resolve_remote_addr(&format!("ws://{}", listener.local_addr()?))?, - }; - let (request_tx, request_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel::<()>(); - let codex_home = app.config.codex_home.display().to_string(); - let recorder = tokio::spawn(async move { - let (socket, _) = listener.accept().await?; - let mut websocket = accept_async(socket).await?; - while let Some(frame) = websocket.next().await { - let Message::Text(text) = frame? else { - continue; - }; - match serde_json::from_str::(&text)? { - JSONRPCMessage::Request(request) if request.method == "initialize" => { - let response = JSONRPCMessage::Response(JSONRPCResponse { - id: request.id, - result: serde_json::json!({ - "userAgent": "codex-tui-test", - "codexHome": codex_home, - }), - }); - websocket - .send(Message::Text(serde_json::to_string(&response)?.into())) - .await?; - } - JSONRPCMessage::Request(request) if request.method == "thread/list" => { - let _ = request_tx.send(request.params.unwrap_or_default()); - let _ = release_rx.await; - return Ok::<(), color_eyre::Report>(()); - } - _ => {} - } - } - Err(color_eyre::eyre::eyre!( - "picker disconnected before requesting its thread list" - )) - }); + // Provider A's conversations must remain visible after switching to provider B. + app_server.model_provider_override = Some("provider-b".to_string()); + let mut backends = crate::app::backend_sessions::BackendSessions::new(app_server); + let displayed = backends.slot(backend_id).expect("registered backend"); let mut tui = crate::tui::test_support::make_test_tui()?; tui.set_alt_screen_enabled(/*enabled*/ false); tui.pause_events(); - let mut picker = Box::pin(app.open_resume_picker(&mut tui, &mut app_server)); + let mut picker = + Box::pin(app.open_resume_picker_for_backends(&mut tui, &mut backends, displayed)); let params = tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 5), async { - tokio::select! { - biased; - params = request_rx => params.map_err(|_| { - color_eyre::eyre::eyre!("picker request recorder closed") - }), - result = &mut picker => { - result?; - Err(color_eyre::eyre::eyre!( - "picker exited before requesting its thread list" - )) + loop { + if let Some(params) = recorded_params(&requests, "thread/list").pop() { + break Ok::<_, color_eyre::Report>(params); + } + tokio::select! { + result = &mut picker => { + result?; + break Err(color_eyre::eyre::eyre!( + "picker exited before requesting its thread list" + )); + } + _ = tokio::task::yield_now() => {} } } }) .await??; drop(picker); - let _ = release_tx.send(()); - recorder.await??; - app_server.shutdown().await?; - active_proxy.await??; + backends.shutdown().await?; + proxy.await??; assert_eq!( params["cwd"], serde_json::Value::String(remote_cwd.display().to_string()) ); + assert_eq!(params["modelProviders"], serde_json::json!([])); Ok(()) } diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 104bbaec09..ae79c9ad3d 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -327,7 +327,11 @@ pub(crate) struct AppServerBootstrap { pub(crate) struct AppServerSession { backend_id: crate::backend_id::BackendId, - client: AppServerClient, + /// Owns events and the connection lifetime for live sessions. Picker views borrow only the + /// request handle below and deliberately have no connection owner. + connection_owner: Option, + request_handle: AppServerRequestHandle, + request_id_namespace: Option, next_request_id: i64, history_pagination: HashMap, pub(crate) focused_conversation: Option, @@ -473,21 +477,37 @@ pub(crate) struct UnsupportedLegacyPermissionProfile; impl AppServerSession { /// Platform of the app-server process, not necessarily its executor. pub(crate) fn app_server_platform_family(&self) -> Option<&str> { - self.client.platform_family() + self.connection_owner + .as_ref() + .and_then(AppServerClient::platform_family) } pub(crate) fn app_server_platform_os(&self) -> Option<&str> { - self.client.platform_os() + self.connection_owner + .as_ref() + .and_then(AppServerClient::platform_os) } pub(crate) fn new( backend_id: crate::backend_id::BackendId, client: AppServerClient, thread_params_mode: ThreadParamsMode, + ) -> Self { + let request_handle = client.request_handle(); + Self::from_parts(backend_id, Some(client), request_handle, thread_params_mode) + } + + fn from_parts( + backend_id: crate::backend_id::BackendId, + connection_owner: Option, + request_handle: AppServerRequestHandle, + thread_params_mode: ThreadParamsMode, ) -> Self { Self { backend_id, - client, + connection_owner, + request_handle, + request_id_namespace: None, next_request_id: 1, history_pagination: HashMap::new(), focused_conversation: None, @@ -516,6 +536,24 @@ impl AppServerSession { self.backend_id } + /// Creates a request-only picker view without transferring the live connection's events. + pub(crate) fn picker_view(&self) -> Self { + let mut view = Self::from_parts( + self.backend_id, + /*connection_owner*/ None, + self.request_handle(), + self.thread_params_mode, + ); + view.request_id_namespace = Some(Uuid::new_v4()); + view.remote_cwd_override = self.remote_cwd_override.clone(); + view.model_provider_override = self.model_provider_override.clone(); + view.history_support = self.history_support; + view.thread_settings_update_supported = self.thread_settings_update_supported; + view.default_model = self.default_model.clone(); + view.available_models = self.available_models.clone(); + view + } + pub(crate) async fn start_dynamic_tool_mcp( &mut self, config: Config, @@ -658,7 +696,7 @@ impl AppServerSession { } pub(crate) fn uses_embedded_app_server(&self) -> bool { - matches!(&self.client, AppServerClient::InProcess(_)) + matches!(&self.connection_owner, Some(AppServerClient::InProcess(_))) } /// Carry capabilities that may exist only in memory when the optional cache is unwritable. @@ -692,18 +730,20 @@ impl AppServerSession { &self, local_codex_home: &AbsolutePathBuf, ) -> Option { - self.client.codex_home(local_codex_home) + self.connection_owner + .as_ref() + .and_then(|client| client.codex_home(local_codex_home)) } pub(crate) fn server_version(&self) -> Option<&str> { - let AppServerClient::Remote(client) = &self.client else { + let Some(AppServerClient::Remote(client)) = &self.connection_owner else { return None; }; client.server_version() } pub(crate) fn server_codex_home(&self) -> Option<&str> { - let AppServerClient::Remote(client) = &self.client else { + let Some(AppServerClient::Remote(client)) = &self.connection_owner else { return None; }; client.codex_home() @@ -755,7 +795,7 @@ impl AppServerSession { let requirements_request_id = self.next_request_id(); let (models, requirements, collaboration_modes) = tokio::try_join!( async { - self.client + self.request_handle .request_typed::(ClientRequest::ModelList { request_id: model_request_id, params: ModelListParams { @@ -771,7 +811,7 @@ impl AppServerSession { }, async { match self - .client + .request_handle .request_typed::( ClientRequest::ConfigRequirementsRead { request_id: requirements_request_id, @@ -928,7 +968,7 @@ impl AppServerSession { /// (to check auth mode without the overhead of a full bootstrap). pub(crate) async fn read_account(&mut self) -> Result { let account_request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::GetAccount { request_id: account_request_id, params: GetAccountParams { @@ -940,7 +980,7 @@ impl AppServerSession { } pub(crate) async fn next_event(&mut self) -> Option { - let event = self.client.next_event().await?; + let event = self.connection_owner.as_mut()?.next_event().await?; if let codex_app_server_client::AppServerEvent::ServerNotification(notification) = &event && let codex_app_server_protocol::ServerNotification::ThreadReadStateChanged(change) = notification.as_ref() @@ -1219,7 +1259,7 @@ impl AppServerSession { self.thread_tool_transport() .configure_mcp(&mut params.config); let response: ThreadForkResponse = match self - .client + .request_handle .request_typed(ClientRequest::ThreadFork { request_id, params: params.clone(), @@ -1233,7 +1273,7 @@ impl AppServerSession { self.history_support = ThreadHistorySupport::LegacyOnly; params.exclude_turns = false; let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadFork { request_id, params }) .await .map_err(|err| { @@ -1362,7 +1402,7 @@ impl AppServerSession { params: ThreadLoadedListParams, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadLoadedList { request_id, params }) .await .wrap_err("failed to list loaded threads from app server") @@ -1375,7 +1415,7 @@ impl AppServerSession { ) -> Result { let request_id = self.next_request_id(); let response = self - .client + .request_handle .request_typed::(ClientRequest::ThreadRead { request_id, params: ThreadReadParams { @@ -1398,7 +1438,7 @@ impl AppServerSession { == "paginated threads do not support thread/read(includeTurns=true)" => { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadRead { request_id, params: ThreadReadParams { @@ -1426,7 +1466,7 @@ impl AppServerSession { pub(crate) async fn thread_archive(&mut self, thread_id: ThreadId) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadArchiveResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadArchive { request_id, params: ThreadArchiveParams { @@ -1441,7 +1481,7 @@ impl AppServerSession { pub(crate) async fn thread_delete(&mut self, thread_id: ThreadId) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadDeleteResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadDelete { request_id, params: ThreadDeleteParams { @@ -1456,7 +1496,7 @@ impl AppServerSession { pub(crate) async fn thread_unarchive(&mut self, thread_id: ThreadId) -> Result { let request_id = self.next_request_id(); let response: ThreadUnarchiveResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadUnarchive { request_id, params: ThreadUnarchiveParams { @@ -1474,7 +1514,7 @@ impl AppServerSession { branch: String, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadMetadataUpdate { request_id, params: ThreadMetadataUpdateParams { @@ -1501,7 +1541,7 @@ impl AppServerSession { } let request_id = self.next_request_id(); match self - .client + .request_handle .request_typed::(ClientRequest::ThreadSettingsUpdate { request_id, params, @@ -1535,7 +1575,7 @@ impl AppServerSession { .collect::, _>>() .wrap_err("failed to encode thread/inject_items payload")?; let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadInjectItems { request_id, params: ThreadInjectItemsParams { @@ -1592,7 +1632,7 @@ impl AppServerSession { params.summary = summary; params.collaboration_mode = collaboration_mode; } - self.client + self.request_handle .request_typed(ClientRequest::TurnStart { request_id, params }) .await .wrap_err("turn/start failed in TUI") @@ -1605,7 +1645,7 @@ impl AppServerSession { ) -> std::result::Result<(), TypedRequestError> { let request_id = self.next_request_id(); let _: TurnInterruptResponse = self - .client + .request_handle .request_typed(ClientRequest::TurnInterrupt { request_id, params: TurnInterruptParams { @@ -1632,7 +1672,7 @@ impl AppServerSession { items: Vec, ) -> std::result::Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::TurnSteer { request_id, params: TurnSteerParams { @@ -1655,7 +1695,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadSetNameResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadSetName { request_id, params: ThreadSetNameParams { @@ -1675,7 +1715,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadMemoryModeSetResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadMemoryModeSet { request_id, params: ThreadMemoryModeSetParams { @@ -1691,7 +1731,7 @@ impl AppServerSession { pub(crate) async fn memory_reset(&mut self) -> Result<()> { let request_id = self.next_request_id(); let _: MemoryResetResponse = self - .client + .request_handle .request_typed(ClientRequest::MemoryReset { request_id, params: None, @@ -1706,7 +1746,7 @@ impl AppServerSession { thread_id: ThreadId, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadGoalGet { request_id, params: ThreadGoalGetParams { @@ -1725,7 +1765,7 @@ impl AppServerSession { token_budget: Option>, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadGoalSet { request_id, params: ThreadGoalSetParams { @@ -1745,7 +1785,7 @@ impl AppServerSession { thread_id: ThreadId, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadGoalClear { request_id, params: ThreadGoalClearParams { @@ -1760,7 +1800,7 @@ impl AppServerSession { pub(crate) async fn logout_account(&mut self) -> Result<()> { let request_id = self.next_request_id(); let _: LogoutAccountResponse = self - .client + .request_handle .request_typed(ClientRequest::LogoutAccount { request_id, params: None, @@ -1773,7 +1813,7 @@ impl AppServerSession { pub(crate) async fn thread_unsubscribe(&mut self, thread_id: ThreadId) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadUnsubscribeResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadUnsubscribe { request_id, params: ThreadUnsubscribeParams { @@ -1788,7 +1828,7 @@ impl AppServerSession { pub(crate) async fn thread_compact_start(&mut self, thread_id: ThreadId) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadCompactStartResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadCompactStart { request_id, params: ThreadCompactStartParams { @@ -1807,7 +1847,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadShellCommandResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadShellCommand { request_id, params: ThreadShellCommandParams { @@ -1828,7 +1868,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadApproveGuardianDeniedActionResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadApproveGuardianDeniedAction { request_id, params: ThreadApproveGuardianDeniedActionParams { @@ -1848,7 +1888,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadBackgroundTerminalsCleanResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadBackgroundTerminalsClean { request_id, params: ThreadBackgroundTerminalsCleanParams { @@ -1866,7 +1906,7 @@ impl AppServerSession { target: ReviewTarget, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ReviewStart { request_id, params: ReviewStartParams { @@ -1884,7 +1924,7 @@ impl AppServerSession { params: SkillsListParams, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::SkillsList { request_id, params }) .await .wrap_err("skills/list failed in TUI") @@ -1893,7 +1933,7 @@ impl AppServerSession { pub(crate) async fn reload_user_config(&mut self) -> Result<()> { let request_id = self.next_request_id(); let _: ConfigWriteResponse = self - .client + .request_handle .request_typed(ClientRequest::ConfigBatchWrite { request_id, params: ConfigBatchWriteParams { @@ -1913,21 +1953,30 @@ impl AppServerSession { request_id: RequestId, result: serde_json::Value, ) -> std::io::Result<()> { - self.client.resolve_server_request(request_id, result).await + let client = self.connection_owner.as_ref().ok_or_else(|| { + std::io::Error::other("request-only session cannot resolve server requests") + })?; + client.resolve_server_request(request_id, result).await } pub(crate) async fn shutdown(self) -> std::io::Result<()> { - self.client.shutdown().await + match self.connection_owner { + Some(client) => client.shutdown().await, + None => Ok(()), + } } pub(crate) fn request_handle(&self) -> AppServerRequestHandle { - self.client.request_handle() + self.request_handle.clone() } pub(crate) fn next_request_id(&mut self) -> RequestId { let request_id = self.next_request_id; self.next_request_id += 1; - RequestId::Integer(request_id) + self.request_id_namespace.map_or_else( + || RequestId::Integer(request_id), + |namespace| RequestId::String(format!("picker-{namespace}-{request_id}")), + ) } } @@ -2825,6 +2874,49 @@ mod tests { .expect("config should build") } + #[test] + fn picker_views_use_distinct_request_id_namespaces() { + let (client, _commands, _events) = + codex_app_server_client::CloudClient::channel(/*capacity*/ 1); + let mut session = AppServerSession::new( + crate::backend_id::BackendId::new(), + AppServerClient::Cloud(client), + ThreadParamsMode::Cloud, + ); + let mut first = session.picker_view(); + let mut second = session.picker_view(); + + let first_id = first.next_request_id(); + assert_ne!(session.next_request_id(), first_id); + assert_ne!(first_id, second.next_request_id()); + } + + #[tokio::test] + async fn picker_view_shutdown_preserves_live_event_stream() { + let (client, _commands, events) = + codex_app_server_client::CloudClient::channel(/*capacity*/ 1); + let backend_id = crate::backend_id::BackendId::new(); + let mut session = AppServerSession::new( + backend_id, + AppServerClient::Cloud(client), + ThreadParamsMode::Cloud, + ); + + session.picker_view().shutdown().await.unwrap(); + events + .send(codex_app_server_client::AppServerEvent::Disconnected { + message: "still live".to_string(), + }) + .unwrap(); + + let event = session.next_event().await.unwrap(); + assert_eq!(event.backend_id, backend_id); + assert!(matches!( + event.event, + codex_app_server_client::AppServerEvent::Disconnected { .. } + )); + } + #[tokio::test] async fn bootstrap_reuses_prefetched_account_without_another_account_read() -> Result<()> { let codex_home = tempfile::tempdir()?; diff --git a/codex-rs/tui/src/app_server_session/external_agent_config.rs b/codex-rs/tui/src/app_server_session/external_agent_config.rs index 424490ccda..5e2c075226 100644 --- a/codex-rs/tui/src/app_server_session/external_agent_config.rs +++ b/codex-rs/tui/src/app_server_session/external_agent_config.rs @@ -13,7 +13,7 @@ impl AppServerSession { params: ExternalAgentConfigDetectParams, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ExternalAgentConfigDetect { request_id, params }) .await .wrap_err("externalAgentConfig/detect failed during external agent import") @@ -29,7 +29,7 @@ impl AppServerSession { } let request_id = self.next_request_id(); let response: ExternalAgentConfigImportResponse = self - .client + .request_handle .request_typed(ClientRequest::ExternalAgentConfigImport { request_id, params: ExternalAgentConfigImportParams { diff --git a/codex-rs/tui/src/app_server_session/history.rs b/codex-rs/tui/src/app_server_session/history.rs index 956bb96372..8febf1c028 100644 --- a/codex-rs/tui/src/app_server_session/history.rs +++ b/codex-rs/tui/src/app_server_session/history.rs @@ -145,7 +145,7 @@ impl AppServerSession { ) -> std::result::Result { let request_id = self.next_request_id(); let response: ThreadRevertResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadRevert { request_id, params: ThreadRevertParams { @@ -232,7 +232,7 @@ impl AppServerSession { limit: u32, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadItemsList { request_id, params: thread_items_page_params(thread_id, turn_id, cursor, limit), @@ -248,7 +248,7 @@ impl AppServerSession { limit: u32, ) -> Result { let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadTurnsList { request_id, params: ThreadTurnsListParams { diff --git a/codex-rs/tui/src/app_server_session/realtime.rs b/codex-rs/tui/src/app_server_session/realtime.rs index 1ad46db83a..14f62f7a13 100644 --- a/codex-rs/tui/src/app_server_session/realtime.rs +++ b/codex-rs/tui/src/app_server_session/realtime.rs @@ -26,7 +26,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadRealtimeStartResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadRealtimeStart { request_id, params: ThreadRealtimeStartParams { @@ -61,7 +61,7 @@ impl AppServerSession { pub(crate) async fn thread_realtime_stop(&mut self, thread_id: ThreadId) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadRealtimeStopResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadRealtimeStop { request_id, params: ThreadRealtimeStopParams { @@ -80,7 +80,7 @@ impl AppServerSession { ) -> Result<()> { let request_id = self.next_request_id(); let _: ThreadRealtimeAppendSpeechResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadRealtimeAppendSpeech { request_id, params: ThreadRealtimeAppendSpeechParams { 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 07d22565fd..b072d8f907 100644 --- a/codex-rs/tui/src/app_server_session/rollout_history.rs +++ b/codex-rs/tui/src/app_server_session/rollout_history.rs @@ -58,7 +58,7 @@ impl AppServerSession { tracing::warn!("Failed to page read-only thread history; trying summary pages"); let request_id = self.next_request_id(); let page: ThreadTurnsListResponse = self - .client + .request_handle .request_typed(ClientRequest::ThreadTurnsList { request_id, params: ThreadTurnsListParams { @@ -215,7 +215,7 @@ impl AppServerSession { } let request_id = self.next_request_id(); let resume_response = self - .client + .request_handle .request_typed(ClientRequest::ThreadResume { request_id, params: params.clone(), @@ -230,7 +230,7 @@ impl AppServerSession { self.history_support = ThreadHistorySupport::LegacyOnly; params.exclude_turns = false; let request_id = self.next_request_id(); - self.client + self.request_handle .request_typed(ClientRequest::ThreadResume { request_id, params }) .await .map_err(|err| { diff --git a/codex-rs/tui/src/app_server_session/thread_list.rs b/codex-rs/tui/src/app_server_session/thread_list.rs index 17a45f0280..cc048f9d4a 100644 --- a/codex-rs/tui/src/app_server_session/thread_list.rs +++ b/codex-rs/tui/src/app_server_session/thread_list.rs @@ -11,7 +11,7 @@ impl AppServerSession { loop { let request_id = self.next_request_id(); let response = self - .client + .request_handle .request_typed(ClientRequest::ThreadList { request_id, params: params.clone(), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 50b07118cd..e1bf6223ed 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -633,6 +633,7 @@ async fn start_app_server( .map(AppServerClient::InProcess) } +#[cfg(test)] pub(crate) async fn start_app_server_for_picker( backend_id: crate::backend_id::BackendId, config: &Config, diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index b9dc24dd2a..8b8401ccb7 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -390,6 +390,71 @@ enum PageCursor { SourcesPending, } +pub(crate) struct ResumePickerAppServers { + ordered: Vec, + displayed_owner: crate::backend_id::BackendId, +} + +impl ResumePickerAppServers { + pub(crate) fn new( + ordered: Vec, + displayed_owner: crate::backend_id::BackendId, + ) -> Self { + assert!( + ordered + .iter() + .any(|session| session.backend_id() == displayed_owner) + ); + Self { + ordered, + displayed_owner, + } + } + + fn singleton(session: AppServerSession) -> Self { + let displayed_owner = session.backend_id(); + Self::new(vec![session], displayed_owner) + } +} + +struct AppServerPickerBackend { + uses_remote_filesystem: bool, + session: AppServerSession, + archive_request_handle: AppServerRequestHandle, + include_non_interactive: bool, + raw_reasoning_visibility: RawReasoningVisibility, + config: Option, +} + +fn resume_picker_backends( + uses_remote_filesystem: bool, + config: &Config, + include_non_interactive: bool, + app_servers: ResumePickerAppServers, +) -> (Vec, crate::backend_id::BackendId) { + let displayed_owner = app_servers.displayed_owner; + let backends = app_servers + .ordered + .into_iter() + .map(|session| { + let is_displayed_owner = session.backend_id() == displayed_owner; + AppServerPickerBackend { + uses_remote_filesystem: if is_displayed_owner { + uses_remote_filesystem + } else { + session.uses_remote_workspace() + }, + archive_request_handle: session.request_handle(), + config: (!session.uses_remote_workspace()).then(|| config.clone()), + session, + include_non_interactive, + raw_reasoning_visibility: raw_reasoning_visibility(config), + } + }) + .collect(); + (backends, displayed_owner) +} + struct PickerPage { rows: Vec, read_states: HashMap, @@ -447,16 +512,41 @@ pub async fn run_resume_picker_with_app_server( include_non_interactive: bool, app_server: AppServerSession, ) -> Result { - let archive_request_handle = app_server.request_handle(); - run_resume_picker_with_launch_context( + run_resume_picker_with_app_servers( uses_remote_filesystem, tui, config, local_settings, show_all, include_non_interactive, - app_server, - archive_request_handle, + ResumePickerAppServers::singleton(app_server), + ) + .await +} + +pub(crate) async fn run_resume_picker_with_app_servers( + uses_remote_filesystem: bool, + tui: &mut Tui, + config: &Config, + local_settings: &crate::local_settings::LocalSettings, + show_all: bool, + include_non_interactive: bool, + app_servers: ResumePickerAppServers, +) -> Result { + let (backends, displayed_owner) = resume_picker_backends( + uses_remote_filesystem, + config, + include_non_interactive, + app_servers, + ); + run_resume_picker_with_launch_context( + uses_remote_filesystem, + tui, + config, + local_settings, + show_all, + backends, + displayed_owner, SessionPickerLaunchContext::Startup, ) .await @@ -464,28 +554,32 @@ pub async fn run_resume_picker_with_app_server( #[expect( clippy::too_many_arguments, - reason = "keep local preferences separate while the legacy Config parameter is still required" + reason = "picker configuration follows the existing session entry point" )] -pub async fn run_resume_picker_from_existing_session_with_app_server( +pub(crate) async fn run_resume_picker_from_existing_session_with_app_servers( uses_remote_filesystem: bool, tui: &mut Tui, config: &Config, local_settings: &crate::local_settings::LocalSettings, show_all: bool, include_non_interactive: bool, - app_server: AppServerSession, - archive_request_handle: AppServerRequestHandle, + app_servers: ResumePickerAppServers, current_thread_id: Option, ) -> Result { + let (backends, displayed_owner) = resume_picker_backends( + uses_remote_filesystem, + config, + include_non_interactive, + app_servers, + ); run_resume_picker_with_launch_context( uses_remote_filesystem, tui, config, local_settings, show_all, - include_non_interactive, - app_server, - archive_request_handle, + backends, + displayed_owner, SessionPickerLaunchContext::ExistingSession { current_thread_id }, ) .await @@ -501,24 +595,36 @@ async fn run_resume_picker_with_launch_context( config: &Config, local_settings: &crate::local_settings::LocalSettings, show_all: bool, - include_non_interactive: bool, - app_server: AppServerSession, - archive_request_handle: AppServerRequestHandle, + app_servers: Vec, + displayed_owner: crate::backend_id::BackendId, launch_context: SessionPickerLaunchContext, ) -> Result { + let has_multiple_backends = app_servers.len() > 1; + let Some(app_server) = app_servers + .iter() + .find(|backend| backend.session.backend_id() == displayed_owner) + else { + unreachable!("displayed owner must belong to the resume picker's nonempty backend set"); + }; let (bg_tx, bg_rx) = mpsc::unbounded_channel(); - let uses_remote_workspace = app_server.uses_remote_workspace(); - let cwd_filter = picker_cwd_filter( - config.cwd.as_path(), - /*show_all*/ false, - uses_remote_workspace, - app_server.remote_cwd_override(), - ); + let uses_remote_workspace = app_server.session.uses_remote_workspace(); + // A single path cannot describe both local and remote workspaces. Until the + // picker owns per-backend filters, a mixed listing starts unfiltered. + let cwd_filter = (!has_multiple_backends) + .then(|| { + picker_cwd_filter( + config.cwd.as_path(), + /*show_all*/ false, + uses_remote_workspace, + app_server.session.remote_cwd_override(), + ) + }) + .flatten(); let local_filter_cwd = local_picker_cwd_filter(&cwd_filter, uses_remote_filesystem); let provider_filter = ProviderFilter::All; let runtime_keymap = picker_runtime_keymap(local_settings)?; let options = SessionPickerRunOptions { - backend_id: app_server.backend_id(), + backend_id: app_server.session.backend_id(), use_theme_colors: local_settings.tui.status_line_use_colors, copy_on_select: local_settings.copy_on_select(&codex_terminal_detection::terminal_info()), copy_on_select_clear_selection: local_settings.tui.copy_on_select_clear_selection, @@ -536,7 +642,7 @@ async fn run_resume_picker_with_launch_context( codex_home: local_settings.codex_home.to_path_buf(), }), keymap: runtime_keymap, - initial_page_mode: if uses_remote_workspace { + initial_page_mode: if uses_remote_workspace || has_multiple_backends { PageLoadMode::StoreDefault } else { PageLoadMode::StateDbOnly @@ -545,15 +651,7 @@ async fn run_resume_picker_with_launch_context( run_session_picker_with_loader( tui, options, - spawn_app_server_page_loader( - uses_remote_filesystem, - app_server, - archive_request_handle, - include_non_interactive, - raw_reasoning_visibility(config), - (!uses_remote_workspace).then(|| config.clone()), - bg_tx, - ), + spawn_app_server_page_loaders(app_servers, bg_tx), bg_rx, ) .await @@ -807,6 +905,39 @@ fn spawn_app_server_page_loader( config: Option, output: mpsc::UnboundedSender, ) -> PickerBackground { + spawn_app_server_page_loaders( + vec![AppServerPickerBackend { + uses_remote_filesystem, + session: app_server, + archive_request_handle, + include_non_interactive, + raw_reasoning_visibility, + config, + }], + output, + ) +} + +fn spawn_app_server_page_loaders( + backends: Vec, + output: mpsc::UnboundedSender, +) -> PickerBackground { + let mut sources = Vec::with_capacity(backends.len()); + for backend in backends { + sources.push(spawn_app_server_source(backend)); + } + sources::spawn(sources, output) +} + +fn spawn_app_server_source(backend: AppServerPickerBackend) -> sources::Source { + let AppServerPickerBackend { + uses_remote_filesystem, + session: app_server, + archive_request_handle, + include_non_interactive, + raw_reasoning_visibility, + config, + } = backend; let backend_id = app_server.backend_id(); let (bg_tx, events) = mpsc::unbounded_channel(); let (request_tx, mut request_rx) = mpsc::unbounded_channel::(); @@ -818,6 +949,7 @@ fn spawn_app_server_page_loader( match request { PickerLoadRequest::Page(request) => { let mode = request.mode; + let provider_filter = request.provider_filter; let cwd_filter = page_cwd_filter.for_request( request.cursor.as_ref(), request.cwd_filter.as_deref(), @@ -834,7 +966,7 @@ fn spawn_app_server_page_loader( cursor, cwd_filter, request.status, - request.provider_filter, + provider_filter, request.sort_key, include_non_interactive, matches!(request.mode, PageLoadMode::StateDbOnly), @@ -941,15 +1073,14 @@ fn spawn_app_server_page_loader( } }); - let source = sources::Source { + sources::Source { backend_id, loader: Arc::new(move |request: PickerLoadRequest| { let _ = request_tx.send(request); }), events, task, - }; - sources::spawn(vec![source], output) + } } /// Returns the human-readable column header for the given sort key. diff --git a/codex-rs/tui/src/resume_picker/sources_tests.rs b/codex-rs/tui/src/resume_picker/sources_tests.rs index ac3dd836a9..f53c2e7188 100644 --- a/codex-rs/tui/src/resume_picker/sources_tests.rs +++ b/codex-rs/tui/src/resume_picker/sources_tests.rs @@ -226,6 +226,51 @@ async fn sources_stream_rows_and_route_actions_to_each_source() { assert!(first_requests.recv().await.is_none()); assert!(second_requests.recv().await.is_none()); } +#[tokio::test] +async fn late_source_failure_clears_pending_cursor() { + let (first, mut first_requests, first_events) = source(); + let (second, mut second_requests, second_events) = source(); + let (output, mut events) = mpsc::unbounded_channel(); + let PickerBackground { loader, task } = spawn(vec![first, second], output); + loader(request(/*token*/ 1, /*cursor*/ None)); + first_requests.recv().await.unwrap(); + second_requests.recv().await.unwrap(); + + first_events + .send(reply( + /*token*/ 1, + ThreadId::new(), + /*cursor*/ None, + )) + .unwrap(); + let BackgroundEvent::PartialPage { page, .. } = next(&mut events).await else { + panic!("expected first source page") + }; + assert!(matches!(page.next_cursor, Some(PageCursor::SourcesPending))); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::Page { .. } + )); + + second_events + .send(BackgroundEvent::Page { + request_token: 1, + search_token: None, + mode: PageLoadMode::StoreDefault, + page: Err(std::io::Error::other("source failed")), + }) + .unwrap(); + assert!(matches!( + next(&mut events).await, + BackgroundEvent::SourceError { + next_cursor: None, + .. + } + )); + + drop(loader); + task.await.unwrap(); +} #[tokio::test] async fn late_source_failure_finishes_catalog_progress() {