feat: route cli commands to saved ssh machines (#3918)

2 files changed, 138 insertions, 10 deletions

diff --git a/src/remote/attach.rs b/src/remote/attach.rs
index 2aee5efe..042676d0 100644
--- a/src/remote/attach.rs
+++ b/src/remote/attach.rs
@@ -957,6 +957,35 @@ fn prepare_windows_remote_herdr(
     })
 }
 
+pub(super) fn find_installed_remote_api_herdr(
+    ssh: &RemoteSsh,
+    session: &str,
+) -> io::Result<RemoteHerdr> {
+    let platform = detect_remote_platform(ssh)?;
+    let remote_herdr = RemoteHerdr::for_platform(platform);
+    let candidates = if remote_herdr.platform.is_windows() {
+        vec![remote_herdr]
+    } else {
+        remote_binary_candidates(ssh, &remote_herdr)?
+    };
+    for candidate in candidates {
+        let probe =
+            ssh.framed_user_shell_output(&remote_api_bridge_command(&candidate, session, true))?;
+        if probe.status.code() == Some(255) {
+            return Err(command_failed("remote SSH connection failed", &probe));
+        }
+        if probe.status.success()
+            && String::from_utf8_lossy(&probe.stdout).trim() == "herdr-api-bridge-v1"
+        {
+            return Ok(candidate);
+        }
+    }
+    Err(io::Error::new(
+        io::ErrorKind::Unsupported,
+        "remote Herdr does not support machine API forwarding; update Herdr on this machine",
+    ))
+}
+
 fn detect_remote_platform(ssh: &RemoteSsh) -> io::Result<RemotePlatform> {
     let output = ssh.sh_output("uname -s\nuname -m\n")?;
     let mut windows_uname_hint = false;
@@ -2015,6 +2044,24 @@ fn confirm_remote_install(
     Ok(())
 }
 
+pub(super) fn remote_api_bridge_command(
+    remote_herdr: &RemoteHerdr,
+    session_name: &str,
+    check: bool,
+) -> String {
+    let mut args = vec!["--session", session_name, "remote-api-bridge"];
+    if check {
+        args.push("--check");
+    }
+    match &remote_herdr.executable {
+        RemoteExecutable::PosixShellPath(_) => {
+            posix_remote_output_command(&format!("exec {}", remote_herdr.executable.command(&args)))
+        }
+        RemoteExecutable::WindowsPath(path) => {
+            windows_powershell_streaming_application_command(path, &args)
+        }
+    }
+}
 fn reattach_command(
     program: &str,
     target: &str,
@@ -2066,6 +2113,22 @@ impl SshStdioBridge {
         ssh_options: Option<&ManagedSshOptions>,
         noninteractive: bool,
     ) -> io::Result<Self> {
+        Self::start_command(
+            target,
+            remote_herdr.executable.bridge_command(&session_name),
+            local_socket,
+            ssh_options,
+            noninteractive,
+        )
+    }
+
+    pub(super) fn start_command(
+        target: String,
+        remote_command: String,
+        local_socket: PathBuf,
+        ssh_options: Option<&ManagedSshOptions>,
+        noninteractive: bool,
+    ) -> io::Result<Self> {
         crate::ipc::prepare_socket_path(&local_socket, |path| {
             format!("remote bridge is already listening at {}", path.display())
         })?;
@@ -2103,8 +2166,7 @@ impl SshStdioBridge {
                         if let Err(err) = bridge_connection(
                             stream,
                             &target,
-                            &remote_herdr,
-                            &session_name,
+                            &remote_command,
                             thread_ssh_options.as_ref(),
                             noninteractive,
                             &thread_stop,
@@ -2142,7 +2204,7 @@ impl SshStdioBridge {
         })
     }
 
-    fn reported_failure(&self) -> Option<io::Error> {
+    pub(super) fn reported_failure(&self) -> Option<io::Error> {
         self.failure_rx
             .recv_timeout(BRIDGE_FAILURE_REPORT_TIMEOUT)
             .ok()
@@ -2285,8 +2347,7 @@ pub(crate) fn bridge_upload_cancellation_for_test(
 fn bridge_connection(
     mut stream: crate::ipc::LocalStream,
     target: &str,
-    remote_herdr: &RemoteHerdr,
-    session_name: &str,
+    remote_command: &str,
     ssh_options: Option<&ManagedSshOptions>,
     noninteractive: bool,
     bridge_stop: &Arc<AtomicBool>,
@@ -2300,7 +2361,7 @@ fn bridge_connection(
     command
         .arg("-T")
         .arg(target)
-        .arg(remote_herdr.executable.bridge_command(session_name))
+        .arg(remote_command)
         .stdin(Stdio::piped())
         .stdout(Stdio::piped())
         .stderr(if noninteractive {
@@ -3551,6 +3612,16 @@ mod tests {
                 "$process = Start-Process -FilePath herdr.exe -ArgumentList '--session agents remote-client-bridge' -NoNewWindow -Wait -PassThru -ErrorAction Stop; exit $process.ExitCode",
             ),
             (
+                "API bridge with explicit default session",
+                remote_api_bridge_command(&RemoteHerdr::for_platform(RemotePlatform { os: "windows", arch: "x86_64" }), "default", false),
+                "$process = Start-Process -FilePath herdr.exe -ArgumentList '--session default remote-api-bridge' -NoNewWindow -Wait -PassThru -ErrorAction Stop; exit $process.ExitCode",
+            ),
+            (
+                "API bridge capability probe",
+                remote_api_bridge_command(&RemoteHerdr::for_platform(RemotePlatform { os: "windows", arch: "x86_64" }), "agents", true),
+                "$process = Start-Process -FilePath herdr.exe -ArgumentList '--session agents remote-api-bridge --check' -NoNewWindow -Wait -PassThru -ErrorAction Stop; exit $process.ExitCode",
+            ),
+            (
                 "saved bridge with closed stdin",
                 executable.saved_bridge_command("agents"),
                 "$process = Start-Process -FilePath herdr.exe -ArgumentList '--session agents remote-client-bridge' -NoNewWindow -Wait -PassThru -ErrorAction Stop; exit $process.ExitCode",
@@ -3665,6 +3736,22 @@ mod tests {
     }
 
     #[test]
+    fn remote_api_bridge_always_selects_the_saved_session() {
+        let remote_herdr = RemoteHerdr::for_platform(RemotePlatform {
+            os: "linux",
+            arch: "x86_64",
+        });
+        for session in ["default", "agents"] {
+            assert_eq!(
+                remote_api_bridge_command(&remote_herdr, session, false),
+                posix_remote_output_command(&format!(
+                    "exec \"$HOME/.local/bin/herdr\" --session {session} remote-api-bridge"
+                ))
+            );
+        }
+    }
+
+    #[test]
     fn remote_bridge_command_uses_installed_binary() {
         let remote_herdr = RemoteHerdr::for_platform(RemotePlatform {
             os: "linux",
diff --git a/src/remote/saved.rs b/src/remote/saved.rs
index 5b570af4..561ee1a0 100644
--- a/src/remote/saved.rs
+++ b/src/remote/saved.rs
@@ -17,10 +17,7 @@ pub(crate) fn connect_saved_ssh(
     target: &str,
     session: &str,
 ) -> io::Result<SavedSshStream> {
-    validate_profile_path_id(profile_id)?;
-    crate::session::validate_name(session)
-        .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
-    let ssh = RemoteSsh::new_noninteractive(target.to_owned());
+    let ssh = validated_saved_ssh(profile_id, target, session)?;
     let remote_herdr = find_installed_remote_herdr(&ssh)?;
     let path = saved_bridge_path(profile_id);
     let bridge = SshStdioBridge::start(
@@ -38,6 +35,43 @@ pub(crate) fn connect_saved_ssh(
     })
 }
 
+pub(crate) struct SavedSshApiBridge {
+    path: PathBuf,
+    bridge: SshStdioBridge,
+}
+
+impl SavedSshApiBridge {
+    pub(crate) fn start(profile_id: &str, target: &str, session: &str) -> io::Result<Self> {
+        let ssh = validated_saved_ssh(profile_id, target, session)?;
+        let remote_herdr = super::attach::find_installed_remote_api_herdr(&ssh, session)?;
+        let command = super::attach::remote_api_bridge_command(&remote_herdr, session, false);
+        let path = crate::platform::remote_bridge_endpoint_path(
+            &format!("herdr-api-ssh-{}-{profile_id}.sock", std::process::id()),
+            &format!(
+                "herdr-api-{}-{}.sock",
+                std::process::id(),
+                &profile_id[..16]
+            ),
+        );
+        let bridge = SshStdioBridge::start_command(
+            target.to_owned(),
+            command,
+            path.clone(),
+            ssh.options(),
+            true,
+        )?;
+        Ok(Self { path, bridge })
+    }
+
+    pub(crate) fn socket_path(&self) -> &std::path::Path {
+        &self.path
+    }
+
+    pub(crate) fn reported_failure(&self) -> Option<io::Error> {
+        self.bridge.reported_failure()
+    }
+}
+
 pub(crate) fn saved_ssh_bootstrap_command(target: &str, session: &str) -> String {
     format!(
         "herdr --remote {} --session {}",
@@ -80,6 +114,13 @@ fn saved_bridge_path(profile_id: &str) -> PathBuf {
     crate::platform::remote_bridge_endpoint_path(&readable, &short)
 }
 
+fn validated_saved_ssh(profile_id: &str, target: &str, session: &str) -> io::Result<RemoteSsh> {
+    validate_profile_path_id(profile_id)?;
+    crate::session::validate_name(session)
+        .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
+    Ok(RemoteSsh::new_noninteractive(target.to_owned()))
+}
+
 fn validate_profile_path_id(profile_id: &str) -> io::Result<()> {
     if profile_id.len() == 32
         && profile_id