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

//! End-to-end tests for the SessionManager: create, kill, recover, broadcast.

mod common;

use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;

use common::{wait_for_async, DaemonFixture};
use maestro_daemon::SessionEvent;
use maestro_protocol::CreateSessionRequest;
use tempfile::TempDir;

#[tokio::test]
async fn create_then_kill_round_trip() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();

    let session = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: cwd.path().to_path_buf(),
            command: "sleep 30".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .expect("create");

    assert!(fx.manager.has(session.id).await);
    let snap = fx.manager.snapshot().await;
    assert_eq!(snap.len(), 1);
    assert_eq!(snap[0].id, session.id);

    // sidecar should be on disk
    assert!(fx
        .sidecar_dir()
        .join(format!("{}.json", session.id))
        .exists());
    // FIFO should also exist
    assert!(fx.fifo_dir().join(format!("{}.fifo", session.id)).exists());

    fx.manager.kill(session.id).await.expect("kill");
    assert!(!fx.manager.has(session.id).await);
    assert!(!fx
        .sidecar_dir()
        .join(format!("{}.json", session.id))
        .exists());
}

#[tokio::test]
async fn rejects_wrong_daemon_in_request() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();
    let err = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "dev-vm".into(),
            cwd: cwd.path().to_path_buf(),
            command: "sleep 5".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .unwrap_err();
    assert!(format!("{err}").contains("targets daemon"), "{err}");
}

#[tokio::test]
async fn recover_adopts_existing_tmux_sessions() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();

    // create a session, then "throw away" the SessionManager (the tmux session
    // and sidecar persist) and build a fresh manager that recovers them.
    let s = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: cwd.path().to_path_buf(),
            command: "sleep 30".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .expect("create");

    // build a parallel manager pointing at the same dirs + same tmux server
    let sidecar = maestro_daemon::sidecar::SidecarStore::new(fx.sidecar_dir());
    let worktree_mgr =
        maestro_daemon::worktree::WorktreeManager::new(fx.worktree_root.path().to_path_buf());
    let revived = maestro_daemon::SessionManager::new(
        "laptop".to_string(),
        fx.tmux.tmux.clone(),
        sidecar,
        worktree_mgr,
        fx.fifo_dir(),
    );
    let adopted = revived.recover().await.expect("recover");
    assert_eq!(adopted, 1);
    assert!(revived.has(s.id).await);
    let snap = revived.snapshot().await;
    assert_eq!(snap[0].id, s.id);
    assert_eq!(snap[0].command, "sleep 30");
}

#[tokio::test]
async fn pty_broadcast_delivers_bytes_to_subscribers() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();

    let s = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: cwd.path().to_path_buf(),
            command: "bash --noprofile --norc".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .unwrap();

    let mut a = fx.manager.subscribe_pty(s.id).await.unwrap();
    let mut b = fx.manager.subscribe_pty(s.id).await.unwrap();

    fx.manager
        .send_keys(s.id, b"echo MAESTRO_BROADCAST_OK\n")
        .await
        .unwrap();

    let want = b"MAESTRO_BROADCAST_OK";
    let mut buf_a: Vec<u8> = Vec::new();
    let mut buf_b: Vec<u8> = Vec::new();

    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    while std::time::Instant::now() < deadline {
        if let Ok(Ok(bytes)) = tokio::time::timeout(Duration::from_millis(50), a.recv()).await {
            buf_a.extend_from_slice(&bytes);
        }
        if let Ok(Ok(bytes)) = tokio::time::timeout(Duration::from_millis(50), b.recv()).await {
            buf_b.extend_from_slice(&bytes);
        }
        if buf_a.windows(want.len()).any(|w| w == want)
            && buf_b.windows(want.len()).any(|w| w == want)
        {
            break;
        }
    }
    assert!(buf_a.windows(want.len()).any(|w| w == want), "A missed");
    assert!(buf_b.windows(want.len()).any(|w| w == want), "B missed");

    fx.manager.kill(s.id).await.unwrap();
    let _ = (cwd, PathBuf::from("/"));
}

#[tokio::test]
async fn liveness_watcher_reaps_dead_sessions_and_emits_removed() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();

    let s = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: cwd.path().to_path_buf(),
            command: "sleep 30".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .expect("create");

    // Subscribe *after* the create so we don't have to filter the Added event.
    let mut events = fx.manager.events();
    fx.manager.start_liveness_watcher(Duration::from_millis(50));

    // Kill the tmux session out from under the daemon. This is the scenario
    // the watcher exists for: command exits, tmux destroys the session, the
    // daemon never finds out through its normal channels.
    fx.tmux.tmux.kill_session(s.id).expect("kill tmux directly");

    // Within a few watcher ticks the daemon should drop the session.
    wait_for_async(Duration::from_secs(3), || async {
        !fx.manager.has(s.id).await
    })
    .await;

    // Sidecar + FIFO should be cleaned up too.
    assert!(!fx.sidecar_dir().join(format!("{}.json", s.id)).exists());
    assert!(!fx.fifo_dir().join(format!("{}.fifo", s.id)).exists());

    // And we should have seen a Removed event for this id.
    let mut saw_removed = false;
    let deadline = std::time::Instant::now() + Duration::from_secs(1);
    while std::time::Instant::now() < deadline && !saw_removed {
        match tokio::time::timeout(Duration::from_millis(50), events.recv()).await {
            Ok(Ok(SessionEvent::Removed(id))) if id == s.id => saw_removed = true,
            Ok(Ok(_)) => continue,
            _ => break,
        }
    }
    assert!(
        saw_removed,
        "expected SessionEvent::Removed for dead session"
    );
}

#[tokio::test]
async fn liveness_watcher_leaves_running_sessions_alone() {
    let fx = DaemonFixture::new("laptop");
    let cwd = TempDir::new().unwrap();

    let s = fx
        .manager
        .create(CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: cwd.path().to_path_buf(),
            command: "sleep 30".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .expect("create");

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

    // Give the watcher enough time to run a handful of ticks.
    tokio::time::sleep(Duration::from_millis(400)).await;
    assert!(fx.manager.has(s.id).await, "live session got reaped");

    fx.manager.kill(s.id).await.unwrap();
}