omni_dev/

worktrees.rs

//! The cross-window worktree registry engine.
//!
//! Maintains the live, authoritative set of repos/worktrees open across *every*
//! VS Code window, fed by a first-party companion extension that reports from
//! each window over the daemon's control socket. The resident daemon is the
//! rendezvous point the per-window extension sandbox cannot replace: each window
//! can see only its own `workspace.workspaceFolders`, so a single process
//! aggregating those registrations is the only cross-window source of truth.
//! See ADR-0040.
//!
//! This is the standalone engine, analogous to [`crate::browser`] and
//! [`crate::snowflake`]; the daemon adapter lives in
//! [`crate::daemon::services::worktrees`].
//!
//! Like the Snowflake engine this is cheap and in-memory — no async setup, no
//! secret persisted. The registry lives behind a [`std::sync::Mutex`] that is
//! **never held across an `.await`** (the Snowflake rule); every op is pure CPU
//! under the lock, so liveness reaping happens inline on each read rather than
//! from a background task.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Mutex, MutexGuard, PoisonError};
use std::time::Duration;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// How long a window may go silent before it ages out of the registry. Three
/// missed ~10s heartbeats; a window that crashed without firing `unregister`
/// disappears on the next read. The resident process is what makes this
/// liveness correct — a flat shared file could not reap stale entries.
const DEFAULT_TTL: Duration = Duration::from_secs(30);

/// Ceiling on live registry entries, so a misbehaving companion flooding
/// `register` with distinct keys cannot grow daemon memory faster than the TTL
/// reaps it (#1140). Far above any real window count; when a new key would
/// exceed it, the longest-silent entry is evicted instead of rejecting the
/// request — an evicted live window self-heals via the `heartbeat` →
/// `{known: false}` → re-register path, so `register` stays infallible for the
/// companion.
const MAX_WINDOWS: usize = 256;

/// A `register` request from a companion extension.
///
/// The companion owns its `key` (a per-`activate()` UUID) so the registry never
/// has to reason about whether `vscode.env.sessionId` is unique per window;
/// everything else is best-effort metadata.
#[derive(Debug, Clone, Deserialize)]
pub struct RegisterRequest {
    /// Stable per-window identity, generated by the companion on activation.
    pub key: String,
    /// Absolute paths of the window's workspace folders.
    #[serde(default)]
    pub folders: Vec<PathBuf>,
    /// Repository root or name, when the window has one.
    #[serde(default)]
    pub repo: Option<String>,
    /// The window title, for display.
    #[serde(default)]
    pub title: Option<String>,
    /// The reporting extension-host process id.
    #[serde(default)]
    pub pid: Option<u32>,
}

/// One open window's live registration. Serialized verbatim into `list` /
/// `status` payloads; consumers compute age from `last_seen` (RFC 3339).
#[derive(Debug, Clone, Serialize)]
pub struct WindowEntry {
    /// The companion-owned per-window key.
    pub key: String,
    /// Absolute workspace-folder paths.
    pub folders: Vec<PathBuf>,
    /// Repository root or name, if reported.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub repo: Option<String>,
    /// Window title, if reported.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub title: Option<String>,
    /// Reporting extension-host pid, if reported.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pid: Option<u32>,
    /// When the registry last heard from this window (register or heartbeat).
    pub last_seen: DateTime<Utc>,
}

/// The cross-window worktree registry: the in-memory, TTL-reaped set of open
/// windows. Hosted by
/// [`WorktreesService`](crate::daemon::services::worktrees::WorktreesService).
pub struct WorktreesRegistry {
    /// Open windows keyed by their companion-owned `key`.
    windows: Mutex<HashMap<String, WindowEntry>>,
    /// How long an entry survives without a heartbeat.
    ttl: Duration,
}

