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}