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
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
+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),
})
})
}
}