impl WorktreesRegistry {
    /// Creates the registry with the default liveness TTL. Cheap — no I/O.
    #[must_use]
    pub fn new() -> Self {
        Self {
            windows: Mutex::new(HashMap::new()),
            ttl: DEFAULT_TTL,
        }
    }

    /// Locks the registry, recovering from a poisoned mutex (a panic in a prior
    /// critical section must not wedge the whole registry).
    fn lock(&self) -> MutexGuard<'_, HashMap<String, WindowEntry>> {
        self.windows.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// Records (upserts) a window registration. Reaps stale entries first, then
    /// — only when a genuinely new key would grow the map past [`MAX_WINDOWS`] —
    /// evicts the longest-silent entry. Infallible: an upsert never evicts, and
    /// callers validate the `key` before reaching here.
    pub fn register(&self, req: RegisterRequest) {
        let now = Utc::now();
        let mut windows = self.lock();
        reap(&mut windows, self.ttl, now);
        // Upserts never evict; only a genuinely new key can grow the map, and
        // never past MAX_WINDOWS.
        if !windows.contains_key(&req.key) && windows.len() >= MAX_WINDOWS {
            evict_oldest(&mut windows);
        }
        windows.insert(
            req.key.clone(),
            WindowEntry {
                key: req.key,
                folders: req.folders,
                repo: req.repo,
                title: req.title,
                pid: req.pid,
                last_seen: now,
            },
        );
    }

    /// Refreshes a window's liveness. Returns whether the key was known: a
    /// `false` tells a window that started before the daemon — or survived a
    /// daemon restart — to re-`register`, since the registry is in-memory and
    /// has no record of it.
    pub fn heartbeat(&self, key: &str) -> bool {
        let now = Utc::now();
        let mut windows = self.lock();
        reap(&mut windows, self.ttl, now);
        match windows.get_mut(key) {
            Some(entry) => {
                entry.last_seen = now;
                true
            }
            None => false,
        }
    }

    /// Drops a window's registration. Returns whether an entry was present.
    pub fn unregister(&self, key: &str) -> bool {
        let now = Utc::now();
        let mut windows = self.lock();
        let removed = windows.remove(key).is_some();
        reap(&mut windows, self.ttl, now);
        removed
    }

    /// Reaps stale entries, then returns the live set sorted for deterministic
    /// output. Holds the lock only for pure-CPU work.
    pub fn list(&self) -> Vec<WindowEntry> {
        let now = Utc::now();
        let mut windows = self.lock();
        reap(&mut windows, self.ttl, now);
        sorted_entries(&windows)
    }

    /// The first workspace folder of a still-live window, if it has one. Used by
    /// the tray "focus" action to resolve a key to a folder to open. Does not
    /// reap — a menu action races the reaper either way, and the caller handles
    /// a `None` (the window may have closed).
    pub fn first_folder(&self, key: &str) -> Option<PathBuf> {
        let windows = self.lock();
        windows.get(key).and_then(|e| e.folders.first().cloned())
    }
}

impl Default for WorktreesRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Removes entries last seen longer than `ttl` ago. Pure CPU; the caller holds
/// the registry lock but never `.await`s while holding it.
fn reap(windows: &mut HashMap<String, WindowEntry>, ttl: Duration, now: DateTime<Utc>) {
    let max_age = ttl.as_secs() as i64;
    windows.retain(|_, e| (now - e.last_seen).num_seconds() <= max_age);
}

/// Removes the entry with the oldest `last_seen` (ties broken by lowest key
/// for determinism). Called when a `register` of a new key would grow the
/// registry past [`MAX_WINDOWS`]. Pure CPU under the registry lock, like
/// [`reap`].
fn evict_oldest(windows: &mut HashMap<String, WindowEntry>) {
    let oldest = windows
        .values()
        .min_by(|a, b| {
            a.last_seen
                .cmp(&b.last_seen)
                .then_with(|| a.key.cmp(&b.key))
        })
        .map(|e| e.key.clone());
    if let Some(key) = oldest {
        windows.remove(&key);
    }
}

