maestro/crates/maestro-tui/src/catalog.rs

//! `CatalogStore` — the TUI's view of the session list.
//!
//! Backed entirely by the daemon's `WS /api/events` stream:
//!
//! 1. The first frame after connect is a [`CatalogEvent::Snapshot`] — that's
//!    the bootstrap. No separate REST round-trip needed.
//! 2. Subsequent frames are deltas (`SessionAdded`, `SessionUpdated`, …) that we
//!    apply atomically and bump a monotonic version counter.
//! 3. On WS disconnect: exponential backoff (1s → 30s, ±25% jitter); the next
//!    connect resends a fresh `Snapshot` so we converge without needing
//!    server-side replay.
//! 4. `Lagged { skipped }` from the broadcast queue is logged at warn level
//!    and ignored — the server will re-snapshot on its own bookkeeping; for
//!    safety we trigger a self-reconnect so the next Snapshot reconciles us.
//!
//! Cloneable and `Send + Sync` — pass it to the render loop and to a Tokio
//! task that drives [`CatalogStore::run_events_loop`].

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Duration;

use futures_util::StreamExt;
use maestro_protocol::{CatalogEvent, Session};
use tokio::sync::Notify;
use tokio_tungstenite::tungstenite::Message;
use uuid::Uuid;

/// Snapshot returned by [`CatalogStore::snapshot`]. `version` is monotone — a
/// renderer can cheap-compare it across frames to know whether re-render is
/// needed.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CatalogSnapshot {
    pub sessions: Vec<Session>,
    pub version: u64,
}

#[derive(Debug, Default, Clone)]
pub struct CatalogStore {
    inner: Arc<Inner>,
}

#[derive(Debug, Default)]
struct Inner {
    state: RwLock<State>,
    notify: Notify,
}

#[derive(Debug, Default)]
struct State {
    sessions: HashMap<Uuid, Session>,
    /// Bumped on every applied event. Wraps at u64::MAX, which is fine
    /// because comparisons are equality, not ordering.
    version: u64,
    /// True once we've applied at least one Snapshot. Used by
    /// [`CatalogStore::wait_for_initial_snapshot`].
    bootstrapped: bool,
}

impl CatalogStore {
    /// Apply one event atomically. Used by tests directly and by the WS loop.
    pub fn apply_event(&self, ev: CatalogEvent) {
        let mut s = self.inner.state.write().expect("catalog rwlock poisoned");
        match ev {
            CatalogEvent::Snapshot { sessions } => {
                s.sessions = sessions.into_iter().map(|x| (x.id, x)).collect();
                s.bootstrapped = true;
            }
            CatalogEvent::SessionAdded { session } | CatalogEvent::SessionUpdated { session } => {
                s.sessions.insert(session.id, session);
            }
            CatalogEvent::SessionRemoved { id } => {
                s.sessions.remove(&id);
            }
            CatalogEvent::Lagged { skipped } => {
                tracing::warn!(
                    skipped,
                    "catalog events stream lagged; resync via reconnect"
                );
                // We don't reset state here — the WS loop notices the warn and
                // can choose to drop+reconnect. Keep state until that happens.
            }
        }
        s.version = s.version.wrapping_add(1);
        drop(s);
        self.inner.notify.notify_waiters();
    }

    /// Cheap O(sessions) read of current state.
    pub fn snapshot(&self) -> CatalogSnapshot {
        let s = self.inner.state.read().expect("catalog rwlock poisoned");
        let mut sessions: Vec<Session> = s.sessions.values().cloned().collect();
        sessions.sort_by_key(|x| x.id);
        CatalogSnapshot {
            sessions,
            version: s.version,
        }
    }

    /// Returns true once we've applied at least one Snapshot frame.
    pub fn is_bootstrapped(&self) -> bool {
        self.inner
            .state
            .read()
            .expect("catalog rwlock poisoned")
            .bootstrapped
    }

    /// Wait until the catalog has been changed by something — applied event
    /// or reconnect Snapshot. The render loop awaits this between frames.
    pub async fn changed(&self) {
        self.inner.notify.notified().await;
    }

    /// Wait until the first Snapshot has landed. Resolves immediately if
    /// already bootstrapped.
    pub async fn wait_for_initial_snapshot(&self) {
        loop {
            if self.is_bootstrapped() {
                return;
            }
            self.inner.notify.notified().await;
        }
    }

    /// Long-running task: connect to `ws_url`, parse [`CatalogEvent`] frames,
    /// apply them. Reconnects with exponential backoff on disconnect. Never
    /// returns — spawn it on a Tokio task.
    pub async fn run_events_loop(self, ws_url: String) {
        self.run_events_loop_reporting(ws_url, crate::daemon_link::DaemonLink::default())
            .await
    }

    /// [`Self::run_events_loop`], reporting every connect and disconnect to
    /// `link` — the socket dying is the first sign the daemon went down or was
    /// upgraded under the client (`docs/RESTART.md`).
    pub async fn run_events_loop_reporting(
        self,
        ws_url: String,
        link: crate::daemon_link::DaemonLink,
    ) {
        let mut backoff = Duration::from_secs(1);
        loop {
            match tokio_tungstenite::connect_async(&ws_url).await {
                Ok((ws, _resp)) => {
                    backoff = Duration::from_secs(1);
                    link.note_connected();
                    let (_sink, mut stream) = ws.split();
                    while let Some(frame) = stream.next().await {
                        match frame {
                            Ok(Message::Text(text)) => {
                                match serde_json::from_str::<CatalogEvent>(&text) {
                                    Ok(ev) => self.apply_event(ev),
                                    Err(e) => {
                                        tracing::warn!(error = ?e, "catalog: bad event frame");
                                    }
                                }
                            }
                            Ok(Message::Close(_)) | Err(_) => break,
                            _ => {}
                        }
                    }
                }
                Err(e) => {
                    tracing::debug!(error = ?e, "catalog: connect failed");
                }
            }
            link.note_disconnected(std::time::Instant::now());
            tracing::debug!(?backoff, "catalog: sleeping before reconnect");
            tokio::time::sleep(jitter(backoff)).await;
            backoff = std::cmp::min(backoff.saturating_mul(2), Duration::from_secs(30));
        }
    }
}

