maestro/crates/maestro-daemon/src/registry_client.rs

//! Outbound WebSocket client to the always-on registry.
//!
//! Holds a single persistent connection. On disconnect, reconnects with
//! exponential backoff. Per connection:
//!
//! 1. Send `Hello { daemon_id, hostname, version, token }`.
//! 2. Send `SessionsSnapshot { sessions: snapshot() }`.
//! 3. Drive a `select!` loop that pumps:
//!    - inbound `CreateSession` / `KillSession` → `SessionManager` → reply
//!    - outbound `SessionAdded/Removed/Updated` from the local event bus
//!    - a 10s `Heartbeat` ticker.

use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{DaemonInbound, DaemonOutbound, RpcResult, PROTOCOL_VERSION};
use tokio::sync::broadcast::error::RecvError;
use tokio_tungstenite::tungstenite::Message;

use crate::session::{SessionEvent, SessionManager};

/// Per-message heartbeat interval. Less than the registry's 30s timeout, with
/// headroom for jitter. Spec §"Wire protocol (D↔R)".
pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10);
/// Bounds on the reconnect backoff.
const BACKOFF_MIN: Duration = Duration::from_millis(250);
const BACKOFF_MAX: Duration = Duration::from_secs(30);

#[derive(Debug, thiserror::Error)]
pub enum ClientError {
    #[error("ws error: {0}")]
    Ws(#[from] tokio_tungstenite::tungstenite::Error),
    #[error("json: {0}")]
    Json(#[from] serde_json::Error),
    #[error("registry handshake failed: {0}")]
    Handshake(String),
}

#[derive(Debug, Clone)]
pub struct RegistryClient {
    registry_url: String,
    daemon_id: String,
    hostname: String,
    daemon_url: String,
    token: String,
    manager: SessionManager,
}

impl RegistryClient {
    pub fn new(
        registry_url: impl Into<String>,
        daemon_id: impl Into<String>,
        hostname: impl Into<String>,
        daemon_url: impl Into<String>,
        token: impl Into<String>,
        manager: SessionManager,
    ) -> Self {
        Self {
            registry_url: registry_url.into(),
            daemon_id: daemon_id.into(),
            hostname: hostname.into(),
            daemon_url: daemon_url.into(),
            token: token.into(),
            manager,
        }
    }

    /// Run forever, reconnecting with exponential backoff between attempts.
    /// Returns only if cancelled by the caller dropping the task.
    pub async fn run(self) {
        let mut backoff = BACKOFF_MIN;
        loop {
            match self.connect_once().await {
                Ok(()) => {
                    tracing::info!("registry connection ended cleanly");
                    backoff = BACKOFF_MIN;
                }
                Err(e) => {
                    tracing::warn!(error = %e, ?backoff, "registry connection failed");
                }
            }
            tokio::time::sleep(backoff).await;
            backoff = (backoff * 2).min(BACKOFF_MAX);
        }
    }

    async fn connect_once(&self) -> Result<(), ClientError> {
        let (ws, _) = tokio_tungstenite::connect_async(&self.registry_url).await?;
        let (mut sink, mut stream) = ws.split();

        // Hello + wait for Ack
        let hello = DaemonOutbound::Hello {
            daemon_id: self.daemon_id.clone(),
            hostname: self.hostname.clone(),
            version: PROTOCOL_VERSION.into(),
            token: self.token.clone(),
            daemon_url: self.daemon_url.clone(),
        };
        sink.send(Message::Text(serde_json::to_string(&hello)?))
            .await?;

        match stream.next().await {
            Some(Ok(Message::Text(t))) => match serde_json::from_str::<DaemonInbound>(&t)? {
                DaemonInbound::Ack => {}
                DaemonInbound::Error(e) => {
                    return Err(ClientError::Handshake(format!("{}: {}", e.code, e.message)));
                }
                other => {
                    return Err(ClientError::Handshake(format!(
                        "unexpected first frame {other:?}"
                    )))
                }
            },
            other => {
                return Err(ClientError::Handshake(format!(
                    "no hello-ack frame: {other:?}"
                )))
            }
        }

        // Initial snapshot
        let snapshot = DaemonOutbound::SessionsSnapshot {
            sessions: self.manager.snapshot().await,
        };
        sink.send(Message::Text(serde_json::to_string(&snapshot)?))
            .await?;

        let mut events = self.manager.events();
        let mut hb = tokio::time::interval(HEARTBEAT_INTERVAL);
        hb.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        let _ = hb.tick().await; // consume the immediate first tick

        loop {
            tokio::select! {
                msg = stream.next() => {
                    match msg {
                        None => break,
                        Some(Err(e)) => return Err(e.into()),
                        Some(Ok(Message::Close(_))) => break,
                        Some(Ok(Message::Text(text))) => {
                            self.handle_inbound(&text, &mut sink).await?;
                        }
                        Some(Ok(_)) => {}
                    }
                }
                ev = events.recv() => {
                    match ev {
                        Ok(SessionEvent::Added(s)) => {
                            sink.send(Message::Text(serde_json::to_string(
                                &DaemonOutbound::SessionAdded { session: s },
                            )?)).await?;
                        }
                        Ok(SessionEvent::Removed(id)) => {
                            sink.send(Message::Text(serde_json::to_string(
                                &DaemonOutbound::SessionRemoved { id },
                            )?)).await?;
                        }
                        Ok(SessionEvent::Updated(s)) => {
                            sink.send(Message::Text(serde_json::to_string(
                                &DaemonOutbound::SessionUpdated { session: s },
                            )?)).await?;
                        }
                        Err(RecvError::Lagged(_)) => continue,
                        Err(RecvError::Closed) => break,
                    }
                }
                _ = hb.tick() => {
                    // perimeter_id / private_addr are wired from daemon config
                    // in NAP-243; NAP-242 lands the protocol field only.
                    sink.send(Message::Text(serde_json::to_string(
                        &DaemonOutbound::Heartbeat {
                            perimeter_id: None,
                            private_addr: None,
                        },
                    )?)).await?;
                }
            }
        }
        Ok(())
    }

    async fn handle_inbound<S>(&self, text: &str, sink: &mut S) -> Result<(), ClientError>
    where
        S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
    {
        let parsed: DaemonInbound = serde_json::from_str(text)?;
        match parsed {
            DaemonInbound::Ack => {}
            DaemonInbound::Error(e) => {
                tracing::warn!(code = %e.code, message = %e.message, "registry-reported error");
            }
            DaemonInbound::CreateSession { reply_to, request } => {
                let result = match self.manager.create(request).await {
                    Ok(s) => RpcResult::Ok { value: s },
                    Err(e) => RpcResult::Err {
                        error: e.to_error_info(),
                    },
                };
                sink.send(Message::Text(serde_json::to_string(
                    &DaemonOutbound::CreateSessionResult { reply_to, result },
                )?))
                .await?;
            }
            DaemonInbound::KillSession { reply_to, id } => {
                let result = match self.manager.kill(id).await {
                    Ok(()) => RpcResult::Ok { value: () },
                    Err(e) => RpcResult::Err {
                        error: e.to_error_info(),
                    },
                };
                sink.send(Message::Text(serde_json::to_string(
                    &DaemonOutbound::KillSessionResult { reply_to, result },
                )?))
                .await?;
            }
            DaemonInbound::ResolveAttachResult { .. } | DaemonInbound::RelaySetup { .. } => {
                // Wired by NAP-243 (broker mode). Lands here only so registries
                // emitting these messages don't crash older daemons on parse.
                tracing::debug!("broker message received but unhandled (NAP-243 pending)");
            }
        }
        Ok(())
    }
}