feat: hand off ptys during updates

1 files changed, 233 insertions, 11 deletions

diff --git a/src/remote.rs b/src/remote.rs
index 237f06de..990ffb7b 100644
--- a/src/remote.rs
+++ b/src/remote.rs
@@ -322,11 +322,15 @@ fn prepare_remote_herdr(target: &str) -> io::Result<PreparedRemoteHerdr> {
     let platform = detect_remote_platform(target)?;
     let remote_herdr = RemoteHerdr::for_platform(platform);
     let override_binary = remote_binary_override_path()?;
+    let path_remote_herdr = remote_binary_on_path_any(target, &remote_herdr)?;
 
     if override_binary.is_none() {
-        if let Some(path_remote_herdr) = remote_binary_on_path(target, &remote_herdr)? {
+        if let Some(path_remote_herdr) = path_remote_herdr
+            .as_ref()
+            .filter(|candidate| remote_binary_matches(target, candidate).unwrap_or(false))
+        {
             return Ok(PreparedRemoteHerdr {
-                remote_herdr: path_remote_herdr,
+                remote_herdr: path_remote_herdr.clone(),
                 installed_or_replaced: false,
             });
         }
@@ -338,6 +342,13 @@ fn prepare_remote_herdr(target: &str) -> io::Result<PreparedRemoteHerdr> {
         }
     }
 
+    if let Some(status_probe_herdr) = path_remote_herdr.as_ref().or_else(|| {
+        remote_binary_exists(target, &remote_herdr)
+            .ok()
+            .and_then(|exists| exists.then_some(&remote_herdr))
+    }) {
+        confirm_remote_install_with_running_server(target, status_probe_herdr)?;
+    }
     confirm_remote_install(
         target,
         &remote_herdr,
@@ -381,28 +392,27 @@ fn detect_remote_platform(target: &str) -> io::Result<RemotePlatform> {
     })
 }
 
-fn remote_binary_on_path(
+fn remote_binary_on_path_any(
     target: &str,
     remote_herdr: &RemoteHerdr,
 ) -> io::Result<Option<RemoteHerdr>> {
-    let output = ssh_output(target, remote_path_probe_command())?;
+    let output = ssh_output(target, remote_path_probe_any_command())?;
     if !output.status.success() {
         return Ok(None);
     }
 
     let stdout = String::from_utf8_lossy(&output.stdout);
-    Ok(remote_herdr_from_path_probe(remote_herdr, &stdout))
+    Ok(remote_herdr_from_path_probe_any(remote_herdr, &stdout))
 }
 
-fn remote_path_probe_command() -> &'static str {
+fn remote_path_probe_any_command() -> &'static str {
     r#"path=$(command -v herdr) || exit 1
 test -n "$path" || exit 1
-version=$("$path" --version) || exit 1
-status=$("$path" status client --json) || exit 1
-printf '%s\n%s\n%s\n' "$path" "$version" "$status"
+printf '%s\n' "$path"
 "#
 }
 
+#[cfg(test)]
 fn remote_herdr_from_path_probe(remote_herdr: &RemoteHerdr, stdout: &str) -> Option<RemoteHerdr> {
     let mut lines = stdout.lines();
     let path = lines.next()?;
@@ -419,6 +429,18 @@ fn remote_herdr_from_path_probe(remote_herdr: &RemoteHerdr, stdout: &str) -> Opt
     Some(remote_herdr.clone().with_shell_path(shell_quote(path)))
 }
 
+fn remote_herdr_from_path_probe_any(
+    remote_herdr: &RemoteHerdr,
+    stdout: &str,
+) -> Option<RemoteHerdr> {
+    let mut lines = stdout.lines();
+    let path = lines.next()?;
+    if !path.starts_with('/') {
+        return None;
+    }
+    Some(remote_herdr.clone().with_shell_path(shell_quote(path)))
+}
+
 fn remote_binary_matches(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<bool> {
     let command = format!(
         "test -x {0} && {0} --version && {0} status client --json",
@@ -439,6 +461,11 @@ fn remote_binary_matches(target: &str, remote_herdr: &RemoteHerdr) -> io::Result
             .unwrap_or(false))
 }
 
+fn remote_binary_exists(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<bool> {
+    let command = format!("test -x {}", remote_herdr.shell_path);
+    Ok(ssh_output(target, &command)?.status.success())
+}
+
 fn remote_binary_override_path() -> io::Result<Option<PathBuf>> {
     let Some(value) = std::env::var_os(REMOTE_BINARY_ENV_VAR) else {
         return Ok(None);
@@ -509,6 +536,7 @@ enum RemoteServerStatus {
     Running {
         version: Option<String>,
         protocol: Option<u32>,
+        live_handoff: bool,
     },
     NotRunning,
 }
@@ -526,7 +554,12 @@ fn ensure_remote_server_ready(
     remote_binary_changed: bool,
 ) -> io::Result<()> {
     let status = remote_server_status(target, remote_herdr)?;
-    let RemoteServerStatus::Running { version, protocol } = status else {
+    let RemoteServerStatus::Running {
+        version,
+        protocol,
+        live_handoff,
+    } = status
+    else {
         return Ok(());
     };
 
@@ -536,6 +569,17 @@ fn ensure_remote_server_ready(
         return Ok(());
     };
 
+    if live_handoff && confirm_remote_server_handoff(target, version.as_deref(), protocol, reason)?
+    {
+        match live_handoff_remote_server(target, remote_herdr) {
+            Ok(()) => return Ok(()),
+            Err(err) => {
+                eprintln!("remote live handoff failed: {err}");
+                eprintln!("falling back to remote server restart.");
+            }
+        }
+    }
+
     if confirm_remote_server_stop(target, version.as_deref(), protocol, reason)? {
         stop_remote_server(target, remote_herdr)?;
     }
@@ -559,6 +603,85 @@ fn remote_server_restart_reason(
     None
 }
 
+fn confirm_remote_install_with_running_server(
+    target: &str,
+    remote_herdr: &RemoteHerdr,
+) -> io::Result<()> {
+    let status = match remote_server_status(target, remote_herdr) {
+        Ok(status) => status,
+        Err(err) => {
+            if !io::stdin().is_terminal() {
+                return Err(io::Error::other(format!(
+                    "could not inspect the running remote herdr server on {target} before installing: {err}; run from an interactive terminal to approve updating the remote binary"
+                )));
+            }
+            eprintln!(
+                "could not inspect the running remote herdr server on {target} before installing: {err}"
+            );
+            eprint!("continue installing the remote herdr binary? [Y/n] ");
+            io::stderr().flush()?;
+
+            let mut answer = String::new();
+            io::stdin().read_line(&mut answer)?;
+            let answer = answer.trim().to_ascii_lowercase();
+            if answer == "n" || answer == "no" {
+                return Err(io::Error::new(
+                    io::ErrorKind::Interrupted,
+                    "remote herdr install cancelled",
+                ));
+            }
+            return Ok(());
+        }
+    };
+    let RemoteServerStatus::Running {
+        version,
+        protocol,
+        live_handoff,
+    } = status
+    else {
+        return Ok(());
+    };
+    if live_handoff {
+        return Ok(());
+    }
+
+    if !io::stdin().is_terminal() {
+        return Err(io::Error::other(format!(
+            "remote herdr server on {target} is running v{} protocol {}, but it does not advertise live handoff; run from an interactive terminal to approve updating the remote binary",
+            version_label(version.as_deref()),
+            protocol_label(protocol)
+        )));
+    }
+
+    eprintln!("remote herdr server on {target} is currently running:");
+    eprintln!(
+        "  server: v{} protocol {}",
+        version_label(version.as_deref()),
+        protocol_label(protocol)
+    );
+    eprintln!(
+        "this server does not advertise live handoff, so this attach cannot preserve its running panes during the update."
+    );
+    eprintln!(
+        "future remote updates can preserve panes after the remote server has run a handoff-capable version once."
+    );
+    eprintln!();
+    eprint!("continue installing the remote herdr binary? [Y/n] ");
+    io::stderr().flush()?;
+
+    let mut answer = String::new();
+    io::stdin().read_line(&mut answer)?;
+    let answer = answer.trim().to_ascii_lowercase();
+    if answer == "n" || answer == "no" {
+        return Err(io::Error::new(
+            io::ErrorKind::Interrupted,
+            "remote herdr install cancelled",
+        ));
+    }
+
+    Ok(())
+}
+
 fn remote_server_status(
     target: &str,
     remote_herdr: &RemoteHerdr,
@@ -583,6 +706,12 @@ struct RemoteServerStatusJson {
     running: bool,
     version: Option<String>,
     protocol: Option<u32>,
+    capabilities: Option<RemoteServerCapabilitiesJson>,
+}
+
+#[derive(Debug, Deserialize)]
+struct RemoteServerCapabilitiesJson {
+    live_handoff: bool,
 }
 
 fn parse_client_status_json(status: &str) -> Option<RemoteClientStatusJson> {
@@ -602,6 +731,9 @@ fn parse_remote_server_status_json(status: &str) -> io::Result<RemoteServerStatu
     Ok(RemoteServerStatus::Running {
         version: parsed.version,
         protocol: parsed.protocol,
+        live_handoff: parsed
+            .capabilities
+            .is_some_and(|capabilities| capabilities.live_handoff),
     })
 }
 
@@ -677,6 +809,80 @@ fn confirm_remote_server_stop(
     Ok(true)
 }
 
+fn confirm_remote_server_handoff(
+    target: &str,
+    version: Option<&str>,
+    protocol: Option<u32>,
+    reason: RemoteServerRestartReason,
+) -> io::Result<bool> {
+    if !io::stdin().is_terminal() {
+        if reason == RemoteServerRestartReason::ProtocolMismatch {
+            return Err(io::Error::other(format!(
+                "remote herdr server on {target} is running with protocol {}, but this client needs protocol {CURRENT_PROTOCOL}; run from an interactive terminal to approve live handoff or stopping it",
+                protocol_label(protocol)
+            )));
+        }
+
+        eprintln!(
+            "remote herdr server on {target} is still running v{}; it will use v{CURRENT_VERSION} after it restarts.",
+            version_label(version)
+        );
+        return Ok(false);
+    }
+
+    eprintln!("remote herdr server on {target} is currently running:");
+    eprintln!(
+        "  server: v{} protocol {}",
+        version_label(version),
+        protocol_label(protocol)
+    );
+    eprintln!("  prepared binary: v{CURRENT_VERSION} protocol {CURRENT_PROTOCOL}");
+    eprintln!();
+
+    match reason {
+        RemoteServerRestartReason::ProtocolMismatch => {
+            eprintln!(
+                "the remote server protocol does not match this client. herdr will try to hand off live pane processes to the prepared remote server before the old server exits."
+            );
+        }
+        RemoteServerRestartReason::BinaryUpdated => {
+            eprintln!(
+                "the remote herdr binary was installed or replaced. herdr will try to hand off live pane processes to the prepared remote server."
+            );
+        }
+        RemoteServerRestartReason::VersionMismatch => {
+            eprintln!(
+                "the remote server is still running a different herdr version. herdr will try to hand off live pane processes to the prepared remote server."
+            );
+        }
+    }
+
+    eprint!("live-handoff remote panes to the prepared server? [Y/n] ");
+    io::stderr().flush()?;
+
+    let mut answer = String::new();
+    io::stdin().read_line(&mut answer)?;
+    let answer = answer.trim().to_ascii_lowercase();
+    Ok(answer != "n" && answer != "no")
+}
+
+fn live_handoff_remote_server(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<()> {
+    let command = format!(
+        "{} server live-handoff --import-exe {} --expected-protocol {CURRENT_PROTOCOL} --expected-version {CURRENT_VERSION}",
+        remote_herdr.shell_path,
+        remote_herdr.shell_path
+    );
+    let output = ssh_output(target, &command)?;
+    if !output.status.success() {
+        return Err(command_failed("remote server live handoff failed", &output));
+    }
+
+    eprintln!(
+        "handed off the remote herdr server on {target}; reconnecting to the prepared server."
+    );
+    Ok(())
+}
+
 fn stop_remote_server(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<()> {
     let command = format!("{} server stop", remote_herdr.shell_path);
     let output = ssh_output(target, &command)?;
@@ -1497,12 +1703,28 @@ mod tests {
     fn parse_remote_server_status_json_reads_running_server() {
         assert_eq!(
             parse_remote_server_status_json(
+                r#"{"status":"running","running":true,"version":"0.6.0","protocol":8,"capabilities":{"live_handoff":true}}"#
+            )
+            .unwrap(),
+            RemoteServerStatus::Running {
+                version: Some("0.6.0".into()),
+                protocol: Some(8),
+                live_handoff: true
+            }
+        );
+    }
+
+    #[test]
+    fn parse_remote_server_status_json_treats_missing_capability_as_no_handoff() {
+        assert_eq!(
+            parse_remote_server_status_json(
                 r#"{"status":"running","running":true,"version":"0.6.0","protocol":8}"#
             )
             .unwrap(),
             RemoteServerStatus::Running {
                 version: Some("0.6.0".into()),
-                protocol: Some(8)
+                protocol: Some(8),
+                live_handoff: false
             }
         );
     }