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

//! End-to-end tests for the per-session attach WS endpoint:
//! - bad / good JWT
//! - 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 chrono::{Duration as ChronoDuration, Utc};
use common::DaemonFixture;
use futures_util::{SinkExt, StreamExt};
use maestro_daemon::{attach, issue_jwt, JwtVerifier, PairAttachVerifier, SessionManager};
use maestro_protocol::{CreateSessionRequest, JwtClaims};
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;

const TOKEN: &str = "shared-secret-for-tests";

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

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

async fn spawn_daemon_ws(daemon_id: &str) -> Harness {
    let fx = DaemonFixture::new(daemon_id);
    let manager = fx.manager.clone();
    let verifier = JwtVerifier::new(TOKEN, daemon_id);
    let pair_verifier = PairAttachVerifier::new(TOKEN, daemon_id);
    let router = attach::router(manager.clone(), verifier, pair_verifier);
    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,
    }
}

fn token_for(session_id: Uuid, daemon_id: &str, valid_seconds: i64) -> String {
    issue_jwt(
        TOKEN,
        &JwtClaims {
            session_id,
            daemon_id: daemon_id.into(),
            exp: (Utc::now() + ChronoDuration::seconds(valid_seconds)).timestamp(),
            iat: Some(Utc::now().timestamp()),
        },
    )
    .unwrap()
}

async fn connect_attach(
    addr: SocketAddr,
    session: Uuid,
    token: &str,
) -> Result<WsClient, tokio_tungstenite::tungstenite::Error> {
    let url = format!("ws://{addr}/sessions/{session}/stream?token={token}");
    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 {
            daemon_id: "laptop".into(),
            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_bad_token() {
    let h = spawn_daemon_ws("laptop").await;
    let id = create_bash_session(&h).await;
    let err = connect_attach(h.addr, id, "not-a-jwt").await.unwrap_err();
    let msg = format!("{err}");
    assert!(msg.contains("401") || msg.contains("403"), "{msg}");
    let _ = h.manager.kill(id).await;
}

#[tokio::test]
async fn rejects_token_for_different_session() {
    let h = spawn_daemon_ws("laptop").await;
    let real_id = create_bash_session(&h).await;
    let other_id = Uuid::new_v4();
    let token = token_for(other_id, "laptop", 300);
    let err = connect_attach(h.addr, real_id, &token).await.unwrap_err();
    assert!(format!("{err}").contains("403"));
    let _ = h.manager.kill(real_id).await;
}

#[tokio::test]
async fn rejects_token_for_different_daemon() {
    let h = spawn_daemon_ws("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "dev-vm", 300);
    let err = connect_attach(h.addr, id, &token).await.unwrap_err();
    assert!(format!("{err}").contains("401"));
    let _ = h.manager.kill(id).await;
}

#[tokio::test]
async fn pty_round_trip_through_real_tmux() {
    let h = spawn_daemon_ws("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "laptop", 300);
    let mut ws = connect_attach(h.addr, id, &token).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).await;
}

#[tokio::test]
async fn multi_attach_broadcasts_to_every_client() {
    let h = spawn_daemon_ws("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "laptop", 300);

    let mut ws_a = connect_attach(h.addr, id, &token).await.unwrap();
    let mut ws_b = connect_attach(h.addr, id, &token).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).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("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "laptop", 300);
    let mut ws = connect_attach(h.addr, id, &token).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("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "laptop", 300);
    let mut ws = connect_attach(h.addr, id, &token).await.unwrap();

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

    h.manager.kill(id).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("laptop").await;
    let id = create_bash_session(&h).await;
    let token = token_for(id, "laptop", 300);
    let mut ws = connect_attach(h.addr, id, &token).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).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("laptop").await;
    let session = create_bash_session(&h).await;
    let token = token_for(session, "laptop", 60);

    let mut a = connect_attach(h.addr, session, &token).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, &token).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"
    );
}