fix: restart remote server after binary update

1 files changed, 214 insertions, 35 deletions

diff --git a/src/remote.rs b/src/remote.rs
index 089772b4..66b9b6e5 100644
--- a/src/remote.rs
+++ b/src/remote.rs
@@ -13,10 +13,12 @@ use std::sync::{
     Arc,
 };
 use std::thread::{self, JoinHandle};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 const BRIDGE_ACCEPT_POLL: Duration = Duration::from_millis(50);
 const BRIDGE_SOCKET_PERMISSION_MODE: u32 = 0o600;
+const REMOTE_SERVER_SHUTDOWN_CONFIRM_TIMEOUT: Duration = Duration::from_secs(5);
+const REMOTE_SERVER_SHUTDOWN_POLL_INTERVAL: Duration = Duration::from_millis(100);
 const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
 const CURRENT_PROTOCOL: u32 = crate::server::protocol::PROTOCOL_VERSION;
 const UPDATE_MANIFEST_URL: &str = "https://herdr.dev/latest.json";
@@ -89,12 +91,16 @@ pub(crate) fn run_remote(remote: RemoteLaunch) -> io::Result<()> {
         .next()
         .unwrap_or_else(|| "herdr".to_string());
     let reattach_command = reattach_command(&program, &remote.target, &session_name);
-    let remote_herdr = prepare_remote_herdr(&remote.target)?;
-    ensure_remote_server_compatible(&remote.target, &remote_herdr)?;
+    let prepared_remote = prepare_remote_herdr(&remote.target)?;
+    ensure_remote_server_ready(
+        &remote.target,
+        &prepared_remote.remote_herdr,
+        prepared_remote.installed_or_replaced,
+    )?;
 
     let _bridge = SshStdioBridge::start(
         remote.target,
-        remote_herdr,
+        prepared_remote.remote_herdr,
         local_socket.clone(),
         session_name,
     )?;
@@ -141,7 +147,7 @@ fn ensure_remote_server_running() -> io::Result<()> {
             return Ok(());
         }
         return Err(io::Error::other(format!(
-            "remote herdr server is running with protocol {}, but this bridge needs protocol {CURRENT_PROTOCOL}; rerun `herdr --remote` from an interactive terminal to approve replacing it",
+            "remote herdr server is running with protocol {}, but this bridge needs protocol {CURRENT_PROTOCOL}; rerun `herdr --remote` from an interactive terminal to approve stopping it",
             protocol_label(status.protocol)
         )));
     }
@@ -231,6 +237,11 @@ struct InstallSource {
     temporary_dir: Option<PathBuf>,
 }
 
