maestro/crates/maestro-registry/tests/daemon_ws.rs

//! End-to-end tests for the `WS /daemon` endpoint, driven by a real
//! `tokio-tungstenite` client against an axum server bound to `:0`.

use std::collections::HashMap;
use std::future::Future;
use std::net::{Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use std::time::{Duration, Instant};

use chrono::{DateTime, Utc};
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{
    CreateSessionRequest, DaemonInbound, DaemonOutbound, RpcResult, Session, SessionStatus,
};
use maestro_registry::{router_with_timeout, Registry};
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 = "test-token-please-allow";
const TEST_HEARTBEAT: Duration = Duration::from_millis(500);

struct Harness {
    addr: SocketAddr,
    registry: Registry,
    _shutdown_tx: oneshot::Sender<()>,
}

async fn spawn_registry() -> Harness {
    spawn_registry_with_timeout(TEST_HEARTBEAT).await
}

async fn spawn_registry_with_timeout(heartbeat: Duration) -> Harness {
    let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
        .await
        .unwrap();
    let addr = listener.local_addr().unwrap();
    let registry = Registry::new();
    let app = router_with_timeout(registry.clone(), TOKEN.to_string(), heartbeat);
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = shutdown_rx.await;
            })
            .await;
    });
    Harness {
        addr,
        registry,
        _shutdown_tx: shutdown_tx,
    }
}

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

async fn connect(addr: SocketAddr) -> WsClient {
    let url = format!("ws://{addr}/daemon");
    let (ws, _) = connect_async(url).await.expect("ws connect");
    ws
}

async fn send_json(ws: &mut WsClient, msg: &DaemonOutbound) {
    let text = serde_json::to_string(msg).unwrap();
    ws.send(WsMessage::Text(text)).await.unwrap();
}

async fn recv_inbound(ws: &mut WsClient) -> Option<DaemonInbound> {
    while let Some(Ok(msg)) = ws.next().await {
        if let WsMessage::Text(t) = msg {
            let parsed: DaemonInbound = serde_json::from_str(&t).unwrap();
            // NAP-246: skip catalog frames — these are informational and not
            // what most tests are asserting on.
            if matches!(parsed, DaemonInbound::CatalogEvent { .. }) {
                continue;
            }
            return Some(parsed);
        }
    }
    None
}

fn hello(token: &str) -> DaemonOutbound {
    DaemonOutbound::Hello {
        daemon_id: "laptop".into(),
        hostname: "naps-laptop".into(),
        version: "0.1.0".into(),
        token: token.into(),
        daemon_url: "ws://127.0.0.1:8081".into(),
        home_dir: String::new(),
    }
}

fn sample_session(id: Uuid, daemon_id: &str) -> Session {
    Session {
        id,
        daemon_id: daemon_id.to_string(),
        cwd: PathBuf::from("/tmp"),
        command: "bash".into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        status: SessionStatus::Running,
        created_at: DateTime::parse_from_rfc3339("2026-05-24T12:00:00Z")
            .unwrap()
            .with_timezone(&Utc),
        foreground: None,
        activity: Default::default(),
        last_activity_at: chrono::Utc::now(),

        current_cwd: None,
    }
}

async fn wait_for<F, Fut>(timeout: Duration, mut cond: F)
where
    F: FnMut() -> Fut,
    Fut: Future<Output = bool>,
{
    let deadline = Instant::now() + timeout;
    while Instant::now() < deadline {
        if cond().await {
            return;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
    panic!("timed out after {timeout:?} waiting for condition");
}

// --------------------------------------------------------------------------
// Tests
// --------------------------------------------------------------------------

#[tokio::test]
async fn rejects_bad_token() {
    let h = spawn_registry().await;
    let mut ws = connect(h.addr).await;

    send_json(&mut ws, &hello("wrong")).await;
    let reply = recv_inbound(&mut ws).await.expect("error frame");
    match reply {
        DaemonInbound::Error(e) => assert_eq!(e.code, "auth_failed"),
        other => panic!("expected Error, got {other:?}"),
    }
    let _ = ws.next().await;
    assert!(h.registry.daemons().await.is_empty());
}

#[tokio::test]
async fn accepts_good_token_and_snapshot() {
    let h = spawn_registry().await;
    let mut ws = connect(h.addr).await;

    send_json(&mut ws, &hello(TOKEN)).await;
    let ack = recv_inbound(&mut ws).await.expect("ack");
    assert!(matches!(ack, DaemonInbound::Ack));

    let sessions = vec![
        sample_session(Uuid::new_v4(), "laptop"),
        sample_session(Uuid::new_v4(), "laptop"),
    ];
    send_json(
        &mut ws,
        &DaemonOutbound::SessionsSnapshot {
            sessions: sessions.clone(),
        },
    )
    .await;

    let expected = sessions.len();
    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.sessions().await.len() == expected }
    })
    .await;

    let daemons = h.registry.daemons().await;
    assert_eq!(daemons.len(), 1);
    assert_eq!(daemons[0].daemon_id, "laptop");
}

