Prepare TUI startup to accept multiple backend sessions (#56097)

## What changed

- Add `InitialBackendSessions` to carry the initial session and additional sessions into the TUI event loop. Existing launches continue to supply a single session.
- Shut down all supplied sessions on startup errors, cancellation, and early exits.
- Register the under-development `tui_multi_backend` feature flag, disabled by default, and expose it in the configuration schema.

## Testing

Add a regression test that fails bootstrap and verifies that both initial backends receive shutdown requests. Update multi-source event polling coverage to use the new initialization path.

GitOrigin-RevId: e5db439552de97594e6f13360b90be71256ab2f0
This commit is contained in:
Benjamin Carlsson
2026-10-22 06:32:01 +00:00
committed by copyberry
parent c2d5ea52e1
commit 86e17e98c7
9 changed files with 243 additions and 72 deletions
+7 -1
View File
@@ -1264,6 +1264,9 @@
"tui_app_server": {
"type": "boolean"
},
"tui_multi_backend": {
"type": "boolean"
},
"ultrafast_mode": {
"type": "boolean"
},
@@ -7776,6 +7779,9 @@
"tui_app_server": {
"type": "boolean"
},
"tui_multi_backend": {
"type": "boolean"
},
"ultrafast_mode": {
"type": "boolean"
},
@@ -8328,4 +8334,4 @@
},
"title": "ConfigToml",
"type": "object"
}
}
+8
View File
@@ -120,6 +120,8 @@ pub enum Feature {
/// Automatically start the shared local daemon for eligible interactive launches.
DaemonAutoStart,
/// Allow the TUI to discover and select among multiple conversation backends.
TuiMultiBackend,
/// Enable negotiated zstd compression for app-server WebSocket and UDS connections.
AppServerWebSocketCompression,
@@ -1057,6 +1059,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::Stable,
default_enabled: true,
},
FeatureSpec {
id: Feature::TuiMultiBackend,
key: "tui_multi_backend",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::AppServerWebSocketCompression,
key: "app_server_websocket_compression",
+1 -1
View File
@@ -216,7 +216,7 @@ pub(crate) mod app_server_requests;
mod app_server_thread_ownership;
mod backend_banner_fallback;
mod backend_events;
mod backend_sessions;
pub(crate) mod backend_sessions;
mod background_requests;
mod composer_hints;
mod config_panel;
+40
View File
@@ -7,6 +7,38 @@ use crate::backend_id::BackendId;
use std::future::Future;
use std::task::Poll;
/// The initial conventional session and any additional sessions registered at launch.
pub(crate) struct InitialBackendSessions {
initial: AppServerSession,
others: Vec<AppServerSession>,
}
impl InitialBackendSessions {
pub(crate) fn singleton(initial: AppServerSession) -> Self {
Self {
initial,
others: Vec::new(),
}
}
#[allow(
dead_code,
reason = "used by the under-development backend registration path"
)]
pub(crate) fn multiple(initial: AppServerSession, others: Vec<AppServerSession>) -> Self {
assert!(
others
.iter()
.all(|session| session.backend_id() != initial.backend_id())
);
Self { initial, others }
}
pub(super) fn into_parts(self) -> (AppServerSession, Vec<AppServerSession>) {
(self.initial, self.others)
}
}
struct SessionEntry {
session: AppServerSession,
listening: bool,
@@ -22,6 +54,14 @@ pub(super) struct BackendSessions {
}
impl BackendSessions {
pub(super) fn from_initial(initial: AppServerSession, others: Vec<AppServerSession>) -> Self {
let mut sessions = Self::new(initial);
for session in others {
sessions.insert(session);
}
sessions
}
pub(super) fn new(session: AppServerSession) -> Self {
let mut sessions = Self {
entries: Vec::new(),
+117 -49
View File
@@ -229,7 +229,7 @@ impl App {
#[allow(clippy::too_many_arguments)]
pub async fn run(
tui: &mut tui::Tui,
mut app_server: AppServerSession,
initial_backends: super::backend_sessions::InitialBackendSessions,
mut config: Config,
launch_cwd: PathBuf,
cli_kv_overrides: Vec<(String, TomlValue)>,
@@ -256,24 +256,58 @@ impl App {
) -> Result<AppExitInfo> {
use tokio_stream::StreamExt;
async fn shutdown_on_startup_error(
app_server: AppServerSession,
error: impl Into<color_eyre::eyre::Report>,
) -> Result<AppExitInfo> {
if let Err(shutdown_error) = app_server.shutdown().await {
async fn shutdown_startup_sessions(sessions: impl IntoIterator<Item = AppServerSession>) {
let shutdowns = sessions.into_iter().map(AppServerSession::shutdown);
let shutdown_result: std::io::Result<Vec<()>> = futures::future::join_all(shutdowns)
.await
.into_iter()
.collect();
if let Err(shutdown_error) = shutdown_result {
tracing::warn!("app-server shutdown failed: {shutdown_error}");
}
}
async fn shutdown_on_startup_error(
app_server: AppServerSession,
other_backends: &mut Vec<AppServerSession>,
error: impl Into<color_eyre::eyre::Report>,
) -> Result<AppExitInfo> {
shutdown_startup_sessions(std::iter::once(app_server).chain(other_backends.drain(..)))
.await;
Err(error.into())
}
async fn cancel_startup_sessions(
app_server: AppServerSession,
other_backends: &mut Vec<AppServerSession>,
) -> AppExitInfo {
let exit_info = cancel_session_start(app_server).await;
shutdown_startup_sessions(other_backends.drain(..)).await;
exit_info
}
let (mut app_server, mut other_backends) = initial_backends.into_parts();
// Preserve singleton error propagation while releasing all sessions on multi-backend exits.
macro_rules! startup_try {
($result:expr) => {
match $result {
Ok(value) => value,
Err(error) if other_backends.is_empty() => return Err(error.into()),
Err(error) => {
return shutdown_on_startup_error(app_server, &mut other_backends, error)
.await;
}
}
};
}
// Adopt actual launch ownership before constructing session-local preferences.
tui.prepare_owned_screen(config.tui_fullscreen_transcript)?;
startup_try!(tui.prepare_owned_screen(config.tui_fullscreen_transcript));
let mut local_settings = crate::local_settings::LocalSettings::for_tui(&config, tui);
let startup_started_at = Instant::now();
if let SessionSelection::Resume(target) | SessionSelection::Fork(target) =
&session_selection
{
target.ensure_backend(app_server.backend_id())?;
startup_try!(target.ensure_backend(app_server.backend_id()));
}
let (app_event_tx, mut app_event_rx) = unbounded_channel();
let app_event_tx = AppEventSender::new(app_event_tx);
@@ -288,8 +322,10 @@ impl App {
local_settings.tui.notification_settings.condition,
);
let harness_overrides =
normalize_harness_overrides_for_cwd(harness_overrides, &config.cwd)?;
let harness_overrides = startup_try!(normalize_harness_overrides_for_cwd(
harness_overrides,
&config.cwd
));
app_server.model_provider_override = harness_overrides.model_provider.clone();
let fresh_start = matches!(
&session_selection,
@@ -307,8 +343,11 @@ impl App {
.run_until(tui, app_server.bootstrap(&config))
.await
{
Ok(bootstrap) => bootstrap?,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Ok(bootstrap) => startup_try!(bootstrap),
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
},
);
}
@@ -342,8 +381,12 @@ impl App {
.await
{
Ok(Ok(defaults)) => defaults,
Ok(Err(err)) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Ok(Err(err)) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
}
} else {
FreshStartupDefaults::default()
@@ -369,8 +412,10 @@ impl App {
)
.await
{
Ok(bootstrap) => bootstrap?,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Ok(bootstrap) => startup_try!(bootstrap),
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
},
};
tracing::debug!(
@@ -413,13 +458,13 @@ impl App {
None
};
if let Err(err) = startup_draft.flush_pending_events(tui).await {
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
let exit_info =
if matches!(&session_selection, SessionSelection::Fork(_)) && config.model.is_none() {
None
} else {
handle_model_migration_prompt_if_needed(
let result = handle_model_migration_prompt_if_needed(
tui,
&mut config,
&local_settings,
@@ -428,16 +473,11 @@ impl App {
&app_event_tx,
&available_models,
)
.await?
.await;
startup_try!(result)
};
if let Some(exit_info) = exit_info {
app_server
.shutdown()
.await
.inspect_err(|err| {
tracing::warn!("app-server shutdown failed: {err}");
})
.ok();
shutdown_startup_sessions(std::iter::once(app_server).chain(other_backends)).await;
return Ok(exit_info);
}
if let Some(updated_model) = config.model.clone() {
@@ -556,7 +596,10 @@ impl App {
.await
{
Ok(tooltip_override) => tooltip_override,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
}
};
let init = crate::chatwidget::ChatWidgetInit {
@@ -603,6 +646,7 @@ impl App {
{
return shutdown_on_startup_error(
app_server,
&mut other_backends,
color_eyre::eyre::eyre!(
"Permission overrides are not supported when resuming a remote task."
),
@@ -650,16 +694,29 @@ impl App {
history_notice = notice;
thread
}),
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => {
return shutdown_on_startup_error(
app_server,
&mut other_backends,
err,
)
.await;
}
}
}
Ok(resumed) => resumed,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
};
let resumed = if read_only_thread {
match resumed {
Ok(resumed) => Some(resumed),
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
}
} else {
let action = SessionStartAction::Resume(
@@ -684,12 +741,14 @@ impl App {
run_unarchive_prompt(tui, target_session.thread_id, action).await
},
)
.await?;
match resumed {
.await;
match startup_try!(resumed) {
SessionStartOutcome::Started(started) => Some(*started),
SessionStartOutcome::CommandCenter => None,
SessionStartOutcome::Exit => {
return Ok(cancel_session_start(app_server).await);
return Ok(
cancel_startup_sessions(app_server, &mut other_backends).await
);
}
}
};
@@ -736,6 +795,7 @@ impl App {
{
return shutdown_on_startup_error(
app_server,
&mut other_backends,
color_eyre::eyre::eyre!(
"Permission overrides are not supported when forking a remote task."
),
@@ -765,7 +825,10 @@ impl App {
.await
{
Ok(forked) => forked,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
};
let action = SessionStartAction::Fork(permission_mode);
let forked = complete_session_start(
@@ -783,19 +846,20 @@ impl App {
run_unarchive_prompt(tui, target_session.thread_id, action).await
},
)
.await?;
let forked = match forked {
.await;
let forked = match startup_try!(forked) {
SessionStartOutcome::Started(started) => Some(*started),
SessionStartOutcome::CommandCenter => None,
SessionStartOutcome::Exit => {
return Ok(cancel_session_start(app_server).await);
return Ok(cancel_startup_sessions(app_server, &mut other_backends).await);
}
};
if let Some(forked) = &forked {
if let Some(worktree) = managed_worktree.as_ref()
&& let Err(err) = worktree.bind(forked.session.thread_id)
{
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err)
.await;
}
if config.model_reasoning_effort.is_none() {
config.model_reasoning_effort = forked.session.reasoning_effort.clone();
@@ -882,7 +946,8 @@ impl App {
Fix the config and retry.\n\
See the Codex keymap documentation for supported actions and examples."
)
})?;
});
let runtime_keymap = startup_try!(runtime_keymap);
#[cfg(not(debug_assertions))]
let upgrade_version = crate::updates::get_upgrade_version(&config);
@@ -1057,8 +1122,10 @@ See the Codex keymap documentation for supported actions and examples."
)
.await
{
Ok(result) => result?,
Err(err) => return shutdown_on_startup_error(app_server, err).await,
Ok(result) => startup_try!(result),
Err(err) => {
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
}
if read_only_thread {
app.ensure_thread_channel(thread_id).mark_external_writer();
@@ -1080,7 +1147,7 @@ See the Codex keymap documentation for supported actions and examples."
)
.await
{
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
}
if !start_in_agents_overview
@@ -1094,14 +1161,14 @@ See the Codex keymap documentation for supported actions and examples."
let initial_session_ms = initial_session_started_at.elapsed().as_millis();
if let Err(err) = startup_draft.flush_pending_events(tui).await {
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
if Self::should_handle_active_thread_events(
wait_for_initial_session_configured,
app.active_thread_rx.is_some(),
) && let Err(err) = app.drain_active_thread_events(tui).await
{
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
if app_event_rx.is_empty()
&& !app.has_queued_startup_protected_request()
@@ -1110,7 +1177,7 @@ See the Codex keymap documentation for supported actions and examples."
&& !app.chat_widget.has_pending_protected_request()
&& let Err(err) = startup_draft.flush_pending_paste_newline(tui).await
{
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
// Keep cancelled resume/fork text editable, but never carry confirmation to another session.
let mut pending_startup_submission =
@@ -1134,7 +1201,7 @@ See the Codex keymap documentation for supported actions and examples."
let event_stream_started_at = Instant::now();
tui.schedule_screen_size_recheck(Duration::ZERO);
if let Err(err) = app.render_startup_frame(tui, &app_event_rx) {
return shutdown_on_startup_error(app_server, err).await;
return shutdown_on_startup_error(app_server, &mut other_backends, err).await;
}
let tui_events = tui.event_stream();
tokio::pin!(tui_events);
@@ -1183,8 +1250,8 @@ See the Codex keymap documentation for supported actions and examples."
crate::update_action::get_update_action(),
))),
))
.await?;
match control {
.await;
match startup_try!(control) {
AppRunControl::Continue => None,
AppRunControl::Exit(exit_reason) => Some(exit_reason),
}
@@ -1195,7 +1262,8 @@ See the Codex keymap documentation for supported actions and examples."
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 mut backends =
super::backend_sessions::BackendSessions::from_initial(app_server, other_backends);
let exit_reason_result = if let Some(exit_reason) = pre_loop_exit_reason {
Ok(exit_reason)
@@ -255,9 +255,9 @@ async fn startup_rejects_foreign_resume_and_fork_before_requests() -> Result<()>
let mut tui = crate::tui::test_support::make_test_tui()?;
let result = super::startup_defaults_tests::run_startup_for_test(
&mut tui,
server,
crate::app::backend_sessions::InitialBackendSessions::singleton(server),
app.config.clone(),
bootstrap,
Some(bootstrap),
selection,
app.app_server_target.clone(),
)
@@ -34,10 +34,7 @@ async fn backend_sessions_poll_three_sources_and_isolate_closed_or_paused_stream
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 backends = BackendSessions::from_initial(sessions.next().unwrap(), sessions.collect());
let mut unpaused = HashSet::new();
for _ in 0..2 {
let (owner, event) =
@@ -54,7 +51,7 @@ async fn backend_sessions_poll_three_sources_and_isolate_closed_or_paused_stream
.await?;
assert_eq!(owner, owners[0]);
assert_eq!(event.unwrap().backend_id, owner);
// Cancel a pending poll, then close only the selected transport.
// Cancel a pending poll, then close only the initial transport.
assert!(
backends
.next_event(/*paused_backend*/ None)
@@ -1,6 +1,7 @@
//! Request-level coverage for the fresh-startup server defaults overlay.
//! Startup boundary coverage for server defaults and connection ownership.
use super::*;
use crate::app::backend_sessions::InitialBackendSessions;
use crate::app::startup::prepare_fresh_startup_config;
use crate::app::startup::startup_model;
use crate::app::startup_bootstrap::bootstrap_server_owned_start;
@@ -17,15 +18,15 @@ fn remote_target() -> AppServerTarget {
pub(super) async fn run_startup_for_test(
tui: &mut crate::tui::Tui,
server: AppServerSession,
initial_backends: InitialBackendSessions,
config: Config,
bootstrap: AppServerBootstrap,
bootstrap: Option<AppServerBootstrap>,
selection: SessionSelection,
app_server_target: AppServerTarget,
) -> Result<AppExitInfo> {
App::run(
tui,
server,
initial_backends,
config.clone(),
config.cwd.to_path_buf(),
Vec::new(),
@@ -42,7 +43,7 @@ pub(super) async fn run_startup_for_test(
/*state_db*/ None,
Arc::new(EnvironmentManager::default_for_tests()),
Duration::ZERO,
Some(bootstrap),
bootstrap,
/*startup_hooks_browser*/ None,
/*daemon_startup_warning*/ None,
/*managed_daemon*/ false,
@@ -72,9 +73,9 @@ async fn agents_overview_startup_restores_saved_grouping() -> Result<()> {
let endpoint = crate::resolve_remote_addr("ws://127.0.0.1:4500")?;
let mut run = Box::pin(run_startup_for_test(
&mut tui,
server,
InitialBackendSessions::singleton(server),
config,
bootstrap,
Some(bootstrap),
SessionSelection::AgentsOverview,
AppServerTarget::LocalDaemon {
allow_embedded_fallback: true,
@@ -179,9 +180,9 @@ async fn cli_fork_omits_implicit_model_and_effort() -> Result<()> {
let mut tui = crate::tui::test_support::make_test_tui()?;
let mut run = Box::pin(run_startup_for_test(
&mut tui,
server,
InitialBackendSessions::singleton(server),
config,
bootstrap,
Some(bootstrap),
SessionSelection::Fork(crate::resume_picker::SessionTarget {
backend_id,
path: None,
@@ -227,9 +228,9 @@ async fn run_until_thread_start(
tui.pause_events();
let mut run = Box::pin(run_startup_for_test(
&mut tui,
server,
InitialBackendSessions::singleton(server),
config,
bootstrap,
Some(bootstrap),
SessionSelection::StartFresh,
AppServerTarget::Embedded,
));
@@ -757,9 +758,9 @@ async fn startup_read_failure_exits_before_thread_creation() -> Result<()> {
let mut tui = crate::tui::test_support::make_test_tui()?;
let result = run_startup_for_test(
&mut tui,
server,
InitialBackendSessions::singleton(server),
config,
bootstrap,
Some(bootstrap),
SessionSelection::StartFresh,
AppServerTarget::Embedded,
)
@@ -770,3 +771,54 @@ async fn startup_read_failure_exits_before_thread_creation() -> Result<()> {
proxy.await??;
Ok(())
}
#[tokio::test]
async fn bootstrap_failure_shuts_down_every_initial_backend() -> Result<()> {
use codex_app_server_client::AppServerClient;
use codex_app_server_client::CloudClient;
use codex_app_server_client::CloudOperation;
let app = make_test_app().await;
let (first, mut first_commands, _first_events) = CloudClient::channel(/*capacity*/ 1);
let (second, mut second_commands, _second_events) = CloudClient::channel(/*capacity*/ 1);
let [first, second] = [first, second].map(|client| {
AppServerSession::new(
crate::backend_id::BackendId::new(),
AppServerClient::Cloud(client),
crate::app_server_session::ThreadParamsMode::Embedded,
)
});
let mut tui = crate::tui::test_support::make_test_tui()?;
let (result, ()) = tokio::join!(
run_startup_for_test(
&mut tui,
InitialBackendSessions::multiple(first, vec![second]),
app.config.clone(),
/*bootstrap*/ None,
SessionSelection::StartFresh,
AppServerTarget::Embedded,
),
async {
let request = first_commands.recv().await.expect("initial account read");
assert!(matches!(
request.operation,
CloudOperation::Request(request)
if matches!(*request, codex_app_server_protocol::ClientRequest::GetAccount { .. })
));
request
.response
.send(Err(std::io::Error::other("bootstrap unavailable")))
.unwrap();
for commands in [&mut first_commands, &mut second_commands] {
let command = commands.recv().await.expect("startup shutdown");
assert!(matches!(command.operation, CloudOperation::Shutdown));
command
.response
.send(Ok(Ok(serde_json::Value::Null)))
.unwrap();
}
}
);
assert!(result.unwrap_err().to_string().contains("account/read"));
Ok(())
}
+1 -1
View File
@@ -2181,7 +2181,7 @@ async fn run_ratatui_app(
// transitions have enough stack headroom to rebuild configuration and the chat widget.
let app_result = Box::pin(App::run(
&mut tui,
app_server,
crate::app::backend_sessions::InitialBackendSessions::singleton(app_server),
config,
current_cwd.to_path_buf(),
cli_kv_overrides.clone(),