use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Mutex, MutexGuard, PoisonError};
use std::time::Duration;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
const DEFAULT_TTL: Duration = Duration::from_secs(30);
const DEFAULT_POLL_LEASE: Duration = Duration::from_secs(15 * 60);
const MAX_WINDOWS: usize = 256;
#[derive(Debug, Clone, Deserialize)]
pub struct RegisterRequest {
pub key: String,
#[serde(default)]
pub folders: Vec<PathBuf>,
#[serde(default)]
pub repo: Option<String>,
#[serde(default)]
pub title: Option<String>,
#[serde(default)]
pub pid: Option<u32>,
}
#[derive(Debug, Clone, Serialize)]
pub struct WindowEntry {
pub key: String,
pub folders: Vec<PathBuf>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pid: Option<u32>,
pub last_seen: DateTime<Utc>,
}
pub struct WorktreesRegistry {
windows: Mutex<HashMap<String, WindowEntry>>,
ttl: Duration,
changes: watch::Sender<u64>,
close_pending: Mutex<HashSet<String>>,
show_closed: AtomicBool,
polling_enabled: Mutex<HashMap<String, DateTime<Utc>>>,
poll_ttl: Duration,
}
impl WorktreesRegistry {
#[must_use]
pub fn new() -> Self {
Self {
windows: Mutex::new(HashMap::new()),
ttl: DEFAULT_TTL,
changes: watch::channel(0).0,
close_pending: Mutex::new(HashSet::new()),
show_closed: AtomicBool::new(true),
polling_enabled: Mutex::new(HashMap::new()),
poll_ttl: DEFAULT_POLL_LEASE,
}
}
#[cfg(test)]
#[must_use]
pub fn with_poll_ttl(poll_ttl: Duration) -> Self {
Self {
poll_ttl,
..Self::new()
}
}
#[must_use]
pub fn subscribe_changes(&self) -> watch::Receiver<u64> {
self.changes.subscribe()
}
pub(crate) fn bump(&self) {
self.changes.send_modify(|v| *v = v.wrapping_add(1));
}
#[must_use]
pub fn change_generation(&self) -> u64 {
*self.changes.borrow()
}
fn lock(&self) -> MutexGuard<'_, HashMap<String, WindowEntry>> {
self.windows.lock().unwrap_or_else(PoisonError::into_inner)
}
pub fn register(&self, req: RegisterRequest) {
let now = Utc::now();
{
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
if !windows.contains_key(&req.key) && windows.len() >= MAX_WINDOWS {
evict_oldest(&mut windows);
}
windows.insert(
req.key.clone(),
WindowEntry {
key: req.key,
folders: req.folders,
repo: req.repo,
title: req.title,
pid: req.pid,
last_seen: now,
},
);
}
self.bump();
}
pub fn heartbeat(&self, key: &str) -> bool {
let now = Utc::now();
let (known, reaped) = {
let mut windows = self.lock();
let reaped = reap(&mut windows, self.ttl, now);
let known = match windows.get_mut(key) {
Some(entry) => {
entry.last_seen = now;
true
}
None => false,
};
(known, reaped)
};
if reaped > 0 {
self.bump();
}
known
}
pub fn unregister(&self, key: &str) -> bool {
let now = Utc::now();
let (removed, reaped) = {
let mut windows = self.lock();
let removed = windows.remove(key).is_some();
let reaped = reap(&mut windows, self.ttl, now);
(removed, reaped)
};
self.take_close_pending(key);
if removed || reaped > 0 {
self.bump();
}
removed
}
pub fn mark_close_pending(&self, key: &str) {
self.close_pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(key.to_string());
}
pub fn take_close_pending(&self, key: &str) -> bool {
self.close_pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(key)
}
#[must_use]
pub fn show_closed(&self) -> bool {
self.show_closed.load(Ordering::Relaxed)
}
pub fn set_show_closed(&self, show_closed: bool) -> bool {
let changed = self.show_closed.swap(show_closed, Ordering::Relaxed) != show_closed;
if changed {
self.bump();
}
changed
}
fn polling_lock(&self) -> MutexGuard<'_, HashMap<String, DateTime<Utc>>> {
self.polling_enabled
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
#[must_use]
pub fn is_polling_enabled(&self, owner: &str, name: &str) -> bool {
let now = Utc::now();
self.polling_lock()
.get(&polling_key(owner, name))
.is_some_and(|expiry| *expiry > now)
}
#[must_use]
pub fn enabled_polling_repos(&self) -> HashSet<String> {
let now = Utc::now();
let mut map = self.polling_lock();
map.retain(|_, expiry| *expiry > now);
map.keys().cloned().collect()
}
#[must_use]
pub fn polling_snapshot(&self) -> Vec<(String, DateTime<Utc>)> {
let now = Utc::now();
let mut map = self.polling_lock();
map.retain(|_, expiry| *expiry > now);
let mut entries: Vec<(String, DateTime<Utc>)> =
map.iter().map(|(k, v)| (k.clone(), *v)).collect();
entries.sort_by(|a, b| a.0.cmp(&b.0));
entries
}
pub fn set_polling(&self, owner: &str, name: &str, enabled: bool) -> bool {
let key = polling_key(owner, name);
let now = Utc::now();
let (changed, flipped) = {
let mut map = self.polling_lock();
let was_enabled = map.get(&key).is_some_and(|expiry| *expiry > now);
if enabled {
let expiry = now
+ ChronoDuration::from_std(self.poll_ttl).unwrap_or_else(|_| {
ChronoDuration::seconds(DEFAULT_POLL_LEASE.as_secs() as i64)
});
let changed = map.insert(key, expiry) != Some(expiry);
(changed, !was_enabled)
} else {
let removed = map.remove(&key).is_some();
(removed, was_enabled)
}
};
if flipped {
self.bump();
}
changed
}
pub fn seed_polling(&self, leases: impl IntoIterator<Item = (String, DateTime<Utc>)>) {
let now = Utc::now();
*self.polling_lock() = leases
.into_iter()
.filter(|(_, expiry)| *expiry > now)
.collect();
}
#[cfg(test)]
pub fn set_polling_expiry(&self, owner: &str, name: &str, expiry: DateTime<Utc>) {
self.polling_lock().insert(polling_key(owner, name), expiry);
}
pub fn list(&self) -> Vec<WindowEntry> {
let now = Utc::now();
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
sorted_entries(&windows)
}
pub fn first_folder(&self, key: &str) -> Option<PathBuf> {
let windows = self.lock();
windows.get(key).and_then(|e| e.folders.first().cloned())
}
pub fn open_folders(&self) -> Vec<PathBuf> {
let now = Utc::now();
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
let mut folders: Vec<PathBuf> = windows
.values()
.flat_map(|e| e.folders.iter().cloned())
.collect();
folders.sort();
folders.dedup();
folders
}
}
impl Default for WorktreesRegistry {
fn default() -> Self {
Self::new()
}
}
fn reap(windows: &mut HashMap<String, WindowEntry>, ttl: Duration, now: DateTime<Utc>) -> usize {
let max_age = ttl.as_secs() as i64;
let before = windows.len();
windows.retain(|_, e| (now - e.last_seen).num_seconds() <= max_age);
before - windows.len()
}
fn evict_oldest(windows: &mut HashMap<String, WindowEntry>) {
let oldest = windows
.values()
.min_by(|a, b| {
a.last_seen
.cmp(&b.last_seen)
.then_with(|| a.key.cmp(&b.key))
})
.map(|e| e.key.clone());
if let Some(key) = oldest {
windows.remove(&key);
}
}
fn polling_key(owner: &str, name: &str) -> String {
format!("{owner}/{name}")
}
fn sorted_entries(windows: &HashMap<String, WindowEntry>) -> Vec<WindowEntry> {
let mut entries: Vec<WindowEntry> = windows.values().cloned().collect();
entries.sort_by(|a, b| a.repo.cmp(&b.repo).then_with(|| a.key.cmp(&b.key)));
entries
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
fn register_request(key: &str, repo: Option<&str>, folder: &str) -> RegisterRequest {
RegisterRequest {
key: key.to_string(),
folders: vec![PathBuf::from(folder)],
repo: repo.map(str::to_string),
title: Some(format!("{key}-title")),
pid: Some(1234),
}
}
#[test]
fn list_is_empty_initially() {
let reg = WorktreesRegistry::new();
assert!(reg.list().is_empty());
}
#[test]
fn register_then_list_round_trips() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
let windows = reg.list();
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].key, "w1");
assert_eq!(windows[0].repo.as_deref(), Some("repo-a"));
}
#[test]
fn register_is_idempotent_upsert() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
reg.register(register_request("w1", Some("repo-b"), "/tmp/b"));
let windows = reg.list();
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].repo.as_deref(), Some("repo-b"));
}
#[test]
fn heartbeat_reports_known_and_unknown() {
let reg = WorktreesRegistry::new();
assert!(!reg.heartbeat("w1"));
reg.register(register_request("w1", None, "/tmp/a"));
assert!(reg.heartbeat("w1"));
}
#[test]
fn unregister_removes() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
assert!(reg.unregister("w1"));
assert!(!reg.unregister("w1"));
}
#[test]
fn first_folder_returns_first_folder_or_none() {
let reg = WorktreesRegistry::new();
assert!(reg.first_folder("missing").is_none());
reg.register(register_request("w1", None, "/tmp/a"));
assert_eq!(reg.first_folder("w1"), Some(PathBuf::from("/tmp/a")));
reg.register(RegisterRequest {
key: "w2".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
});
assert!(reg.first_folder("w2").is_none());
}
#[test]
fn open_folders_dedups_and_sorts_across_windows() {
let reg = WorktreesRegistry::new();
assert!(reg.open_folders().is_empty());
reg.register(register_request("w1", Some("repo-a"), "/tmp/shared"));
reg.register(RegisterRequest {
key: "w2".to_string(),
folders: vec![PathBuf::from("/tmp/shared"), PathBuf::from("/tmp/b")],
repo: Some("repo-a".to_string()),
title: None,
pid: None,
});
reg.register(register_request("w3", Some("repo-b"), "/tmp/a"));
assert_eq!(
reg.open_folders(),
vec![
PathBuf::from("/tmp/a"),
PathBuf::from("/tmp/b"),
PathBuf::from("/tmp/shared"),
]
);
}
#[test]
fn open_folders_reaps_stale_windows() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
windows.insert(
"fresh".to_string(),
WindowEntry {
key: "fresh".to_string(),
folders: vec![PathBuf::from("/tmp/fresh")],
repo: None,
title: None,
pid: None,
last_seen: Utc::now(),
},
);
windows.insert(
"stale".to_string(),
WindowEntry {
key: "stale".to_string(),
folders: vec![PathBuf::from("/tmp/stale")],
repo: None,
title: None,
pid: None,
last_seen: Utc::now() - chrono::Duration::seconds(120),
},
);
}
assert_eq!(reg.open_folders(), vec![PathBuf::from("/tmp/fresh")]);
}
#[test]
fn reap_evicts_only_stale_entries() {
let now = Utc::now();
let mut windows = HashMap::new();
windows.insert(
"fresh".to_string(),
WindowEntry {
key: "fresh".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(5),
},
);
windows.insert(
"stale".to_string(),
WindowEntry {
key: "stale".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(120),
},
);
reap(&mut windows, DEFAULT_TTL, now);
assert!(windows.contains_key("fresh"));
assert!(!windows.contains_key("stale"));
}
fn entry_at(key: &str, last_seen: DateTime<Utc>) -> WindowEntry {
WindowEntry {
key: key.to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen,
}
}
#[test]
fn evict_oldest_removes_oldest_with_key_tiebreak() {
let now = Utc::now();
let mut windows = HashMap::new();
windows.insert("young".to_string(), entry_at("young", now));
windows.insert(
"old-b".to_string(),
entry_at("old-b", now - chrono::Duration::seconds(10)),
);
windows.insert(
"old-a".to_string(),
entry_at("old-a", now - chrono::Duration::seconds(10)),
);
evict_oldest(&mut windows);
assert!(!windows.contains_key("old-a"));
assert!(windows.contains_key("old-b"));
assert!(windows.contains_key("young"));
let mut empty: HashMap<String, WindowEntry> = HashMap::new();
evict_oldest(&mut empty);
assert!(empty.is_empty());
}
#[test]
fn register_at_cap_evicts_only_the_oldest() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
let base = Utc::now();
for i in 0..MAX_WINDOWS {
let key = format!("w{i:03}");
windows.insert(
key.clone(),
entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
);
}
}
reg.register(register_request("fresh", None, "/tmp/f"));
let windows = reg.lock();
assert_eq!(windows.len(), MAX_WINDOWS);
assert!(windows.contains_key("fresh"));
assert!(!windows.contains_key(&format!("w{:03}", MAX_WINDOWS - 1)));
assert!(windows.contains_key("w000"));
}
#[test]
fn register_upsert_at_cap_does_not_evict() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
let base = Utc::now();
for i in 0..MAX_WINDOWS {
let key = format!("w{i:03}");
windows.insert(
key.clone(),
entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
);
}
}
let oldest = format!("w{:03}", MAX_WINDOWS - 1);
reg.register(register_request(&oldest, Some("r"), "/tmp/a"));
let windows = reg.lock();
assert_eq!(windows.len(), MAX_WINDOWS);
assert!(windows.contains_key(&oldest));
assert!(windows.contains_key("w000"));
}
#[test]
fn sorted_entries_orders_by_repo_then_key() {
let now = Utc::now();
let mut windows = HashMap::new();
for (key, repo) in [("z", "repo-a"), ("a", "repo-b"), ("m", "repo-a")] {
windows.insert(
key.to_string(),
WindowEntry {
key: key.to_string(),
folders: vec![],
repo: Some(repo.to_string()),
title: None,
pid: None,
last_seen: now,
},
);
}
let entries = sorted_entries(&windows);
let ordered: Vec<(&str, &str)> = entries
.iter()
.map(|e| (e.key.as_str(), e.repo.as_deref().unwrap()))
.collect();
assert_eq!(
ordered,
vec![("m", "repo-a"), ("z", "repo-a"), ("a", "repo-b")]
);
}
#[test]
fn default_constructs_an_empty_registry() {
let reg = WorktreesRegistry::default();
assert!(reg.lock().is_empty());
}
#[test]
fn subscribe_changes_starts_seen_and_register_bumps() {
let reg = WorktreesRegistry::new();
let mut rx = reg.subscribe_changes();
assert!(!rx.has_changed().unwrap());
reg.register(register_request("w1", None, "/tmp/a"));
assert!(rx.has_changed().unwrap(), "register should bump");
rx.borrow_and_update();
assert!(!rx.has_changed().unwrap());
}
#[test]
fn unregister_bumps_only_when_it_removes() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
let rx = reg.subscribe_changes();
assert!(!reg.unregister("ghost"));
assert!(
!rx.has_changed().unwrap(),
"a no-op unregister must not bump"
);
assert!(reg.unregister("w1"));
assert!(
rx.has_changed().unwrap(),
"a removing unregister should bump"
);
}
#[test]
fn change_generation_advances_only_on_a_visible_change() {
let reg = WorktreesRegistry::new();
let g0 = reg.change_generation();
assert!(!reg.heartbeat("ghost"));
assert_eq!(
reg.change_generation(),
g0,
"a no-op must not advance the generation"
);
reg.register(register_request("w1", None, "/tmp/a"));
assert_ne!(
reg.change_generation(),
g0,
"a register should advance the generation"
);
}
#[test]
fn heartbeat_bumps_only_when_it_reaps() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
let rx = reg.subscribe_changes();
assert!(reg.heartbeat("w1"));
assert!(!rx.has_changed().unwrap(), "a pure heartbeat must not bump");
{
let mut windows = reg.lock();
windows.insert(
"stale".to_string(),
entry_at("stale", Utc::now() - chrono::Duration::seconds(120)),
);
}
assert!(reg.heartbeat("w1"));
assert!(
rx.has_changed().unwrap(),
"a heartbeat that reaps a stale sibling should bump"
);
}
#[test]
fn close_pending_is_taken_once_then_cleared() {
let reg = WorktreesRegistry::new();
assert!(!reg.take_close_pending("w1"));
reg.mark_close_pending("w1");
assert!(reg.take_close_pending("w1"));
assert!(!reg.take_close_pending("w1"));
}
#[test]
fn unregister_clears_a_pending_close_directive() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
reg.mark_close_pending("w1");
assert!(reg.unregister("w1"));
assert!(!reg.take_close_pending("w1"));
}
#[test]
fn show_closed_defaults_to_true() {
let reg = WorktreesRegistry::new();
assert!(reg.show_closed(), "default is show all");
}
#[test]
fn set_show_closed_reports_change_and_is_idempotent() {
let reg = WorktreesRegistry::new();
assert!(reg.set_show_closed(false));
assert!(!reg.show_closed());
assert!(!reg.set_show_closed(false));
assert!(reg.set_show_closed(true));
assert!(reg.show_closed());
}
#[test]
fn set_show_closed_bumps_only_on_change() {
let reg = WorktreesRegistry::new();
let rx = reg.subscribe_changes();
assert!(!reg.set_show_closed(true));
assert!(
!rx.has_changed().unwrap(),
"a no-op toggle must not bump the change-notify"
);
assert!(reg.set_show_closed(false));
assert!(
rx.has_changed().unwrap(),
"flipping the toggle should bump the change-notify"
);
}
#[test]
fn polling_defaults_off_for_an_untoggled_repo() {
let reg = WorktreesRegistry::new();
assert!(!reg.is_polling_enabled("rust-works", "omni-dev"));
assert!(reg.enabled_polling_repos().is_empty());
assert!(reg.polling_snapshot().is_empty());
}
#[test]
fn set_polling_leases_and_disables() {
let reg = WorktreesRegistry::new();
assert!(reg.set_polling("rust-works", "omni-dev", true));
assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
assert!(reg.set_polling("rust-works", "omni-dev", true));
assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
assert!(reg.set_polling("rust-works", "omni-dev", false));
assert!(!reg.is_polling_enabled("rust-works", "omni-dev"));
assert!(!reg.set_polling("rust-works", "omni-dev", false));
}
#[test]
fn polling_lease_auto_expires() {
let reg = WorktreesRegistry::new();
reg.set_polling("rust-works", "omni-dev", true);
assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
reg.set_polling_expiry(
"rust-works",
"omni-dev",
Utc::now() - ChronoDuration::seconds(1),
);
assert!(
!reg.is_polling_enabled("rust-works", "omni-dev"),
"an expired lease reads as disabled"
);
assert!(reg.enabled_polling_repos().is_empty());
assert!(reg.polling_snapshot().is_empty());
let short = WorktreesRegistry::with_poll_ttl(Duration::from_millis(30));
short.set_polling("o", "n", true);
assert!(short.is_polling_enabled("o", "n"));
std::thread::sleep(Duration::from_millis(60));
assert!(!short.is_polling_enabled("o", "n"));
}
#[test]
fn set_polling_bumps_only_on_an_effective_flip() {
let reg = WorktreesRegistry::new();
let rx = reg.subscribe_changes();
assert!(!reg.set_polling("o", "n", false));
assert!(
!rx.has_changed().unwrap(),
"a no-op poll toggle must not bump the change-notify"
);
assert!(reg.set_polling("o", "n", true));
assert!(
rx.has_changed().unwrap(),
"enabling a repo should bump the change-notify"
);
}
#[test]
fn refreshing_a_live_lease_does_not_bump() {
let reg = WorktreesRegistry::new();
reg.set_polling("o", "n", true);
let rx = reg.subscribe_changes();
assert!(
reg.set_polling("o", "n", true),
"re-enabling refreshes the lease (map changed → persist)"
);
assert!(
!rx.has_changed().unwrap(),
"refreshing a live lease must not bump — the visible state is unchanged"
);
}
#[test]
fn seed_polling_loads_leases_and_drops_expired() {
let reg = WorktreesRegistry::new();
reg.set_polling("a", "z", true);
let future = Utc::now() + ChronoDuration::minutes(10);
let past = Utc::now() - ChronoDuration::minutes(1);
reg.seed_polling([
("rust-works/omni-dev".to_string(), future),
("acme/widgets".to_string(), future),
("stale/repo".to_string(), past),
]);
assert!(!reg.is_polling_enabled("a", "z"));
assert!(reg.is_polling_enabled("rust-works", "omni-dev"));
assert!(
!reg.is_polling_enabled("stale", "repo"),
"expired lease dropped"
);
let snap = reg.polling_snapshot();
assert_eq!(
snap.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>(),
vec![
"acme/widgets".to_string(),
"rust-works/omni-dev".to_string()
]
);
assert!(snap.iter().all(|(_, expiry)| *expiry == future));
}
}