Skip to main content

omni_dev/
worktrees.rs

1//! The cross-window worktree registry engine.
2//!
3//! Maintains the live, authoritative set of repos/worktrees open across *every*
4//! VS Code window, fed by a first-party companion extension that reports from
5//! each window over the daemon's control socket. The resident daemon is the
6//! rendezvous point the per-window extension sandbox cannot replace: each window
7//! can see only its own `workspace.workspaceFolders`, so a single process
8//! aggregating those registrations is the only cross-window source of truth.
9//! See ADR-0040.
10//!
11//! This is the standalone engine, analogous to [`crate::browser`] and
12//! [`crate::snowflake`]; the daemon adapter lives in
13//! [`crate::daemon::services::worktrees`].
14//!
15//! Like the Snowflake engine this is cheap and in-memory — no async setup, no
16//! secret persisted. The registry lives behind a [`std::sync::Mutex`] that is
17//! **never held across an `.await`** (the Snowflake rule); every op is pure CPU
18//! under the lock, so liveness reaping happens inline on each read rather than
19//! from a background task.
20
21use std::collections::{HashMap, HashSet};
22use std::path::PathBuf;
23use std::sync::atomic::{AtomicBool, Ordering};
24use std::sync::{Mutex, MutexGuard, PoisonError};
25use std::time::Duration;
26
27use chrono::{DateTime, Duration as ChronoDuration, Utc};
28use serde::{Deserialize, Serialize};
29use tokio::sync::watch;
30
31/// How long a window may go silent before it ages out of the registry. Three
32/// missed ~10s heartbeats; a window that crashed without firing `unregister`
33/// disappears on the next read. The resident process is what makes this
34/// liveness correct — a flat shared file could not reap stale entries.
35const DEFAULT_TTL: Duration = Duration::from_secs(30);
36
37/// How long a per-repository PR-poll lease lasts before it auto-expires (#1376).
38/// Enabling polling for a repo is deliberately **temporary** — 15 minutes — so an
39/// idle repo stops costing GitHub budget without the user remembering to disable
40/// it; re-enabling refreshes the lease.
41const DEFAULT_POLL_LEASE: Duration = Duration::from_secs(15 * 60);
42
43/// Ceiling on live registry entries, so a misbehaving companion flooding
44/// `register` with distinct keys cannot grow daemon memory faster than the TTL
45/// reaps it (#1140). Far above any real window count; when a new key would
46/// exceed it, the longest-silent entry is evicted instead of rejecting the
47/// request — an evicted live window self-heals via the `heartbeat` →
48/// `{known: false}` → re-register path, so `register` stays infallible for the
49/// companion.
50const MAX_WINDOWS: usize = 256;
51
52/// A `register` request from a companion extension.
53///
54/// The companion owns its `key` (a per-`activate()` UUID) so the registry never
55/// has to reason about whether `vscode.env.sessionId` is unique per window;
56/// everything else is best-effort metadata.
57#[derive(Debug, Clone, Deserialize)]
58pub struct RegisterRequest {
59    /// Stable per-window identity, generated by the companion on activation.
60    pub key: String,
61    /// Absolute paths of the window's workspace folders.
62    #[serde(default)]
63    pub folders: Vec<PathBuf>,
64    /// Repository root or name, when the window has one.
65    #[serde(default)]
66    pub repo: Option<String>,
67    /// The window title, for display.
68    #[serde(default)]
69    pub title: Option<String>,
70    /// The reporting extension-host process id.
71    #[serde(default)]
72    pub pid: Option<u32>,
73}
74
75/// One open window's live registration. Serialized verbatim into `list` /
76/// `status` payloads; consumers compute age from `last_seen` (RFC 3339).
77#[derive(Debug, Clone, Serialize)]
78pub struct WindowEntry {
79    /// The companion-owned per-window key.
80    pub key: String,
81    /// Absolute workspace-folder paths.
82    pub folders: Vec<PathBuf>,
83    /// Repository root or name, if reported.
84    #[serde(skip_serializing_if = "Option::is_none")]
85    pub repo: Option<String>,
86    /// Window title, if reported.
87    #[serde(skip_serializing_if = "Option::is_none")]
88    pub title: Option<String>,
89    /// Reporting extension-host pid, if reported.
90    #[serde(skip_serializing_if = "Option::is_none")]
91    pub pid: Option<u32>,
92    /// When the registry last heard from this window (register or heartbeat).
93    pub last_seen: DateTime<Utc>,
94}
95
96/// The cross-window worktree registry: the in-memory, TTL-reaped set of open
97/// windows. Hosted by
98/// [`WorktreesService`](crate::daemon::services::worktrees::WorktreesService).
99pub struct WorktreesRegistry {
100    /// Open windows keyed by their companion-owned `key`.
101    windows: Mutex<HashMap<String, WindowEntry>>,
102    /// How long an entry survives without a heartbeat.
103    ttl: Duration,
104    /// A monotonically-bumped version counter, incremented whenever the visible
105    /// set of windows changes (a `register`, a removing `unregister`, or a
106    /// mutation-driven reap that drops a stale entry). A push-subscription
107    /// consumer holds a [`watch::Receiver`] from
108    /// [`subscribe_changes`](Self::subscribe_changes) and wakes on each bump to
109    /// re-snapshot (#1267). The counter's *value* is immaterial —
110    /// only that it changed — so a burst coalesces into one wake and the
111    /// subscriber diffs the resulting snapshot to suppress duplicate frames.
112    ///
113    /// `watch` needs no runtime and never blocks, so it fits this engine's
114    /// no-async-setup posture; every [`bump`](Self::bump) happens *after* the map
115    /// guard is dropped, so the `std::Mutex`-never-across-`.await` rule is intact
116    /// (and the watch's own internal lock is never nested under the map lock).
117    changes: watch::Sender<u64>,
118    /// Window keys with a pending "close yourself" directive, set by the
119    /// `close` op (#1277) when a cross-window close must reach a window the
120    /// daemon can only *reply* to, never call. Each key is surfaced — and
121    /// cleared — on that window's next `heartbeat` (the `known:false →
122    /// re-register` precedent, riding the same reply). In-memory only, like the
123    /// window map: a daemon restart drops any pending directive (the close op
124    /// aborts and the user retries — an accepted failure mode). Behind its own
125    /// `Mutex`, taken independently of the window map's, so neither nests.
126    close_pending: Mutex<HashSet<String>>,
127    /// The daemon-backed **show/hide-closed** toggle (#1301): whether the
128    /// companion's tree view shows worktrees with no open window. A single
129    /// cross-window value carried in every `tree`/`subscribe` snapshot so all
130    /// windows read (and live-sync) the same state — `context.globalState` could
131    /// not, being read-once with no cross-window change event. Defaults to `true`
132    /// (show all, the original behavior). A lock-free [`AtomicBool`] rather than a
133    /// `Mutex`, so it is never a `.await`-holding-a-lock hazard; a flip
134    /// [`bump`](Self::bump)s the change-notify so subscribers re-push. In-memory
135    /// like the window map: a daemon restart resets it to the default, which the
136    /// next snapshot propagates to every window.
137    show_closed: AtomicBool,
138    /// The **per-repository PR-poll** enable set (#1376): the GitHub repos
139    /// (`"owner/name"`) whose PR badges the daemon polls. Polling defaults
140    /// **off** — a repo not in this set issues zero `gh` — so the user enables
141    /// only the handful of repos they are actively working on, rather than the
142    /// daemon polling all 29 open repos and exhausting the GitHub budget. A
143    /// cross-window value like [`show_closed`](Self::show_closed): the daemon
144    /// stamps each repo's state onto the `tree` snapshot (`polling_enabled`) and
145    /// a [`set_polling`](Self::set_polling) flip [`bump`](Self::bump)s the
146    /// change-notify so every window recolors and drops/keeps badges in sync.
147    ///
148    /// Unlike `show_closed` this survives a daemon restart: the adapter seeds it
149    /// from a `0600` file on startup ([`seed_polling`](Self::seed_polling)) and
150    /// persists it on each change — otherwise a restart would silently re-disable
151    /// every repo. Behind its **own** `Mutex`, taken independently of the window
152    /// map's (neither nests) and never held across an `.await`.
153    ///
154    /// Each enable is a **time-boxed lease**, not a permanent flag (#1376): the
155    /// value is the wall-clock instant the lease **expires** ([`poll_ttl`] after
156    /// it was enabled), so an idle repo auto-disables and stops costing `gh`
157    /// without the user remembering to turn it back off. Expired entries are
158    /// reaped on read — the window-TTL precedent — so the icon greys, badges
159    /// drop, and the poller stops within one snapshot tick of expiry.
160    ///
161    /// [`poll_ttl`]: Self::poll_ttl
162    polling_enabled: Mutex<HashMap<String, DateTime<Utc>>>,
163    /// How long a repo's PR-poll lease lasts before it auto-expires (#1376).
164    /// [`DEFAULT_POLL_LEASE`] in production; tests inject a short value via the
165    /// `#[cfg(test)]` `with_poll_ttl` constructor (not linked — it does not exist
166    /// in a non-test doc build).
167    poll_ttl: Duration,
168}
169
170impl WorktreesRegistry {
171    /// Creates the registry with the default liveness TTL. Cheap — no I/O.
172    #[must_use]
173    pub fn new() -> Self {
174        Self {
175            windows: Mutex::new(HashMap::new()),
176            ttl: DEFAULT_TTL,
177            changes: watch::channel(0).0,
178            close_pending: Mutex::new(HashSet::new()),
179            show_closed: AtomicBool::new(true),
180            polling_enabled: Mutex::new(HashMap::new()),
181            poll_ttl: DEFAULT_POLL_LEASE,
182        }
183    }
184
185    /// Creates a registry with a custom PR-poll lease duration, for tests that
186    /// exercise auto-expiry without waiting the full 15 minutes.
187    #[cfg(test)]
188    #[must_use]
189    pub fn with_poll_ttl(poll_ttl: Duration) -> Self {
190        Self {
191            poll_ttl,
192            ..Self::new()
193        }
194    }
195
196    /// A change-notification receiver for the push subscription: it observes a
197    /// new value each time the visible window set changes (see [`changes`] and
198    /// [`bump`]). Created with the current version already marked seen, so the
199    /// first [`watch::Receiver::changed`] resolves on the *next* change — the
200    /// subscriber sends its own initial snapshot up front and then waits for
201    /// deltas (#1267).
202    ///
203    /// [`changes`]: Self::changes
204    /// [`bump`]: Self::bump
205    #[must_use]
206    pub fn subscribe_changes(&self) -> watch::Receiver<u64> {
207        self.changes.subscribe()
208    }
209
210    /// Signals subscribers that the visible state changed. Non-blocking and
211    /// runtime-free; called only *after* the map guard is released so the two
212    /// locks never nest. A send never fails here (the sender is owned by the
213    /// registry, which outlives every receiver, and `send_modify` bumps even
214    /// with no receivers).
215    ///
216    /// Visible outside the registry so the daemon's PR badge poller can signal a
217    /// changed CI verdict (#1337) — the tree snapshot carries more than the window
218    /// set. Callers must bump **only on a real change**: an unconditional bump
219    /// defeats the server's snapshot diff and re-pushes to every window on every
220    /// tick.
221    pub(crate) fn bump(&self) {
222        self.changes.send_modify(|v| *v = v.wrapping_add(1));
223    }
224
225    /// The current change-notify generation — the counter [`bump`](Self::bump)
226    /// increments on every visible-set change. Read (never subscribed) so a
227    /// coalescing consumer — the service's shared tree-snapshot cache (#1303) —
228    /// can tell whether the registry has changed since it last computed and
229    /// rebuild only then. The value itself is immaterial; only whether it
230    /// differs between two reads matters, so `wrapping_add` overflow is benign.
231    #[must_use]
232    pub fn change_generation(&self) -> u64 {
233        *self.changes.borrow()
234    }
235
236    /// Locks the registry, recovering from a poisoned mutex (a panic in a prior
237    /// critical section must not wedge the whole registry).
238    fn lock(&self) -> MutexGuard<'_, HashMap<String, WindowEntry>> {
239        self.windows.lock().unwrap_or_else(PoisonError::into_inner)
240    }
241
242    /// Records (upserts) a window registration. Reaps stale entries first, then
243    /// — only when a genuinely new key would grow the map past [`MAX_WINDOWS`] —
244    /// evicts the longest-silent entry. Infallible: an upsert never evicts, and
245    /// callers validate the `key` before reaching here.
246    pub fn register(&self, req: RegisterRequest) {
247        let now = Utc::now();
248        {
249            let mut windows = self.lock();
250            reap(&mut windows, self.ttl, now);
251            // Upserts never evict; only a genuinely new key can grow the map, and
252            // never past MAX_WINDOWS.
253            if !windows.contains_key(&req.key) && windows.len() >= MAX_WINDOWS {
254                evict_oldest(&mut windows);
255            }
256            windows.insert(
257                req.key.clone(),
258                WindowEntry {
259                    key: req.key,
260                    folders: req.folders,
261                    repo: req.repo,
262                    title: req.title,
263                    pid: req.pid,
264                    last_seen: now,
265                },
266            );
267        }
268        // Always bump: a register is infrequent (once per companion `activate()`,
269        // not per heartbeat) and may add or alter a window's folders/repo. A
270        // no-op re-register with identical data is harmless — the subscriber
271        // diffs the snapshot and suppresses the duplicate frame.
272        self.bump();
273    }
274
275    /// Refreshes a window's liveness. Returns whether the key was known: a
276    /// `false` tells a window that started before the daemon — or survived a
277    /// daemon restart — to re-`register`, since the registry is in-memory and
278    /// has no record of it.
279    pub fn heartbeat(&self, key: &str) -> bool {
280        let now = Utc::now();
281        let (known, reaped) = {
282            let mut windows = self.lock();
283            let reaped = reap(&mut windows, self.ttl, now);
284            let known = match windows.get_mut(key) {
285                Some(entry) => {
286                    entry.last_seen = now;
287                    true
288                }
289                None => false,
290            };
291            (known, reaped)
292        };
293        // A heartbeat is frequent (~every 10 s per window); a pure liveness
294        // refresh does not change the visible set, so bump *only* when this
295        // heartbeat's inline reap actually aged a stale sibling out.
296        if reaped > 0 {
297            self.bump();
298        }
299        known
300    }
301
302    /// Drops a window's registration. Returns whether an entry was present.
303    pub fn unregister(&self, key: &str) -> bool {
304        let now = Utc::now();
305        let (removed, reaped) = {
306            let mut windows = self.lock();
307            let removed = windows.remove(key).is_some();
308            let reaped = reap(&mut windows, self.ttl, now);
309            (removed, reaped)
310        };
311        // The window is gone; any close directive for it is fulfilled or moot.
312        // (Keys are per-`activate()` UUIDs, never reused, so a stale directive
313        // would only ever leak a little memory — but clearing keeps it tidy.)
314        self.take_close_pending(key);
315        if removed || reaped > 0 {
316            self.bump();
317        }
318        removed
319    }
320
321    /// Records a pending "close yourself" directive for `key`, to be surfaced on
322    /// that window's next `heartbeat`. Set by the `close` op when signalling a
323    /// window it can only reply to. Idempotent; infallible.
324    pub fn mark_close_pending(&self, key: &str) {
325        self.close_pending
326            .lock()
327            .unwrap_or_else(PoisonError::into_inner)
328            .insert(key.to_string());
329    }
330
331    /// Takes (returns and clears) `key`'s pending close directive. Called on
332    /// each `heartbeat` so the directive fires exactly once; a `false` means no
333    /// close is pending.
334    pub fn take_close_pending(&self, key: &str) -> bool {
335        self.close_pending
336            .lock()
337            .unwrap_or_else(PoisonError::into_inner)
338            .remove(key)
339    }
340
341    /// The current show/hide-closed toggle: whether the tree view shows
342    /// worktrees with no open window (#1301). Read into every `tree`/`subscribe`
343    /// snapshot so every window renders the same, live-synced state.
344    #[must_use]
345    pub fn show_closed(&self) -> bool {
346        self.show_closed.load(Ordering::Relaxed)
347    }
348
349    /// Sets the show/hide-closed toggle, returning whether the value actually
350    /// changed. A real change [`bump`](Self::bump)s the change-notify so every
351    /// subscriber re-pushes a snapshot carrying the new value — the reliable
352    /// cross-window sync `context.globalState` could not provide. A no-op set
353    /// (same value) neither bumps nor wakes anyone.
354    pub fn set_show_closed(&self, show_closed: bool) -> bool {
355        let changed = self.show_closed.swap(show_closed, Ordering::Relaxed) != show_closed;
356        if changed {
357            self.bump();
358        }
359        changed
360    }
361
362    /// Locks the per-repo PR-poll lease map (`"owner/name"` → lease-expiry
363    /// instant), recovering from a poisoned mutex (a panic in a prior critical
364    /// section must not wedge polling for the whole daemon).
365    fn polling_lock(&self) -> MutexGuard<'_, HashMap<String, DateTime<Utc>>> {
366        self.polling_enabled
367            .lock()
368            .unwrap_or_else(PoisonError::into_inner)
369    }
370
371    /// Whether PR polling is currently leased for the GitHub repo `owner/name`
372    /// (#1376) — the entry exists **and** its lease has not expired. Defaults
373    /// **false** (a never-toggled repo does not poll), so only repos the user has
374    /// explicitly enabled, within the last [`poll_ttl`](Self::poll_ttl), poll.
375    #[must_use]
376    pub fn is_polling_enabled(&self, owner: &str, name: &str) -> bool {
377        let now = Utc::now();
378        self.polling_lock()
379            .get(&polling_key(owner, name))
380            .is_some_and(|expiry| *expiry > now)
381    }
382
383    /// The repos with a **live** (unexpired) lease, as a set of `"owner/name"`
384    /// keys — what stamps `polling_enabled` onto the `tree` snapshot. Reaps
385    /// expired leases first (the window-TTL reap-on-read precedent), so an idle
386    /// repo drops out on the next snapshot build without a background timer.
387    /// Cloned out so the lock is never held across the (blocking-thread) tree
388    /// build that reads it.
389    #[must_use]
390    pub fn enabled_polling_repos(&self) -> HashSet<String> {
391        let now = Utc::now();
392        let mut map = self.polling_lock();
393        map.retain(|_, expiry| *expiry > now);
394        map.keys().cloned().collect()
395    }
396
397    /// The live leases as `(repo, expiry)` pairs sorted by repo, for
398    /// deterministic persistence to the `0600` prefs file (#1376) — the expiry is
399    /// stored so a daemon restart within the lease window keeps the *remaining*
400    /// time rather than resetting the clock. Reaps expired leases first, so a
401    /// stale entry is never written back.
402    #[must_use]
403    pub fn polling_snapshot(&self) -> Vec<(String, DateTime<Utc>)> {
404        let now = Utc::now();
405        let mut map = self.polling_lock();
406        map.retain(|_, expiry| *expiry > now);
407        let mut entries: Vec<(String, DateTime<Utc>)> =
408            map.iter().map(|(k, v)| (k.clone(), *v)).collect();
409        entries.sort_by(|a, b| a.0.cmp(&b.0));
410        entries
411    }
412
413    /// Enables (leases for [`poll_ttl`](Self::poll_ttl)) or disables PR polling
414    /// for `owner/name`, returning whether the stored map changed — which the
415    /// adapter uses to decide whether to persist. Enabling an already-leased repo
416    /// **refreshes** the lease (a new expiry), which is a change worth persisting.
417    ///
418    /// [`bump`](Self::bump)s the change-notify only when the repo's **effective**
419    /// enabled state flips (off→on or on→off), so every subscribed window
420    /// re-pushes a snapshot that recolours the icon and drops/keeps badges — the
421    /// [`set_show_closed`](Self::set_show_closed) precedent. A lease *refresh*
422    /// (already on, still on) changes the expiry but not the visible state, so it
423    /// persists without waking anyone.
424    pub fn set_polling(&self, owner: &str, name: &str, enabled: bool) -> bool {
425        let key = polling_key(owner, name);
426        let now = Utc::now();
427        let (changed, flipped) = {
428            let mut map = self.polling_lock();
429            let was_enabled = map.get(&key).is_some_and(|expiry| *expiry > now);
430            if enabled {
431                let expiry = now
432                    + ChronoDuration::from_std(self.poll_ttl).unwrap_or_else(|_| {
433                        ChronoDuration::seconds(DEFAULT_POLL_LEASE.as_secs() as i64)
434                    });
435                let changed = map.insert(key, expiry) != Some(expiry);
436                (changed, !was_enabled)
437            } else {
438                let removed = map.remove(&key).is_some();
439                (removed, was_enabled)
440            }
441        };
442        if flipped {
443            self.bump();
444        }
445        changed
446    }
447
448    /// Replaces the lease map wholesale from the persisted `0600` prefs file
449    /// (#1376), dropping any lease that already expired while the daemon was down.
450    /// Does **not** [`bump`](Self::bump): it runs before any window subscribes, so
451    /// there is no one to notify, and each window's first snapshot already
452    /// reflects the seeded leases.
453    pub fn seed_polling(&self, leases: impl IntoIterator<Item = (String, DateTime<Utc>)>) {
454        let now = Utc::now();
455        *self.polling_lock() = leases
456            .into_iter()
457            .filter(|(_, expiry)| *expiry > now)
458            .collect();
459    }
460
461    /// Test-only: forces `owner/name`'s lease to `expiry`, so a test can simulate
462    /// an elapsed lease deterministically without sleeping for [`poll_ttl`].
463    #[cfg(test)]
464    pub fn set_polling_expiry(&self, owner: &str, name: &str, expiry: DateTime<Utc>) {
465        self.polling_lock().insert(polling_key(owner, name), expiry);
466    }
467
468    /// Reaps stale entries, then returns the live set sorted for deterministic
469    /// output. Holds the lock only for pure-CPU work.
470    ///
471    /// Like the other reads ([`open_folders`](Self::open_folders),
472    /// [`first_folder`](Self::first_folder)) this reaps but never
473    /// [`bump`](Self::bump)s: the only observer of a read-path reap is the push
474    /// subscription's own re-snapshot (or `status`/`menu`), and the
475    /// subscription's periodic tick already re-samples read-only staleness — so
476    /// bumping here would only make the subscription wake itself (#1267).
477    pub fn list(&self) -> Vec<WindowEntry> {
478        let now = Utc::now();
479        let mut windows = self.lock();
480        reap(&mut windows, self.ttl, now);
481        sorted_entries(&windows)
482    }
483
484    /// The first workspace folder of a still-live window, if it has one. Used by
485    /// the tray "focus" action to resolve a key to a folder to open. Does not
486    /// reap — a menu action races the reaper either way, and the caller handles
487    /// a `None` (the window may have closed).
488    pub fn first_folder(&self, key: &str) -> Option<PathBuf> {
489        let windows = self.lock();
490        windows.get(key).and_then(|e| e.folders.first().cloned())
491    }
492
493    /// Snapshots the distinct workspace folders across all live windows — the
494    /// seed set the adapter resolves to repositories (each folder → its git
495    /// common dir → repo root) to enumerate every worktree per repo (#1265).
496    ///
497    /// Reaps stale entries first, then returns the folders sorted and
498    /// deduplicated. Like [`list`](Self::list) it is pure CPU under the lock:
499    /// the git resolution the "distinct repos" derivation needs is disk I/O and
500    /// stays in the adapter, off the registry lock, honouring the
501    /// `Mutex`-never-across-`.await` invariant.
502    pub fn open_folders(&self) -> Vec<PathBuf> {
503        let now = Utc::now();
504        let mut windows = self.lock();
505        reap(&mut windows, self.ttl, now);
506        let mut folders: Vec<PathBuf> = windows
507            .values()
508            .flat_map(|e| e.folders.iter().cloned())
509            .collect();
510        folders.sort();
511        folders.dedup();
512        folders
513    }
514}
515
516impl Default for WorktreesRegistry {
517    fn default() -> Self {
518        Self::new()
519    }
520}
521
522/// Removes entries last seen longer than `ttl` ago, returning how many were
523/// dropped. Pure CPU; the caller holds the registry lock but never `.await`s
524/// while holding it. The count lets a *mutation* path
525/// ([`register`](WorktreesRegistry::register) et al.) decide whether to
526/// [`bump`](WorktreesRegistry::bump) the change-notify; read paths ignore it (see
527/// [`list`](WorktreesRegistry::list)).
528fn reap(windows: &mut HashMap<String, WindowEntry>, ttl: Duration, now: DateTime<Utc>) -> usize {
529    let max_age = ttl.as_secs() as i64;
530    let before = windows.len();
531    windows.retain(|_, e| (now - e.last_seen).num_seconds() <= max_age);
532    before - windows.len()
533}
534
535/// Removes the entry with the oldest `last_seen` (ties broken by lowest key
536/// for determinism). Called when a `register` of a new key would grow the
537/// registry past [`MAX_WINDOWS`]. Pure CPU under the registry lock, like
538/// [`reap`].
539fn evict_oldest(windows: &mut HashMap<String, WindowEntry>) {
540    let oldest = windows
541        .values()
542        .min_by(|a, b| {
543            a.last_seen
544                .cmp(&b.last_seen)
545                .then_with(|| a.key.cmp(&b.key))
546        })
547        .map(|e| e.key.clone());
548    if let Some(key) = oldest {
549        windows.remove(&key);
550    }
551}
552
553/// The canonical key for a GitHub repo in the per-repo PR-poll set: `owner/name`
554/// (#1376). One place so the registry's set, the snapshot stamp, and the poller
555/// filter all agree on the exact string.
556fn polling_key(owner: &str, name: &str) -> String {
557    format!("{owner}/{name}")
558}
559
560/// Snapshots the registry into a stably-ordered vector (by repo, then key) so
561/// `list`/`status`/`menu` output is deterministic despite `HashMap` ordering.
562fn sorted_entries(windows: &HashMap<String, WindowEntry>) -> Vec<WindowEntry> {
563    let mut entries: Vec<WindowEntry> = windows.values().cloned().collect();
564    entries.sort_by(|a, b| a.repo.cmp(&b.repo).then_with(|| a.key.cmp(&b.key)));
565    entries
566}
567
568#[cfg(test)]
569#[allow(clippy::unwrap_used, clippy::expect_used)]
570mod tests {
571    use super::*;
572
573    fn register_request(key: &str, repo: Option<&str>, folder: &str) -> RegisterRequest {
574        RegisterRequest {
575            key: key.to_string(),
576            folders: vec![PathBuf::from(folder)],
577            repo: repo.map(str::to_string),
578            title: Some(format!("{key}-title")),
579            pid: Some(1234),
580        }
581    }
582
583    #[test]
584    fn list_is_empty_initially() {
585        let reg = WorktreesRegistry::new();
586        assert!(reg.list().is_empty());
587    }
588
589    #[test]
590    fn register_then_list_round_trips() {
591        let reg = WorktreesRegistry::new();
592        reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
593        let windows = reg.list();
594        assert_eq!(windows.len(), 1);
595        assert_eq!(windows[0].key, "w1");
596        assert_eq!(windows[0].repo.as_deref(), Some("repo-a"));
597    }
598
599    #[test]
600    fn register_is_idempotent_upsert() {
601        let reg = WorktreesRegistry::new();
602        reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
603        // Re-registering the same key updates rather than duplicates.
604        reg.register(register_request("w1", Some("repo-b"), "/tmp/b"));
605        let windows = reg.list();
606        assert_eq!(windows.len(), 1);
607        assert_eq!(windows[0].repo.as_deref(), Some("repo-b"));
608    }
609
610    #[test]
611    fn heartbeat_reports_known_and_unknown() {
612        let reg = WorktreesRegistry::new();
613        // Unknown before registration: the window must re-register.
614        assert!(!reg.heartbeat("w1"));
615        reg.register(register_request("w1", None, "/tmp/a"));
616        assert!(reg.heartbeat("w1"));
617    }
618
619    #[test]
620    fn unregister_removes() {
621        let reg = WorktreesRegistry::new();
622        reg.register(register_request("w1", None, "/tmp/a"));
623        assert!(reg.unregister("w1"));
624        // Removing again is a no-op.
625        assert!(!reg.unregister("w1"));
626    }
627
628    #[test]
629    fn first_folder_returns_first_folder_or_none() {
630        let reg = WorktreesRegistry::new();
631        // No such key.
632        assert!(reg.first_folder("missing").is_none());
633        reg.register(register_request("w1", None, "/tmp/a"));
634        assert_eq!(reg.first_folder("w1"), Some(PathBuf::from("/tmp/a")));
635        // A folderless window resolves to None rather than a folder.
636        reg.register(RegisterRequest {
637            key: "w2".to_string(),
638            folders: vec![],
639            repo: None,
640            title: None,
641            pid: None,
642        });
643        assert!(reg.first_folder("w2").is_none());
644    }
645
646    #[test]
647    fn open_folders_dedups_and_sorts_across_windows() {
648        let reg = WorktreesRegistry::new();
649        assert!(reg.open_folders().is_empty());
650        // Two windows sharing a folder, plus a multi-folder window: the shared
651        // path collapses and the result is sorted.
652        reg.register(register_request("w1", Some("repo-a"), "/tmp/shared"));
653        reg.register(RegisterRequest {
654            key: "w2".to_string(),
655            folders: vec![PathBuf::from("/tmp/shared"), PathBuf::from("/tmp/b")],
656            repo: Some("repo-a".to_string()),
657            title: None,
658            pid: None,
659        });
660        reg.register(register_request("w3", Some("repo-b"), "/tmp/a"));
661        assert_eq!(
662            reg.open_folders(),
663            vec![
664                PathBuf::from("/tmp/a"),
665                PathBuf::from("/tmp/b"),
666                PathBuf::from("/tmp/shared"),
667            ]
668        );
669    }
670
671    #[test]
672    fn open_folders_reaps_stale_windows() {
673        let reg = WorktreesRegistry::new();
674        {
675            let mut windows = reg.lock();
676            windows.insert(
677                "fresh".to_string(),
678                WindowEntry {
679                    key: "fresh".to_string(),
680                    folders: vec![PathBuf::from("/tmp/fresh")],
681                    repo: None,
682                    title: None,
683                    pid: None,
684                    last_seen: Utc::now(),
685                },
686            );
687            windows.insert(
688                "stale".to_string(),
689                WindowEntry {
690                    key: "stale".to_string(),
691                    folders: vec![PathBuf::from("/tmp/stale")],
692                    repo: None,
693                    title: None,
694                    pid: None,
695                    last_seen: Utc::now() - chrono::Duration::seconds(120),
696                },
697            );
698        }
699        // The stale window's folder is reaped out of the snapshot.
700        assert_eq!(reg.open_folders(), vec![PathBuf::from("/tmp/fresh")]);
701    }
702
703    #[test]
704    fn reap_evicts_only_stale_entries() {
705        let now = Utc::now();
706        let mut windows = HashMap::new();
707        windows.insert(
708            "fresh".to_string(),
709            WindowEntry {
710                key: "fresh".to_string(),
711                folders: vec![],
712                repo: None,
713                title: None,
714                pid: None,
715                last_seen: now - chrono::Duration::seconds(5),
716            },
717        );
718        windows.insert(
719            "stale".to_string(),
720            WindowEntry {
721                key: "stale".to_string(),
722                folders: vec![],
723                repo: None,
724                title: None,
725                pid: None,
726                last_seen: now - chrono::Duration::seconds(120),
727            },
728        );
729        reap(&mut windows, DEFAULT_TTL, now);
730        assert!(windows.contains_key("fresh"));
731        assert!(!windows.contains_key("stale"));
732    }
733
734    /// A minimal entry for cap/eviction tests; only `key` and `last_seen`
735    /// participate in eviction order.
736    fn entry_at(key: &str, last_seen: DateTime<Utc>) -> WindowEntry {
737        WindowEntry {
738            key: key.to_string(),
739            folders: vec![],
740            repo: None,
741            title: None,
742            pid: None,
743            last_seen,
744        }
745    }
746
747    #[test]
748    fn evict_oldest_removes_oldest_with_key_tiebreak() {
749        let now = Utc::now();
750        let mut windows = HashMap::new();
751        windows.insert("young".to_string(), entry_at("young", now));
752        windows.insert(
753            "old-b".to_string(),
754            entry_at("old-b", now - chrono::Duration::seconds(10)),
755        );
756        windows.insert(
757            "old-a".to_string(),
758            entry_at("old-a", now - chrono::Duration::seconds(10)),
759        );
760        // Oldest `last_seen` is shared by two entries; the lowest key loses.
761        evict_oldest(&mut windows);
762        assert!(!windows.contains_key("old-a"));
763        assert!(windows.contains_key("old-b"));
764        assert!(windows.contains_key("young"));
765        // Empty map is a no-op rather than a panic.
766        let mut empty: HashMap<String, WindowEntry> = HashMap::new();
767        evict_oldest(&mut empty);
768        assert!(empty.is_empty());
769    }
770
771    #[test]
772    fn register_at_cap_evicts_only_the_oldest() {
773        let reg = WorktreesRegistry::new();
774        // Seed a full registry directly (registering 256 times would work too,
775        // but sub-second timestamps may tie; explicit timestamps make the
776        // highest-numbered key unambiguously the oldest).
777        {
778            let mut windows = reg.lock();
779            let base = Utc::now();
780            for i in 0..MAX_WINDOWS {
781                let key = format!("w{i:03}");
782                windows.insert(
783                    key.clone(),
784                    entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
785                );
786            }
787        }
788        // A new key at the cap displaces exactly the longest-silent entry.
789        reg.register(register_request("fresh", None, "/tmp/f"));
790        let windows = reg.lock();
791        assert_eq!(windows.len(), MAX_WINDOWS);
792        assert!(windows.contains_key("fresh"));
793        assert!(!windows.contains_key(&format!("w{:03}", MAX_WINDOWS - 1)));
794        assert!(windows.contains_key("w000"));
795    }
796
797    #[test]
798    fn register_upsert_at_cap_does_not_evict() {
799        let reg = WorktreesRegistry::new();
800        {
801            let mut windows = reg.lock();
802            let base = Utc::now();
803            for i in 0..MAX_WINDOWS {
804                let key = format!("w{i:03}");
805                windows.insert(
806                    key.clone(),
807                    entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
808                );
809            }
810        }
811        // Re-registering an existing key is an upsert: nothing is displaced,
812        // not even the oldest entry.
813        let oldest = format!("w{:03}", MAX_WINDOWS - 1);
814        reg.register(register_request(&oldest, Some("r"), "/tmp/a"));
815        let windows = reg.lock();
816        assert_eq!(windows.len(), MAX_WINDOWS);
817        assert!(windows.contains_key(&oldest));
818        assert!(windows.contains_key("w000"));
819    }
820
821    #[test]
822    fn sorted_entries_orders_by_repo_then_key() {
823        let now = Utc::now();
824        let mut windows = HashMap::new();
825        for (key, repo) in [("z", "repo-a"), ("a", "repo-b"), ("m", "repo-a")] {
826            windows.insert(
827                key.to_string(),
828                WindowEntry {
829                    key: key.to_string(),
830                    folders: vec![],
831                    repo: Some(repo.to_string()),
832                    title: None,
833                    pid: None,
834                    last_seen: now,
835                },
836            );
837        }
838        let entries = sorted_entries(&windows);
839        let ordered: Vec<(&str, &str)> = entries
840            .iter()
841            .map(|e| (e.key.as_str(), e.repo.as_deref().unwrap()))
842            .collect();
843        assert_eq!(
844            ordered,
845            vec![("m", "repo-a"), ("z", "repo-a"), ("a", "repo-b")]
846        );
847    }
848
849    #[test]
850    fn default_constructs_an_empty_registry() {
851        let reg = WorktreesRegistry::default();
852        assert!(reg.lock().is_empty());
853    }
854
855    // --- Change-notify for the push subscription (#1267) --------------------
856
857    #[test]
858    fn subscribe_changes_starts_seen_and_register_bumps() {
859        let reg = WorktreesRegistry::new();
860        let mut rx = reg.subscribe_changes();
861        // A fresh receiver has the current version already marked seen.
862        assert!(!rx.has_changed().unwrap());
863        // A register changes the visible set → the receiver observes a new value.
864        reg.register(register_request("w1", None, "/tmp/a"));
865        assert!(rx.has_changed().unwrap(), "register should bump");
866        // Marking it seen clears the pending change.
867        rx.borrow_and_update();
868        assert!(!rx.has_changed().unwrap());
869    }
870
871    #[test]
872    fn unregister_bumps_only_when_it_removes() {
873        let reg = WorktreesRegistry::new();
874        reg.register(register_request("w1", None, "/tmp/a"));
875        // Subscribe *after* the register so its bump is already seen.
876        let rx = reg.subscribe_changes();
877        // Removing a missing key changes nothing (and reaps nothing) → no bump.
878        assert!(!reg.unregister("ghost"));
879        assert!(
880            !rx.has_changed().unwrap(),
881            "a no-op unregister must not bump"
882        );
883        // Removing a present key bumps.
884        assert!(reg.unregister("w1"));
885        assert!(
886            rx.has_changed().unwrap(),
887            "a removing unregister should bump"
888        );
889    }
890
891    #[test]
892    fn change_generation_advances_only_on_a_visible_change() {
893        let reg = WorktreesRegistry::new();
894        let g0 = reg.change_generation();
895        // A no-op (heartbeat of an unknown key) leaves the generation untouched.
896        assert!(!reg.heartbeat("ghost"));
897        assert_eq!(
898            reg.change_generation(),
899            g0,
900            "a no-op must not advance the generation"
901        );
902        // A register changes the visible set → the generation advances, so a
903        // cache keyed on it rebuilds.
904        reg.register(register_request("w1", None, "/tmp/a"));
905        assert_ne!(
906            reg.change_generation(),
907            g0,
908            "a register should advance the generation"
909        );
910    }
911
912    #[test]
913    fn heartbeat_bumps_only_when_it_reaps() {
914        let reg = WorktreesRegistry::new();
915        reg.register(register_request("w1", None, "/tmp/a"));
916        let rx = reg.subscribe_changes();
917        // A plain heartbeat refreshes liveness but changes no visible state.
918        assert!(reg.heartbeat("w1"));
919        assert!(!rx.has_changed().unwrap(), "a pure heartbeat must not bump");
920        // Seed a stale sibling directly; a heartbeat that reaps it *does* bump.
921        {
922            let mut windows = reg.lock();
923            windows.insert(
924                "stale".to_string(),
925                entry_at("stale", Utc::now() - chrono::Duration::seconds(120)),
926            );
927        }
928        assert!(reg.heartbeat("w1"));
929        assert!(
930            rx.has_changed().unwrap(),
931            "a heartbeat that reaps a stale sibling should bump"
932        );
933    }
934
935    // --- Close-pending directive (#1277) -----------------------------------
936
937    #[test]
938    fn close_pending_is_taken_once_then_cleared() {
939        let reg = WorktreesRegistry::new();
940        // No directive by default.
941        assert!(!reg.take_close_pending("w1"));
942        // Marked → the first take observes it, the next does not (fires once).
943        reg.mark_close_pending("w1");
944        assert!(reg.take_close_pending("w1"));
945        assert!(!reg.take_close_pending("w1"));
946    }
947
948    #[test]
949    fn unregister_clears_a_pending_close_directive() {
950        let reg = WorktreesRegistry::new();
951        reg.register(register_request("w1", None, "/tmp/a"));
952        reg.mark_close_pending("w1");
953        // Unregistering the window drops any pending directive with it.
954        assert!(reg.unregister("w1"));
955        assert!(!reg.take_close_pending("w1"));
956    }
957
958    // --- Show/hide-closed toggle (#1301) -----------------------------------
959
960    #[test]
961    fn show_closed_defaults_to_true() {
962        let reg = WorktreesRegistry::new();
963        assert!(reg.show_closed(), "default is show all");
964    }
965
966    #[test]
967    fn set_show_closed_reports_change_and_is_idempotent() {
968        let reg = WorktreesRegistry::new();
969        // Flipping to a new value reports a change and is observable.
970        assert!(reg.set_show_closed(false));
971        assert!(!reg.show_closed());
972        // Setting the same value again is a no-op (no change reported).
973        assert!(!reg.set_show_closed(false));
974        // Flipping back reports a change again.
975        assert!(reg.set_show_closed(true));
976        assert!(reg.show_closed());
977    }
978
979    #[test]
980    fn set_show_closed_bumps_only_on_change() {
981        let reg = WorktreesRegistry::new();
982        let rx = reg.subscribe_changes();
983        // A no-op set (already the default) does not wake subscribers.
984        assert!(!reg.set_show_closed(true));
985        assert!(
986            !rx.has_changed().unwrap(),
987            "a no-op toggle must not bump the change-notify"
988        );
989        // A real flip bumps so subscribers re-push a snapshot with the new value.
990        assert!(reg.set_show_closed(false));
991        assert!(
992            rx.has_changed().unwrap(),
993            "flipping the toggle should bump the change-notify"
994        );
995    }
996
997    #[test]
998    fn polling_defaults_off_for_an_untoggled_repo() {
999        let reg = WorktreesRegistry::new();
1000        // #1376: the whole point — a repo the user has never enabled is not polled.
1001        assert!(!reg.is_polling_enabled("rust-works", "omni-dev"));
1002        assert!(reg.enabled_polling_repos().is_empty());
1003        assert!(reg.polling_snapshot().is_empty());
1004    }
1005
1006    #[test]
1007    fn set_polling_leases_and_disables() {
1008        let reg = WorktreesRegistry::new();
1009        // Enabling a fresh repo reports a change and leases it live.
1010        assert!(reg.set_polling("rust-works", "omni-dev", true));
1011        assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
1012        // Re-enabling refreshes the lease — the map changes (new expiry).
1013        assert!(reg.set_polling("rust-works", "omni-dev", true));
1014        assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
1015        // Disabling reports a change and clears it.
1016        assert!(reg.set_polling("rust-works", "omni-dev", false));
1017        assert!(!reg.is_polling_enabled("rust-works", "omni-dev"));
1018        // Disabling an already-disabled repo is a no-op.
1019        assert!(!reg.set_polling("rust-works", "omni-dev", false));
1020    }
1021
1022    #[test]
1023    fn polling_lease_auto_expires() {
1024        // The 15-minute lease (#1376): an enabled repo drops out once its lease
1025        // elapses, reaped on the next read — no background timer.
1026        let reg = WorktreesRegistry::new();
1027        reg.set_polling("rust-works", "omni-dev", true);
1028        assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
1029        // Force the lease into the past (as if 15 min elapsed).
1030        reg.set_polling_expiry(
1031            "rust-works",
1032            "omni-dev",
1033            Utc::now() - ChronoDuration::seconds(1),
1034        );
1035        assert!(
1036            !reg.is_polling_enabled("rust-works", "omni-dev"),
1037            "an expired lease reads as disabled"
1038        );
1039        // The read paths reap it, so it is gone from the snapshot entirely.
1040        assert!(reg.enabled_polling_repos().is_empty());
1041        assert!(reg.polling_snapshot().is_empty());
1042        // A tiny real TTL expires on its own after a short sleep.
1043        let short = WorktreesRegistry::with_poll_ttl(Duration::from_millis(30));
1044        short.set_polling("o", "n", true);
1045        assert!(short.is_polling_enabled("o", "n"));
1046        std::thread::sleep(Duration::from_millis(60));
1047        assert!(!short.is_polling_enabled("o", "n"));
1048    }
1049
1050    #[test]
1051    fn set_polling_bumps_only_on_an_effective_flip() {
1052        let reg = WorktreesRegistry::new();
1053        let rx = reg.subscribe_changes();
1054        // A no-op (disabling an already-off repo) does not wake subscribers.
1055        assert!(!reg.set_polling("o", "n", false));
1056        assert!(
1057            !rx.has_changed().unwrap(),
1058            "a no-op poll toggle must not bump the change-notify"
1059        );
1060        // A real enable flips off→on and bumps so every window recolors.
1061        assert!(reg.set_polling("o", "n", true));
1062        assert!(
1063            rx.has_changed().unwrap(),
1064            "enabling a repo should bump the change-notify"
1065        );
1066    }
1067
1068    #[test]
1069    fn refreshing_a_live_lease_does_not_bump() {
1070        // A lease refresh (already on, still on) persists a new expiry but does
1071        // not change the visible state, so it must not wake every window.
1072        let reg = WorktreesRegistry::new();
1073        reg.set_polling("o", "n", true);
1074        let rx = reg.subscribe_changes();
1075        assert!(
1076            reg.set_polling("o", "n", true),
1077            "re-enabling refreshes the lease (map changed → persist)"
1078        );
1079        assert!(
1080            !rx.has_changed().unwrap(),
1081            "refreshing a live lease must not bump — the visible state is unchanged"
1082        );
1083    }
1084
1085    #[test]
1086    fn seed_polling_loads_leases_and_drops_expired() {
1087        let reg = WorktreesRegistry::new();
1088        reg.set_polling("a", "z", true);
1089        let future = Utc::now() + ChronoDuration::minutes(10);
1090        let past = Utc::now() - ChronoDuration::minutes(1);
1091        // Seeding (the startup load) replaces wholesale, dropping the prior entry
1092        // and any already-expired lease from the file.
1093        reg.seed_polling([
1094            ("rust-works/omni-dev".to_string(), future),
1095            ("acme/widgets".to_string(), future),
1096            ("stale/repo".to_string(), past),
1097        ]);
1098        assert!(!reg.is_polling_enabled("a", "z"));
1099        assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
1100        assert!(
1101            !reg.is_polling_enabled("stale", "repo"),
1102            "expired lease dropped"
1103        );
1104        // The persisted form is deterministic (sorted) and carries the expiries.
1105        let snap = reg.polling_snapshot();
1106        assert_eq!(
1107            snap.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>(),
1108            vec![
1109                "acme/widgets".to_string(),
1110                "rust-works/omni-dev".to_string()
1111            ]
1112        );
1113        assert!(snap.iter().all(|(_, expiry)| *expiry == future));
1114    }
1115}