+struct PreparedRemoteHerdr {
+    remote_herdr: RemoteHerdr,
+    installed_or_replaced: bool,
+}
+
 impl InstallSource {
     fn persistent(path: PathBuf) -> Self {
         Self {
@@ -253,17 +264,23 @@ impl InstallSource {
     }
 }
 
-fn prepare_remote_herdr(target: &str) -> io::Result<RemoteHerdr> {
+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()?;
 
     if override_binary.is_none() {
         if let Some(path_remote_herdr) = remote_binary_on_path(target, &remote_herdr)? {
-            return Ok(path_remote_herdr);
+            return Ok(PreparedRemoteHerdr {
+                remote_herdr: path_remote_herdr,
+                installed_or_replaced: false,
+            });
         }
         if remote_binary_matches(target, &remote_herdr)? {
-            return Ok(remote_herdr);
+            return Ok(PreparedRemoteHerdr {
+                remote_herdr,
+                installed_or_replaced: false,
+            });
         }
     }
 
@@ -286,7 +303,10 @@ fn prepare_remote_herdr(target: &str) -> io::Result<RemoteHerdr> {
     warn_if_remote_bin_not_on_path(target)?;
     maybe_copy_local_keybindings_to_remote(target, &remote_herdr)?;
 
-    Ok(remote_herdr)
+    Ok(PreparedRemoteHerdr {
+        remote_herdr,
+        installed_or_replaced: true,
+    })
 }
 
 fn detect_remote_platform(target: &str) -> io::Result<RemotePlatform> {
@@ -429,24 +449,59 @@ fn resolve_install_source(
     download_release_asset(platform)
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 enum RemoteServerStatus {
-    Running { protocol: Option<u32> },
+    Running {
+        version: Option<String>,
+        protocol: Option<u32>,
+    },
     NotRunning,
 }
 
-fn ensure_remote_server_compatible(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<()> {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum RemoteServerRestartReason {
+    ProtocolMismatch,
+    BinaryUpdated,
+    VersionMismatch,
+}
+
+fn ensure_remote_server_ready(
+    target: &str,
+    remote_herdr: &RemoteHerdr,
+    remote_binary_changed: bool,
+) -> io::Result<()> {
     let status = remote_server_status(target, remote_herdr)?;
-    let RemoteServerStatus::Running { protocol } = status else {
+    let RemoteServerStatus::Running { version, protocol } = status else {
         return Ok(());
     };
 
-    if protocol == Some(CURRENT_PROTOCOL) {
+    let Some(reason) =
+        remote_server_restart_reason(version.as_deref(), protocol, remote_binary_changed)
+    else {
         return Ok(());
+    };
+
+    if confirm_remote_server_stop(target, version.as_deref(), protocol, reason)? {
+        stop_remote_server(target, remote_herdr)?;
     }
+    Ok(())
+}
 
-    confirm_remote_server_stop(target, protocol)?;
-    stop_remote_server(target, remote_herdr)
+fn remote_server_restart_reason(
+    version: Option<&str>,
+    protocol: Option<u32>,
+    remote_binary_changed: bool,
+) -> Option<RemoteServerRestartReason> {
+    if protocol != Some(CURRENT_PROTOCOL) {
+        return Some(RemoteServerRestartReason::ProtocolMismatch);
+    }
+    if remote_binary_changed {
+        return Some(RemoteServerRestartReason::BinaryUpdated);
+    }
+    if version != Some(CURRENT_VERSION) {
+        return Some(RemoteServerRestartReason::VersionMismatch);
+    }
+    None
 }
 
 fn remote_server_status(
@@ -469,6 +524,7 @@ fn remote_server_status(
 
     if stdout.lines().any(|line| line.trim() == "status: running") {
         return Ok(RemoteServerStatus::Running {
+            version: parse_status_version(&stdout).map(str::to_owned),
             protocol: parse_status_protocol(&stdout),
         });
     }
@@ -479,49 +535,127 @@ fn remote_server_status(
     )))
 }
 
-fn parse_status_protocol(status: &str) -> Option<u32> {
+fn parse_status_value<'a>(status: &'a str, key: &str) -> Option<&'a str> {
     status.lines().find_map(|line| {
-        let (_, value) = line.trim().split_once(':')?;
-        (line.trim_start().starts_with("protocol:")).then(|| value.trim().parse().ok())?
+        let line = line.trim();
+        let (found_key, value) = line.split_once(':')?;
+        (found_key == key).then_some(value.trim())
     })
 }
 
-fn confirm_remote_server_stop(target: &str, protocol: Option<u32>) -> io::Result<()> {
+fn parse_status_version(status: &str) -> Option<&str> {
+    parse_status_value(status, "version")
+}
+
+fn parse_status_protocol(status: &str) -> Option<u32> {
+    parse_status_value(status, "protocol")?.parse().ok()
+}
+
+fn confirm_remote_server_stop(
+    target: &str,
+    version: Option<&str>,
+    protocol: Option<u32>,
+    reason: RemoteServerRestartReason,
+) -> io::Result<bool> {
     if !io::stdin().is_terminal() {
-        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 stopping it",
-            protocol_label(protocol)
-        )));
+        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 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!(
-        "remote herdr server on {target} is running with protocol {}, but this client needs protocol {CURRENT_PROTOCOL}.",
+        "  server: v{} protocol {}",
+        version_label(version),
         protocol_label(protocol)
     );
-    eprint!("Install/replace the remote binary if needed and stop the running remote server now? This will kill the remote Herdr server. [Y/n] ");
+    eprintln!("  prepared binary: v{CURRENT_VERSION} protocol {CURRENT_PROTOCOL}");
+    eprintln!();
+
+    match reason {
+        RemoteServerRestartReason::ProtocolMismatch => {
+            eprintln!(
+                "the remote server protocol does not match this client. the remote server must be stopped before attaching."
+            );
+        }
+        RemoteServerRestartReason::BinaryUpdated => {
+            eprintln!(
+                "the remote herdr binary was installed or replaced. restart the remote server so it uses the prepared binary."
+            );
+        }
+        RemoteServerRestartReason::VersionMismatch => {
+            eprintln!(
+                "the remote server is still running a different herdr version. restart it so it uses the prepared binary."
+            );
+        }
+    }
+
+    let prompt = if reason == RemoteServerRestartReason::ProtocolMismatch {
+        "stop the remote server and continue attaching? [Y/n] "
+    } else {
+        "restart the remote server now? [Y/n] "
+    };
+    eprint!("{prompt}");
     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 server replacement cancelled",
-        ));
+        if reason == RemoteServerRestartReason::ProtocolMismatch {
+            return Err(io::Error::new(
+                io::ErrorKind::Interrupted,
+                "remote herdr server stop cancelled",
+            ));
+        }
+        return Ok(false);
     }
 
-    Ok(())
+    Ok(true)
 }
 
 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)?;
-    if output.status.success() {
-        Ok(())
-    } else {
-        Err(command_failed("remote server stop failed", &output))
+    if !output.status.success() {
+        return Err(command_failed("remote server stop failed", &output));
     }
+
+    wait_for_remote_server_shutdown(target, remote_herdr)?;
+    eprintln!("stopped the remote herdr server on {target}; it will restart when the remote client bridge attaches.");
+    Ok(())
+}
+
+fn wait_for_remote_server_shutdown(target: &str, remote_herdr: &RemoteHerdr) -> io::Result<()> {
+    let deadline = Instant::now() + REMOTE_SERVER_SHUTDOWN_CONFIRM_TIMEOUT;
+    loop {
+        if remote_server_status(target, remote_herdr)? == RemoteServerStatus::NotRunning {
+            return Ok(());
+        }
+        if Instant::now() >= deadline {
+            return Err(io::Error::new(
+                io::ErrorKind::TimedOut,
+                format!(
+                    "shutdown was requested, but the old remote herdr server on {target} is still responding after {} seconds",
+                    REMOTE_SERVER_SHUTDOWN_CONFIRM_TIMEOUT.as_secs()
+                ),
+            ));
+        }
+        thread::sleep(REMOTE_SERVER_SHUTDOWN_POLL_INTERVAL);
+    }
+}
+
+fn version_label(version: Option<&str>) -> &str {
+    version.unwrap_or("unknown")
 }
 
 fn protocol_label(protocol: Option<u32>) -> String {
@@ -1387,6 +1521,51 @@ command = "lazygit"
     }
 
     #[test]
+    fn parse_status_version_reads_version_line() {
+        assert_eq!(
+            parse_status_version("status: running\nversion: 0.6.0\nprotocol: 8"),
+            Some("0.6.0")
+        );
+        assert_eq!(parse_status_version("status: not running"), None);
+    }
+
+    #[test]
+    fn remote_server_restart_reason_requires_stop_for_protocol_mismatch() {
+        assert_eq!(
+            remote_server_restart_reason(Some(CURRENT_VERSION), Some(0), false),
+            Some(RemoteServerRestartReason::ProtocolMismatch)
+        );
+    }
+
+    #[test]
+    fn remote_server_restart_reason_offers_restart_after_binary_update() {
+        assert_eq!(
+            remote_server_restart_reason(Some(CURRENT_VERSION), Some(CURRENT_PROTOCOL), true),
+            Some(RemoteServerRestartReason::BinaryUpdated)
+        );
+    }
+
+    #[test]
+    fn remote_server_restart_reason_offers_restart_for_version_mismatch() {
+        assert_eq!(
+            remote_server_restart_reason(Some("0.0.0"), Some(CURRENT_PROTOCOL), false),
+            Some(RemoteServerRestartReason::VersionMismatch)
+        );
+        assert_eq!(
+            remote_server_restart_reason(None, Some(CURRENT_PROTOCOL), false),
+            Some(RemoteServerRestartReason::VersionMismatch)
+        );
+    }
+
+    #[test]
+    fn remote_server_restart_reason_allows_current_server() {
+        assert_eq!(
+            remote_server_restart_reason(Some(CURRENT_VERSION), Some(CURRENT_PROTOCOL), false),
+            None
+        );
+    }
+
+    #[test]
     fn install_source_description_uses_override_binary() {
         let platform = RemotePlatform {
             os: "linux",