feat: manage multiple ssh machines from one client (#3670)

3 files changed, 270 insertions, 0 deletions

diff --git a/tests/client_mode.rs b/tests/client_mode.rs
index 351c5fa2..06bd1123 100644
--- a/tests/client_mode.rs
+++ b/tests/client_mode.rs
@@ -125,6 +125,7 @@ fn spawn_client_process_with_args_and_env(
     let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
     cmd.args(args);
     cmd.env("HERDR_DISABLE_SOUND", "1");
+    cmd.env("XDG_STATE_HOME", runtime_dir.join("state"));
     cmd.env("XDG_CONFIG_HOME", config_home);
     cmd.env("XDG_RUNTIME_DIR", runtime_dir);
     cmd.env("HERDR_SOCKET_PATH", api_socket_path);
@@ -492,6 +493,7 @@ fn server_unreachable_shows_clear_error() {
         .env("HERDR_DISABLE_SOUND", "1")
         .env("XDG_CONFIG_HOME", &config_home)
         .env("XDG_RUNTIME_DIR", &runtime_dir)
+        .env("XDG_STATE_HOME", runtime_dir.join("state"))
         .env("HERDR_SOCKET_PATH", &api_socket)
         .env_remove("HERDR_CLIENT_SOCKET_PATH")
         .env_remove("HERDR_ENV")
@@ -753,6 +755,271 @@ fn attach_thin_client_with_config(
 }
 
 #[test]
+fn federated_launch_opens_local_directly_while_saved_ssh_is_unavailable() {
+    use std::os::unix::fs::PermissionsExt;
+
+    let _lock = test_lock();
+    for select_remote in [false, true] {
+        let base = unique_test_dir();
+        let config_home = base.join("config");
+        let runtime_dir = base.join("runtime");
+        let api_socket = runtime_dir.join("herdr.sock");
+        fs::create_dir_all(config_home.join(app_dir_name())).unwrap();
+        fs::write(
+            config_home.join(app_dir_name()).join("config.toml"),
+            "onboarding = false\n",
+        )
+        .unwrap();
+        let catalog_dir = runtime_dir
+            .join("state")
+            .join(app_dir_name())
+            .join("client");
+        fs::create_dir_all(&catalog_dir).unwrap();
+        let profile = "0123456789abcdef0123456789abcdef";
+        fs::write(catalog_dir.join("endpoints.json"), serde_json::json!({
+            "version": 1, "selected_profile": select_remote.then_some(profile),
+            "ssh": [{"id": profile, "label": "Unavailable remote", "target": "test-only", "session": "default", "enabled": true}],
+        }).to_string()).unwrap();
+        let bin = base.join("bin");
+        fs::create_dir_all(&bin).unwrap();
+        fs::write(bin.join("ssh"), "#!/bin/sh\nexit 255\n").unwrap();
+        fs::set_permissions(bin.join("ssh"), fs::Permissions::from_mode(0o700)).unwrap();
+        let path = format!(
+            "{}:{}",
+            bin.display(),
+            std::env::var("PATH").unwrap_or_default()
+        );
+
+        // Exercise both auto-start and a subsequent attach to the healthy Local server.
+        for args in [&[][..], &["client"][..]] {
+            let client = spawn_client_process_with_args_and_env(
+                &config_home,
+                &runtime_dir,
+                &api_socket,
+                args,
+                &[("PATH", &path)],
+            );
+            let output =
+                spawn_pty_drain(client._master.as_ref().unwrap().try_clone_reader().unwrap());
+            wait_for_socket(&api_socket, Duration::from_secs(10));
+            assert!(wait_until(
+                Duration::from_secs(10),
+                Duration::from_millis(20),
+                || { read_output(&output).contains("Local") }
+            ));
+            let mut input = client._master.as_ref().unwrap().take_writer().unwrap();
+            input
+                .write_all(b"printf 'LOCAL_%s\\n' DIRECT_READY\r")
+                .unwrap();
+            assert!(wait_until(Duration::from_secs(10), Duration::from_millis(20), || {
+                read_output(&output).contains("LOCAL_DIRECT_READY")
+            }), "Local must accept input without waiting for SSH (remote selected: {select_remote}): {}", read_output(&output));
+            let text = read_output(&output);
+            assert!(!text.contains("Local: connecting"), "{text}");
+            assert!(!text.contains("Local: reconnecting"), "{text}");
+            drop(input);
+            drop(client);
+        }
+        let _ = send_json_request(
+            &api_socket,
+            r#"{"id":"stop","method":"server.stop","params":{}}"#,
+        );
+        cleanup_test_base(&base);
+    }
+}
+
+#[test]
+fn federated_client_starts_without_local_and_survives_its_restart() {
+    use std::os::unix::fs::PermissionsExt;
+
+    let _lock = test_lock();
+    let base = unique_test_dir();
+    let config_home = base.join("config");
+    let runtime_dir = base.join("runtime");
+    let api_socket = runtime_dir.join("herdr.sock");
+    let client_socket = runtime_dir.join("herdr-client.sock");
+    let remote_config = base.join("remote-config");
+    let remote_runtime = base.join("remote-runtime");
+    let remote_api = remote_runtime.join("herdr.sock");
+    let remote_client = remote_runtime.join("herdr-client.sock");
+    let mut remote_server =
+        spawn_server(&remote_config, &remote_runtime, &remote_api, &remote_client);
+    wait_for_socket(&remote_api, Duration::from_secs(10));
+    wait_for_socket(&remote_client, Duration::from_secs(10));
+    let created = send_json_request(
+        &remote_api,
+        &serde_json::json!({
+            "id": "remote-workspace", "method": "workspace.create",
+            "params": {"cwd": base, "focus": true, "label": "remote-ready"},
+        })
+        .to_string(),
+    );
+    let remote_pane = created["result"]["root_pane"]["pane_id"].as_str().unwrap();
+    send_pane_shell_command(&remote_api, remote_pane, "printf 'REMOTE_INITIAL_FRAME\\n'");
+
+    fs::create_dir_all(config_home.join(app_dir_name())).unwrap();
+    fs::write(
+        config_home.join(app_dir_name()).join("config.toml"),
+        "onboarding = false\n",
+    )
+    .unwrap();
+    let catalog_dir = runtime_dir
+        .join("state")
+        .join(app_dir_name())
+        .join("client");
+    fs::create_dir_all(&catalog_dir).unwrap();
+    let profile = "0123456789abcdef0123456789abcdef";
+    fs::write(catalog_dir.join("endpoints.json"), serde_json::json!({
+        "version": 1, "selected_profile": profile,
+        "ssh": [{"id": profile, "label": "Test remote", "target": "test-only", "session": "default", "enabled": true}],
+    }).to_string()).unwrap();
+
+    // The SSH executable is private to this client. Discovery and the stdio bridge run the real
+    // binary against a second disposable local server, never the developer's saved hosts.
+    let bin = base.join("bin");
+    fs::create_dir_all(&bin).unwrap();
+    fs::create_dir_all(base.join("home")).unwrap();
+    std::os::unix::fs::symlink(env!("CARGO_BIN_EXE_herdr"), bin.join("herdr")).unwrap();
+    let quote =
+        |path: &std::path::Path| format!("'{}'", path.display().to_string().replace('\'', "'\\''"));
+    fs::write(bin.join("ssh"), format!(
+        "#!/bin/sh\nexport HOME={} XDG_CONFIG_HOME={} XDG_RUNTIME_DIR={} HERDR_SOCKET_PATH={}\nunset HERDR_CLIENT_SOCKET_PATH HERDR_SESSION\nfor arg do last=\"$arg\"; done\nexec /bin/sh -c \"$last\"\n",
+        quote(&base.join("home")), quote(&remote_config), quote(&remote_runtime), quote(&remote_api),
+    )).unwrap();
+    fs::set_permissions(bin.join("ssh"), fs::Permissions::from_mode(0o700)).unwrap();
+    let path = format!(
+        "{}:{}",
+        bin.display(),
+        std::env::var("PATH").unwrap_or_default()
+    );
+    let mut client = spawn_client_process_with_args_and_env(
+        &config_home,
+        &runtime_dir,
+        &api_socket,
+        &["client"],
+        &[("PATH", &path)],
+    );
+    let output = spawn_pty_drain(client._master.as_ref().unwrap().try_clone_reader().unwrap());
+    assert!(
+        wait_until(Duration::from_secs(12), Duration::from_millis(20), || {
+            read_output(&output).contains("REMOTE_INITIAL_FRAME")
+        }),
+        "remote must be usable before Local exists: {}",
+        read_output(&output)
+    );
+
+    let mut local = spawn_server(&config_home, &runtime_dir, &api_socket, &client_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    let created = send_json_request(
+        &api_socket,
+        &serde_json::json!({
+            "id": "local-workspace", "method": "workspace.create",
+            "params": {"cwd": base, "focus": true, "label": "local-online"},
+        })
+        .to_string(),
+    );
+    assert_eq!(created["result"]["type"], "workspace_created");
+    assert!(wait_until(
+        Duration::from_secs(10),
+        Duration::from_millis(20),
+        || read_output(&output).contains("local-online")
+    ));
+
+    local.child.kill().unwrap();
+    local.close_master();
+    drop(local);
+    let watermark = output_len(&output);
+    let mut input = client._master.as_ref().unwrap().take_writer().unwrap();
+    input
+        .write_all(b"printf 'REMOTE_%s\\n' SURVIVED\r")
+        .unwrap();
+    assert!(
+        wait_until(Duration::from_secs(8), Duration::from_millis(20), || {
+            read_output(&output)[watermark..].contains("REMOTE_SURVIVED")
+        }),
+        "Local loss must not interrupt remote input or output"
+    );
+    assert!(client.child.try_wait().unwrap().is_none());
+
+    let restarted = spawn_server(&config_home, &runtime_dir, &api_socket, &client_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    let created = send_json_request(
+        &api_socket,
+        &serde_json::json!({
+            "id": "local-returned", "method": "workspace.create",
+            "params": {"cwd": base, "focus": true, "label": "local-returned"},
+        })
+        .to_string(),
+    );
+    assert_eq!(created["result"]["type"], "workspace_created");
+    assert!(
+        wait_until(Duration::from_secs(12), Duration::from_millis(20), || {
+            read_output(&output).contains("local-returned")
+        }),
+        "Local must reconnect with fresh metadata"
+    );
+    let watermark = output_len(&output);
+    input
+        .write_all(b"printf 'REMOTE_%s\\n' STILL_SELECTED\r")
+        .unwrap();
+    assert!(
+        wait_until(Duration::from_secs(8), Duration::from_millis(20), || {
+            read_output(&output)[watermark..].contains("REMOTE_STILL_SELECTED")
+        }),
+        "Local recovery must not steal selection"
+    );
+    let watermark = output_len(&output);
+    remote_server.child.kill().unwrap();
+    assert!(
+        wait_until(Duration::from_secs(10), Duration::from_millis(20), || {
+            read_output(&output)[watermark..].contains("reconnecting")
+        }),
+        "the selected remote must be marked disconnected"
+    );
+    let text = read_output(&output);
+    assert!(
+        text.rfind("\x1b[?1000h") > text.rfind("\x1b[?1000l"),
+        "losing the selected remote must keep host mouse reporting enabled"
+    );
+
+    let local_pane = created["result"]["root_pane"]["pane_id"].as_str().unwrap();
+    send_pane_shell_command(
+        &api_socket,
+        local_pane,
+        "printf 'LOCAL_RECOVERED_SURFACE\\n'",
+    );
+    let watermark = output_len(&output);
+    // Select the fresh workspace below Local's restored workspace.
+    input.write_all(b"\x1b[<0;7;5M\x1b[<0;7;5m").unwrap();
+    assert!(
+        wait_until(Duration::from_secs(10), Duration::from_millis(20), || {
+            read_output(&output)[watermark..].contains("LOCAL_RECOVERED_SURFACE")
+        }),
+        "recovered Local must be selectable: {}",
+        read_output(&output)
+    );
+    // A coherent frame precedes the final host-effects fence; input stays gated until then.
+    assert!(
+        wait_until(Duration::from_secs(8), Duration::from_millis(100), || {
+            if read_output(&output)[watermark..].contains("LOCAL_INPUT_RECOVERED") {
+                return true;
+            }
+            input
+                .write_all(b"printf 'LOCAL_%s\\n' INPUT_RECOVERED\r")
+                .unwrap();
+            false
+        }),
+        "recovered Local must accept input: {}",
+        read_output(&output)
+    );
+    drop(input);
+    drop(client);
+    drop(restarted);
+    drop(remote_server);
+    cleanup_test_base(&base);
+}
+
+#[test]
 fn client_shell_detaches_restores_and_freshly_reattaches_to_current_state() {
     let _lock = test_lock();
     let base = unique_test_dir();
diff --git a/tests/cross_area.rs b/tests/cross_area.rs
index 01c15c68..3e933c95 100644
--- a/tests/cross_area.rs
+++ b/tests/cross_area.rs
@@ -119,6 +119,7 @@ fn spawn_server_with_path(
 
     let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
     cmd.arg("server");
+    cmd.env("XDG_STATE_HOME", runtime_dir.join("state"));
     cmd.env("XDG_CONFIG_HOME", config_home);
     cmd.env("XDG_RUNTIME_DIR", runtime_dir);
     cmd.env("HERDR_SOCKET_PATH", api_socket_path);
@@ -157,6 +158,7 @@ fn spawn_client_process(
     let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
     cmd.arg("client");
     cmd.env("HERDR_DISABLE_SOUND", "1");
+    cmd.env("XDG_STATE_HOME", runtime_dir.join("state"));
     cmd.env("XDG_CONFIG_HOME", config_home);
     cmd.env("XDG_RUNTIME_DIR", runtime_dir);
     cmd.env("HERDR_SOCKET_PATH", api_socket_path);
diff --git a/tests/fixtures/endpoint-method-shapes-v1.json b/tests/fixtures/endpoint-method-shapes-v1.json
index b473869f..f8ea6444 100644
--- a/tests/fixtures/endpoint-method-shapes-v1.json
+++ b/tests/fixtures/endpoint-method-shapes-v1.json
@@ -1,4 +1,5 @@
 {
+  "client_shell.surface.set": "919df2ba99fe7d5df003a051dbbaf9929998387c53ccb652b6262bd8e4650330",
   "command.invoke": "b16817652442b1ea2405e60a0443826a55baca461e938f524e79b16a675c5899",
   "integration.install": "913055b838166650fee6713332c1cc436efd4f8fae9a2064d8f9b6616a7a66f8",
   "integration.list": "60c57189adf9c281958062f0e11680bda0ece83307c506fa18bff8b01ebfb4e8",