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");
}

/// Regression for NAP-253: the very first bytes a freshly-spawned command
/// writes must reach the replay buffer. Before the fix, `pipe-pane` was a
/// separate tmux invocation that ran after `new-session` + an async FIFO
/// setup, so the child's first output was buffered by tmux but never tee'd
/// into the FIFO.
#[tokio::test]
async fn create_captures_initial_output_in_replay_buffer() {
    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(),
            // The shell prints the marker before doing anything else and
            // then sleeps so the session stays alive. Before NAP-253 this
            // marker was the canonical case that got swallowed.
            command: "sh -c 'echo NAP_253_FIRST_LINE_MARKER; exec sleep 30'".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .await
        .expect("create");

    let want = b"NAP_253_FIRST_LINE_MARKER";
    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    let mut found = false;
    while std::time::Instant::now() < deadline {
        let (replay, _rx) = fx
            .manager
            .subscribe_pty_with_replay(s.id)
            .await
            .expect("subscribe");
        if replay.windows(want.len()).any(|w| w == want) {
            found = true;
            break;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    assert!(
        found,
        "first-line marker never reached the replay buffer — pipe-pane raced past the child's initial output",
    );

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

/// Regression for NAP-228: after `recover`, PTY output must reach new
/// subscribers. Before the fix tmux's stale `pipe-pane` (with `-o`) silently
/// kept feeding the old FIFO and the revived manager's subscribers saw
/// nothing.
#[tokio::test]
async fn recover_re_pipes_pane_so_output_reaches_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
        .expect("create");

    // Capture paths + tmux handle before dropping the original manager —
    // `drop(fx.manager)` partially moves `fx` so we can't call its helpers
    // afterwards.
    let sidecar_dir = fx.sidecar_dir();
    let fifo_dir = fx.fifo_dir();
    let tmux = fx.tmux.tmux.clone();
    let worktree_root = fx.worktree_root.path().to_path_buf();

    // Drop the original manager. Its PtyBridge closes the FIFO's only
    // reader, so tmux's `cat >> fifo` writer hits SIGPIPE and exits — but
    // tmux's pipe-pane bookkeeping still thinks the pane is piped. That's
    // exactly the stale state recover() has to overwrite.
    drop(fx.manager);

    let sidecar = maestro_daemon::sidecar::SidecarStore::new(sidecar_dir);
    let worktree_mgr = maestro_daemon::worktree::WorktreeManager::new(worktree_root);
    let revived = maestro_daemon::SessionManager::new(
        "laptop".to_string(),
        tmux,
        sidecar,
        worktree_mgr,
        fifo_dir,
    );
    assert_eq!(revived.recover().await.expect("recover"), 1);

    let mut rx = revived.subscribe_pty(s.id).await.expect("subscribe");
    revived
        .send_keys(s.id, b"echo NAP_228_RECOVERED\n")
        .await
        .unwrap();

    let want = b"NAP_228_RECOVERED";
    let mut buf: 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), rx.recv()).await {
            buf.extend_from_slice(&bytes);
            if buf.windows(want.len()).any(|w| w == want) {
                break;
            }
        }
    }
    assert!(
        buf.windows(want.len()).any(|w| w == want),
        "recover() did not restore pty output flow; buf={:?}",
        String::from_utf8_lossy(&buf),
    );

    revived.kill(s.id).await.unwrap();
}

#[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();
}