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

//! Integration tests for the client-facing HTTP API (NAP-215).
//!
//! Spins up a registry on `:0`, plus a tiny mock-daemon WS client that exists
//! purely to satisfy `CreateSession` / `KillSession` round-trips. All
//! HTTP/WS calls go over the wire via `reqwest` and `tokio-tungstenite`.

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

use chrono::Utc;
use futures_util::{SinkExt, StreamExt};
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use maestro_protocol::{
    ConnectInfo, CreateSessionRequest, DaemonInbound, DaemonInfo, DaemonOutbound, JwtClaims,
    RegistryEvent, RpcResult, Session, SessionStatus,
};
use maestro_registry::{router_full, Registry, COOKIE_NAME};
use reqwest::Client;
use serde_json::Value;
use tokio::net::TcpListener;
use tokio::sync::{oneshot, Mutex};
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_tungstenite::{
    connect_async, tungstenite::client::IntoClientRequest, MaybeTlsStream, WebSocketStream,
};
use uuid::Uuid;

const TOKEN: &str = "http-api-test-token";
const TEST_HEARTBEAT: Duration = Duration::from_secs(30);

// ---------------------------------------------------------------------------
// Harness
// ---------------------------------------------------------------------------

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

impl Harness {
    fn base(&self) -> String {
        format!("http://{}", self.addr)
    }

    fn ws_base(&self) -> String {
        format!("ws://{}", self.addr)
    }
}

async fn spawn(rpc_timeout: 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_full(
        registry.clone(),
        TOKEN.to_string(),
        TEST_HEARTBEAT,
        rpc_timeout,
    );
    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,
    }
}

// ---------------------------------------------------------------------------
// Mock daemon
// ---------------------------------------------------------------------------

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

#[derive(Clone, Copy)]
enum DaemonBehavior {
    /// Replies immediately with the session the request asked for.
    ReplyOk,
    /// Never replies — exercises the registry's 504 timeout path.
    Hang,
}

#[derive(Clone)]
struct MockDaemon {
    daemon_id: String,
    daemon_url: String,
    captured: Arc<Mutex<Vec<DaemonInbound>>>,
}

impl MockDaemon {
    fn new(daemon_id: &str, daemon_url: &str) -> Self {
        Self {
            daemon_id: daemon_id.into(),
            daemon_url: daemon_url.into(),
            captured: Arc::new(Mutex::new(Vec::new())),
        }
    }

    async fn spawn(self, addr: SocketAddr, behavior: DaemonBehavior) {
        let url = format!("ws://{addr}/daemon");
        let (mut ws, _) = connect_async(url).await.expect("ws connect");
        let hello = DaemonOutbound::Hello {
            daemon_id: self.daemon_id.clone(),
            hostname: "mock-host".into(),
            version: "0.1.0".into(),
            token: TOKEN.into(),
            daemon_url: self.daemon_url.clone(),
            home_dir: String::new(),
        };
        ws.send(WsMessage::Text(serde_json::to_string(&hello).unwrap()))
            .await
            .unwrap();
        // consume the Ack
        let _ = ws.next().await;
        // send an empty initial snapshot so the registry's view is consistent
        ws.send(WsMessage::Text(
            serde_json::to_string(&DaemonOutbound::SessionsSnapshot { sessions: vec![] }).unwrap(),
        ))
        .await
        .unwrap();

        let captured = self.captured.clone();
        let daemon_id = self.daemon_id.clone();
        tokio::spawn(async move {
            run_loop(ws, behavior, captured, daemon_id).await;
        });
    }

    async fn captured(&self) -> Vec<DaemonInbound> {
        self.captured.lock().await.clone()
    }
}

