maestro/crates/maestro-daemon/tests/attach.rs

//! End-to-end tests for the per-session attach WS endpoint:
//! - unknown session is rejected
//! - PTY bytes round-trip through real tmux
//! - multi-attach: two clients see the same output

mod common;

use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

use common::DaemonFixture;
use futures_util::{SinkExt, StreamExt};
use maestro_daemon::{local_api, SessionManager};
use maestro_protocol::CreateSessionRequest;
use tempfile::TempDir;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use uuid::Uuid;

type WsClient = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;

struct Harness {
    addr: SocketAddr,
    manager: SessionManager,
    _shutdown: oneshot::Sender<()>,
    _fx: DaemonFixture,
}

async fn spawn_daemon_ws() -> Harness {
    let fx = DaemonFixture::new();
    let manager = fx.manager.clone();
    let router = local_api::router(local_api::state_for(
        manager.clone(),
        "test-host",
        "/home/test",
        Vec::new(),
        std::sync::Arc::new(maestro_daemon::projects::ProjectScanner::disabled()),
        maestro_config::DEFAULT_LEADER,
    ));
    let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
        .await
        .unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, router)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    Harness {
        addr,
        manager,
        _shutdown: tx,
        _fx: fx,
    }
}

#[allow(clippy::result_large_err)]
async fn connect_attach(
    addr: SocketAddr,
    session: Uuid,
) -> Result<WsClient, tokio_tungstenite::tungstenite::Error> {
    let url = format!("ws://{addr}/api/attach/{session}");
    connect_async(url).await.map(|(ws, _)| ws)
}

async fn create_bash_session(h: &Harness) -> Uuid {
    let cwd = TempDir::new().unwrap();
    let cwd_path = cwd.path().to_path_buf();
    let s = h
        .manager
        .create(CreateSessionRequest {
            cwd: cwd_path,
            command: "bash --noprofile --norc".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .unwrap();
    std::mem::forget(cwd); // keep dir for the test's duration
    s.id
}

#[tokio::test]
async fn rejects_unknown_session() {
    let h = spawn_daemon_ws().await;
    let err = connect_attach(h.addr, Uuid::new_v4()).await.unwrap_err();
    assert!(format!("{err}").contains("404"));
}

#[tokio::test]
async fn pty_round_trip_through_real_tmux() {
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;
    let mut ws = connect_attach(h.addr, id).await.unwrap();

    ws.send(WsMessage::Binary(b"echo MAESTRO_RT_OK\n".to_vec()))
        .await
        .unwrap();

    let want = b"MAESTRO_RT_OK";
    let mut got: Vec<u8> = Vec::new();
    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    while std::time::Instant::now() < deadline {
        if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
            tokio::time::timeout(Duration::from_millis(50), ws.next()).await
        {
            got.extend_from_slice(&bytes);
            if got.windows(want.len()).any(|w| w == want) {
                break;
            }
        }
    }
    assert!(
        got.windows(want.len()).any(|w| w == want),
        "marker not seen in {} bytes",
        got.len()
    );
    let _ = h.manager.kill(id, false).await;
}

#[tokio::test]
async fn multi_attach_broadcasts_to_every_client() {
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;

    let mut ws_a = connect_attach(h.addr, id).await.unwrap();
    let mut ws_b = connect_attach(h.addr, id).await.unwrap();

    // give the broadcast subscribers a moment to register
    tokio::time::sleep(Duration::from_millis(100)).await;

    ws_a.send(WsMessage::Binary(b"echo MAESTRO_MULTI_OK\n".to_vec()))
        .await
        .unwrap();

    let want = b"MAESTRO_MULTI_OK";
    let mut got_a: Vec<u8> = Vec::new();
    let mut got_b: Vec<u8> = Vec::new();
    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    while std::time::Instant::now() < deadline {
        if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
            tokio::time::timeout(Duration::from_millis(50), ws_a.next()).await
        {
            got_a.extend_from_slice(&bytes);
        }
        if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
            tokio::time::timeout(Duration::from_millis(50), ws_b.next()).await
        {
            got_b.extend_from_slice(&bytes);
        }
        if got_a.windows(want.len()).any(|w| w == want)
            && got_b.windows(want.len()).any(|w| w == want)
        {
            break;
        }
    }
    assert!(
        got_a.windows(want.len()).any(|w| w == want),
        "client A missed marker"
    );
    assert!(
        got_b.windows(want.len()).any(|w| w == want),
        "client B missed marker"
    );

    let _ = h.manager.kill(id, false).await;
}

#[tokio::test]
async fn sends_close_frame_when_tmux_session_dies_under_liveness_watcher() {
    // Bug NAP-226: when bash exits, tmux destroys the session, the liveness
    // watcher reaps it, and the broadcast channel closes — but the WS used to
    // be dropped silently, leaving the client hung in raw mode.
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;
    let mut ws = connect_attach(h.addr, id).await.unwrap();

    tokio::time::sleep(Duration::from_millis(100)).await;

    h.manager.start_liveness_watcher(Duration::from_millis(50));

    // Kill the tmux session out-of-band; the watcher should reap and the
    // broadcast channel should close, triggering serve_attach to send Close.
    h._fx.tmux.tmux.kill_session(id).expect("kill tmux");

    let mut saw_close = false;
    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    while std::time::Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_millis(100), ws.next()).await {
            Ok(Some(Ok(WsMessage::Close(_)))) => {
                saw_close = true;
                break;
            }
            Ok(None) => break,
            Ok(Some(Err(_))) => break,
            Ok(_) => continue,
            Err(_) => continue,
        }
    }
    assert!(
        saw_close,
        "daemon must send a Close frame when the liveness watcher reaps the session"
    );
}

