feat: simplify ssh machine setup discovery (#4204)

2 files changed, 250 insertions, 23 deletions

diff --git a/tests/machine_setup.rs b/tests/machine_setup.rs
index 3b577f4e..649ba7c0 100644
--- a/tests/machine_setup.rs
+++ b/tests/machine_setup.rs
@@ -21,7 +21,7 @@ if [ "$FAKE_STRICT_HOST_KEY_FAILURE" = yes ] && [ "$strict_host_key_check" = yes
     exit 255
 fi
 if [ "$last" = 'command -v herdr' ]; then
-    echo /home/remote/.local/bin/herdr
+    if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi
     exit 0
 fi
 case "$last" in
@@ -33,8 +33,10 @@ case "$last" in
 esac
 printf '\n%s\n' 'herdr-remote-output-ready:1'
 case "$script" in
-    *'uname -s'*) uname -s; uname -m ;;
-    *'version='*) echo /home/remote/.local/bin/herdr ;;
+    *'uname -s'*) echo platform >>"$FAKE_ROOT/probes"; uname -s; uname -m ;;
+    *'version='*)
+        echo candidates >>"$FAKE_ROOT/probes"
+        if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi ;;
     *'status client --json'*)
         if [ "$FAKE_INSTALLED" = new ] || [ -f "$FAKE_ROOT/installed" ]; then
             printf '%s\n' "$FAKE_CLIENT_STATUS"
@@ -51,10 +53,17 @@ case "$script" in
         fi ;;
     *'server live-handoff'*) echo handoff >>"$FAKE_ROOT/actions"; exit 1 ;;
     *'server stop'*) echo stop >>"$FAKE_ROOT/actions"; touch "$FAKE_ROOT/stopped" ;;
-    *'remote-client-bridge'*) echo start >>"$FAKE_ROOT/actions"; echo 'test startup failure' >&2; exit 1 ;;
+    *'session list --json'*)
+        echo sessions >>"$FAKE_ROOT/probes"
+        if [ "$FAKE_SESSIONS" = failed ]; then echo 'session query denied' >&2; exit 1; fi
+        printf '%s\n' "$FAKE_SESSIONS" ;;
+    *'remote-client-bridge'*)
+        printf '%s\n' "$script" >"$FAKE_ROOT/start-command"
+        echo start >>"$FAKE_ROOT/actions"; echo 'test startup failure' >&2; exit 1 ;;
     *'mkdir -p'*) printf '/fake/tmp\000/fake/herdr\000' ;;
     *'chmod 755'*) echo install >>"$FAKE_ROOT/actions"; touch "$FAKE_ROOT/installed" ;;
-    *'command -v herdr'*) echo /home/remote/.local/bin/herdr ;;
+    *'command -v herdr'*) if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi ;;
+    *'test -x '*) exit 1 ;;
     *) echo "unexpected fake SSH script: $script" >&2; exit 1 ;;
 esac
 "#;
