feat: hand off ptys during updates

2 files changed, 839 insertions, 8 deletions

diff --git a/tests/live_handoff.rs b/tests/live_handoff.rs
new file mode 100644
index 00000000..b8e9ad76
--- /dev/null
+++ b/tests/live_handoff.rs
@@ -0,0 +1,817 @@
+mod support;
+
+use std::fs;
+use std::io::{BufRead, BufReader, Read, Write};
+use std::net::{TcpListener, TcpStream};
+use std::os::unix::net::UnixStream;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Mutex, MutexGuard, OnceLock};
+use std::thread;
+use std::time::{Duration, Instant};
+
+use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize};
+use support::{
+    cleanup_test_base, client_handshake, register_runtime_dir, register_spawned_herdr_pid,
+    unregister_spawned_herdr_pid, wait_for_disconnect, wait_for_socket,
+};
+
+struct SpawnedHerdr {
+    _master: Box<dyn MasterPty + Send>,
+    child: Box<dyn Child + Send + Sync>,
+}
+
+impl Drop for SpawnedHerdr {
+    fn drop(&mut self) {
+        let pid = self.child.process_id();
+        let _ = self.child.kill();
+        unregister_spawned_herdr_pid(pid);
+    }
+}
+
+fn test_lock() -> MutexGuard<'static, ()> {
+    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
+    LOCK.get_or_init(|| Mutex::new(()))
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner())
+}
+
+fn unique_test_dir() -> PathBuf {
+    static COUNTER: AtomicUsize = AtomicUsize::new(0);
+    let n = COUNTER.fetch_add(1, Ordering::Relaxed);
+    PathBuf::from(format!("/tmp/hlh-{}-{n}", std::process::id()))
+}
+
+fn spawn_server(config_home: &Path, runtime_dir: &Path, api_socket: &Path) -> SpawnedHerdr {
+    spawn_server_with_env(config_home, runtime_dir, api_socket, &[])
+}
+
+fn spawn_server_with_env(
+    config_home: &Path,
+    runtime_dir: &Path,
+    api_socket: &Path,
+    extra_env: &[(&str, &str)],
+) -> SpawnedHerdr {
+    fs::create_dir_all(config_home.join("herdr")).unwrap();
+    fs::create_dir_all(runtime_dir).unwrap();
+    fs::write(
+        config_home.join("herdr/config.toml"),
+        "onboarding = false\n",
+    )
+    .unwrap();
+
+    let pair = native_pty_system()
+        .openpty(PtySize {
+            rows: 24,
+            cols: 80,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .unwrap();
+    let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
+    cmd.arg("server");
+    cmd.env("XDG_CONFIG_HOME", config_home);
+    cmd.env("XDG_RUNTIME_DIR", runtime_dir);
+    cmd.env("HERDR_SOCKET_PATH", api_socket);
+    cmd.env(
+        "HERDR_CLIENT_SOCKET_PATH",
+        runtime_dir.join("herdr-client.sock"),
+    );
+    cmd.env("SHELL", "/bin/sh");
+    for (key, value) in extra_env {
+        cmd.env(key, value);
+    }
+
+    let child = pair.slave.spawn_command(cmd).unwrap();
+    register_spawned_herdr_pid(child.process_id());
+    SpawnedHerdr {
+        _master: pair.master,
+        child,
+    }
+}
+
+fn spawn_named_session_server(
+    config_home: &Path,
+    runtime_dir: &Path,
+    session_name: &str,
+) -> SpawnedHerdr {
+    fs::create_dir_all(config_home.join("herdr-dev")).unwrap();
+    fs::create_dir_all(runtime_dir).unwrap();
+    fs::write(
+        config_home.join("herdr-dev/config.toml"),
+        "onboarding = false\n",
+    )
+    .unwrap();
+
+    let pair = native_pty_system()
+        .openpty(PtySize {
+            rows: 24,
+            cols: 80,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .unwrap();
+    let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
+    cmd.arg("server");
+    cmd.env("XDG_CONFIG_HOME", config_home);
+    cmd.env("XDG_RUNTIME_DIR", runtime_dir);
+    cmd.env("HERDR_SESSION", session_name);
+    cmd.env_remove("HERDR_SOCKET_PATH");
+    cmd.env_remove("HERDR_CLIENT_SOCKET_PATH");
+    cmd.env("SHELL", "/bin/sh");
+
+    let child = pair.slave.spawn_command(cmd).unwrap();
+    register_spawned_herdr_pid(child.process_id());
+    SpawnedHerdr {
+        _master: pair.master,
+        child,
+    }
+}
+
+fn spawn_default_session_server(config_home: &Path, runtime_dir: &Path) -> SpawnedHerdr {
+    fs::create_dir_all(config_home.join("herdr-dev")).unwrap();
+    fs::create_dir_all(runtime_dir).unwrap();
+    fs::write(
+        config_home.join("herdr-dev/config.toml"),
+        "onboarding = false\n",
+    )
+    .unwrap();
+
+    let pair = native_pty_system()
+        .openpty(PtySize {
+            rows: 24,
+            cols: 80,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .unwrap();
+    let mut cmd = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
+    cmd.arg("server");
+    cmd.env("XDG_CONFIG_HOME", config_home);
+    cmd.env("XDG_RUNTIME_DIR", runtime_dir);
+    cmd.env_remove("HERDR_SESSION");
+    cmd.env_remove("HERDR_SOCKET_PATH");
+    cmd.env_remove("HERDR_CLIENT_SOCKET_PATH");
+    cmd.env("SHELL", "/bin/sh");
+
+    let child = pair.slave.spawn_command(cmd).unwrap();
+    register_spawned_herdr_pid(child.process_id());
+    SpawnedHerdr {
+        _master: pair.master,
+        child,
+    }
+}
+
+fn request(socket_path: &Path, request: serde_json::Value) -> serde_json::Value {
+    let mut stream = UnixStream::connect(socket_path).unwrap();
+    stream.write_all(request.to_string().as_bytes()).unwrap();
+    stream.write_all(b"\n").unwrap();
+    stream.flush().unwrap();
+    let mut line = String::new();
+    BufReader::new(stream).read_line(&mut line).unwrap();
+    serde_json::from_str(&line).unwrap()
+}
+
+fn assert_ok(response: serde_json::Value) {
+    assert!(
+        response.get("result").is_some(),
+        "api request failed: {response}"
+    );
+}
+
+fn wait_for_api(socket_path: &Path, timeout: Duration) {
+    let deadline = Instant::now() + timeout;
+    while Instant::now() < deadline {
+        if UnixStream::connect(socket_path).is_ok() {
+            let response = request(
+                socket_path,
+                serde_json::json!({"id":"test:ping","method":"ping","params":{}}),
+            );
+            if response.get("result").is_some() {
+                return;
+            }
+        }
+        thread::sleep(Duration::from_millis(25));
+    }
+    panic!("api did not become ready at {}", socket_path.display());
+}
+
+fn wait_for_output(socket_path: &Path, pane_id: &str, needle: &str) {
+    let deadline = Instant::now() + Duration::from_secs(5);
+    let mut last_text = String::new();
+    let mut last_response = serde_json::Value::Null;
+    while Instant::now() < deadline {
+        let response = request(
+            socket_path,
+            serde_json::json!({
+                "id": "test:pane:read",
+                "method": "pane.read",
+                "params": {
+                    "pane_id": pane_id,
+                    "source": "visible",
+                    "lines": 20,
+                    "format": "text",
+                    "strip_ansi": true
+                }
+            }),
+        );
+        last_response = response.clone();
+        let text = response["result"]["read"]["text"]
+            .as_str()
+            .unwrap_or_default();
+        last_text = text.to_string();
+        if text.contains(needle) {
+            return;
+        }
+        thread::sleep(Duration::from_millis(50));
+    }
+    panic!(
+        "pane output did not contain {needle:?}; last text was {last_text:?}; last response was {last_response}"
+    );
+}
+
+fn wait_for_file_contains(path: &Path, needle: &str, timeout: Duration) -> String {
+    let deadline = Instant::now() + timeout;
+    let mut last_text = String::new();
+    while Instant::now() < deadline {
+        if let Ok(text) = fs::read_to_string(path) {
+            last_text = text;
+            if last_text.contains(needle) {
+                return last_text;
+            }
+        }
+        thread::sleep(Duration::from_millis(50));
+    }
+    panic!(
+        "{} did not contain {needle:?}; last text was {last_text:?}",
+        path.display()
+    );
+}
+
+fn unused_local_port() -> u16 {
+    TcpListener::bind("127.0.0.1:0")
+        .unwrap()
+        .local_addr()
+        .unwrap()
+        .port()
+}
+
+fn wait_for_http_contains(port: u16, needle: &str, timeout: Duration) -> String {
+    let deadline = Instant::now() + timeout;
+    let mut last_response = String::new();
+    while Instant::now() < deadline {
+        if let Ok(mut stream) = TcpStream::connect(("127.0.0.1", port)) {
+            let _ =
+                stream.write_all(b"GET / HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n");
+            let mut response = String::new();
+            let _ = stream.read_to_string(&mut response);
+            last_response = response;
+            if last_response.contains(needle) {
+                return last_response;
+            }
+        }
+        thread::sleep(Duration::from_millis(50));
+    }
+    panic!(
+        "http server on port {port} did not return {needle:?}; last response was {last_response:?}"
+    );
+}
+
+#[test]
+fn live_handoff_preserves_named_session_socket_paths() {
+    let _lock = test_lock();
+    let base = unique_test_dir();
+    let config_home = base.join("config");
+    let runtime_dir = base.join("runtime");
+    let session_dir = config_home.join("herdr-dev/sessions/work");
+    let api_socket = session_dir.join("herdr.sock");
+    let client_socket = session_dir.join("herdr-client.sock");
+
+    let spawned = spawn_named_session_server(&config_home, &runtime_dir, "work");
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    register_runtime_dir(&runtime_dir);
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({"id":"test:handoff","method":"server.live_handoff","params":{}}),
+    ));
+    drop(spawned);
+    wait_for_api(&api_socket, Duration::from_secs(10));
+    wait_for_socket(&client_socket, Duration::from_secs(5));
+    assert!(
+        !config_home.join("herdr-dev/herdr.sock").exists(),
+        "named handoff unexpectedly bound the default session API socket"
+    );
+
+    let _ = request(
+        &api_socket,
+        serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+    );
+    cleanup_test_base(&base);
+}
+
+#[test]
+fn live_handoff_preserves_pane_process_io() {
+    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 marker = base.join("child.pid");
+    let second_marker = base.join("second-child.pid");
+    let hup_marker = base.join("hup");
+    let second_hup_marker = base.join("second-hup");
+    let received_marker = base.join("received");
+    let second_received_marker = base.join("second-received");
+
+    let spawned = spawn_server(&config_home, &runtime_dir, &api_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    register_runtime_dir(&runtime_dir);
+
+    let created = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:workspace:create",
+            "method": "workspace.create",
+            "params": {"cwd": "/tmp", "focus": true}
+        }),
+    );
+    let pane_id = created["result"]["root_pane"]["pane_id"]
+        .as_str()
+        .unwrap()
+        .to_string();
+    let split = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:split",
+            "method": "pane.split",
+            "params": {
+                "target_pane_id": pane_id,
+                "direction": "right",
+                "focus": false
+            }
+        }),
+    );
+    assert_ok(split.clone());
+    let second_pane_id = split["result"]["pane"]["pane_id"]
+        .as_str()
+        .unwrap()
+        .to_string();
+
+    let command = format!(
+        "sh -c 'echo READY $$ > {}; trap \"echo HUP >> {}\" HUP; while read line; do echo got:$line; echo got:$line >> {}; done'",
+        marker.display(),
+        hup_marker.display(),
+        received_marker.display()
+    );
+    let second_command = format!(
+        "sh -c 'echo SECOND_READY $$ > {}; trap \"echo HUP >> {}\" HUP; while read line; do echo second:$line; echo second:$line >> {}; done'",
+        second_marker.display(),
+        second_hup_marker.display(),
+        second_received_marker.display()
+    );
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:run",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": command, "keys": ["Enter"]}
+        }),
+    ));
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:second-pane:run",
+            "method": "pane.send_input",
+            "params": {"pane_id": second_pane_id, "text": second_command, "keys": ["Enter"]}
+        }),
+    ));
+    support::wait_for_file(&marker, Duration::from_secs(5));
+    support::wait_for_file(&second_marker, Duration::from_secs(5));
+    let pid_text = fs::read_to_string(&marker).unwrap();
+    let child_pid: u32 = pid_text.split_whitespace().last().unwrap().parse().unwrap();
+    let second_pid_text = fs::read_to_string(&second_marker).unwrap();
+    let second_child_pid: u32 = second_pid_text
+        .split_whitespace()
+        .last()
+        .unwrap()
+        .parse()
+        .unwrap();
+    assert_eq!(unsafe { libc::kill(child_pid as libc::pid_t, 0) }, 0);
+    assert_eq!(unsafe { libc::kill(second_child_pid as libc::pid_t, 0) }, 0);
+
+    let protocol = request(
+        &api_socket,
+        serde_json::json!({"id":"test:protocol","method":"ping","params":{}}),
+    )["result"]["protocol"]
+        .as_u64()
+        .unwrap() as u32;
+    let mut client_stream = UnixStream::connect(&client_socket).unwrap();
+    let (server_protocol, error) = client_handshake(&mut client_stream, protocol, 80, 24).unwrap();
+    assert_eq!(server_protocol, protocol);
+    assert!(error.is_none(), "client handshake failed: {error:?}");
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:before-log",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": "before_replay", "keys": ["Enter"]}
+        }),
+    ));
+    wait_for_output(&api_socket, &pane_id, "got:before_replay");
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({"id":"test:handoff","method":"server.live_handoff","params":{}}),
+    ));
+    drop(spawned);
+    assert!(
+        wait_for_disconnect(&mut client_stream, Duration::from_secs(5)).unwrap(),
+        "connected clients should disconnect during live handoff"
+    );
+    thread::sleep(Duration::from_millis(300));
+    wait_for_api(&api_socket, Duration::from_secs(10));
+    wait_for_socket(&client_socket, Duration::from_secs(5));
+    assert_eq!(unsafe { libc::kill(child_pid as libc::pid_t, 0) }, 0);
+    assert_eq!(unsafe { libc::kill(second_child_pid as libc::pid_t, 0) }, 0);
+    assert!(
+        !hup_marker.exists(),
+        "pane process received HUP during handoff"
+    );
+    assert!(
+        !second_hup_marker.exists(),
+        "second pane process received HUP during handoff"
+    );
+    wait_for_output(&api_socket, &pane_id, "got:before_replay");
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:send",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": "after-handoff", "keys": ["Enter"]}
+        }),
+    ));
+    wait_for_file_contains(
+        &received_marker,
+        "got:after-handoff",
+        Duration::from_secs(5),
+    );
+    wait_for_output(&api_socket, &pane_id, "got:after-handoff");
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:second-pane:send",
+            "method": "pane.send_input",
+            "params": {"pane_id": second_pane_id, "text": "after-handoff-second", "keys": ["Enter"]}
+        }),
+    ));
+    wait_for_file_contains(
+        &second_received_marker,
+        "second:after-handoff-second",
+        Duration::from_secs(5),
+    );
+    wait_for_output(&api_socket, &second_pane_id, "second:after-handoff-sec");
+
+    let _ = request(
+        &api_socket,
+        serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+    );
+    let _ = client_socket;
+    cleanup_test_base(&base);
+}
+
+#[test]
+fn live_handoff_preserves_python_http_server() {
+    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 web_root = base.join("web");
+    fs::create_dir_all(&web_root).unwrap();
+    fs::write(
+        web_root.join("index.html"),
+        "hello-from-python-before-and-after",
+    )
+    .unwrap();
+    let port = unused_local_port();
+
+    let spawned = spawn_server(&config_home, &runtime_dir, &api_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    register_runtime_dir(&runtime_dir);
+
+    let created = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:workspace:create",
+            "method": "workspace.create",
+            "params": {"cwd": web_root, "focus": true}
+        }),
+    );
+    let pane_id = created["result"]["root_pane"]["pane_id"]
+        .as_str()
+        .unwrap()
+        .to_string();
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:run-python",
+            "method": "pane.send_input",
+            "params": {
+                "pane_id": pane_id,
+                "text": format!("python3 -m http.server {port} --bind 127.0.0.1"),
+                "keys": ["Enter"]
+            }
+        }),
+    ));
+    wait_for_http_contains(
+        port,
+        "hello-from-python-before-and-after",
+        Duration::from_secs(10),
+    );
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({"id":"test:handoff","method":"server.live_handoff","params":{}}),
+    ));
+    drop(spawned);
+    wait_for_api(&api_socket, Duration::from_secs(10));
+    wait_for_http_contains(
+        port,
+        "hello-from-python-before-and-after",
+        Duration::from_secs(10),
+    );
+
+    let _ = request(
+        &api_socket,
+        serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+    );
+    let _ = client_socket;
+    cleanup_test_base(&base);
+}
+
+#[test]
+fn live_handoff_preserves_http_servers_across_multiple_sessions() {
+    let _lock = test_lock();
+    let base = unique_test_dir();
+    let config_home = base.join("config");
+    let runtime_dir = base.join("runtime");
+    let sessions = [
+        (None, config_home.join("herdr-dev/herdr.sock")),
+        (
+            Some("work"),
+            config_home.join("herdr-dev/sessions/work/herdr.sock"),
+        ),
+    ];
+    let mut spawned = Vec::new();
+    let mut ports = Vec::new();
+
+    for (session_name, api_socket) in &sessions {
+        let web_root = base.join(format!("web-{}", session_name.unwrap_or("default")));
+        fs::create_dir_all(&web_root).unwrap();
+        fs::write(
+            web_root.join("index.html"),
+            format!("hello-from-{}", session_name.unwrap_or("default")),
+        )
+        .unwrap();
+        let port = unused_local_port();
+        let server = if let Some(session_name) = session_name {
+            spawn_named_session_server(&config_home, &runtime_dir, session_name)
+        } else {
+            spawn_default_session_server(&config_home, &runtime_dir)
+        };
+        wait_for_socket(api_socket, Duration::from_secs(10));
+        let created = request(
+            api_socket,
+            serde_json::json!({
+                "id": "test:workspace:create",
+                "method": "workspace.create",
+                "params": {"cwd": web_root, "focus": true}
+            }),
+        );
+        let pane_id = created["result"]["root_pane"]["pane_id"]
+            .as_str()
+            .unwrap()
+            .to_string();
+        assert_ok(request(
+            api_socket,
+            serde_json::json!({
+                "id": "test:pane:run-python",
+                "method": "pane.send_input",
+                "params": {
+                    "pane_id": pane_id,
+                    "text": format!("python3 -m http.server {port} --bind 127.0.0.1"),
+                    "keys": ["Enter"]
+                }
+            }),
+        ));
+        wait_for_http_contains(
+            port,
+            &format!("hello-from-{}", session_name.unwrap_or("default")),
+            Duration::from_secs(10),
+        );
+        spawned.push(server);
+        ports.push((port, session_name.unwrap_or("default").to_string()));
+    }
+    register_runtime_dir(&runtime_dir);
+
+    for (_session_name, api_socket) in &sessions {
+        assert_ok(request(
+            api_socket,
+            serde_json::json!({"id":"test:handoff","method":"server.live_handoff","params":{}}),
+        ));
+    }
+    drop(spawned);
+
+    for (_session_name, api_socket) in &sessions {
+        wait_for_api(api_socket, Duration::from_secs(10));
+    }
+    for (port, label) in ports {
+        wait_for_http_contains(
+            port,
+            &format!("hello-from-{label}"),
+            Duration::from_secs(10),
+        );
+    }
+
+    for (_session_name, api_socket) in &sessions {
+        let _ = request(
+            api_socket,
+            serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+        );
+    }
+    cleanup_test_base(&base);
+}
+
+#[test]
+fn live_handoff_bad_expected_protocol_rolls_back_old_server() {
+    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 marker = base.join("child.pid");
+    let received_marker = base.join("received");
+
+    let spawned = spawn_server(&config_home, &runtime_dir, &api_socket);
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    register_runtime_dir(&runtime_dir);
+
+    let created = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:workspace:create",
+            "method": "workspace.create",
+            "params": {"cwd": "/tmp", "focus": true}
+        }),
+    );
+    let pane_id = created["result"]["root_pane"]["pane_id"]
+        .as_str()
+        .unwrap()
+        .to_string();
+    let command = format!(
+        "sh -c 'echo READY $$ > {}; while read line; do echo got:$line; echo got:$line >> {}; done'",
+        marker.display(),
+        received_marker.display()
+    );
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:run",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": command, "keys": ["Enter"]}
+        }),
+    ));
+    support::wait_for_file(&marker, Duration::from_secs(5));
+    let pid_text = fs::read_to_string(&marker).unwrap();
+    let child_pid: u32 = pid_text.split_whitespace().last().unwrap().parse().unwrap();
+
+    let failed = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:bad-handoff",
+            "method": "server.live_handoff",
+            "params": {"expected_protocol": 999999}
+        }),
+    );
+    assert!(
+        failed.get("error").is_some(),
+        "bad protocol handoff should fail: {failed}"
+    );
+    wait_for_api(&api_socket, Duration::from_secs(5));
+    assert_eq!(unsafe { libc::kill(child_pid as libc::pid_t, 0) }, 0);
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:send-after-failed-handoff",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": "after-failed-handoff", "keys": ["Enter"]}
+        }),
+    ));
+    wait_for_file_contains(
+        &received_marker,
+        "got:after-failed-handoff",
+        Duration::from_secs(5),
+    );
+    wait_for_output(&api_socket, &pane_id, "got:after-failed-handoff");
+
+    let _ = request(
+        &api_socket,
+        serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+    );
+    drop(spawned);
+    cleanup_test_base(&base);
+}
+
+fn live_handoff_import_failure_rolls_back_old_server_at(failure_point: &str) {
+    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 marker = base.join("child.pid");
+    let received_marker = base.join("received");
+
+    let spawned = spawn_server_with_env(
+        &config_home,
+        &runtime_dir,
+        &api_socket,
+        &[("HERDR_TEST_HANDOFF_IMPORT_FAIL", failure_point)],
+    );
+    wait_for_socket(&api_socket, Duration::from_secs(10));
+    register_runtime_dir(&runtime_dir);
+
+    let created = request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:workspace:create",
+            "method": "workspace.create",
+            "params": {"cwd": "/tmp", "focus": true}
+        }),
+    );
+    let pane_id = created["result"]["root_pane"]["pane_id"]
+        .as_str()
+        .unwrap()
+        .to_string();
+    let command = format!(
+        "sh -c 'echo READY $$ > {}; while read line; do echo got:$line; echo got:$line >> {}; done'",
+        marker.display(),
+        received_marker.display()
+    );
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:run",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": command, "keys": ["Enter"]}
+        }),
+    ));
+    support::wait_for_file(&marker, Duration::from_secs(5));
+    let pid_text = fs::read_to_string(&marker).unwrap();
+    let child_pid: u32 = pid_text.split_whitespace().last().unwrap().parse().unwrap();
+
+    let failed = request(
+        &api_socket,
+        serde_json::json!({"id":"test:handoff-fail","method":"server.live_handoff","params":{}}),
+    );
+    assert!(
+        failed.get("error").is_some(),
+        "{failure_point} handoff should fail: {failed}"
+    );
+    wait_for_api(&api_socket, Duration::from_secs(10));
+    wait_for_socket(&client_socket, Duration::from_secs(5));
+    assert_eq!(unsafe { libc::kill(child_pid as libc::pid_t, 0) }, 0);
+
+    assert_ok(request(
+        &api_socket,
+        serde_json::json!({
+            "id": "test:pane:send-after-import-failure",
+            "method": "pane.send_input",
+            "params": {"pane_id": pane_id, "text": failure_point, "keys": ["Enter"]}
+        }),
+    ));
+    wait_for_file_contains(
+        &received_marker,
+        &format!("got:{failure_point}"),
+        Duration::from_secs(5),
+    );
+
+    let _ = request(
+        &api_socket,
+        serde_json::json!({"id":"test:stop","method":"server.stop","params":{}}),
+    );
+    drop(spawned);
+    cleanup_test_base(&base);
+}
+
+#[test]
+fn live_handoff_after_restored_failure_rolls_back_old_server() {
+    live_handoff_import_failure_rolls_back_old_server_at("after_restored");
+}
diff --git a/tests/support/mod.rs b/tests/support/mod.rs
index e810b843..cae140c9 100644
--- a/tests/support/mod.rs
+++ b/tests/support/mod.rs
@@ -328,16 +328,30 @@ pub fn wait_for_message_variant(
 }
 
 pub fn wait_for_disconnect(stream: &mut UnixStream, timeout: Duration) -> Result<bool, String> {
-    stream
-        .set_read_timeout(Some(Duration::from_millis(200)))
-        .map_err(|e| e.to_string())?;
+    stream.set_nonblocking(true).map_err(|e| e.to_string())?;
     let deadline = Instant::now() + timeout;
-    while Instant::now() < deadline {
-        if read_server_message(stream).is_err() {
-            return Ok(true);
+    let mut idle_since = None;
+    let result = loop {
+        match read_server_message(stream) {
+            Ok(_) => idle_since = None,
+            Err(err)
+                if err.to_ascii_lowercase().contains("would block")
+                    || err.contains("Resource temporarily unavailable") =>
+            {
+                let idle_started = *idle_since.get_or_insert_with(Instant::now);
+                if idle_started.elapsed() >= Duration::from_millis(200) {
+                    break Ok(true);
+                }
+            }
+            Err(_) => break Ok(true),
         }
-    }
-    Ok(false)
+        if Instant::now() >= deadline {
+            break Ok(false);
+        }
+        thread::sleep(Duration::from_millis(25));
+    };
+    let _ = stream.set_nonblocking(false);
+    result
 }
 
 pub fn cleanup_registered_herdr_pids() {