public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add reconnection logic to remote server (#9289)

Browse Source 
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
This commit is contained in:
Yunfan Yang
2026-04-28 18:02:55 -04:00
committed by GitHub
parent 00df35b5dc
commit 2d0d88fea6
8 changed files with 743 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/remote_server/ssh_transport.rs
+105 -111
View File
@@ -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=<path>`). 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<RemotePlatform, String> {
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<bool, String> {
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<Box<dyn Future<Output = Result<RemotePlatform, String>> + 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<Connection> {
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<Box<dyn Future<Output = Result<bool, String>> + 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<RemoteServerClient>` 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<Box<dyn Future<Output = Result<(), String>> + 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<executor::Background>,
) -> Pin<Box<dyn Future<Output = Result<Connection>> + 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),
})
})
}
}
app/src/terminal/model/session.rs
+43 -13
View File
@@ -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<BootstrapSessionType> for SessionType {
pub struct Session {
info: SessionInfo,
external_commands: Arc<OnceCell<HashSet<SmolStr>>>,
command_executor: Arc<dyn CommandExecutor>,
/// 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<Arc<dyn CommandExecutor>>,
load_external_commands_future: OnceCell<Shared<BoxFuture<'static, ()>>>,
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>`.
session_type: parking_lot::Mutex<SessionType>,
session_type: Mutex<SessionType>,
}
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<dyn CommandExecutor>) {
*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::<InBandCommandExecutor>()
@@ -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<HashMap<String, String>>,
execute_command_options: ExecuteCommandOptions,
) -> Result<CommandOutput> {
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<String> {
@@ -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),
}
}
app/src/terminal/view.rs
+1
View File
@@ -4376,6 +4376,7 @@ impl TerminalView {
}
}
RemoteServerManagerEvent::SessionConnecting { .. }
| RemoteServerManagerEvent::SessionReconnected { .. }
| RemoteServerManagerEvent::HostConnected { .. }
| RemoteServerManagerEvent::HostDisconnected { .. }
| RemoteServerManagerEvent::RepoMetadataSnapshot { .. }