@@ -64,6 +73,16 @@ struct SetupResult {
     actions: String,
     prompts: usize,
     success: bool,
+    probes: String,
+    start_command: String,
+}
+
+#[derive(Default)]
+struct SetupOptions<'a> {
+    sessions: Option<&'a str>,
+    selection: &'a [u8],
+    explicit_session: Option<&'a str>,
+    noninteractive: bool,
 }
 
 fn setup(installed: &str, answer: &str, handoff: bool) -> SetupResult {
@@ -76,13 +95,33 @@ fn setup_with_strict_host_key_failure(
     handoff: bool,
     strict_host_key_failure: bool,
 ) -> SetupResult {
+    setup_options(
+        installed,
+        answer,
+        handoff,
+        strict_host_key_failure,
+        SetupOptions::default(),
+    )
+}
+
+fn setup_options(
+    installed: &str,
+    answer: &str,
+    handoff: bool,
+    strict_host_key_failure: bool,
+    options: SetupOptions<'_>,
+) -> SetupResult {
     let root = std::path::PathBuf::from(format!(
-        "/var/tmp/herdr-machine-setup-{}-{}-{}-{}-{}",
+        "/var/tmp/herdr-machine-setup-{}-{}-{}-{}-{}-{}",
         std::process::id(),
         installed,
         answer.trim().is_empty(),
         handoff,
-        strict_host_key_failure
+        strict_host_key_failure,
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos()
     ));
     let app = if cfg!(debug_assertions) {
         "herdr-dev"
@@ -110,6 +149,9 @@ fn setup_with_strict_host_key_failure(
         command.args(["--remote", "fake-host", "--handoff"]);
     } else {
         command.args(["machine", "add", "fake-host", "--label", "VPS"]);
+        if let Some(session) = options.explicit_session {
+            command.args(["--remote-session", session]);
+        }
     }
     command.env(
         "PATH",
@@ -122,6 +164,12 @@ fn setup_with_strict_host_key_failure(
     command.env("FAKE_ROOT", &root);
     command.env("FAKE_INSTALLED", installed);
     command.env(
+        "FAKE_SESSIONS",
+        options
+            .sessions
+            .unwrap_or(r#"{"sessions":[{"name":"default","running":true}]}"#),
+    );
+    command.env(
         "FAKE_STRICT_HOST_KEY_FAILURE",
         if strict_host_key_failure { "yes" } else { "no" },
     );
@@ -139,6 +187,18 @@ fn setup_with_strict_host_key_failure(
     ] {
         command.env_remove(name);
     }
+    if options.noninteractive {
+        // A redirected stdin makes machine add non-interactive while retaining
+        // the same isolated SSH fixture and captured terminal diagnostics.
+        let mut shell = CommandBuilder::new("/bin/sh");
+        shell.args(["-c", "exec \"$@\" </dev/null", "sh"]);
+        shell.args(command.get_argv());
+        shell.env_clear();
+        for (key, value) in command.iter_full_env_as_str() {
+            shell.env(key, value);
+        }
+        command = shell;
+    }
     let mut child = pair.slave.spawn_command(command).unwrap();
     drop(pair.slave);
     let mut reader = pair.master.try_clone_reader().unwrap();
@@ -155,10 +215,16 @@ fn setup_with_strict_host_key_failure(
     let mut output = String::new();
     let mut prompts = 0;
     let mut timed_out = false;
+    let mut picked = false;
     loop {
         match rx.recv_timeout(Duration::from_secs(20)) {
             Ok(bytes) => {
                 output.push_str(&String::from_utf8_lossy(&bytes));
+                if !picked && output.contains("Esc cancel") {
+                    writer.write_all(options.selection).unwrap();
+                    writer.flush().unwrap();
+                    picked = true;
+                }
                 let count = output.matches("[y/N]").count() + output.matches("[Y/n]").count();
                 if count > prompts {
                     writer
@@ -185,6 +251,8 @@ fn setup_with_strict_host_key_failure(
     drop(pair.master);
     reading.join().unwrap();
     let actions = fs::read_to_string(root.join("actions")).unwrap_or_default();
+    let probes = fs::read_to_string(root.join("probes")).unwrap_or_default();
+    let start_command = fs::read_to_string(root.join("start-command")).unwrap_or_default();
     // Every scenario either cancels or reaches the intentionally failed startup.
     let saved = root
         .join("state")
@@ -202,6 +270,8 @@ fn setup_with_strict_host_key_failure(
         actions,
         prompts,
         success: status.success(),
+        probes,
+        start_command,
     }
 }
 
@@ -259,6 +329,155 @@ fn machine_add_accepts_help_argument_order() {
 }
 
 #[test]
+fn machine_add_fresh_host_offers_installation_and_cancellation_saves_nothing() {
+    let result = setup("missing", "n\n", false);
+    assert_eq!(result.probes, "platform\ncandidates\n", "{}", result.output);
+    assert_eq!(result.prompts, 1, "{}", result.output);
+    assert!(
+        result.output.contains("installation cancelled"),
+        "{}",
+        result.output
+    );
+    assert!(result.actions.is_empty(), "{}", result.output);
+}
+
+#[test]
+fn machine_add_discovers_sessions_once_and_prepares_the_selected_session() {
+    for (sessions, selection, expected) in [
+        (r#"{"sessions":[]}"#, &b""[..], "default"),
+        (
+            r#"{"extra":1,"sessions":[{"name":"stopped","running":false},{"name":"bad/name","running":true},{"name":"agents","running":true,"future":true}]}"#,
+            &b""[..],
+            "agents",
+        ),
+        (
+            r#"{"sessions":[{"name":"default","running":true},{"name":"agents","running":true}]}"#,
+            &b"\x1b[B\r"[..],
+            "agents",
+        ),
+    ] {
+        let result = setup_options(
+            "new",
+            "y\n",
+            false,
+            false,
+            SetupOptions {
+                sessions: Some(sessions),
+                selection,
+                ..Default::default()
+            },
+        );
+        assert_eq!(
+            result.probes, "platform\ncandidates\nsessions\n",
+            "{}",
+            result.output
+        );
+        assert_eq!(result.actions, "stop\nstart\n", "{}", result.output);
+        if expected == "default" {
+            assert!(
+                !result.start_command.contains("--session"),
+                "{}",
+                result.start_command
+            );
+        } else {
+            assert!(
+                result.start_command.contains("--session agents"),
+                "{}",
+                result.start_command
+            );
+        }
+        if !selection.is_empty() {
+            let picker = result
+                .output
+                .split("Running sessions on fake-host:")
+                .nth(1)
+                .unwrap();
+            assert!(
+                picker.starts_with("\r\n"),
+                "picker did not return to column zero: {picker:?}"
+            );
+            assert!(picker.contains("default\u{1b}[0m\r\n"), "{picker:?}");
+            assert!(picker.contains("  agents\r\n"), "{picker:?}");
+        }
+    }
+}
+
+#[test]
+fn machine_add_failed_discovery_and_cancel_never_prepare_or_save() {
+    for (sessions, selection, expected) in [
+        ("failed", &b""[..], "session query denied"),
+        ("not json", &b""[..], "--remote-session"),
+        (
+            r#"{"sessions":[{"name":"agents"}]}"#,
+            &b""[..],
+            "--remote-session",
+        ),
+        (
+            r#"{"sessions":[{"name":"default","running":true},{"name":"agents","running":true}]}"#,
+            &b"\x1b"[..],
+            "selection cancelled",
+        ),
+    ] {
+        let result = setup_options(
+            "new",
+            "y\n",
+            false,
+            false,
+            SetupOptions {
+                sessions: Some(sessions),
+                selection,
+                ..Default::default()
+            },
+        );
+        assert!(!result.success, "{}", result.output);
+        assert!(result.actions.is_empty(), "{}", result.output);
+        assert!(result.output.contains(expected), "{}", result.output);
+        assert!(
+            result.output.contains("machine was not saved"),
+            "{}",
+            result.output
+        );
+    }
+}
+
+#[test]
+fn machine_add_skips_discovery_for_scripts_and_explicit_sessions() {
+    for (noninteractive, explicit_session) in [
+        (true, None),
+        (true, Some("agents")),
+        (false, Some("agents")),
+    ] {
+        let result = setup_options(
+            "new",
+            "y\n",
+            false,
+            false,
+            SetupOptions {
+                noninteractive,
+                explicit_session,
+                sessions: Some("failed"),
+                ..Default::default()
+            },
+        );
+        assert_eq!(result.probes, "platform\ncandidates\n", "{}", result.output);
+        assert!(
+            result.output.contains(if explicit_session.is_some() {
+                "session agents"
+            } else {
+                "session default"
+            }),
+            "{}",
+            result.output
+        );
+        if noninteractive {
+            assert!(result.actions.is_empty(), "{}", result.output);
+        } else {
+            assert_eq!(result.actions, "stop\nstart\n", "{}", result.output);
+        }
+    }
+}
+
+#[test]
 fn machine_add_reports_strict_host_key_failure() {
     let result = setup_with_strict_host_key_failure("new", "", false, true);
     assert!(!result.success, "{}", result.output);
diff --git a/tests/remote_attach.rs b/tests/remote_attach.rs
index b99e4335..2af61c77 100644
--- a/tests/remote_attach.rs
+++ b/tests/remote_attach.rs
@@ -1,7 +1,7 @@
 #![cfg(unix)]
 
 use std::fs;
-use std::io::{BufRead, BufReader};
+use std::io::Read;
 use std::os::unix::fs::PermissionsExt;
 use std::os::unix::process::CommandExt;
 use std::path::{Path, PathBuf};
@@ -77,7 +77,7 @@ fn check_authentication_output(framed_shell: bool) {
             r#"#!/bin/sh
 authenticate() {{
     : > "$FAKE_SSH_STARTED"
-    printf '%s\n%s\n' '{CHECK_NOTICE}' '{CHECK_URL}' >&2
+    printf '/bin/sh: expected probe noise\n%s\n%s' '{CHECK_NOTICE}' '{CHECK_URL}' >&2
     while [ ! -e "$FAKE_SSH_APPROVED" ]; do
         /bin/sleep 0.01
     done
@@ -94,7 +94,7 @@ if [ "$FAKE_SSH_FRAMED" = 1 ] && [ ! -e "$FAKE_SSH_STARTED" ]; then
 fi
 /bin/cat >/dev/null
 : > "$FAKE_SSH_ADVANCED"
-printf '%s\n' '{LATER_FAILURE}' >&2
+printf '\n%s\n' '{LATER_FAILURE}' >&2
 exit 255
 "#
         ),
@@ -128,22 +128,25 @@ exit 255
 
     cleanup.child = Some(child);
     let child = cleanup.child.as_mut().expect("registered child");
-    let stderr = child.stderr.take().expect("remote attach stderr");
+    let mut stderr = child.stderr.take().expect("remote attach stderr");
     let (line_tx, line_rx) = mpsc::channel();
     cleanup.reader = Some(thread::spawn(move || {
-        for line in BufReader::new(stderr).lines() {
-            if line_tx
-                .send(line.expect("read remote attach stderr"))
-                .is_err()
-            {
+        let mut buffer = [0; 4096];
+        while let Ok(len) = stderr.read(&mut buffer) {
+            if len == 0 || line_tx.send(buffer[..len].to_vec()).is_err() {
                 break;
             }
         }
     }));
 
     wait_for_file(&started_path, Duration::from_secs(2));
-    let notice = line_rx.recv_timeout(Duration::from_secs(2));
-    let url = line_rx.recv_timeout(Duration::from_secs(2));
+    let mut before_approval = Vec::new();
+    while !String::from_utf8_lossy(&before_approval).contains(CHECK_URL) {
+        match line_rx.recv_timeout(Duration::from_secs(2)) {
+            Ok(bytes) => before_approval.extend(bytes),
+            Err(_) => break,
+        }
+    }
     fs::write(&approval_path, b"approved").expect("release fake ssh approval");
     wait_for_file(&advanced_path, Duration::from_secs(2));
 
@@ -155,16 +158,21 @@ exit 255
         .expect("registered stderr reader")
         .join()
         .expect("join stderr reader");
-    let later_lines = line_rx.try_iter().collect::<Vec<_>>();
+    let later_bytes = line_rx.try_iter().flatten().collect::<Vec<_>>();
+    let later_lines = String::from_utf8_lossy(&later_bytes);
 
-    assert_eq!(notice.as_deref(), Ok(CHECK_NOTICE));
-    assert_eq!(url.as_deref(), Ok(CHECK_URL));
+    let before_approval = String::from_utf8_lossy(&before_approval);
+    assert!(before_approval.contains(CHECK_NOTICE), "{before_approval}");
+    assert!(
+        before_approval.contains(CHECK_URL),
+        "partial prompt was hidden: {before_approval}"
+    );
     assert!(
-        later_lines.iter().any(|line| line == LATER_FAILURE),
+        later_lines.contains(LATER_FAILURE),
         "later SSH stderr should also be visible: {later_lines:?}"
     );
     assert!(
-        later_lines.iter().any(|line| {
+        later_lines.lines().any(|line| {
             line.contains("error: remote binary discovery failed") && line.contains(LATER_FAILURE)
         }),
         "failed SSH stderr should remain in the contextual error: {later_lines:?}"