//! 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, //! except the event that an in-flight RPC reply already conveys, which is //! swallowed so the registry writes each change once (see [`RpcCovered`]). //! - a 10s `Heartbeat` ticker. use std::collections::HashSet; use std::time::Duration; use futures_util::{SinkExt, StreamExt}; use maestro_protocol::{ DaemonInbound, DaemonOutbound, PerimeterCandidate, RpcResult, PROTOCOL_VERSION, }; use tokio::sync::broadcast::error::RecvError; use tokio_tungstenite::tungstenite::Message; use uuid::Uuid; use tokio::sync::mpsc; use crate::broker::BrokerHandle; use crate::catalog_mirror::CatalogMirror; 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); /// Session ids whose next lifecycle event the registry already learns about via /// an RPC reply, so forwarding the matching async `SessionAdded` / /// `SessionRemoved` would be a redundant *second* write to the registry catalog /// (NAP-268). /// /// A registry-driven `CreateSession` / `KillSession` is reported back to the /// registry in two ways: the `CreateSessionResult` / `KillSessionResult` reply /// (which the registry applies eagerly), *and* the async lifecycle event this /// daemon emits on its own `SessionManager` stream. Both used to land, so the /// registry's catalog write had to be idempotent to absorb the duplicate. We /// drop the duplicate at the source instead: when an RPC creates or kills a /// session, its id is parked here and the matching async event is swallowed /// (consumed exactly once) before it reaches the registry. Spontaneous changes /// — sessions created via the daemon's local HTTP API, or reaped on their own — /// are not parked and still forward normally, since the registry has no other /// way to learn about them. /// /// Tracked per-connection: a reconnect starts fresh (the post-handshake /// `SessionsSnapshot` reconciles the registry's full set anyway). #[derive(Debug, Default)] struct RpcCovered { added: HashSet, removed: HashSet, } #[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)] pub struct RegistryClient { registry_url: String, daemon_id: String, hostname: String, daemon_url: String, home_dir: String, token: String, /// Operator-set perimeter (id + reachable address) advertised in heartbeats /// so peer daemons can dial this one over T2 when they share the same /// perimeter. Typed so "both or neither" holds by construction (NAP-272). perimeter: Option, manager: SessionManager, /// Local mirror of the global catalog (peer daemons + peer sessions), /// populated from `DaemonInbound::CatalogEvent` frames the registry /// fans out. catalog: CatalogMirror, /// Shared with HTTP handlers (NAP-247): when one calls `resolve_attach`, /// we route the request through this client's control link and route the /// reply back via the handle's pending map. broker: BrokerHandle, /// Outbound channel from `BrokerHandle::resolve_attach` (and future broker /// requests). Drained inside `connect_once` and forwarded into the WS. broker_rx: mpsc::Receiver, } impl RegistryClient { #[allow(clippy::too_many_arguments)] pub fn new( registry_url: impl Into, daemon_id: impl Into, hostname: impl Into, daemon_url: impl Into, home_dir: impl Into, token: impl Into, manager: SessionManager, ) -> Self { Self::new_with_bus( registry_url, daemon_id, hostname, daemon_url, home_dir, token, manager, None, ) } /// Variant that wires a [`LocalEventBus`](crate::LocalEventBus) into the /// catalog mirror so peer deltas reach `WS /api/events` (NAP-248). #[allow(clippy::too_many_arguments)] pub fn new_with_bus( registry_url: impl Into, daemon_id: impl Into, hostname: impl Into, daemon_url: impl Into, home_dir: impl Into, token: impl Into, manager: SessionManager, bus: Option, ) -> Self { let (broker, broker_rx) = BrokerHandle::new(); let catalog = match &bus { Some(b) => CatalogMirror::new().with_bus(b.clone()), None => CatalogMirror::new(), }; Self { registry_url: registry_url.into(), daemon_id: daemon_id.into(), hostname: hostname.into(), daemon_url: daemon_url.into(), home_dir: home_dir.into(), token: token.into(), perimeter: None, manager, catalog, broker, broker_rx, } } /// Cheap clone of the broker handle for HTTP handlers (NAP-247). Sends /// here are pumped through this client's control link. pub fn broker(&self) -> BrokerHandle { self.broker.clone() } /// Cheap handle to the daemon's peer-catalog mirror. Clone freely; /// internal state is `Arc>`. pub fn catalog(&self) -> CatalogMirror { self.catalog.clone() } /// Set the perimeter advertised in heartbeats. The typed argument encodes /// "both id and addr, or neither" — a private address without a perimeter /// id is no longer representable (NAP-272). pub fn with_perimeter(mut self, perimeter: Option) -> Self { self.perimeter = perimeter; self } /// Run forever, reconnecting with exponential backoff between attempts. /// Returns only if cancelled by the caller dropping the task. pub async fn run(mut self) { let mut backoff = BACKOFF_MIN; loop { let outcome = self.connect_once().await; // Whatever happened, the link is down now — make the broker fail // fast until the next successful handshake re-marks it up (NAP-263). self.broker.set_connected(false); match outcome { 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(&mut 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(), home_dir: self.home_dir.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::(&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:?}" ))) } } // Link is up — let the broker queue ResolveAttach frames again // (NAP-263). Cleared in `run` once this connection ends. self.broker.set_connected(true); // Subscribe to the session-event stream BEFORE taking the snapshot. A // reap or kill that races the handshake would otherwise slip between // the snapshot and the subscription — present in neither — leaving the // registry advertising a session this daemon no longer has until the // next reconnect. `local_api::pump_events` subscribes-then-snapshots // for the browser/CLI event WS for exactly this reason. let mut events = self.manager.events(); let snapshot = DaemonOutbound::SessionsSnapshot { sessions: self.manager.snapshot().await, }; sink.send(Message::Text(serde_json::to_string(&snapshot)?)) .await?; // The first interval tick fires immediately, so the registry gets a // perimeter advertisement (NAP-247) right after the handshake instead // of waiting a full HEARTBEAT_INTERVAL for `update_perimeter` to land. let mut hb = tokio::time::interval(HEARTBEAT_INTERVAL); hb.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); // Session ids whose async lifecycle event the in-flight RPC reply // already conveys to the registry — swallowed here so the registry // writes each change exactly once (NAP-268). let mut covered = RpcCovered::default(); 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, &mut covered).await?; } Some(Ok(_)) => {} } } ev = events.recv() => { if !forward_session_event(&self.manager, ev, &mut sink, &mut covered).await? { break; } } _ = hb.tick() => { sink.send(Message::Text(serde_json::to_string( &DaemonOutbound::Heartbeat { perimeter: self.perimeter.clone(), }, )?)).await?; } Some(broker_msg) = self.broker_rx.recv() => { // Outbound broker requests (NAP-247) — e.g. ResolveAttach. sink.send(Message::Text(serde_json::to_string(&broker_msg)?)) .await?; } } } Ok(()) } async fn handle_inbound( &self, text: &str, sink: &mut S, covered: &mut RpcCovered, ) -> Result<(), ClientError> where S: SinkExt + Unpin, { let parsed: DaemonInbound = match serde_json::from_str(text) { Ok(p) => p, Err(e) => { // A single malformed frame must not drop the registry link — // log and skip it (NAP-269). tracing::warn!(error = %e, "malformed registry frame — skipping"); return Ok(()); } }; 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) => { // The reply carries the session, so the registry applies // it on receipt; swallow the async `SessionAdded` this // create just queued on our event stream (NAP-268). covered.added.insert(s.id); 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(()) => { // The reply confirms the removal; swallow the async // `SessionRemoved` it queued so the registry drops the // session exactly once (NAP-268). covered.removed.insert(id); 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 { reply_to, result } => { self.broker.complete_resolve(reply_to, result).await; } DaemonInbound::ForwardCreateResult { reply_to, result } => { self.broker.complete_forward_create(reply_to, result).await; } DaemonInbound::ForwardKillResult { reply_to, result } => { self.broker.complete_forward_kill(reply_to, result).await; } DaemonInbound::ListDir { reply_to, path } => { // List a directory on this host for a client's spawn-form // autocomplete (NAP-329). Resolved against this daemon's home. let result = match crate::dirlist::list_dir(&self.home_dir, &path) { Ok(listing) => RpcResult::Ok { value: listing }, Err(error) => RpcResult::Err { error }, }; sink.send(Message::Text(serde_json::to_string( &DaemonOutbound::ListDirResult { reply_to, result }, )?)) .await?; } DaemonInbound::ForwardListDirResult { reply_to, result } => { self.broker .complete_forward_list_dir(reply_to, result) .await; } DaemonInbound::RecentDirs { reply_to } => { // Answer a peer's recent-dirs query from this host's spawn // history (#26). Daemon-wide, so no path to resolve. let value = crate::spawn_history::to_wire(self.manager.recent_dirs()); sink.send(Message::Text(serde_json::to_string( &DaemonOutbound::RecentDirsResult { reply_to, result: RpcResult::Ok { value }, }, )?)) .await?; } DaemonInbound::ForwardRecentDirsResult { reply_to, result } => { self.broker .complete_forward_recent_dirs(reply_to, result) .await; } DaemonInbound::GitRoot { reply_to, path } => { // Probe the enclosing repo of `path` on this host (#26), // resolving the abbreviated path against this daemon's home the // same way `ListDir` does. let dir = crate::dirlist::resolve(&self.home_dir, &path); let value = maestro_protocol::GitRoot { root: crate::gitroot::find_root(&dir), }; sink.send(Message::Text(serde_json::to_string( &DaemonOutbound::GitRootResult { reply_to, result: RpcResult::Ok { value }, }, )?)) .await?; } DaemonInbound::ForwardGitRootResult { reply_to, result } => { self.broker .complete_forward_git_root(reply_to, result) .await; } DaemonInbound::RelaySetup { .. } => { // T3 (relay) is gated on NAP-238. Lands here only so newer // registries sending it don't crash this daemon on parse. tracing::debug!("RelaySetup received but T3 not implemented (NAP-238 pending)"); } DaemonInbound::CatalogEvent { event } => { // NAP-246: peer-filtered catalog delta from the registry. self.catalog.apply(event).await; } DaemonInbound::Unknown => { // A newer registry emitted a frame type this build doesn't know. // Skip it rather than dropping the link (NAP-269). tracing::debug!("unknown registry frame — skipping (forward-compat, NAP-269)"); } } Ok(()) } } /// Translate one item from the session-event stream into an outbound frame on /// `sink`. Returns `Ok(false)` once the stream has closed and the caller /// should stop the loop. /// /// An `Added` / `Removed` whose id is parked in `covered` was driven by an RPC /// whose reply already reports the change to the registry; it is consumed and /// *not* forwarded, so the registry writes each change exactly once (NAP-268). /// /// On `Lagged` the daemon's own view of its sessions skipped frames — which /// may include a `Removed`. Emitting nothing (the previous behaviour) left the /// registry advertising a now-dead session until the next reconnect, since the /// registry only reconciles a daemon's full session set on `apply_snapshot`. /// So we re-send a full `SessionsSnapshot`, which the registry applies by /// replacing this daemon's set — healing the missed removal immediately. The /// snapshot is the full truth, so any still-parked `covered` ids are stale and /// cleared. async fn forward_session_event( manager: &SessionManager, ev: Result, sink: &mut S, covered: &mut RpcCovered, ) -> Result where S: SinkExt + Unpin, { let frame = match ev { Ok(SessionEvent::Added(session)) => { if covered.added.remove(&session.id) { return Ok(true); } DaemonOutbound::SessionAdded { session } } Ok(SessionEvent::Removed(id)) => { if covered.removed.remove(&id) { return Ok(true); } DaemonOutbound::SessionRemoved { id } } Ok(SessionEvent::Updated(session)) => DaemonOutbound::SessionUpdated { session }, Err(RecvError::Lagged(skipped)) => { tracing::warn!( skipped, "session-event stream lagged — resyncing registry with a full snapshot" ); covered.added.clear(); covered.removed.clear(); DaemonOutbound::SessionsSnapshot { sessions: manager.snapshot().await, } } Err(RecvError::Closed) => return Ok(false), }; sink.send(Message::Text(serde_json::to_string(&frame)?)) .await?; Ok(true) } #[cfg(test)] mod tests { use super::*; use std::pin::Pin; use std::task::{Context, Poll}; use uuid::Uuid; /// Minimal in-memory sink that records every frame sent to it, so we can /// assert exactly what `forward_session_event` emits. #[derive(Default)] struct VecSink(Vec); impl futures_util::Sink for VecSink { type Error = tokio_tungstenite::tungstenite::Error; fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { self.get_mut().0.push(item); Ok(()) } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } /// A `SessionManager` backed by temp dirs. It never spawns tmux here — the /// tests only drive `forward_session_event`, which at most reads the /// (empty) in-memory session map via `snapshot()`. fn test_manager() -> (SessionManager, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let mgr = SessionManager::new( "test-daemon", crate::tmux::Tmux::new(), crate::sidecar::SidecarStore::new(dir.path().join("sessions")), crate::worktree::WorktreeManager::new(dir.path().join("worktrees")), dir.path().join("fifos"), ); (mgr, dir) } fn decode(sink: &VecSink) -> Vec { sink.0 .iter() .map(|m| match m { Message::Text(t) => serde_json::from_str(t).unwrap(), other => panic!("unexpected frame {other:?}"), }) .collect() } fn sample_session(id: Uuid) -> maestro_protocol::Session { maestro_protocol::Session { id, daemon_id: "test-daemon".into(), cwd: std::path::PathBuf::from("/tmp"), command: "bash".into(), env: Default::default(), metadata: Default::default(), status: maestro_protocol::SessionStatus::Running, created_at: chrono::Utc::now(), foreground: None, activity: Default::default(), last_activity_at: chrono::Utc::now(), current_cwd: None, } } // Regression test for NAP-291 (defect B): a lagged event stream may have // dropped a `Removed`, so the daemon must resync the registry's full set // rather than silently skip — otherwise a dead session lingers in the // catalog until the next reconnect. The previous `continue` sent nothing. #[tokio::test] async fn lagged_stream_resyncs_with_full_snapshot() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let cont = forward_session_event( &mgr, Err(RecvError::Lagged(9)), &mut sink, &mut RpcCovered::default(), ) .await .unwrap(); assert!(cont, "lag is recoverable; the loop must continue"); let frames = decode(&sink); assert!( matches!(frames.as_slice(), [DaemonOutbound::SessionsSnapshot { .. }]), "a lagged stream must resync via SessionsSnapshot; got {frames:?}", ); } #[tokio::test] async fn removed_event_forwards_session_removed() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let id = Uuid::new_v4(); let cont = forward_session_event( &mgr, Ok(SessionEvent::Removed(id)), &mut sink, &mut RpcCovered::default(), ) .await .unwrap(); assert!(cont); let frames = decode(&sink); assert!( matches!(frames.as_slice(), [DaemonOutbound::SessionRemoved { id: got }] if *got == id), "got {frames:?}", ); } #[tokio::test] async fn closed_stream_signals_break_without_sending() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let cont = forward_session_event( &mgr, Err(RecvError::Closed), &mut sink, &mut RpcCovered::default(), ) .await .unwrap(); assert!(!cont, "a closed stream must signal the loop to break"); assert!(sink.0.is_empty(), "nothing should be sent on close"); } // NAP-268: an RPC-driven create/kill reports the change to the registry via // its reply, so the matching async lifecycle event must be swallowed — never // forwarded — to keep the registry's catalog write single (not idempotent). #[tokio::test] async fn rpc_covered_added_is_swallowed_once() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let id = Uuid::new_v4(); let mut covered = RpcCovered::default(); covered.added.insert(id); let cont = forward_session_event( &mgr, Ok(SessionEvent::Added(sample_session(id))), &mut sink, &mut covered, ) .await .unwrap(); assert!(cont, "swallowing an event must not stop the loop"); assert!( sink.0.is_empty(), "an RPC-covered SessionAdded must not reach the registry; got {:?}", decode(&sink), ); assert!( !covered.added.contains(&id), "the parked id must be consumed exactly once", ); } #[tokio::test] async fn rpc_covered_removed_is_swallowed_once() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let id = Uuid::new_v4(); let mut covered = RpcCovered::default(); covered.removed.insert(id); let cont = forward_session_event(&mgr, Ok(SessionEvent::Removed(id)), &mut sink, &mut covered) .await .unwrap(); assert!(cont); assert!( sink.0.is_empty(), "an RPC-covered SessionRemoved must not reach the registry; got {:?}", decode(&sink), ); assert!(!covered.removed.contains(&id)); } // A spontaneous change (e.g. a session created via the daemon's local HTTP // API, or reaped on its own) is the registry's only signal, so it must still // forward even while *other* ids are parked. #[tokio::test] async fn uncovered_added_still_forwards() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let parked = Uuid::new_v4(); let spontaneous = Uuid::new_v4(); let mut covered = RpcCovered::default(); covered.added.insert(parked); let cont = forward_session_event( &mgr, Ok(SessionEvent::Added(sample_session(spontaneous))), &mut sink, &mut covered, ) .await .unwrap(); assert!(cont); let frames = decode(&sink); assert!( matches!( frames.as_slice(), [DaemonOutbound::SessionAdded { session }] if session.id == spontaneous ), "an unparked create must forward; got {frames:?}", ); assert!( covered.added.contains(&parked), "an unrelated parked id must be left intact", ); } // A second event for the same id (e.g. the RPC was swallowed, then the // session is reaped) is no longer covered and forwards normally. #[tokio::test] async fn covered_consumed_then_subsequent_event_forwards() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let id = Uuid::new_v4(); let mut covered = RpcCovered::default(); covered.added.insert(id); // First event: the RPC-covered create — swallowed. forward_session_event( &mgr, Ok(SessionEvent::Added(sample_session(id))), &mut sink, &mut covered, ) .await .unwrap(); assert!(sink.0.is_empty()); // Second event: a spontaneous removal of that same session — forwarded, // since only the create was parked. forward_session_event(&mgr, Ok(SessionEvent::Removed(id)), &mut sink, &mut covered) .await .unwrap(); let frames = decode(&sink); assert!( matches!(frames.as_slice(), [DaemonOutbound::SessionRemoved { id: got }] if *got == id), "got {frames:?}", ); } // A lag resync is the registry's full truth, so any still-parked ids are // stale and must be cleared (else a later unrelated event could be wrongly // swallowed). #[tokio::test] async fn lag_resync_clears_parked_ids() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let mut covered = RpcCovered::default(); covered.added.insert(Uuid::new_v4()); covered.removed.insert(Uuid::new_v4()); forward_session_event(&mgr, Err(RecvError::Lagged(3)), &mut sink, &mut covered) .await .unwrap(); assert!( covered.added.is_empty() && covered.removed.is_empty(), "a full-snapshot resync must clear stale parked ids", ); } }