/// ±25 % jitter on the backoff. Cheap LCG seeded from the duration itself —
/// avoids touching `rand` for one knob.
fn jitter(base: Duration) -> Duration {
    let ms = base.as_millis() as u64;
    // Tiny pseudo-random in [0, 1) from system time nanos.
    let now_nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.subsec_nanos())
        .unwrap_or(0) as u64;
    // ±25 %: scale by [-250, 250] of one quarter ms per ms.
    let delta = (now_nanos % 500) as i64 - 250;
    let scaled = (ms as i64).saturating_mul(delta).div_euclid(1000);
    let total = ms.saturating_add_signed(scaled);
    Duration::from_millis(total)
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use maestro_protocol::{Session, SessionStatus};

    fn session(id: u128, command: &str) -> Session {
        Session {
            id: Uuid::from_u128(id),
            command: command.into(),
            cwd: std::path::PathBuf::from("/tmp"),
            env: Default::default(),
            metadata: Default::default(),
            worktree: None,
            status: SessionStatus::Running,
            created_at: Utc::now(),
            foreground: None,
            activity: Default::default(),
            activity_detail: None,
            last_activity_at: chrono::Utc::now(),
            activity_since: chrono::Utc::now(),
            current_cwd: None,
            blocked_work: None,
        }
    }

    #[test]
    fn snapshot_event_replaces_state() {
        let store = CatalogStore::default();
        store.apply_event(CatalogEvent::Snapshot {
            sessions: vec![session(1, "a")],
        });
        let snap = store.snapshot();
        assert_eq!(snap.sessions.len(), 1);
        assert_eq!(snap.version, 1);
        assert!(store.is_bootstrapped());

        // A second snapshot wipes out the old state.
        store.apply_event(CatalogEvent::Snapshot { sessions: vec![] });
        let snap = store.snapshot();
        assert!(snap.sessions.is_empty());
        assert_eq!(snap.version, 2);
    }

    #[test]
    fn session_added_updated_removed() {
        let store = CatalogStore::default();
        store.apply_event(CatalogEvent::Snapshot { sessions: vec![] });
        store.apply_event(CatalogEvent::SessionAdded {
            session: session(7, "vim"),
        });
        assert_eq!(store.snapshot().sessions.len(), 1);
        assert_eq!(store.snapshot().sessions[0].command, "vim");

        let mut updated = session(7, "claude");
        updated.status = SessionStatus::Dead;
        store.apply_event(CatalogEvent::SessionUpdated { session: updated });
        let snap = store.snapshot();
        assert_eq!(snap.sessions[0].command, "claude");
        assert!(matches!(snap.sessions[0].status, SessionStatus::Dead));

        store.apply_event(CatalogEvent::SessionRemoved {
            id: Uuid::from_u128(7),
        });
        assert!(store.snapshot().sessions.is_empty());
    }

    #[test]
    fn lagged_event_doesnt_blow_up_state() {
        let store = CatalogStore::default();
        store.apply_event(CatalogEvent::Snapshot {
            sessions: vec![session(1, "a")],
        });
        let before = store.snapshot();
        store.apply_event(CatalogEvent::Lagged { skipped: 42 });
        let after = store.snapshot();
        assert_eq!(after.sessions, before.sessions);
        assert_eq!(after.version, before.version + 1);
    }

    #[test]
    fn snapshot_is_sorted() {
        let store = CatalogStore::default();
        store.apply_event(CatalogEvent::Snapshot {
            sessions: vec![session(3, "c"), session(1, "a"), session(2, "b")],
        });
        let snap = store.snapshot();
        let session_ids: Vec<_> = snap.sessions.iter().map(|s| s.id.as_u128()).collect();
        assert_eq!(session_ids, vec![1, 2, 3]);
    }

    #[tokio::test]
    async fn wait_for_initial_snapshot_resolves_after_first_event() {
        let store = CatalogStore::default();
        let task_store = store.clone();
        let handle = tokio::spawn(async move {
            task_store.wait_for_initial_snapshot().await;
        });
        tokio::time::sleep(Duration::from_millis(10)).await;
        assert!(!handle.is_finished());
        store.apply_event(CatalogEvent::Snapshot { sessions: vec![] });
        // Give the waker a chance.
        tokio::time::timeout(Duration::from_millis(200), handle)
            .await
            .expect("waker did not fire")
            .unwrap();
    }

    #[tokio::test]
    async fn changed_notifies_on_every_apply() {
        let store = CatalogStore::default();
        let waiter = {
            let s = store.clone();
            tokio::spawn(async move { s.changed().await })
        };
        tokio::time::sleep(Duration::from_millis(10)).await;
        store.apply_event(CatalogEvent::Snapshot { sessions: vec![] });
        tokio::time::timeout(Duration::from_millis(200), waiter)
            .await
            .expect("changed did not fire")
            .unwrap();
    }
}