From 2d0d88fea6b28e7f1f8066a53f93e23ebf9beb88 Mon Sep 17 00:00:00 2001 From: Yunfan Yang Date: Tue, 28 Apr 2026 18:02:55 -0400 Subject: [PATCH] Add reconnection logic to remote server (#9289) Migrated from /Users/kevinyang/Documents/GitHub/warp-internal via `script/migrate-private-to-public`. ## Commits - 823458b Add reconnect logic to remote server - 104aafb tech spec - c4039c3 clippy - c4e3a39 nit --- Cargo.lock | 1 + app/src/remote_server/ssh_transport.rs | 216 +++++---- app/src/terminal/model/session.rs | 56 ++- app/src/terminal/view.rs | 1 + crates/remote_server/Cargo.toml | 1 + crates/remote_server/src/manager.rs | 590 +++++++++++++++++++------ crates/remote_server/src/transport.rs | 28 +- specs/APP-4283/TECH.md | 119 +++++ 8 files changed, 743 insertions(+), 269 deletions(-) create mode 100644 specs/APP-4283/TECH.md diff --git a/Cargo.lock b/Cargo.lock index 9f5999d3..5487866f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10559,6 +10559,7 @@ version = "0.1.0" dependencies = [ "anyhow", "async-channel", + "async-io", "async-process", "command", "dashmap", diff --git a/app/src/remote_server/ssh_transport.rs b/app/src/remote_server/ssh_transport.rs index e3fef712..4a78aa44 100644 --- a/app/src/remote_server/ssh_transport.rs +++ b/app/src/remote_server/ssh_transport.rs @@ -3,13 +3,18 @@ //! [`SshTransport`] uses an existing SSH ControlMaster socket to check/install //! the remote server binary and to launch the `remote-server-proxy` process //! whose stdin/stdout become the protocol channel. +use std::future::Future; use std::path::PathBuf; +use std::pin::Pin; +use std::process::Stdio; +use std::sync::Arc; use anyhow::Result; use warpui::r#async::executor; use remote_server::client::RemoteServerClient; -use remote_server::setup::RemotePlatform; +use remote_server::setup::{self, RemotePlatform, CHECK_TIMEOUT, INSTALL_TIMEOUT}; +use remote_server::ssh::{run_ssh_command, run_ssh_script, ssh_args}; use remote_server::transport::{Connection, RemoteTransport}; /// SSH transport: connects via a ControlMaster socket. @@ -18,7 +23,7 @@ use remote_server::transport::{Connection, RemoteTransport}; /// process (`ssh -N -o ControlMaster=yes -o ControlPath=`). All SSH /// commands (binary check, install, proxy launch) are multiplexed through /// this socket without re-authenticating. -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct SshTransport { socket_path: PathBuf, } @@ -28,123 +33,112 @@ impl SshTransport { Self { socket_path } } } - impl RemoteTransport for SshTransport { - async fn detect_platform(&self) -> Result { - match remote_server::ssh::run_ssh_command( - &self.socket_path, - "uname -sm", - remote_server::setup::CHECK_TIMEOUT, - ) - .await - { - Ok(output) if output.status.success() => { - let stdout = String::from_utf8_lossy(&output.stdout); - remote_server::setup::parse_uname_output(&stdout).map_err(|e| format!("{e:#}")) - } - Ok(output) => { - let code = output.status.code().unwrap_or(-1); - let stderr = String::from_utf8_lossy(&output.stderr); - Err(format!("uname -sm exited with code {code}: {stderr}")) - } - Err(e) => Err(format!("{e:#}")), - } - } - - async fn check_binary(&self) -> Result { - let bin_path = remote_server::setup::remote_server_binary(); - log::info!("Checking for remote server binary at {bin_path}"); - match remote_server::ssh::run_ssh_command( - &self.socket_path, - &remote_server::setup::binary_check_command(), - remote_server::setup::CHECK_TIMEOUT, - ) - .await - { - // `test -x` exits 0 when present, 1 when missing. - // Any other exit code (or None / signal) is treated as a check failure. - Ok(output) => match output.status.code() { - Some(0) => Ok(true), - Some(1) => Ok(false), - Some(code) => { - let stderr = String::from_utf8_lossy(&output.stderr); - Err(format!("binary check exited with code {code}: {stderr}")) + fn detect_platform( + &self, + ) -> Pin> + Send>> { + let socket_path = self.socket_path.clone(); + Box::pin(async move { + match run_ssh_command(&socket_path, "uname -sm", CHECK_TIMEOUT).await { + Ok(output) if output.status.success() => { + let stdout = String::from_utf8_lossy(&output.stdout); + setup::parse_uname_output(&stdout).map_err(|e| format!("{e:#}")) } - None => Err("binary check terminated by signal".into()), - }, - Err(e) => Err(format!("{e:#}")), - } - } - - async fn install_binary(&self) -> Result<(), String> { - let script = remote_server::setup::install_script(); - log::info!( - "Installing remote server binary to {}", - remote_server::setup::remote_server_binary() - ); - match remote_server::ssh::run_ssh_script( - &self.socket_path, - &script, - remote_server::setup::INSTALL_TIMEOUT, - ) - .await - { - Ok(output) if output.status.success() => Ok(()), - Ok(output) => { - let code = output.status.code().unwrap_or(-1); - let stderr = String::from_utf8_lossy(&output.stderr); - Err(format!("install script failed (exit {code}): {stderr}")) + Ok(output) => { + let code = output.status.code().unwrap_or(-1); + let stderr = String::from_utf8_lossy(&output.stderr); + Err(format!("uname -sm exited with code {code}: {stderr}")) + } + Err(e) => Err(format!("{e:#}")), } - Err(e) => Err(format!("{e:#}")), - } + }) } - async fn connect(&self, executor: &executor::Background) -> Result { - let binary = remote_server::setup::remote_server_binary(); - let mut args = remote_server::ssh::ssh_args(&self.socket_path); - args.push(format!("{binary} remote-server-proxy")); + fn check_binary(&self) -> Pin> + Send>> { + let socket_path = self.socket_path.clone(); + Box::pin(async move { + let bin_path = setup::remote_server_binary(); + log::info!("Checking for remote server binary at {bin_path}"); + match run_ssh_command(&socket_path, &setup::binary_check_command(), CHECK_TIMEOUT).await + { + Ok(output) => match output.status.code() { + Some(0) => Ok(true), + Some(1) => Ok(false), + Some(code) => { + let stderr = String::from_utf8_lossy(&output.stderr); + Err(format!("binary check exited with code {code}: {stderr}")) + } + None => Err("binary check terminated by signal".into()), + }, + Err(e) => Err(format!("{e:#}")), + } + }) + } - // `kill_on_drop(true)` pairs with ownership of the `Child` being - // returned in the [`Connection`] below: the - // [`RemoteServerManager`] holds the `Child` on its per-session - // state, and dropping that state (on explicit teardown or - // spontaneous disconnect) sends SIGKILL to this ssh process. - // Without this the ssh child is orphaned and keeps a channel - // open on the ControlMaster socket, blocking the master from - // exiting cleanly when the user logs out. - // - // Note that the child's lifetime is decoupled from any - // `Arc` clones: other owners (e.g. the - // per-session command executor) can keep the client alive for - // their own purposes without pinning the subprocess. - let mut child = command::r#async::Command::new("ssh") - .args(&args) - .stdin(std::process::Stdio::piped()) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .kill_on_drop(true) - .spawn()?; + fn install_binary(&self) -> Pin> + Send>> { + let socket_path = self.socket_path.clone(); + Box::pin(async move { + let script = setup::install_script(); + log::info!( + "Installing remote server binary to {}", + setup::remote_server_binary() + ); + match run_ssh_script(&socket_path, &script, INSTALL_TIMEOUT).await { + Ok(output) if output.status.success() => Ok(()), + Ok(output) => { + let code = output.status.code().unwrap_or(-1); + let stderr = String::from_utf8_lossy(&output.stderr); + Err(format!("install script failed (exit {code}): {stderr}")) + } + Err(e) => Err(format!("{e:#}")), + } + }) + } - let stdin = child - .stdin - .take() - .ok_or_else(|| anyhow::anyhow!("Failed to capture child stdin"))?; - let stdout = child - .stdout - .take() - .ok_or_else(|| anyhow::anyhow!("Failed to capture child stdout"))?; - let stderr = child - .stderr - .take() - .ok_or_else(|| anyhow::anyhow!("Failed to capture child stderr"))?; + fn connect( + &self, + executor: Arc, + ) -> Pin> + Send>> { + let socket_path = self.socket_path.clone(); + Box::pin(async move { + let binary = setup::remote_server_binary(); + let mut args = ssh_args(&socket_path); + args.push(format!("{binary} remote-server-proxy")); - let (client, event_rx) = - RemoteServerClient::from_child_streams(stdin, stdout, stderr, executor); - Ok(Connection { - client, - event_rx, - child, - control_path: Some(self.socket_path.clone()), + // `kill_on_drop(true)` pairs with ownership of the `Child` being + // returned in the [`Connection`] below: the + // [`RemoteServerManager`] holds the `Child` on its per-session + // state, and dropping that state (on explicit teardown or + // spontaneous disconnect) sends SIGKILL to this ssh process. + let mut child = command::r#async::Command::new("ssh") + .args(&args) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + + let stdin = child + .stdin + .take() + .ok_or_else(|| anyhow::anyhow!("Failed to capture child stdin"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("Failed to capture child stdout"))?; + let stderr = child + .stderr + .take() + .ok_or_else(|| anyhow::anyhow!("Failed to capture child stderr"))?; + + let (client, event_rx) = + RemoteServerClient::from_child_streams(stdin, stdout, stderr, &executor); + Ok(Connection { + client, + event_rx, + child, + control_path: Some(socket_path), + }) }) } } diff --git a/app/src/terminal/model/session.rs b/app/src/terminal/model/session.rs index b9911c47..723779c9 100644 --- a/app/src/terminal/model/session.rs +++ b/app/src/terminal/model/session.rs @@ -36,6 +36,9 @@ use crate::server::telemetry::{BootstrappingInfo, TelemetryEvent}; use crate::terminal::event::ExecutedExecutorCommandEvent; use crate::terminal::ShellHost; use crate::terminal::ShellLaunchData; +#[cfg(feature = "local_tty")] +use command_executor::remote_server_executor::RemoteServerCommandExecutor; +use parking_lot::{Mutex, RwLock}; use crate::terminal::shell::{Shell, ShellType}; use crate::terminal::warpify::SubshellSource; @@ -172,6 +175,18 @@ impl Sessions { | RemoteServerManagerEvent::BinaryInstallComplete { .. } | RemoteServerManagerEvent::ClientRequestFailed { .. } | RemoteServerManagerEvent::ServerMessageDecodingError { .. } => {} + RemoteServerManagerEvent::SessionReconnected { + session_id: sid, + client, + .. + } => { + if let Some(session) = sessions.sessions.get(sid) { + let new_executor = + Arc::new(RemoteServerCommandExecutor::new(*sid, client.clone())); + session.set_command_executor(new_executor); + log::info!("Swapped command executor for session {sid:?} after reconnect"); + } + } }); } #[cfg(not(feature = "local_tty"))] @@ -847,14 +862,16 @@ impl From for SessionType { pub struct Session { info: SessionInfo, external_commands: Arc>>, - command_executor: Arc, + /// The command executor for this session. Behind a `RwLock` so it can be + /// swapped after a remote server reconnect (via `set_command_executor`). + command_executor: RwLock>, load_external_commands_future: OnceCell>>, command_case_sensitivity: TopLevelCommandCaseSensitivity, /// The authoritative session type, initially derived from the /// [`BootstrapSessionType`] in `SessionInfo` and updated by [`Sessions`] /// when `RemoteServerManager` reports a connected session (to fill in the /// `host_id`). Interior mutability allows updating through `Arc`. - session_type: parking_lot::Mutex, + session_type: Mutex, } impl Session { @@ -873,10 +890,10 @@ impl Session { Self { info: session_info, external_commands: Arc::new(OnceCell::new()), - command_executor, + command_executor: RwLock::new(command_executor), load_external_commands_future: Default::default(), command_case_sensitivity, - session_type: parking_lot::Mutex::new(session_type), + session_type: Mutex::new(session_type), } } @@ -1025,9 +1042,17 @@ impl Session { &self.info.subshell_info } + /// Replaces the command executor for this session. Used after a remote + /// server reconnect to swap in a new `RemoteServerCommandExecutor` + /// backed by the reconnected client. + pub fn set_command_executor(&self, executor: Arc) { + *self.command_executor.write() = executor; + } + /// Returns true if the session is employing in-band command execution to run generators. pub fn is_using_in_band_command_execution(&self) -> bool { self.command_executor + .read() .as_ref() .as_any() .downcast_ref::() @@ -1089,8 +1114,8 @@ impl Session { .path .as_deref() .map(|path| HashMap::from_iter([("PATH".to_string(), path.to_string())])); - let windows_results = self - .command_executor + let executor = self.command_executor.read().clone(); + let windows_results = executor .execute_command( ShellType::PowerShell.shell_command_to_get_executables(), &Shell::new(ShellType::PowerShell, None, None, Default::default(), None), @@ -1381,7 +1406,10 @@ impl Session { environment_variables: Option>, execute_command_options: ExecuteCommandOptions, ) -> Result { - self.command_executor + // Clone the Arc out of the lock so we don't hold the read guard + // across the await point. + let executor = self.command_executor.read().clone(); + executor .execute_command( command, &self.info.shell, @@ -1394,11 +1422,13 @@ impl Session { /// Whether the backing executor for the session supports execution of commands in parallel. pub fn supports_parallel_command_execution(&self) -> bool { - self.command_executor.supports_parallel_command_execution() + self.command_executor + .read() + .supports_parallel_command_execution() } pub fn cancel_active_commands(&self) { - self.command_executor.cancel_active_commands(); + self.command_executor.read().cancel_active_commands(); } pub async fn git_branches_for_command_corrections(&self, working_dir: &str) -> Vec { @@ -1644,10 +1674,10 @@ pub mod testing { Self { info, external_commands: Default::default(), - command_executor: Arc::new(TestCommandExecutor::default()), + command_executor: RwLock::new(Arc::new(TestCommandExecutor::default())), load_external_commands_future: Default::default(), command_case_sensitivity: TopLevelCommandCaseSensitivity::CaseSensitive, - session_type: parking_lot::Mutex::new(session_type), + session_type: Mutex::new(session_type), } } @@ -1659,10 +1689,10 @@ pub mod testing { Self { info, external_commands: Default::default(), - command_executor: Arc::new(TestCommandExecutor::default()), + command_executor: RwLock::new(Arc::new(TestCommandExecutor::default())), load_external_commands_future: Default::default(), command_case_sensitivity: TopLevelCommandCaseSensitivity::CaseSensitive, - session_type: parking_lot::Mutex::new(session_type), + session_type: Mutex::new(session_type), } } diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index e8a47b0c..559c407c 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -4376,6 +4376,7 @@ impl TerminalView { } } RemoteServerManagerEvent::SessionConnecting { .. } + | RemoteServerManagerEvent::SessionReconnected { .. } | RemoteServerManagerEvent::HostConnected { .. } | RemoteServerManagerEvent::HostDisconnected { .. } | RemoteServerManagerEvent::RepoMetadataSnapshot { .. } diff --git a/crates/remote_server/Cargo.toml b/crates/remote_server/Cargo.toml index 91c6a0a0..791c812b 100644 --- a/crates/remote_server/Cargo.toml +++ b/crates/remote_server/Cargo.toml @@ -24,6 +24,7 @@ warp_util.workspace = true warpui.workspace = true [target.'cfg(not(target_family = "wasm"))'.dependencies] +async-io.workspace = true async-process.workspace = true [target.'cfg(target_family = "wasm")'.dependencies] diff --git a/crates/remote_server/src/manager.rs b/crates/remote_server/src/manager.rs index 96f8999d..541cf68c 100644 --- a/crates/remote_server/src/manager.rs +++ b/crates/remote_server/src/manager.rs @@ -2,6 +2,8 @@ use std::collections::{HashMap, HashSet}; #[cfg(not(target_family = "wasm"))] use std::path::PathBuf; use std::sync::Arc; +#[cfg(not(target_family = "wasm"))] +use std::time::Duration; #[cfg(not(target_family = "wasm"))] use crate::client::ClientEvent; @@ -17,6 +19,47 @@ use serde::Serialize; use warp_core::SessionId; use warpui::{Entity, ModelContext, ModelSpawner, SingletonEntity}; +/// Maximum number of reconnection attempts after a spontaneous disconnect. +#[cfg(not(target_family = "wasm"))] +const MAX_RECONNECT_ATTEMPTS: u32 = 2; +/// Delay between reconnection attempts. +#[cfg(not(target_family = "wasm"))] +const RECONNECT_DELAY: Duration = Duration::from_secs(2); + +/// Parameters that travel together through the reconnection flow. +#[cfg(not(target_family = "wasm"))] +struct ReconnectParams { + attempt: u32, + host_id: HostId, + exit_status: Option, + transport: Arc, + control_path: Option, +} + +/// Error from [`RemoteServerManager::run_connect_and_handshake`] that +/// preserves which phase failed so callers can report accurate telemetry. +#[cfg(not(target_family = "wasm"))] +#[derive(Debug, thiserror::Error)] +enum ConnectAndHandshakeError { + /// `transport.connect()` failed, or the session was deregistered + /// before the connect phase could complete. + #[error("connect: {0:#}")] + Connect(anyhow::Error), + /// `client.initialize()` handshake failed. + #[error("initialize: {0:#}")] + Initialize(anyhow::Error), +} + +#[cfg(not(target_family = "wasm"))] +impl ConnectAndHandshakeError { + fn phase(&self) -> RemoteServerInitPhase { + match self { + Self::Connect(_) => RemoteServerInitPhase::Connect, + Self::Initialize(_) => RemoteServerInitPhase::Initialize, + } + } +} + /// Which phase of the remote server connection flow failed. #[derive(Clone, Copy, Debug, Serialize)] #[serde(rename_all = "snake_case")] @@ -45,6 +88,16 @@ pub enum RemoteServerErrorKind { Other, } +/// Exit status information captured from the remote server subprocess +/// when the connection drops. Used for diagnostics and telemetry. +#[derive(Clone, Debug, Serialize)] +pub struct RemoteServerExitStatus { + /// Process exit code, if the process exited normally. + pub code: Option, + /// True if the process was killed by a signal (Unix only). + pub signal_killed: bool, +} + impl RemoteServerErrorKind { /// Classify a [`ClientError`] into a telemetry error kind. pub fn from_client_error(error: &crate::client::ClientError) -> Self { @@ -104,6 +157,16 @@ pub enum RemoteSessionState { /// See type-level doc. #[cfg(not(target_family = "wasm"))] control_path: Option, + /// Transport stored for reconnection after spontaneous disconnect. + #[cfg(not(target_family = "wasm"))] + transport: Arc, + }, + /// A reconnection attempt is in progress after a spontaneous disconnect. + #[cfg(not(target_family = "wasm"))] + Reconnecting { + attempt: u32, + host_id: HostId, + control_path: Option, }, /// Connection dropped (EOF/error from the reader task). Disconnected, @@ -145,6 +208,19 @@ pub enum RemoteServerManagerEvent { SessionDisconnected { session_id: SessionId, host_id: HostId, + /// Exit status of the remote server subprocess, if available. + /// `None` when the session was explicitly deregistered or when + /// the exit status could not be determined. + exit_status: Option, + }, + /// A reconnection attempt succeeded. Downstream owners (e.g. + /// `RemoteServerCommandExecutor`) should swap their client reference + /// to the new one carried in `client`. + SessionReconnected { + session_id: SessionId, + host_id: HostId, + attempt: u32, + client: Arc, }, /// The manager is no longer tracking this session -- it has been /// removed from the `sessions` map via `deregister_session`. Fires @@ -224,11 +300,13 @@ pub enum RemoteServerManagerEvent { ServerMessageDecodingError { session_id: SessionId }, } -/// Shell info stashed by [`RemoteServerManager::notify_session_bootstrapped`] -/// when the session is not yet in `Connected` state. Flushed automatically -/// when [`RemoteServerManager::mark_session_connected`] fires. +/// Shell info recorded by [`RemoteServerManager::notify_session_bootstrapped`]. +/// +/// Persists for the lifetime of the session (removed only in +/// `deregister_session`) so that `mark_session_connected` can re-send +/// the notification after a reconnect. #[cfg_attr(target_family = "wasm", allow(dead_code))] -struct PendingSessionBootstrappedNotification { +struct SessionBootstrapInfo { shell_type: String, shell_path: Option, } @@ -253,9 +331,10 @@ pub struct RemoteServerManager { /// `navigate_to_directory` calls when `update_active_session` fires /// repeatedly for the same CWD. last_navigated_path: HashMap, - /// Per-session `SessionBootstrapped` notifications that arrived before the - /// session reached `Connected`. Flushed in `mark_session_connected`. - pending_bootstrapped_notifications: HashMap, + /// Per-session shell info recorded at bootstrap time and re-sent to the + /// remote server daemon on every (re)connect. Persists until + /// `deregister_session`. + session_bootstrap_info: HashMap, /// Detected remote platform per session, populated during the binary check /// phase via `detect_platform()`. Used for telemetry. session_platforms: HashMap, @@ -274,7 +353,7 @@ impl RemoteServerManager { host_to_sessions: HashMap::new(), spawner: ctx.spawner(), last_navigated_path: HashMap::new(), - pending_bootstrapped_notifications: HashMap::new(), + session_bootstrap_info: HashMap::new(), session_platforms: HashMap::new(), } } @@ -417,10 +496,7 @@ impl RemoteServerManager { { log::info!("Starting remote server connection for session {session_id:?}"); - // Advance the user-visible setup pipeline. Both callers (binary - // already installed, and binary just installed) enter this - // method right when the Initializing phase begins, so we emit - // the state change from one place. + // Advance the user-visible setup pipeline. ctx.emit(RemoteServerManagerEvent::SetupStateChanged { session_id, state: RemoteServerSetupState::Initializing, @@ -432,99 +508,31 @@ impl RemoteServerManager { let spawner = self.spawner.clone(); let executor = ctx.background_executor().clone(); + // Wrap the transport in an Arc so it can be stored on `Connected` + // for reconnection after a spontaneous disconnect. + let transport: Arc = Arc::new(transport); ctx.background_executor() .spawn(async move { - // ---- Phase 1: Connect (establish streams, create client) ---- - match transport.connect(&executor).await { - Ok(Connection { - client, - event_rx, - child, - control_path, - }) => { - let client = Arc::new(client); - - // Transition to Initializing and start draining - // the event channel for push events and disconnect. - // The `Child` is stashed on the session state so - // its lifetime is controlled by the manager -- on - // teardown the state is dropped, which runs the - // `Child`'s destructor and SIGKILLs the subprocess - // via `kill_on_drop`. `control_path` is stashed - // for explicit teardown's `ssh -O exit` call. - let client_for_state = Arc::clone(&client); + match Self::run_connect_and_handshake( + session_id, + &*transport, + &spawner, + &executor, + ) + .await + { + Ok(host_id) => { let _ = spawner .spawn(move |me, ctx| { - me.sessions.insert( - session_id, - RemoteSessionState::Initializing { - client: client_for_state, - _child: child, - control_path, - }, - ); - - // Drain the event channel on the main thread. - // Each push event is forwarded as a manager - // event in real-time. When the stream closes - // (after Disconnected or channel drop), we - // transition the session to Disconnected. - ctx.spawn_stream_local( - event_rx, - move |me, event, ctx| { - me.forward_client_event(session_id, event, ctx); - }, - move |me, ctx| { - me.mark_session_disconnected(session_id, ctx); - }, - ); + me.mark_session_connected(session_id, host_id, transport, ctx); }) .await; - - // ---- Phase 2: Initialize handshake ---- - match client.initialize().await { - Ok(resp) => { - let host_id = HostId::new(resp.host_id); - let _ = spawner - .spawn(move |me, ctx| { - me.mark_session_connected(session_id, host_id, ctx); - }) - .await; - } - Err(e) => { - log::error!( - "Initialize handshake failed for session {session_id:?}: {e}" - ); - let error = format!("{e:#}"); - let _ = spawner - .spawn(move |me, ctx| { - ctx.emit( - RemoteServerManagerEvent::SetupStateChanged { - session_id, - state: RemoteServerSetupState::Failed { - error: error.clone(), - }, - }, - ); - ctx.emit( - RemoteServerManagerEvent::SessionConnectionFailed { - session_id, - phase: RemoteServerInitPhase::Initialize, - error, - }, - ); - me.mark_session_disconnected(session_id, ctx); - }) - .await; - } - } } Err(e) => { - log::error!( - "Failed to connect remote server for session {session_id:?}: {e:#}" - ); - let error = format!("{e:#}"); + log::error!("Connection failed for session {session_id:?}: {e}"); + let phase = e.phase(); + let error = format!("{e}"); let _ = spawner .spawn(move |me, ctx| { ctx.emit(RemoteServerManagerEvent::SetupStateChanged { @@ -535,7 +543,7 @@ impl RemoteServerManager { }); ctx.emit(RemoteServerManagerEvent::SessionConnectionFailed { session_id, - phase: RemoteServerInitPhase::Connect, + phase, error, }); me.mark_session_disconnected(session_id, ctx); @@ -548,6 +556,81 @@ impl RemoteServerManager { } } + /// Shared connect + handshake logic used by both `connect_session` and + /// `attempt_reconnect`. + /// + /// 1. Calls `transport.connect()` to establish streams. + /// 2. Transitions the session to `Initializing` and starts draining the + /// event channel. + /// 3. Runs the initialize handshake. + /// + /// Returns `Ok(host_id)` on success, or a phase-tagged error. + #[cfg(not(target_family = "wasm"))] + async fn run_connect_and_handshake( + session_id: SessionId, + transport: &dyn RemoteTransport, + spawner: &ModelSpawner, + executor: &Arc, + ) -> Result { + // Phase 1: Connect (establish streams, create client). + let Connection { + client, + event_rx, + child, + control_path, + } = transport + .connect(executor.clone()) + .await + .map_err(ConnectAndHandshakeError::Connect)?; + + let client = Arc::new(client); + let client_for_init = Arc::clone(&client); + + // Transition to Initializing and start draining the event channel. + // Guard: if the session was deregistered during `transport.connect()`, + // the entry will have been removed; don't re-insert it. + let was_inserted = spawner + .spawn(move |me, ctx| { + if !me.sessions.contains_key(&session_id) { + return false; + } + me.sessions.insert( + session_id, + RemoteSessionState::Initializing { + client: client_for_init, + _child: child, + control_path, + }, + ); + + ctx.spawn_stream_local( + event_rx, + move |me, event, ctx| { + me.forward_client_event(session_id, event, ctx); + }, + move |me, ctx| { + me.mark_session_disconnected(session_id, ctx); + }, + ); + true + }) + .await + .unwrap_or(false); + + if !was_inserted { + return Err(ConnectAndHandshakeError::Connect(anyhow::anyhow!( + "Session {session_id:?} was deregistered during connect" + ))); + } + + // Phase 2: Initialize handshake. + let resp = client + .initialize() + .await + .map_err(|e| ConnectAndHandshakeError::Initialize(anyhow::anyhow!("{e:#}")))?; + Ok(HostId::new(resp.host_id)) + } + /// Removes a session from the manager and tears down its connection. /// /// Assumes the caller has already observed that the user's shell @@ -587,7 +670,7 @@ impl RemoteServerManager { /// spontaneous drops -- only for explicit teardown. pub fn deregister_session(&mut self, session_id: SessionId, ctx: &mut ModelContext) { self.last_navigated_path.remove(&session_id); - self.pending_bootstrapped_notifications.remove(&session_id); + self.session_bootstrap_info.remove(&session_id); self.session_platforms.remove(&session_id); // Remove the session entry. Dropping the `RemoteSessionState` @@ -604,19 +687,26 @@ impl RemoteServerManager { let control_path = match &prev { Some(RemoteSessionState::Connected { control_path, .. }) | Some(RemoteSessionState::Initializing { control_path, .. }) => control_path.clone(), + Some(RemoteSessionState::Reconnecting { control_path, .. }) => control_path.clone(), _ => None, }; - if let Some(RemoteSessionState::Connected { host_id, .. }) = prev { + // Extract `host_id` from states that track a host connection. + let host_id = match &prev { + Some(RemoteSessionState::Connected { host_id, .. }) => Some(host_id.clone()), + #[cfg(not(target_family = "wasm"))] + Some(RemoteSessionState::Reconnecting { host_id, .. }) => Some(host_id.clone()), + _ => None, + }; + if let Some(host_id) = host_id { self.remove_from_host_index(&host_id, session_id); ctx.emit(RemoteServerManagerEvent::SessionDisconnected { session_id, host_id: host_id.clone(), + exit_status: None, }); if !self.host_to_sessions.contains_key(&host_id) { - ctx.emit(RemoteServerManagerEvent::HostDisconnected { - host_id: host_id.clone(), - }); + ctx.emit(RemoteServerManagerEvent::HostDisconnected { host_id }); } } ctx.emit(RemoteServerManagerEvent::SessionDeregistered { session_id }); @@ -742,19 +832,21 @@ impl RemoteServerManager { shell_type: &str, shell_path: Option<&str>, ) { + // Always persist so we can re-send after a reconnect. + self.session_bootstrap_info.insert( + session_id, + SessionBootstrapInfo { + shell_type: shell_type.to_owned(), + shell_path: shell_path.map(ToOwned::to_owned), + }, + ); + if let Some(client) = self.client_for_session(session_id) { client.notify_session_bootstrapped(session_id, shell_type, shell_path); } else { log::info!( "notify_session_bootstrapped: session {session_id:?} not yet connected, \ - stashing notification" - ); - self.pending_bootstrapped_notifications.insert( - session_id, - PendingSessionBootstrappedNotification { - shell_type: shell_type.to_owned(), - shell_path: shell_path.map(ToOwned::to_owned), - }, + will send on connect" ); } } @@ -856,18 +948,19 @@ impl RemoteServerManager { } } + /// Transitions a session from `Initializing` to `Connected`. Stores the + /// `transport` for reconnection support after a spontaneous disconnect. #[cfg(not(target_family = "wasm"))] fn mark_session_connected( &mut self, session_id: SessionId, host_id: HostId, + transport: Arc, ctx: &mut ModelContext, ) { log::info!("Remote server connected for session {session_id:?}, host {host_id}"); // Only transition if the session is still in Initializing state. - // Remove first so we can move the client handle (and owned `Child`) - // out. let Some(RemoteSessionState::Initializing { client, _child, @@ -881,10 +974,11 @@ impl RemoteServerManager { self.sessions.insert( session_id, RemoteSessionState::Connected { - client, + client: client.clone(), host_id: host_id.clone(), _child, control_path, + transport, }, ); self.host_to_sessions @@ -905,51 +999,275 @@ impl RemoteServerManager { host_id, }); - // Flush any SessionBootstrapped notification that was stashed before - // the session reached Connected. - if let Some(notif) = self.pending_bootstrapped_notifications.remove(&session_id) { + // (Re-)send the SessionBootstrapped notification so the daemon + // registers an executor for this session. This fires on both the + // initial connect and every reconnect. + if let Some(info) = self.session_bootstrap_info.get(&session_id) { if let Some(client) = self.client_for_session(session_id) { - log::info!( - "Flushing stashed SessionBootstrapped notification for session \ - {session_id:?}" - ); + log::info!("Sending SessionBootstrapped notification for session {session_id:?}"); client.notify_session_bootstrapped( session_id, - ¬if.shell_type, - notif.shell_path.as_deref(), + &info.shell_type, + info.shell_path.as_deref(), ); } } } + /// Captures the exit status from a `Child` process, if available. + #[cfg(not(target_family = "wasm"))] + fn capture_exit_status( + child: &mut async_process::Child, + session_id: SessionId, + ) -> Option { + match child.try_status() { + Ok(Some(status)) => { + let code = status.code(); + #[cfg(unix)] + let signal_killed = { + use std::os::unix::process::ExitStatusExt; + status.signal().is_some() + }; + #[cfg(not(unix))] + let signal_killed = false; + log::warn!( + "Remote server process exited for session {session_id:?}: \ + code={code:?}, signal_killed={signal_killed}" + ); + Some(RemoteServerExitStatus { + code, + signal_killed, + }) + } + Ok(None) => { + log::warn!( + "Remote server process still running for session {session_id:?} \ + despite EOF on reader task" + ); + None + } + Err(e) => { + log::warn!("Failed to read exit status for session {session_id:?}: {e}"); + None + } + } + } + #[cfg(not(target_family = "wasm"))] pub(crate) fn mark_session_disconnected( &mut self, session_id: SessionId, ctx: &mut ModelContext, ) { - self.pending_bootstrapped_notifications.remove(&session_id); let Some(prev) = self.sessions.remove(&session_id) else { return; }; - self.sessions - .insert(session_id, RemoteSessionState::Disconnected); - if let RemoteSessionState::Connected { host_id, .. } = prev { + // Only attempt reconnect for sessions that were in Connected state + // with a transport available, and not being explicitly deregistered. + if let RemoteSessionState::Connected { + host_id, + mut _child, + control_path, + transport, + .. + } = prev + { + let exit_status = Self::capture_exit_status(&mut _child, session_id); + // Drop the old child process explicitly before reconnecting. + drop(_child); + + log::info!( + "Spontaneous disconnect for session {session_id:?}, \ + will attempt reconnect (transport={transport:?})" + ); + + // Clear stale repo metadata and host index so downstream + // models don't hold onto data from the dead server process. self.remove_from_host_index(&host_id, session_id); - // Emit `SessionDisconnected` before `HostDisconnected` so that - // subscribers (e.g. the command executor) drop their - // `Arc` reference before any host-scoped - // teardown runs. This matches the ordering in - // `deregister_session` so both teardown paths look identical - // to subscribers. + if !self.host_to_sessions.contains_key(&host_id) { + ctx.emit(RemoteServerManagerEvent::HostDisconnected { + host_id: host_id.clone(), + }); + } + + // Clear last navigated path so navigate_to_directory + // re-fires after reconnect. + // We need to do this on disconnect because the cached + // navigated path is only deduping for the current _remote server session. + self.last_navigated_path.remove(&session_id); + + self.attempt_reconnect( + session_id, + ReconnectParams { + attempt: 1, + host_id, + exit_status, + transport, + control_path, + }, + ctx, + ); + } else { + // Non-Connected states (Initializing, Connecting, etc.) — + // no reconnect, just mark disconnected. + self.sessions + .insert(session_id, RemoteSessionState::Disconnected); + } + } + + /// Attempt to re-establish the remote server connection. + #[cfg(not(target_family = "wasm"))] + fn attempt_reconnect( + &mut self, + session_id: SessionId, + params: ReconnectParams, + ctx: &mut ModelContext, + ) { + let ReconnectParams { + attempt, + host_id, + exit_status, + transport, + control_path, + } = params; + + log::info!( + "Attempting reconnect for session {session_id:?} \ + (attempt {attempt}/{MAX_RECONNECT_ATTEMPTS})" + ); + + self.sessions.insert( + session_id, + RemoteSessionState::Reconnecting { + attempt, + host_id: host_id.clone(), + control_path: control_path.clone(), + }, + ); + + let spawner = self.spawner.clone(); + let executor = ctx.background_executor().clone(); + let transport_clone = Arc::clone(&transport); + + ctx.background_executor() + .spawn(async move { + async_io::Timer::after(RECONNECT_DELAY).await; + + // Check if the session was deregistered during the delay. + // (Checked via spawner since sessions lives on the main thread.) + let was_removed = spawner + .spawn(move |me, _ctx| !me.sessions.contains_key(&session_id)) + .await + .unwrap_or(true); + if was_removed { + log::info!("Session {session_id:?} removed during reconnect delay, aborting"); + return; + } + + match Self::run_connect_and_handshake( + session_id, + &*transport_clone, + &spawner, + &executor, + ) + .await + { + Ok(new_host_id) => { + let _ = spawner + .spawn(move |me, ctx| { + // If the session was deregistered during the + // handshake, don't resurrect it. + if !me.sessions.contains_key(&session_id) { + log::info!( + "Session {session_id:?} deregistered during \ + reconnect handshake, aborting" + ); + return; + } + me.mark_session_connected( + session_id, + new_host_id.clone(), + transport, + ctx, + ); + if let Some(client) = me.client_for_session(session_id).cloned() { + ctx.emit(RemoteServerManagerEvent::SessionReconnected { + session_id, + host_id: new_host_id, + attempt, + client, + }); + } + }) + .await; + } + Err(e) => { + log::error!( + "Reconnect failed for session {session_id:?} \ + (attempt {attempt}): {e}" + ); + let _ = spawner + .spawn(move |me, ctx| { + // If the session was deregistered during the + // handshake, don't retry or insert Disconnected. + if !me.sessions.contains_key(&session_id) { + log::info!( + "Session {session_id:?} deregistered during \ + reconnect handshake, aborting" + ); + return; + } + me.handle_reconnect_failure( + session_id, + ReconnectParams { + attempt, + host_id, + exit_status, + transport, + control_path, + }, + ctx, + ); + }) + .await; + } + } + }) + .detach(); + } + + /// Handle a failed reconnection attempt: either retry or give up. + #[cfg(not(target_family = "wasm"))] + fn handle_reconnect_failure( + &mut self, + session_id: SessionId, + params: ReconnectParams, + ctx: &mut ModelContext, + ) { + if params.attempt < MAX_RECONNECT_ATTEMPTS { + self.attempt_reconnect( + session_id, + ReconnectParams { + attempt: params.attempt + 1, + ..params + }, + ctx, + ); + } else { + log::warn!( + "Reconnect exhausted for session {session_id:?} after {} attempt(s)", + params.attempt + ); + self.sessions + .insert(session_id, RemoteSessionState::Disconnected); ctx.emit(RemoteServerManagerEvent::SessionDisconnected { session_id, - host_id: host_id.clone(), + host_id: params.host_id, + exit_status: params.exit_status, }); - if !self.host_to_sessions.contains_key(&host_id) { - ctx.emit(RemoteServerManagerEvent::HostDisconnected { host_id }); - } + // Note: HostDisconnected was already emitted by + // mark_session_disconnected when entering the reconnect flow. } } diff --git a/crates/remote_server/src/transport.rs b/crates/remote_server/src/transport.rs index 08b9f450..60e534d2 100644 --- a/crates/remote_server/src/transport.rs +++ b/crates/remote_server/src/transport.rs @@ -6,13 +6,13 @@ //! in-process for tests) implement the same trait without touching the //! manager. //! -//! Methods are async. Callers use the trait via generics -//! (`T: RemoteTransport`) rather than `dyn` dispatch. +//! Returns boxed futures for object safety — the manager stores +//! `Arc` for reconnection. //! //! [`RemoteServerManager`]: crate::manager::RemoteServerManager -use std::future::Future; #[cfg(not(target_family = "wasm"))] use std::path::PathBuf; +use std::pin::Pin; use async_channel::Receiver; use warpui::r#async::executor; @@ -56,12 +56,18 @@ pub struct Connection { pub control_path: Option, } -pub trait RemoteTransport: Send + Sync { +/// Transport abstraction for remote server connections. +/// +/// Object-safe: returns boxed futures so implementations can be stored +/// as `Arc` for reconnection. +pub trait RemoteTransport: Send + Sync + std::fmt::Debug { /// Detects the remote host's OS and architecture by running `uname -sm`. /// /// Returns the parsed [`RemotePlatform`] on success, or an error string /// if the command fails or the output cannot be parsed. - fn detect_platform(&self) -> impl Future> + Send; + fn detect_platform( + &self, + ) -> Pin> + Send>>; /// Checks whether the remote server binary is present on the remote host. /// @@ -72,7 +78,9 @@ pub trait RemoteTransport: Send + Sync { /// Returns `Ok(true)` if the binary is installed and executable, /// `Ok(false)` if it is definitively not installed, and /// `Err(_)` if the check failed (e.g. SSH timeout/unreachable). - fn check_binary(&self) -> impl Future> + Send; + fn check_binary( + &self, + ) -> Pin> + Send>>; /// Installs the remote server binary on the remote host. /// @@ -82,7 +90,9 @@ pub trait RemoteTransport: Send + Sync { /// /// Returns `Ok(())` if the install succeeded, and /// `Err(_)` if the install failed (e.g. SSH timeout, script error). - fn install_binary(&self) -> impl Future> + Send; + fn install_binary( + &self, + ) -> Pin> + Send>>; /// Establish a new connection to the remote server. /// @@ -96,6 +106,6 @@ pub trait RemoteTransport: Send + Sync { /// a socket). Stderr forwarding to local logging should also happen here. fn connect( &self, - executor: &executor::Background, - ) -> impl Future> + Send; + executor: std::sync::Arc, + ) -> Pin> + Send>>; } diff --git a/specs/APP-4283/TECH.md b/specs/APP-4283/TECH.md new file mode 100644 index 00000000..cd6fd7aa --- /dev/null +++ b/specs/APP-4283/TECH.md @@ -0,0 +1,119 @@ +# APP-4283: Remote Server Reconnection on Spontaneous Disconnect + +## Context + +When the SSH remote server connection drops spontaneously (daemon crash, proxy killed, transient network failure), the `RemoteServerClient` reader task hits EOF and calls `mark_session_disconnected`. Before this change the session transitions straight to `Disconnected` and stays there — completions, repo metadata, and all other remote-server-backed features stop working until the user manually exits and re-SSHes. + +The remote server architecture uses a **proxy → daemon** model: + +- `SshTransport` (`app/src/remote_server/ssh_transport.rs`) spawns `ssh … remote-server-proxy` whose stdin/stdout become the protocol channel. +- The proxy (`app/src/remote_server/unix/proxy.rs:34`) connects to a long-lived daemon via a local Unix socket. The daemon survives across proxy lifetimes. + +The core session lifecycle lives in `RemoteServerManager` (`crates/remote_server/src/manager.rs`), a singleton model with per-session state tracked via `RemoteSessionState`. The state machine before this change: + +``` +Connecting → Initializing → Connected → Disconnected +``` + +Key files: + +- `crates/remote_server/src/manager.rs` — `RemoteServerManager`, session state machine, connect/disconnect lifecycle +- `crates/remote_server/src/transport.rs` — `RemoteTransport` trait, `Connection` struct +- `app/src/remote_server/ssh_transport.rs` — SSH `RemoteTransport` implementation +- `app/src/terminal/model/session.rs (135-191)` — `Sessions` model, subscribes to manager events, owns the `Session` and its `CommandExecutor` +- `app/src/terminal/model/session.rs (862-1050)` — `Session` struct, holds the `command_executor` used by completions + +## Proposed changes + +### 1. Object-safe `RemoteTransport` for reconnection + +The `RemoteTransport` trait (`crates/remote_server/src/transport.rs:63`) now returns boxed futures (`Pin + Send>>`) for object safety. This allows `RemoteServerManager` to store `Arc` on the `Connected` state and carry it forward through reconnection without knowing the concrete transport type. + +`SshTransport` (`app/src/remote_server/ssh_transport.rs:36`) implements the trait directly with `Box::pin(async move { … })`. + +### 2. Exit status capture + +New `RemoteServerExitStatus` type (`manager.rs:67`) records `code: Option` and `signal_killed: bool`. + +`capture_exit_status()` (`manager.rs:982`) reads `child.try_status()` before the `Child` is dropped on disconnect. The result is carried on `SessionDisconnected.exit_status` for diagnostics and telemetry. + +### 3. Reconnection state machine + +New state variant `Reconnecting { attempt, host_id, control_path }` (`manager.rs:139`). The state machine becomes: + +``` +Connecting → Initializing → Connected → Reconnecting → Initializing → Connected + ↘ Disconnected (if retries exhausted) +``` + +Constants: `MAX_RECONNECT_ATTEMPTS = 2`, `RECONNECT_DELAY = 2s` (`manager.rs:22-24`). + +`mark_session_disconnected()` (`manager.rs:1020`) checks whether the session was `Connected` with a stored transport. If so it clears stale host-index and repo-metadata state, then calls `attempt_reconnect()`. + +`attempt_reconnect()` (`manager.rs:1100`) transitions to `Reconnecting`, waits `RECONNECT_DELAY` via `async_io::Timer`, then calls the shared `run_connect_and_handshake()`. On success it calls `mark_session_connected()` and emits `SessionReconnected`. On failure, `handle_reconnect_failure()` (`manager.rs:1203`) either increments the attempt and retries, or gives up and emits `SessionDisconnected`. + +The retry parameters are bundled in `ReconnectParams` (`manager.rs:28`) to stay under clippy's argument limit. + +### 4. Shared connect + handshake helper + +`run_connect_and_handshake()` (`manager.rs:546`) is extracted from `connect_session()` so both the initial connect and reconnect share the same two-phase logic: `transport.connect()` → `Initializing` → `client.initialize()` → `HostId`. + +### 5. Session bootstrap info persistence + +Previously, `notify_session_bootstrapped()` stashed shell info in a `pending_bootstrapped_notifications` map that was consumed on the first connect and wiped on disconnect. After reconnect, the daemon had no executor registered for the session (`SessionNotFound` error on completions). + +Fix: renamed to `session_bootstrap_info` (`manager.rs:310`), which persists for the session lifetime (removed only in `deregister_session`). `notify_session_bootstrapped()` (`manager.rs:790`) always stores the info. `mark_session_connected()` (`manager.rs:965`) re-sends the `SessionBootstrapped` notification on every connect/reconnect. + +### 6. Command executor swap on reconnect + +`Session.command_executor` (`session.rs:867`) changed from `Arc` to `RwLock>` for interior mutability through `Arc`. New `Session::set_command_executor()` (`session.rs:1048`). + +All read sites (e.g. `execute_command` at `session.rs:1117`, `load_external_commands` at `session.rs:1083`) clone the `Arc` out of the lock before `.await` to avoid holding the guard across await points. + +The `Sessions` model subscribes to `SessionReconnected` (`session.rs:178-189`) and swaps in a new `RemoteServerCommandExecutor` backed by the reconnected client. + +### 7. Downstream match arm updates + +Added exhaustive match arms for `SessionReconnected` in `Sessions` subscriber (`session.rs:178`) and `terminal/view.rs`. + +## Diagram + +```mermaid +stateDiagram-v2 + [*] --> Connecting: connect_session() + Connecting --> Initializing: transport.connect() ok + Initializing --> Connected: client.initialize() ok + Connected --> Reconnecting: spontaneous EOF + transport available + Reconnecting --> Initializing: transport.connect() ok (after delay) + Reconnecting --> Reconnecting: attempt < MAX (after delay) + Reconnecting --> Disconnected: retries exhausted + Connected --> Disconnected: no transport / explicit deregister + Initializing --> Disconnected: handshake failed +``` + +## Testing and validation + +Manual E2E (verified): + +1. Build Warp from this branch (`cargo run`). +2. SSH into a remote host with the remote server feature flag enabled. +3. Kill the `remote-server-proxy` on the remote side: `ssh "pkill -f remote-server-proxy"`. +4. Observe in `Warp.log`: + - `"Remote server process exited for session …"` — exit status captured + - `"Spontaneous disconnect for session …, will attempt reconnect"` — reconnect triggered + - `"Attempting reconnect for session … (attempt 1/2)"` — delay + retry + - `"Remote server connected for session …"` — reconnect succeeded + - `"Sending SessionBootstrapped notification for session …"` — daemon re-registered + - `"Swapped command executor for session … after reconnect"` — executor swapped +5. Press Tab in the SSH session to trigger completions — should work. +6. Verify `navigate_to_directory` re-fires (repo metadata restored). + +Edge cases verified manually: +- Kill the daemon (`pkill -f remote-server-daemon`) — harder reconnect since daemon must restart; proxy's `run()` re-spawns it. +- `deregister_session` during reconnect delay — reconnect aborts cleanly (checked via `sessions.contains_key`). + +## Follow-ups + +- **User-visible reconnecting indicator**: surface the `Reconnecting` state in the terminal UI so the user knows a retry is in progress. +- **Telemetry**: emit a structured event with `exit_status`, `attempt`, and `reconnect_succeeded` for reconnect outcomes. +- **Exponential backoff**: the current fixed 2s delay works for the proxy-restart case; longer backoffs may be warranted for network-level failures.