Manage TUI backend sessions independently (#55381)

## What changed

Add a `BackendSessions` collection to own app-server sessions and poll their events fairly with cancellation-safe receives. Track listening state per backend so pausing or closing one stream leaves the others available, and resume polling after reconnect. Shut down all registered sessions on exit.

Ignore background backend disconnect, lag, account, and rate-limit events that would otherwise update the displayed conversation's UI.

## Testing

Add coverage for polling three sessions, paused and closed stream isolation, cancellation, and session replacement after reconnect. Extend event ownership coverage to verify that background lifecycle and account events leave the displayed connection and rate-limit state untouched.

GitOrigin-RevId: 9ed82abc77aaf4bfeb86eb72161ad90fe20177ad

codex-rs/tui/src/app.rs

@@ -216,6 +216,7 @@ pub(crate) mod app_server_requests;
 mod app_server_thread_ownership;
 mod backend_banner_fallback;
 mod backend_events;
+mod backend_sessions;
 mod background_requests;
 mod composer_hints;
 mod config_panel;

codex-rs/tui/src/app/backend_events.rs

@@ -1,10 +1,39 @@
 //! Checks incoming event ownership before consulting native-thread caches or UI state.
+//! Connection-wide lifecycle and account events may only update the displayed conversation's UI.
 
 use super::App;
 use crate::app_server_session::AppServerSession;
 use crate::app_server_session::AppServerSessionEvent;
+use crate::backend_id::BackendId;
+use codex_app_server_client::AppServerEvent;
+use codex_app_server_protocol::ServerNotification;
 
 impl App {
+    pub(super) async fn handle_polled_backend_event(
+        &mut self,
+        app_server: &AppServerSession,
+        displayed_owner: BackendId,
+        event: AppServerSessionEvent,
+    ) {
+        if event.backend_id != displayed_owner
+            && match &event.event {
+                AppServerEvent::Disconnected { .. } | AppServerEvent::Lagged { .. } => true,
+                AppServerEvent::ServerNotification(notification) => matches!(
+                    notification.as_ref(),
+                    ServerNotification::AccountUpdated(_)
+                        | ServerNotification::AccountRateLimitsUpdated(_)
+                ),
+                AppServerEvent::ServerRequest(_) => false,
+            }
+        {
+            tracing::debug!(
+                "ignoring a background backend event that targets displayed conversation state"
+            );
+            return;
+        }
+        self.handle_backend_event(app_server, event).await;
+    }
+
     pub(super) async fn handle_backend_event(
         &mut self,
         app_server: &AppServerSession,

codex-rs/tui/src/app/backend_sessions.rs

@@ -0,0 +1,112 @@
+//! Owns backend sessions independently of the displayed conversation's UI state.
+//! Polling is fair and cancellation-safe; closing one stream never closes another.
+
+use crate::app_server_session::AppServerSession;
+use crate::app_server_session::AppServerSessionEvent;
+use crate::backend_id::BackendId;
+use std::future::Future;
+use std::task::Poll;
+
+struct SessionEntry {
+    session: AppServerSession,
+    listening: bool,
+}
+
+/// A session's stable position in an append-only backend collection.
+#[derive(Clone, Copy)]
+pub(super) struct SessionSlot(usize);
+
+pub(super) struct BackendSessions {
+    entries: Vec<SessionEntry>,
+    next_source: usize,
+}
+
+impl BackendSessions {
+    pub(super) fn new(session: AppServerSession) -> Self {
+        let mut sessions = Self {
+            entries: Vec::new(),
+            next_source: 0,
+        };
+        sessions.insert(session);
+        sessions
+    }
+
+    pub(super) fn insert(&mut self, session: AppServerSession) {
+        assert!(self.session(session.backend_id()).is_none());
+        self.entries.push(SessionEntry {
+            session,
+            listening: true,
+        });
+    }
+
+    pub(super) fn session(&self, owner: BackendId) -> Option<&AppServerSession> {
+        self.slot(owner).map(|slot| self.session_at(slot))
+    }
+
+    pub(super) fn slot(&self, owner: BackendId) -> Option<SessionSlot> {
+        self.entries
+            .iter()
+            .position(|entry| entry.session.backend_id() == owner)
+            .map(SessionSlot)
+    }
+
+    pub(super) fn session_at(&self, slot: SessionSlot) -> &AppServerSession {
+        &self.entries[slot.0].session
+    }
+
+    pub(super) fn session_at_mut(&mut self, slot: SessionSlot) -> &mut AppServerSession {
+        &mut self.entries[slot.0].session
+    }
+
+    pub(super) fn is_listening(&self, owner: BackendId) -> bool {
+        self.entries
+            .iter()
+            .any(|entry| entry.session.backend_id() == owner && entry.listening)
+    }
+
+    pub(super) fn resume_events(&mut self, owner: BackendId) {
+        if let Some(entry) = self
+            .entries
+            .iter_mut()
+            .find(|entry| entry.session.backend_id() == owner)
+        {
+            entry.listening = true;
+        }
+    }
+
+    pub(super) async fn next_event(
+        &mut self,
+        paused_backend: Option<BackendId>,
+    ) -> (BackendId, Option<AppServerSessionEvent>) {
+        std::future::poll_fn(|cx| {
+            for offset in 0..self.entries.len() {
+                let index = (self.next_source + offset) % self.entries.len();
+                let entry = &mut self.entries[index];
+                let owner = entry.session.backend_id();
+                if !entry.listening || paused_backend == Some(owner) {
+                    continue;
+                }
+                // Both client transports receive from cancellation-safe queues. Dropping a
+                // pending receive leaves its event available for the next main-loop poll.
+                let event = std::pin::pin!(entry.session.next_event()).poll(cx);
+                if let Poll::Ready(event) = event {
+                    entry.listening = event.is_some();
+                    self.next_source = (index + 1) % self.entries.len();
+                    return Poll::Ready((owner, event));
+                }
+            }
+            Poll::Pending
+        })
+        .await
+    }
+
+    pub(super) async fn shutdown(self) -> std::io::Result<()> {
+        let mut first_error = None;
+        for entry in self.entries {
+            if let Err(error) = entry.session.shutdown().await {
+                first_error.get_or_insert(error);
+            }
+        }
+        first_error.map_or(Ok(()), Err)
+    }
+}

codex-rs/tui/src/app/startup.rs

@@ -1169,7 +1169,6 @@ See the Codex keymap documentation for supported actions and examples."
         let mut last_screen_size = tui.terminal.last_known_screen_size;
         let mut settled_size_tags = crate::terminal_size_telemetry::tags(last_screen_size);
         let mut pending_resize: Option<(ratatui::layout::Size, Instant)> = None;
-        let mut listen_for_app_server_events = true;
         let mut reconnect = None;
         let mut waiting_for_initial_session_configured = wait_for_initial_session_configured;
         let mut waiting_for_initial_session_header = true;
@@ -1195,10 +1194,18 @@ See the Codex keymap documentation for supported actions and examples."
         #[cfg(debug_assertions)]
         let pre_loop_exit_reason: Option<ExitReason> = None;
 
+        let initial_owner = app_server.backend_id();
+        let mut backends = super::backend_sessions::BackendSessions::new(app_server);
+
         let exit_reason_result = if let Some(exit_reason) = pre_loop_exit_reason {
             Ok(exit_reason)
         } else {
             loop {
+                let Some(displayed_session) = backends.slot(initial_owner) else {
+                    break Err(color_eyre::eyre::eyre!(
+                        "the initial backend is not registered"
+                    ));
+                };
                 // Reconnect can dismiss an overlay from the server-event path.
                 if app.overlay.is_none() {
                     if !tui.is_owned_screen() && tui.is_alt_screen_active() {
@@ -1211,7 +1218,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, &mut app_server)).await {
+                    match Box::pin(
+                        app.open_resume_picker(tui, backends.session_at_mut(displayed_session)),
+                    )
+                    .await
+                    {
                         Ok(AppRunControl::Continue) => {}
                         Ok(AppRunControl::Exit(reason)) => break Ok(reason),
                         Err(err) if app.recover_transport_error(&err) => {}
@@ -1220,23 +1231,40 @@ See the Codex keymap documentation for supported actions and examples."
                     continue;
                 }
                 if let Some(pending) = app.pending_working_directory_change.take() {
-                    Box::pin(app.finish_working_directory_change(tui, &mut app_server, pending))
-                        .await;
+                    Box::pin(app.finish_working_directory_change(
+                        tui,
+                        backends.session_at_mut(displayed_session),
+                        pending,
+                    ))
+                    .await;
                     continue;
                 }
                 if let Some((mode, name)) = app.pending_start_managed_worktree.take() {
-                    Box::pin(app.start_managed_worktree(&mut app_server, mode, name)).await;
+                    Box::pin(app.start_managed_worktree(
+                        backends.session_at_mut(displayed_session),
+                        mode,
+                        name,
+                    ))
+                    .await;
                     continue;
                 }
                 // Complete the fork and widget attachment on separate fresh loop iterations.
                 if let Some(attach) = app.pending_managed_worktree_attach.take() {
-                    Box::pin(app.attach_working_directory(tui, &mut app_server, *attach)).await;
+                    Box::pin(app.attach_working_directory(
+                        tui,
+                        backends.session_at_mut(displayed_session),
+                        *attach,
+                    ))
+                    .await;
                     continue;
                 }
                 if let Some(transition) = app.pending_managed_worktree_transition.take() {
-                    if let Err(error) =
-                        Box::pin(app.switch_to_managed_worktree(tui, &mut app_server, *transition))
-                            .await
+                    if let Err(error) = Box::pin(app.switch_to_managed_worktree(
+                        tui,
+                        backends.session_at_mut(displayed_session),
+                        *transition,
+                    ))
+                    .await
                     {
                         app.chat_widget.add_error_message(error.to_string());
                     }
@@ -1248,19 +1276,28 @@ See the Codex keymap documentation for supported actions and examples."
                 }
                 if app.reconnect.offline && !app.reconnect.failed && reconnect.is_none() {
                     reconnect = Some(Box::pin(reconnect::reconnect(
-                        app_server.backend_id(),
+                        initial_owner,
                         app.app_server_target.clone(),
                         app.config.clone(),
                         app.local_settings.clone(),
                         app.current_displayed_thread_id(),
-                        app_server.remote_cwd_override().map(Path::to_path_buf),
-                        app_server.thread_tool_transport(),
+                        backends
+                            .session_at(displayed_session)
+                            .remote_cwd_override()
+                            .map(Path::to_path_buf),
+                        backends
+                            .session_at(displayed_session)
+                            .thread_tool_transport(),
                         app.reconnect.presentation,
                     )));
                 }
                 // Replay queues history and operations. A buffered closure must not switch
                 // widgets before those app events have been applied.
                 let has_pending_app_events = !app_event_rx.is_empty();
+                let paused_backend = (app.reconnect.offline
+                    || (!matches!(app.app_server_target, AppServerTarget::Embedded)
+                        && has_pending_app_events))
+                    .then_some(initial_owner);
                 let initial_session_header_pending = waiting_for_initial_session_header
                     && app.primary_session_configured.is_some()
                     && has_pending_app_events;
@@ -1310,7 +1347,7 @@ See the Codex keymap documentation for supported actions and examples."
                                 if cell.as_any().is::<history_cell::SessionInfoCell>()
                         );
                         let had_active_modal = app.chat_widget.has_active_modal();
-                        match Box::pin(app.handle_event(tui, &mut app_server, event)).await {
+                        match Box::pin(app.handle_event(tui, backends.session_at_mut(displayed_session), event)).await {
                             Ok(AppRunControl::Continue) => {
                                 if is_initial_session_header {
                                     waiting_for_initial_session_header = false;
@@ -1339,7 +1376,7 @@ See the Codex keymap documentation for supported actions and examples."
                         app.active_thread_rx.is_some()
                     ) && !has_pending_app_events && !app.reconnect.offline => {
                         if let Some(event) = active {
-                            if let Err(err) = app.handle_active_thread_event(tui, &mut app_server, event).await {
+                            if let Err(err) = app.handle_active_thread_event(tui, backends.session_at_mut(displayed_session), event).await {
                                 break Err(err);
                             }
                         } else {
@@ -1370,23 +1407,27 @@ See the Codex keymap documentation for supported actions and examples."
                             {
                                 app.startup_protected_input_boundary = false;
                             }
-                            match app.handle_tui_event(tui, &mut app_server, event).await {
+                            match app.handle_tui_event(tui, backends.session_at_mut(displayed_session), event).await {
                                 Ok(control) => control,
                                 Err(err) if app.recover_transport_error(&err) => AppRunControl::Continue,
                             Err(err) => break Err(err),
                             }
                         } else {
                             tracing::warn!("terminal input stream closed; shutting down active thread");
-                            app.handle_exit_mode(&mut app_server, ExitMode::ShutdownFirst).await
+                            app.handle_exit_mode(backends.session_at_mut(displayed_session), ExitMode::ShutdownFirst).await
                         }
                     }
-                    app_server_event = app_server.next_event(), if listen_for_app_server_events && !app.reconnect.offline
-                        && (matches!(app.app_server_target, AppServerTarget::Embedded) || !has_pending_app_events) => {
+                    (owner, app_server_event) = backends.next_event(paused_backend) => {
                         match app_server_event {
-                            Some(event) => app.handle_backend_event(&app_server, event).await,
+                            Some(event) => {
+                                if let Some(session) = backends.session(owner) {
+                                    app.handle_polled_backend_event(session, initial_owner, event).await;
+                                }
+                            }
                             None => {
-                                listen_for_app_server_events = false;
-                                app.begin_reconnect();
+                                if owner == initial_owner {
+                                    app.begin_reconnect();
+                                }
                                 tracing::warn!("app-server event stream closed");
                             }
                         }
@@ -1396,10 +1437,10 @@ See the Codex keymap documentation for supported actions and examples."
                         reconnect = None;
                         match result {
                             Ok(connected) => {
-                                if let Err(err) = app.finish_reconnect(tui, &mut app_server, &mut app_event_rx, connected, CODEX_CLI_VERSION).await {
+                                if let Err(err) = app.finish_reconnect(tui, backends.session_at_mut(displayed_session), &mut app_event_rx, connected, CODEX_CLI_VERSION).await {
                                     break Err(err);
                                 }
-                                listen_for_app_server_events = true;
+                                backends.resume_events(initial_owner);
                                 waiting_for_initial_session_configured = false;
                             }
                             Err(error) => {
@@ -1442,8 +1483,8 @@ See the Codex keymap documentation for supported actions and examples."
                             }
                             None => std::future::pending().await,
                         }
-                    }, if listen_for_app_server_events => {
-                        app.refresh_rate_limits(&app_server, RateLimitRefreshOrigin::Periodic);
+                    }, if backends.is_listening(initial_owner) => {
+                        app.refresh_rate_limits(backends.session_at(displayed_session), RateLimitRefreshOrigin::Periodic);
                         AppRunControl::Continue
                     }
                     () = async {
@@ -1511,7 +1552,7 @@ See the Codex keymap documentation for supported actions and examples."
                 }
             }
         };
-        if let Err(err) = app_server.shutdown().await {
+        if let Err(err) = backends.shutdown().await {
             tracing::warn!(error = %err, "failed to shut down embedded app server");
         }
         let clear_pet_result = tui.clear_ambient_pet_image();

codex-rs/tui/src/app/tests.rs

@@ -71,6 +71,8 @@ mod quick_dials_tests;
 #[path = "tests/unavailable_commands_tests.rs"]
 mod unavailable_commands;
 
+#[path = "tests/backend_sessions_tests.rs"]
+mod backend_sessions;
 #[path = "tests/history_hydration_tests.rs"]
 mod history_hydration_tests;
 #[path = "tests/permission_shortcuts_tests.rs"]

codex-rs/tui/src/app/tests/backend_event_tests.rs

@@ -4,6 +4,8 @@ use super::*;
 use crate::app_server_session::AppServerSessionEvent;
 use crate::backend_id::BackendId;
 use codex_app_server_client::AppServerEvent;
+use codex_app_server_protocol::AccountRateLimitsUpdatedNotification;
+use codex_app_server_protocol::AccountUpdatedNotification;
 use pretty_assertions::assert_eq;
 
 #[tokio::test]
@@ -23,6 +25,7 @@ async fn foreign_backend_events_leave_thread_cache_and_connection_untouched() ->
     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);
     for event in [
@@ -42,7 +45,37 @@ async fn foreign_backend_events_leave_thread_cache_and_connection_untouched() ->
         )
         .await;
     }
+    for event in [
+        AppServerEvent::Lagged { skipped: 1 },
+        AppServerEvent::Disconnected {
+            message: "background disconnected".into(),
+        },
+        AppServerEvent::ServerNotification(Box::new(ServerNotification::AccountUpdated(
+            AccountUpdatedNotification {
+                auth_mode: None,
+                plan_type: None,
+            },
+        ))),
+        AppServerEvent::ServerNotification(Box::new(ServerNotification::AccountRateLimitsUpdated(
+            AccountRateLimitsUpdatedNotification {
+                rate_limits: serde_json::from_value(
+                    serde_json::json!({ "spendControlReached": true }),
+                )?,
+            },
+        ))),
+    ] {
+        app.handle_polled_backend_event(
+            &server,
+            foreign,
+            AppServerSessionEvent {
+                backend_id: server.backend_id(),
+                event,
+            },
+        )
+        .await;
+    }
     assert!(!app.reconnect.offline);
+    assert_eq!(app.rate_limit_hard_stop_generation, hard_stop_generation);
     assert!(
         !app.pending_app_server_requests
             .contains_server_request(&request)

codex-rs/tui/src/app/tests/backend_sessions_tests.rs

@@ -0,0 +1,156 @@
+//! Real session streams remain independent through polling, closure and replacement.
+
+use super::*;
+use crate::app::backend_sessions::BackendSessions;
+use codex_app_server_protocol::RequestId;
+use codex_app_server_protocol::ThreadSettingsUpdateResponse;
+use futures::FutureExt;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn backend_sessions_poll_three_sources_and_isolate_closed_or_paused_streams() -> Result<()> {
+    let mut app = make_test_app().await;
+    app.config.model_reasoning_effort = Some(ReasoningEffortConfig::Low);
+    let mut sessions = Vec::new();
+    let mut proxies = Vec::new();
+    let mut owners = Vec::new();
+    let mut threads = Vec::new();
+    for _ in 0..3 {
+        let (mut session, _, proxy) =
+            session_lifecycle_requests::start_recording_remote_app_server(&app.config).await?;
+        let thread = session.start_thread(&app.config).await?.session.thread_id;
+        assert!(
+            session
+                .thread_settings_update(ThreadSettingsUpdateParams {
+                    thread_id: thread.to_string(),
+                    effort: Some(ReasoningEffortConfig::High),
+                    ..Default::default()
+                })
+                .await?
+        );
+        owners.push(session.backend_id());
+        threads.push(thread);
+        sessions.push(session);
+        proxies.push(proxy);
+    }
+    let mut sessions = sessions.into_iter();
+    let mut backends = BackendSessions::new(sessions.next().unwrap());
+    for session in sessions {
+        backends.insert(session);
+    }
+    let mut unpaused = HashSet::new();
+    for _ in 0..2 {
+        let (owner, event) =
+            time::timeout(Duration::from_secs(5), backends.next_event(Some(owners[0]))).await?;
+        assert_ne!(owner, owners[0]);
+        assert_eq!(event.unwrap().backend_id, owner);
+        unpaused.insert(owner);
+    }
+    assert_eq!(unpaused, HashSet::from([owners[1], owners[2]]));
+    let (owner, event) = time::timeout(
+        Duration::from_secs(5),
+        backends.next_event(/*paused_backend*/ None),
+    )
+    .await?;
+    assert_eq!(owner, owners[0]);
+    assert_eq!(event.unwrap().backend_id, owner);
+    // Cancel a pending poll, then close only the selected transport.
+    assert!(
+        backends
+            .next_event(/*paused_backend*/ None)
+            .now_or_never()
+            .is_none()
+    );
+    proxies[0].abort();
+    loop {
+        let (owner, event) = time::timeout(
+            Duration::from_secs(5),
+            backends.next_event(/*paused_backend*/ None),
+        )
+        .await?;
+        assert_eq!(owner, owners[0]);
+        if event.is_none() {
+            break;
+        }
+    }
+    assert!(!backends.is_listening(owners[0]));
+    for index in 1..3 {
+        assert!(backends.is_listening(owners[index]));
+        let _: ThreadSettingsUpdateResponse = backends
+            .session(owners[index])
+            .unwrap()
+            .request_handle()
+            .request_typed(ClientRequest::ThreadSettingsUpdate {
+                request_id: RequestId::Integer(900),
+                params: ThreadSettingsUpdateParams {
+                    thread_id: threads[index].to_string(),
+                    effort: Some(ReasoningEffortConfig::Low),
+                    ..Default::default()
+                },
+            })
+            .await?;
+    }
+    let mut surviving = HashSet::new();
+    for _ in 0..2 {
+        let (owner, event) = time::timeout(
+            Duration::from_secs(5),
+            backends.next_event(/*paused_backend*/ None),
+        )
+        .await?;
+        assert_eq!(event.unwrap().backend_id, owner);
+        surviving.insert(owner);
+    }
+    assert_eq!(surviving, HashSet::from([owners[1], owners[2]]));
+    // EOF is reported once, and a later cancelled poll does not spin on the closed source.
+    assert!(
+        backends
+            .next_event(/*paused_backend*/ None)
+            .now_or_never()
+            .is_none()
+    );
+
+    let replacement = crate::start_app_server_for_picker(
+        owners[0],
+        &app.config,
+        &AppServerTarget::Embedded,
+        Vec::new(),
+        LoaderOverrides::without_managed_config_for_tests(),
+        /*state_db*/ None,
+        Arc::clone(&app.environment_manager),
+    )
+    .await?;
+    let reconnected_slot = backends.slot(owners[0]).unwrap();
+    let previous = std::mem::replace(backends.session_at_mut(reconnected_slot), replacement);
+    let _ = previous.shutdown().await;
+    backends.resume_events(owners[0]);
+    let thread = backends
+        .session_at_mut(reconnected_slot)
+        .start_thread(&app.config)
+        .await?
+        .session
+        .thread_id;
+    loop {
+        let (owner, event) = time::timeout(
+            Duration::from_secs(5),
+            backends.next_event(/*paused_backend*/ None),
+        )
+        .await?;
+        assert_eq!(owner, owners[0]);
+        if let codex_app_server_client::AppServerEvent::ServerNotification(notification) =
+            event.unwrap().event
+            && let ServerNotification::ThreadStarted(started) = *notification
+        {
+            assert_eq!(started.thread.id, thread.to_string());
+            break;
+        }
+    }
+    backends.shutdown().await?;
+    for (index, proxy) in proxies.into_iter().enumerate() {
+        if index == 0 {
+            assert!(proxy.await.unwrap_err().is_cancelled());
+        } else {
+            proxy.await??;
+        }
+    }
+    Ok(())
+}