#[tokio::test]
async fn applies_individual_diffs() {
    let h = spawn_registry().await;
    let mut ws = connect(h.addr).await;

    send_json(&mut ws, &hello(TOKEN)).await;
    let _ = recv_inbound(&mut ws).await;

    let s1 = sample_session(Uuid::new_v4(), "laptop");
    let s2 = sample_session(Uuid::new_v4(), "laptop");

    send_json(
        &mut ws,
        &DaemonOutbound::SessionAdded {
            session: s1.clone(),
        },
    )
    .await;
    send_json(
        &mut ws,
        &DaemonOutbound::SessionAdded {
            session: s2.clone(),
        },
    )
    .await;

    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.sessions().await.len() == 2 }
    })
    .await;

    let mut s1_updated = s1.clone();
    s1_updated.command = "fish".into();
    send_json(
        &mut ws,
        &DaemonOutbound::SessionUpdated {
            session: s1_updated.clone(),
        },
    )
    .await;
    send_json(&mut ws, &DaemonOutbound::SessionRemoved { id: s2.id }).await;

    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move {
            let s = r.sessions().await;
            s.len() == 1 && s[0].command == "fish"
        }
    })
    .await;
}

#[tokio::test]
async fn heartbeat_timeout_drops_daemon() {
    let h = spawn_registry_with_timeout(Duration::from_millis(200)).await;
    let mut ws = connect(h.addr).await;

    send_json(&mut ws, &hello(TOKEN)).await;
    let _ = recv_inbound(&mut ws).await;
    send_json(
        &mut ws,
        &DaemonOutbound::SessionAdded {
            session: sample_session(Uuid::new_v4(), "laptop"),
        },
    )
    .await;

    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.sessions().await.len() == 1 }
    })
    .await;

    // Hold the WS open but never send another frame. The server's per-read
    // timeout should fire after 200ms and drop the daemon.
    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.daemons().await.is_empty() && r.sessions().await.is_empty() }
    })
    .await;
}

#[tokio::test]
async fn re_registration_replaces_catalog() {
    let h = spawn_registry().await;

    let mut ws1 = connect(h.addr).await;
    send_json(&mut ws1, &hello(TOKEN)).await;
    let _ = recv_inbound(&mut ws1).await;
    let s_old = sample_session(Uuid::new_v4(), "laptop");
    send_json(
        &mut ws1,
        &DaemonOutbound::SessionAdded {
            session: s_old.clone(),
        },
    )
    .await;

    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.sessions().await.len() == 1 }
    })
    .await;

    drop(ws1);

    let registry = h.registry.clone();
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move { r.daemons().await.is_empty() }
    })
    .await;

    let mut ws2 = connect(h.addr).await;
    send_json(&mut ws2, &hello(TOKEN)).await;
    let _ = recv_inbound(&mut ws2).await;
    let s_new = sample_session(Uuid::new_v4(), "laptop");
    send_json(
        &mut ws2,
        &DaemonOutbound::SessionsSnapshot {
            sessions: vec![s_new.clone()],
        },
    )
    .await;

    let registry = h.registry.clone();
    let want_id = s_new.id;
    wait_for(Duration::from_secs(2), || {
        let r = registry.clone();
        async move {
            let s = r.sessions().await;
            s.len() == 1 && s[0].id == want_id
        }
    })
    .await;
}

#[tokio::test]
async fn create_session_rpc_round_trip() {
    let h = spawn_registry().await;
    let mut ws = connect(h.addr).await;

    send_json(&mut ws, &hello(TOKEN)).await;
    let _ = recv_inbound(&mut ws).await;

    let registry = h.registry.clone();
    let request = CreateSessionRequest {
        daemon_id: "laptop".into(),
        cwd: PathBuf::from("/tmp"),
        command: "bash".into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        worktree: None,
    };
    let rpc = tokio::spawn(async move { registry.create_session("laptop", request).await });

    let inbound = recv_inbound(&mut ws).await.expect("CreateSession");
    let (reply_to, _request) = match inbound {
        DaemonInbound::CreateSession { reply_to, request } => (reply_to, request),
        other => panic!("expected CreateSession, got {other:?}"),
    };

    let session = sample_session(Uuid::new_v4(), "laptop");
    send_json(
        &mut ws,
        &DaemonOutbound::CreateSessionResult {
            reply_to,
            result: RpcResult::Ok {
                value: session.clone(),
            },
        },
    )
    .await;

    let outcome = rpc.await.unwrap().expect("rpc transport ok");
    let returned = outcome.expect("rpc daemon ok");
    assert_eq!(returned.id, session.id);
}

#[tokio::test]
async fn create_session_to_unknown_daemon_errors() {
    let h = spawn_registry().await;
    let err = h
        .registry
        .create_session(
            "no-such-daemon",
            CreateSessionRequest {
                daemon_id: "no-such-daemon".into(),
                cwd: PathBuf::from("/tmp"),
                command: "bash".into(),
                env: HashMap::new(),
                metadata: HashMap::new(),
                worktree: None,
            },
        )
        .await
        .unwrap_err();
    assert!(matches!(
        err,
        maestro_registry::RegistryError::DaemonNotConnected(_)
    ));
}