async fn run_loop(
    mut ws: WsClient,
    behavior: DaemonBehavior,
    captured: Arc<Mutex<Vec<DaemonInbound>>>,
    daemon_id: String,
) {
    while let Some(Ok(WsMessage::Text(text))) = ws.next().await {
        let msg: DaemonInbound = match serde_json::from_str(&text) {
            Ok(m) => m,
            Err(_) => continue,
        };
        captured.lock().await.push(msg.clone());

        match (behavior, msg) {
            (DaemonBehavior::ReplyOk, DaemonInbound::CreateSession { reply_to, request }) => {
                let session = mock_session(&daemon_id, &request);
                // Send SessionAdded BEFORE the RPC result so the registry's
                // catalog is updated by the time the HTTP caller resumes.
                // The registry processes inbound frames serially in order,
                // so this ordering guarantee holds end-to-end.
                let _ = ws
                    .send(WsMessage::Text(
                        serde_json::to_string(&DaemonOutbound::SessionAdded {
                            session: session.clone(),
                        })
                        .unwrap(),
                    ))
                    .await;
                let reply = DaemonOutbound::CreateSessionResult {
                    reply_to,
                    result: RpcResult::Ok { value: session },
                };
                let _ = ws
                    .send(WsMessage::Text(serde_json::to_string(&reply).unwrap()))
                    .await;
            }
            (DaemonBehavior::ReplyOk, DaemonInbound::KillSession { reply_to, id }) => {
                let _ = ws
                    .send(WsMessage::Text(
                        serde_json::to_string(&DaemonOutbound::SessionRemoved { id }).unwrap(),
                    ))
                    .await;
                let reply = DaemonOutbound::KillSessionResult {
                    reply_to,
                    result: RpcResult::Ok { value: () },
                };
                let _ = ws
                    .send(WsMessage::Text(serde_json::to_string(&reply).unwrap()))
                    .await;
            }
            (DaemonBehavior::Hang, _) => {
                // intentionally do nothing
            }
            _ => {}
        }
    }
}

fn mock_session(daemon_id: &str, request: &CreateSessionRequest) -> Session {
    // If a worktree spec was attached, pretend the daemon "resolved" the cwd
    // to <base_repo>/.worktrees/<branch> so the test can assert on it.
    let cwd = request
        .worktree
        .as_ref()
        .map(|w| w.base_repo.join(".worktrees").join(&w.branch))
        .unwrap_or_else(|| request.cwd.clone());
    Session {
        id: Uuid::new_v4(),
        daemon_id: daemon_id.into(),
        cwd,
        command: request.command.clone(),
        env: request.env.clone(),
        metadata: request.metadata.clone(),
        status: SessionStatus::Running,
        created_at: Utc::now(),
        foreground: None,
        activity: Default::default(),
        last_activity_at: chrono::Utc::now(),

        current_cwd: None,
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

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

fn client_with_cookies() -> Client {
    reqwest::Client::builder()
        .cookie_store(true)
        .build()
        .unwrap()
}

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

#[tokio::test]
async fn first_request_sets_cookie_and_subsequent_reuses_it() {
    let h = spawn(Duration::from_secs(5)).await;
    let client = client_with_cookies();

    let r1 = client
        .get(format!("{}/api/daemons", h.base()))
        .send()
        .await
        .unwrap();
    assert!(r1.status().is_success());
    let cookies: Vec<_> = r1
        .headers()
        .get_all(reqwest::header::SET_COOKIE)
        .iter()
        .collect();
    assert!(
        !cookies.is_empty(),
        "first request must include a Set-Cookie"
    );
    let raw = cookies[0].to_str().unwrap();
    assert!(raw.contains(COOKIE_NAME));
    assert!(raw.contains("HttpOnly"));
    assert!(raw.contains("SameSite=Strict"));

    let r2 = client
        .get(format!("{}/api/daemons", h.base()))
        .send()
        .await
        .unwrap();
    assert!(r2.status().is_success());
    let new_cookies: Vec<_> = r2
        .headers()
        .get_all(reqwest::header::SET_COOKIE)
        .iter()
        .collect();
    assert!(
        new_cookies.is_empty(),
        "valid existing cookie should not be reissued, got {new_cookies:?}"
    );
}

#[tokio::test]
async fn tampered_cookie_is_replaced_silently() {
    let h = spawn(Duration::from_secs(5)).await;
    let client = reqwest::Client::new();

    let resp = client
        .get(format!("{}/api/daemons", h.base()))
        .header(
            reqwest::header::COOKIE,
            format!("{}=not-a-jwt", COOKIE_NAME),
        )
        .send()
        .await
        .unwrap();
    assert!(resp.status().is_success());
    let cookies: Vec<_> = resp
        .headers()
        .get_all(reqwest::header::SET_COOKIE)
        .iter()
        .collect();
    assert!(
        !cookies.is_empty(),
        "tampered cookie must be replaced with a freshly-signed one"
    );
}

#[tokio::test]
async fn list_daemons_includes_url() {
    let h = spawn(Duration::from_secs(5)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

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

    let client = client_with_cookies();
    let resp = client
        .get(format!("{}/api/daemons", h.base()))
        .send()
        .await
        .unwrap();
    assert!(resp.status().is_success());
    let body: Vec<DaemonInfo> = resp.json().await.unwrap();
    assert_eq!(body.len(), 1);
    assert_eq!(body[0].daemon_id, "laptop");
    assert_eq!(body[0].daemon_url, "ws://10.0.0.1:8081");
}

#[tokio::test]
async fn post_sessions_round_trips_to_daemon() {
    let h = spawn(Duration::from_secs(5)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

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

    let client = client_with_cookies();
    let request = CreateSessionRequest {
        daemon_id: "laptop".into(),
        cwd: PathBuf::from("/tmp"),
        command: "bash".into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        worktree: None,
    };
    let resp = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&request)
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 201);
    let session: Session = resp.json().await.unwrap();
    assert_eq!(session.daemon_id, "laptop");
    assert_eq!(session.command, "bash");

    // GET /api/sessions should now see it.
    let resp = client
        .get(format!("{}/api/sessions", h.base()))
        .send()
        .await
        .unwrap();
    let listed: Vec<Session> = resp.json().await.unwrap();
    assert!(listed.iter().any(|s| s.id == session.id));
}

#[tokio::test]
async fn post_sessions_returns_504_on_daemon_timeout() {
    // Short RPC timeout so the 504 path is fast.
    let h = spawn(Duration::from_millis(200)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::Hang).await;

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

    let client = client_with_cookies();
    let request = CreateSessionRequest {
        daemon_id: "laptop".into(),
        cwd: PathBuf::from("/tmp"),
        command: "bash".into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        worktree: None,
    };
    let resp = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&request)
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 504);

    // The daemon should have received the CreateSession even though it never
    // replied.
    let captured = mock.captured().await;
    assert!(
        captured
            .iter()
            .any(|m| matches!(m, DaemonInbound::CreateSession { .. })),
        "daemon should have seen CreateSession, got {captured:?}"
    );
}

