diff --git a/crates/maestro-daemon/src/registry_client.rs b/crates/maestro-daemon/src/registry_client.rs index 6d3ab7a..96a181f 100644 --- a/crates/maestro-daemon/src/registry_client.rs +++ b/crates/maestro-daemon/src/registry_client.rs @@ -194,14 +194,19 @@ impl RegistryClient { } } - // Initial snapshot + // 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?; - let mut events = self.manager.events(); // 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. @@ -222,24 +227,8 @@ impl RegistryClient { } } 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, + if !forward_session_event(&self.manager, ev, &mut sink).await? { + break; } } _ = hb.tick() => { @@ -310,3 +299,143 @@ impl RegistryClient { 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. +/// +/// 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. +async fn forward_session_event( + manager: &SessionManager, + ev: Result, + sink: &mut S, +) -> Result +where + S: SinkExt + Unpin, +{ + let frame = match ev { + Ok(SessionEvent::Added(session)) => DaemonOutbound::SessionAdded { session }, + Ok(SessionEvent::Removed(id)) => 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" + ); + 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() + } + + // 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) + .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) + .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) + .await + .unwrap(); + assert!(!cont, "a closed stream must signal the loop to break"); + assert!(sink.0.is_empty(), "nothing should be sent on close"); + } +}