/// Snapshots the registry into a stably-ordered vector (by repo, then key) so
/// `list`/`status`/`menu` output is deterministic despite `HashMap` ordering.
fn sorted_entries(windows: &HashMap<String, WindowEntry>) -> Vec<WindowEntry> {
    let mut entries: Vec<WindowEntry> = windows.values().cloned().collect();
    entries.sort_by(|a, b| a.repo.cmp(&b.repo).then_with(|| a.key.cmp(&b.key)));
    entries
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    fn register_request(key: &str, repo: Option<&str>, folder: &str) -> RegisterRequest {
        RegisterRequest {
            key: key.to_string(),
            folders: vec![PathBuf::from(folder)],
            repo: repo.map(str::to_string),
            title: Some(format!("{key}-title")),
            pid: Some(1234),
        }
    }

    #[test]
    fn list_is_empty_initially() {
        let reg = WorktreesRegistry::new();
        assert!(reg.list().is_empty());
    }

    #[test]
    fn register_then_list_round_trips() {
        let reg = WorktreesRegistry::new();
        reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
        let windows = reg.list();
        assert_eq!(windows.len(), 1);
        assert_eq!(windows[0].key, "w1");
        assert_eq!(windows[0].repo.as_deref(), Some("repo-a"));
    }

    #[test]
    fn register_is_idempotent_upsert() {
        let reg = WorktreesRegistry::new();
        reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
        // Re-registering the same key updates rather than duplicates.
        reg.register(register_request("w1", Some("repo-b"), "/tmp/b"));
        let windows = reg.list();
        assert_eq!(windows.len(), 1);
        assert_eq!(windows[0].repo.as_deref(), Some("repo-b"));
    }

    #[test]
    fn heartbeat_reports_known_and_unknown() {
        let reg = WorktreesRegistry::new();
        // Unknown before registration: the window must re-register.
        assert!(!reg.heartbeat("w1"));
        reg.register(register_request("w1", None, "/tmp/a"));
        assert!(reg.heartbeat("w1"));
    }

    #[test]
    fn unregister_removes() {
        let reg = WorktreesRegistry::new();
        reg.register(register_request("w1", None, "/tmp/a"));
        assert!(reg.unregister("w1"));
        // Removing again is a no-op.
        assert!(!reg.unregister("w1"));
    }

    #[test]
    fn first_folder_returns_first_folder_or_none() {
        let reg = WorktreesRegistry::new();
        // No such key.
        assert!(reg.first_folder("missing").is_none());
        reg.register(register_request("w1", None, "/tmp/a"));
        assert_eq!(reg.first_folder("w1"), Some(PathBuf::from("/tmp/a")));
        // A folderless window resolves to None rather than a folder.
        reg.register(RegisterRequest {
            key: "w2".to_string(),
            folders: vec![],
            repo: None,
            title: None,
            pid: None,
        });
        assert!(reg.first_folder("w2").is_none());
    }

    #[test]
    fn reap_evicts_only_stale_entries() {
        let now = Utc::now();
        let mut windows = HashMap::new();
        windows.insert(
            "fresh".to_string(),
            WindowEntry {
                key: "fresh".to_string(),
                folders: vec![],
                repo: None,
                title: None,
                pid: None,
                last_seen: now - chrono::Duration::seconds(5),
            },
        );
        windows.insert(
            "stale".to_string(),
            WindowEntry {
                key: "stale".to_string(),
                folders: vec![],
                repo: None,
                title: None,
                pid: None,
                last_seen: now - chrono::Duration::seconds(120),
            },
        );
        reap(&mut windows, DEFAULT_TTL, now);
        assert!(windows.contains_key("fresh"));
        assert!(!windows.contains_key("stale"));
    }

    /// A minimal entry for cap/eviction tests; only `key` and `last_seen`
    /// participate in eviction order.
    fn entry_at(key: &str, last_seen: DateTime<Utc>) -> WindowEntry {
        WindowEntry {
            key: key.to_string(),
            folders: vec![],
            repo: None,
            title: None,
            pid: None,
            last_seen,
        }
    }

    #[test]
    fn evict_oldest_removes_oldest_with_key_tiebreak() {
        let now = Utc::now();
        let mut windows = HashMap::new();
        windows.insert("young".to_string(), entry_at("young", now));
        windows.insert(
            "old-b".to_string(),
            entry_at("old-b", now - chrono::Duration::seconds(10)),
        );
        windows.insert(
            "old-a".to_string(),
            entry_at("old-a", now - chrono::Duration::seconds(10)),
        );
        // Oldest `last_seen` is shared by two entries; the lowest key loses.
        evict_oldest(&mut windows);
        assert!(!windows.contains_key("old-a"));
        assert!(windows.contains_key("old-b"));
        assert!(windows.contains_key("young"));
        // Empty map is a no-op rather than a panic.
        let mut empty: HashMap<String, WindowEntry> = HashMap::new();
        evict_oldest(&mut empty);
        assert!(empty.is_empty());
    }

    #[test]
    fn register_at_cap_evicts_only_the_oldest() {
        let reg = WorktreesRegistry::new();
        // Seed a full registry directly (registering 256 times would work too,
        // but sub-second timestamps may tie; explicit timestamps make the
        // highest-numbered key unambiguously the oldest).
        {
            let mut windows = reg.lock();
            let base = Utc::now();
            for i in 0..MAX_WINDOWS {
                let key = format!("w{i:03}");
                windows.insert(
                    key.clone(),
                    entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
                );
            }
        }
        // A new key at the cap displaces exactly the longest-silent entry.
        reg.register(register_request("fresh", None, "/tmp/f"));
        let windows = reg.lock();
        assert_eq!(windows.len(), MAX_WINDOWS);
        assert!(windows.contains_key("fresh"));
        assert!(!windows.contains_key(&format!("w{:03}", MAX_WINDOWS - 1)));
        assert!(windows.contains_key("w000"));
    }

    #[test]
    fn register_upsert_at_cap_does_not_evict() {
        let reg = WorktreesRegistry::new();
        {
            let mut windows = reg.lock();
            let base = Utc::now();
            for i in 0..MAX_WINDOWS {
                let key = format!("w{i:03}");
                windows.insert(
                    key.clone(),
                    entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
                );
            }
        }
        // Re-registering an existing key is an upsert: nothing is displaced,
        // not even the oldest entry.
        let oldest = format!("w{:03}", MAX_WINDOWS - 1);
        reg.register(register_request(&oldest, Some("r"), "/tmp/a"));
        let windows = reg.lock();
        assert_eq!(windows.len(), MAX_WINDOWS);
        assert!(windows.contains_key(&oldest));
        assert!(windows.contains_key("w000"));
    }

    #[test]
    fn sorted_entries_orders_by_repo_then_key() {
        let now = Utc::now();
        let mut windows = HashMap::new();
        for (key, repo) in [("z", "repo-a"), ("a", "repo-b"), ("m", "repo-a")] {
            windows.insert(
                key.to_string(),
                WindowEntry {
                    key: key.to_string(),
                    folders: vec![],
                    repo: Some(repo.to_string()),
                    title: None,
                    pid: None,
                    last_seen: now,
                },
            );
        }
        let entries = sorted_entries(&windows);
        let ordered: Vec<(&str, &str)> = entries
            .iter()
            .map(|e| (e.key.as_str(), e.repo.as_deref().unwrap()))
            .collect();
        assert_eq!(
            ordered,
            vec![("m", "repo-a"), ("z", "repo-a"), ("a", "repo-b")]
        );
    }

    #[test]
    fn default_constructs_an_empty_registry() {
        let reg = WorktreesRegistry::default();
        assert!(reg.lock().is_empty());
    }
}