#[tokio::test]
async fn post_sessions_with_worktree_forwards_spec_and_returns_resolved_cwd() {
    let h = spawn(Duration::from_secs(5)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

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

    let client = client_with_cookies();
    let body = serde_json::json!({
        "daemon_id": "laptop",
        "cwd": "/home/naps62/projects/maestro",
        "command": "claude",
        "worktree": {
            "base_repo": "/home/naps62/projects/maestro",
            "branch": "feat/test",
            "cleanup_on_kill": true
        }
    });
    let resp = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&body)
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 201);
    let session: Session = resp.json().await.unwrap();
    assert_eq!(
        session.cwd,
        PathBuf::from("/home/naps62/projects/maestro/.worktrees/feat/test")
    );

    let captured = mock.captured().await;
    let forwarded = captured
        .iter()
        .find_map(|m| match m {
            DaemonInbound::CreateSession { request, .. } => Some(request.clone()),
            _ => None,
        })
        .expect("CreateSession was not forwarded");
    let spec = forwarded.worktree.expect("worktree spec dropped");
    assert_eq!(spec.branch, "feat/test");
    assert!(spec.cleanup_on_kill);
}

#[tokio::test]
async fn delete_session_round_trips() {
    let h = spawn(Duration::from_secs(5)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

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

    let client = client_with_cookies();
    let request = CreateSessionRequest {
        daemon_id: "laptop".into(),
        cwd: PathBuf::from("/tmp"),
        command: "bash".into(),
        env: HashMap::new(),
        metadata: HashMap::new(),
        worktree: None,
    };
    let session: Session = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&request)
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();

    // The mock daemon doesn't actually update its catalog after kill — but the
    // registry's catalog reflects whatever the daemon broadcasts. We're
    // asserting only on the HTTP response.
    let resp = client
        .delete(format!("{}/api/sessions/{}", h.base(), session.id))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 204);

    let captured = mock.captured().await;
    assert!(
        captured
            .iter()
            .any(|m| matches!(m, DaemonInbound::KillSession { id, .. } if *id == session.id)),
        "daemon should have seen KillSession, got {captured:?}"
    );
}

#[tokio::test]
async fn delete_session_returns_404_for_unknown_session() {
    let h = spawn(Duration::from_secs(5)).await;
    let client = client_with_cookies();
    let resp = client
        .delete(format!("{}/api/sessions/{}", h.base(), Uuid::new_v4()))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status().as_u16(), 404);
}

#[tokio::test]
async fn connect_info_issues_valid_jwt() {
    let h = spawn(Duration::from_secs(5)).await;
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

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

    let client = client_with_cookies();
    let session: Session = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: PathBuf::from("/tmp"),
            command: "bash".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();

    let resp = client
        .get(format!(
            "{}/api/sessions/{}/connect-info",
            h.base(),
            session.id
        ))
        .send()
        .await
        .unwrap();
    assert!(resp.status().is_success());
    let info: ConnectInfo = resp.json().await.unwrap();
    assert_eq!(info.daemon_url, "ws://10.0.0.1:8081");

    // Verify the JWT against the same shared secret a daemon would use.
    let decoded = decode::<JwtClaims>(
        &info.jwt,
        &DecodingKey::from_secret(TOKEN.as_bytes()),
        &Validation::new(Algorithm::HS256),
    )
    .expect("jwt decodes");
    assert_eq!(decoded.claims.daemon_id, "laptop");
    assert_eq!(decoded.claims.session_id, session.id);
    assert!(decoded.claims.exp > Utc::now().timestamp());
}