#[tokio::test]
async fn sends_close_frame_when_session_is_killed() {
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;
    let mut ws = connect_attach(h.addr, id).await.unwrap();

    // give the broadcast subscriber a moment to register before killing.
    tokio::time::sleep(Duration::from_millis(100)).await;

    h.manager.kill(id, false).await.unwrap();

    let mut saw_close = false;
    let deadline = std::time::Instant::now() + Duration::from_secs(2);
    while std::time::Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_millis(100), ws.next()).await {
            Ok(Some(Ok(WsMessage::Close(_)))) => {
                saw_close = true;
                break;
            }
            Ok(None) => {
                // socket dropped without a Close frame — that's exactly the
                // bug NAP-226 is fixing.
                break;
            }
            Ok(Some(Err(_))) => break,
            Ok(_) => continue,
            Err(_) => continue,
        }
    }
    assert!(
        saw_close,
        "daemon must send a Close frame when the session dies so the client can exit cleanly"
    );
}

#[tokio::test]
async fn ping_text_frame_gets_pong() {
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;
    let mut ws = connect_attach(h.addr, id).await.unwrap();

    let ping = serde_json::to_string(&maestro_protocol::ClientControl::Ping).unwrap();
    ws.send(WsMessage::Text(ping)).await.unwrap();

    let mut saw_pong = false;
    let deadline = std::time::Instant::now() + Duration::from_secs(2);
    while std::time::Instant::now() < deadline {
        if let Ok(Some(Ok(WsMessage::Text(t)))) =
            tokio::time::timeout(Duration::from_millis(100), ws.next()).await
        {
            if t.contains("Pong") {
                saw_pong = true;
                break;
            }
        }
    }
    assert!(saw_pong, "no Pong frame received");
    let _ = h.manager.kill(id, false).await;
}

#[tokio::test]
async fn second_attach_sees_replay_buffer_of_prior_output() {
    // NAP-229: a fresh attach must receive the buffered scrollback before
    // any live frames, so reconnecting to an idle prompt isn't blank.
    let h = spawn_daemon_ws().await;
    let session = create_bash_session(&h).await;

    let mut a = connect_attach(h.addr, session).await.unwrap();
    // Generate output the bash will echo back.
    a.send(WsMessage::Binary(b"echo MAESTRO_REPLAY_MARK\n".to_vec()))
        .await
        .unwrap();
    // Drain at least until we see the mark on client A.
    let want = b"MAESTRO_REPLAY_MARK";
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
    let mut buf_a = Vec::new();
    while std::time::Instant::now() < deadline {
        tokio::select! {
            msg = a.next() => match msg {
                Some(Ok(WsMessage::Binary(bytes))) => buf_a.extend(bytes),
                Some(Ok(_)) => {}
                _ => break,
            },
            _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {}
        }
        if buf_a.windows(want.len()).any(|w| w == want) {
            break;
        }
    }
    assert!(
        buf_a.windows(want.len()).any(|w| w == want),
        "client A never saw its own echo"
    );
    let _ = a.close(None).await;

    // Now attach a fresh client. It should see the prior MARK in the replay
    // even though no new output has happened since.
    let mut b = connect_attach(h.addr, session).await.unwrap();
    let saw_replay = tokio::time::timeout(std::time::Duration::from_secs(2), async {
        let mut buf = Vec::new();
        while let Some(Ok(msg)) = b.next().await {
            if let WsMessage::Binary(bytes) = msg {
                buf.extend(bytes);
                if buf.windows(want.len()).any(|w| w == want) {
                    return true;
                }
            }
        }
        false
    })
    .await
    .unwrap_or(false);
    assert!(
        saw_replay,
        "fresh attach did not receive replay containing prior output"
    );
}

/// The point of keeping a dead session: attach afterwards and read what it
/// printed before it died. The stream then ends rather than hanging.
#[tokio::test]
async fn attaching_to_a_dead_session_replays_its_final_output() {
    let h = spawn_daemon_ws().await;
    let id = create_bash_session(&h).await;

    // Produce output, then let the session die out-of-band.
    {
        let mut ws = connect_attach(h.addr, id).await.unwrap();
        ws.send(WsMessage::Binary(b"echo last-words\n".to_vec()))
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(300)).await;
    }
    h.manager.start_liveness_watcher(Duration::from_millis(50));
    h._fx.tmux.tmux.kill_session(id).expect("kill tmux");

    // Wait for the daemon to notice.
    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    while std::time::Instant::now() < deadline {
        if h.manager.status(id).await == Some(maestro_protocol::SessionStatus::Dead) {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
    assert_eq!(
        h.manager.status(id).await,
        Some(maestro_protocol::SessionStatus::Dead),
        "session should be marked dead, not removed"
    );

    let mut ws = connect_attach(h.addr, id)
        .await
        .expect("a dead session must still accept an attach");

    let mut seen = Vec::new();
    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    while std::time::Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_millis(200), ws.next()).await {
            Ok(Some(Ok(WsMessage::Binary(b)))) => seen.extend_from_slice(&b),
            Ok(Some(Ok(WsMessage::Close(_)))) | Ok(None) => break,
            Ok(Some(Ok(_))) => continue,
            Ok(Some(Err(_))) => break,
            Err(_) => continue,
        }
    }
    let text = String::from_utf8_lossy(&seen);
    assert!(
        text.contains("last-words"),
        "replay of a dead session lost its final output: {text:?}"
    );
}