#[tokio::test]
async fn events_ws_pushes_snapshot_then_deltas() {
    let h = spawn(Duration::from_secs(5)).await;

    // First obtain a valid cookie via a plain HTTP call.
    let client = client_with_cookies();
    let _ = client
        .get(format!("{}/api/daemons", h.base()))
        .send()
        .await
        .unwrap();
    let cookie_value = harvest_cookie(&client, &h.base()).await;

    // Now open the events WS with that cookie.
    let url = format!("{}/api/events", h.ws_base());
    let mut req = url.into_client_request().unwrap();
    req.headers_mut().insert(
        reqwest::header::COOKIE,
        format!("{}={}", COOKIE_NAME, cookie_value).parse().unwrap(),
    );
    let (mut ws, _) = connect_async(req).await.expect("connect events ws");

    let snapshot = recv_event(&mut ws).await;
    match snapshot {
        RegistryEvent::Snapshot { daemons, sessions } => {
            assert!(daemons.is_empty());
            assert!(sessions.is_empty());
        }
        other => panic!("expected Snapshot first, got {other:?}"),
    }

    // Now connect a daemon — subscriber should see DaemonAdded.
    let mock = MockDaemon::new("laptop", "ws://10.0.0.1:8081");
    mock.clone().spawn(h.addr, DaemonBehavior::ReplyOk).await;

    let event = recv_event(&mut ws).await;
    match event {
        RegistryEvent::DaemonAdded { daemon } => assert_eq!(daemon.daemon_id, "laptop"),
        other => panic!("expected DaemonAdded, got {other:?}"),
    }

    // Spawn a session via HTTP; subscriber should see SessionAdded.
    let _: Session = client
        .post(format!("{}/api/sessions", h.base()))
        .json(&CreateSessionRequest {
            daemon_id: "laptop".into(),
            cwd: PathBuf::from("/tmp"),
            command: "bash".into(),
            env: HashMap::new(),
            metadata: HashMap::new(),
            worktree: None,
        })
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();

    let event = recv_event(&mut ws).await;
    assert!(
        matches!(event, RegistryEvent::SessionAdded { .. }),
        "expected SessionAdded, got {event:?}"
    );
}

#[tokio::test]
async fn root_serves_no_ui() {
    // The registry serves no web UI — that's `maestro-web`'s job (see
    // connection-orchestration spec §"UI hosting"). `/` has no route.
    let h = spawn(Duration::from_secs(5)).await;
    let client = client_with_cookies();
    let resp = client.get(h.base()).send().await.unwrap();
    assert_eq!(resp.status(), reqwest::StatusCode::NOT_FOUND);
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

async fn harvest_cookie(client: &Client, base: &str) -> String {
    // Reqwest's cookie store keeps the cookie, but doesn't expose values
    // directly. Easiest: re-issue a request without the store and read the
    // Set-Cookie that comes back.
    let bare = reqwest::Client::new();
    let resp = bare
        .get(format!("{base}/api/daemons"))
        .send()
        .await
        .unwrap();
    let raw = resp
        .headers()
        .get(reqwest::header::SET_COOKIE)
        .expect("set-cookie")
        .to_str()
        .unwrap();
    // raw looks like "maestro_session=<value>; HttpOnly; ..."
    let prefix = format!("{COOKIE_NAME}=");
    let body = raw.strip_prefix(&prefix).expect("cookie prefix");
    let value = body.split(';').next().unwrap();
    let _ = client; // unused — kept in signature for symmetry
    value.to_string()
}

async fn recv_event(ws: &mut WsClient) -> RegistryEvent {
    let frame = tokio::time::timeout(Duration::from_secs(2), ws.next())
        .await
        .expect("event arrived")
        .expect("ws open")
        .expect("ws frame");
    match frame {
        WsMessage::Text(t) => match serde_json::from_str::<RegistryEvent>(&t) {
            Ok(e) => e,
            Err(e) => {
                let parsed: Value = serde_json::from_str(&t).unwrap();
                panic!("could not parse event {parsed:?}: {e}");
            }
        },
        other => panic!("expected text frame, got {other:?}"),
    }
}