use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, Instant};
use anyhow::{anyhow, bail, Context, Result};
use chrono::{DateTime, Utc};
use crate::github_rate_limit::{
resolve_rate_limit_with, RateLimitCache, RateLimitResource, RateLimitSnapshot,
};
use crate::pr_status::{PrBadge, PrCheckState, PrResolution, PrStatusCache, PrTarget};
use async_trait::async_trait;
use git2::{Repository, RepositoryState, Status, StatusOptions, WorktreeLockStatus};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::watch;
use tokio::sync::Mutex as AsyncMutex;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use crate::daemon::service::{
DaemonService, MenuAction, MenuItem, MenuSnapshot, ServiceStatus, ServiceStream,
};
use crate::worktrees::{RegisterRequest, WindowEntry, WorktreesRegistry};
pub const SERVICE_NAME: &str = "worktrees";
const SUBMENU_TITLE: &str = "Worktrees";
const VSCODE_BIN_ENV: &str = "OMNI_DEV_VSCODE_BIN";
const ENV_MENU_REFRESH_INTERVAL: &str = "OMNI_DEV_DAEMON_MENU_REFRESH";
const DEFAULT_MENU_REFRESH_INTERVAL: Duration = Duration::from_secs(10);
fn menu_refresh_interval() -> Duration {
crate::daemon::server::duration_secs_from_env(
ENV_MENU_REFRESH_INTERVAL,
DEFAULT_MENU_REFRESH_INTERVAL,
)
}
const ENV_PR_POLL_INTERVAL: &str = "OMNI_DEV_DAEMON_PR_POLL";
const DEFAULT_PR_POLL_INTERVAL: Duration = Duration::from_secs(10);
const MAX_PR_POLL_INTERVAL: Duration = Duration::from_secs(30 * 60);
const PENDING_FAST_WINDOW: Duration = Duration::from_secs(2 * 60);
const PENDING_MAX_INTERVAL: Duration = Duration::from_secs(60);
const BUDGET_THROTTLE_INTERVAL: Duration = Duration::from_secs(5 * 60);
const ENV_PR_DEBOUNCE: &str = "OMNI_DEV_DAEMON_PR_DEBOUNCE";
const DEFAULT_PR_DEBOUNCE: Duration = Duration::from_secs(2);
fn pr_poll_interval() -> Duration {
crate::daemon::server::duration_secs_from_env(ENV_PR_POLL_INTERVAL, DEFAULT_PR_POLL_INTERVAL)
}
fn pr_debounce_interval() -> Duration {
crate::daemon::server::duration_secs_from_env(ENV_PR_DEBOUNCE, DEFAULT_PR_DEBOUNCE)
}
const ENV_OPEN_PR_TTL: &str = "OMNI_DEV_DAEMON_OPEN_PR_TTL";
const DEFAULT_OPEN_PR_TTL: Duration = Duration::from_secs(60);
const OPEN_PR_JSON_FIELDS: &str = "number,title,url,headRefName,baseRefName,isDraft,state,author";
const OPEN_PR_LIST_LIMIT: &str = "100";
fn open_pr_ttl() -> Duration {
crate::daemon::server::duration_secs_from_env(ENV_OPEN_PR_TTL, DEFAULT_OPEN_PR_TTL)
}
const ENV_RATE_LIMIT_POLL_INTERVAL: &str = "OMNI_DEV_DAEMON_RATE_LIMIT_POLL";
const DEFAULT_RATE_LIMIT_POLL_INTERVAL: Duration = Duration::from_secs(60);
fn rate_limit_poll_interval() -> Duration {
crate::daemon::server::duration_secs_from_env(
ENV_RATE_LIMIT_POLL_INTERVAL,
DEFAULT_RATE_LIMIT_POLL_INTERVAL,
)
}
struct RefreshTask {
token: CancellationToken,
handle: JoinHandle<()>,
}
fn pr_should_fetch(grew: bool, since_last_fetch: Option<Duration>, backoff: Duration) -> bool {
grew || since_last_fetch.map_or(true, |elapsed| elapsed >= backoff)
}
fn pr_watch_grew(prev: &[PrWatch], next: &[PrWatch]) -> bool {
next.iter().any(|w| !prev.contains(w))
}
fn next_pr_poll_delay(
current: Duration,
base: Duration,
pending: bool,
since_moved: Option<Duration>,
) -> Duration {
if !pending {
return current.saturating_mul(2).min(MAX_PR_POLL_INTERVAL);
}
match since_moved {
Some(elapsed) if elapsed < PENDING_FAST_WINDOW => base,
_ => current.saturating_mul(2).min(PENDING_MAX_INTERVAL),
}
}
fn budget_throttled_delay(delay: Duration, rate_limit: Option<&RateLimitSnapshot>) -> Duration {
if rate_limit.is_some_and(RateLimitSnapshot::over_warn) {
delay.max(BUDGET_THROTTLE_INTERVAL)
} else {
delay
}
}
fn rate_limit_crossed_warn(prev: Option<&RateLimitSnapshot>, next: &RateLimitSnapshot) -> bool {
let over = |res: Option<RateLimitResource>| res.is_some_and(|r| r.over_warn());
[
(prev.and_then(|p| p.graphql), next.graphql),
(prev.and_then(|p| p.core), next.core),
(prev.and_then(|p| p.search), next.search),
]
.into_iter()
.any(|(before, after)| over(after) && !over(before))
}
struct PollerTask {
token: CancellationToken,
handle: JoinHandle<()>,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct PrWatch {
target: PrTarget,
upstream_sha: Option<String>,
}
fn pr_watch_from_snapshot(snapshot: &Value) -> Vec<PrWatch> {
let mut out = Vec::new();
for repo in snapshot
.get("repos")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
if repo.get("polling_enabled").and_then(Value::as_bool) != Some(true) {
continue;
}
let Some(github) = repo.get("github") else {
continue;
};
let (Some(owner), Some(name)) = (
github.get("owner").and_then(Value::as_str),
github.get("name").and_then(Value::as_str),
) else {
continue;
};
for wt in repo
.get("worktrees")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
if let Some(branch) = wt.get("branch").and_then(Value::as_str) {
out.push(PrWatch {
upstream_sha: wt
.get("upstream_sha")
.and_then(Value::as_str)
.map(str::to_string),
target: PrTarget {
owner: owner.to_string(),
name: name.to_string(),
branch: branch.to_string(),
},
});
}
}
}
out.sort();
out.dedup();
out
}
#[cfg(test)]
fn pr_targets_from_snapshot(snapshot: &Value) -> Vec<PrTarget> {
pr_watch_from_snapshot(snapshot)
.into_iter()
.map(|w| w.target)
.collect()
}
pub struct WorktreesService {
registry: Arc<WorktreesRegistry>,
menu_cache: Arc<Mutex<Option<Vec<MenuItem>>>>,
refresh: Mutex<Option<RefreshTask>>,
pr_cache: Arc<PrStatusCache>,
poller: Mutex<Option<PollerTask>>,
rate_limit_cache: Arc<RateLimitCache>,
rate_limit_poller: Mutex<Option<PollerTask>>,
tree_cache: Arc<TreeSnapshotCache>,
prune_lock: tokio::sync::Mutex<()>,
polling_prefs_path: Mutex<Option<PathBuf>>,
pr_cache_path: Mutex<Option<PathBuf>>,
pr_warm_start: Mutex<Option<PrWarmStart>>,
open_pr_cache: Arc<OpenPrCache>,
}
impl WorktreesService {
#[must_use]
pub fn new() -> Self {
let registry = Arc::new(WorktreesRegistry::new());
let pr_cache = Arc::new(PrStatusCache::new());
Self {
registry: registry.clone(),
menu_cache: Arc::new(Mutex::new(None)),
refresh: Mutex::new(None),
pr_cache: pr_cache.clone(),
poller: Mutex::new(None),
rate_limit_cache: Arc::new(RateLimitCache::new()),
rate_limit_poller: Mutex::new(None),
tree_cache: Arc::new(TreeSnapshotCache::new(registry, pr_cache)),
prune_lock: tokio::sync::Mutex::new(()),
polling_prefs_path: Mutex::new(None),
pr_cache_path: Mutex::new(None),
pr_warm_start: Mutex::new(None),
open_pr_cache: Arc::new(OpenPrCache::new(open_pr_ttl())),
}
}
pub fn load_polling_prefs(&self, path: PathBuf) {
match std::fs::read(&path) {
Ok(bytes) => match serde_json::from_slice::<PollingPrefs>(&bytes) {
Ok(prefs) => self
.registry
.seed_polling(prefs.enabled.into_iter().map(|l| (l.repo, l.expires_at))),
Err(err) => tracing::warn!(
"ignoring unreadable worktrees polling prefs at {}: {err:#}",
path.display()
),
},
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => tracing::warn!(
"could not read worktrees polling prefs at {}: {err:#}",
path.display()
),
}
*self
.polling_prefs_path
.lock()
.unwrap_or_else(PoisonError::into_inner) = Some(path);
}
fn persist_polling_prefs(&self) {
let Some(path) = self
.polling_prefs_path
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
else {
return;
};
let prefs = PollingPrefs {
enabled: self
.registry
.polling_snapshot()
.into_iter()
.map(|(repo, expires_at)| PollingLease { repo, expires_at })
.collect(),
};
if let Err(err) = write_polling_prefs(&path, &prefs) {
tracing::warn!(
"could not persist worktrees polling prefs to {}: {err:#}",
path.display()
);
}
}
pub fn load_pr_cache(&self, path: PathBuf) {
match std::fs::read(&path) {
Ok(bytes) => match serde_json::from_slice::<PrCachePrefs>(&bytes) {
Ok(prefs) => {
self.pr_cache.seed(
prefs
.entries
.into_iter()
.map(|e| (e.target, e.resolution.into_resolution())),
);
if let Some(polled_at) = prefs.polled_at {
let watched = prefs
.watched
.into_iter()
.map(|w| PrWatch {
target: w.target,
upstream_sha: w.upstream_sha,
})
.collect();
*self
.pr_warm_start
.lock()
.unwrap_or_else(PoisonError::into_inner) =
Some(PrWarmStart { watched, polled_at });
}
}
Err(err) => {
let at = path.display();
tracing::warn!("ignoring unreadable worktrees PR cache at {at}: {err:#}");
}
},
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => {
let at = path.display();
tracing::warn!("could not read worktrees PR cache at {at}: {err:#}");
}
}
*self
.pr_cache_path
.lock()
.unwrap_or_else(PoisonError::into_inner) = Some(path);
}
async fn open_prs(&self, owner: &str, name: &str) -> Result<Vec<Value>> {
self.open_prs_with(owner, name, crate::pr_status::resolve_gh_binary())
.await
}
async fn open_prs_with(&self, owner: &str, name: &str, bin: PathBuf) -> Result<Vec<Value>> {
let key = format!("{owner}/{name}");
if let Some(prs) = self.open_pr_cache.fresh(&key) {
return Ok(prs);
}
let slug = key.clone();
let prs = tokio::task::spawn_blocking(move || open_pr_list(&bin, &slug))
.await
.unwrap_or_else(|err| Err(anyhow!("blocking open-prs task failed: {err}")))?;
self.open_pr_cache.store(key, prs.clone());
Ok(prs)
}
#[must_use]
pub fn rate_limit_cache(&self) -> Arc<RateLimitCache> {
self.rate_limit_cache.clone()
}
pub fn start_menu_refresh(&self) {
if tokio::runtime::Handle::try_current().is_err() {
tracing::debug!("no tokio runtime; worktrees menu refresh not started");
return;
}
let mut guard = self.refresh.lock().unwrap_or_else(PoisonError::into_inner);
if guard.is_some() {
return;
}
let token = CancellationToken::new();
let loop_token = token.clone();
let registry = self.registry.clone();
let cache = self.menu_cache.clone();
let rate_limit_cache = self.rate_limit_cache.clone();
let interval = menu_refresh_interval();
let handle = tokio::spawn(async move {
loop {
let entries = registry.list();
let rate_limit = rate_limit_cache.get();
if let Ok(items) = tokio::task::spawn_blocking(move || {
menu_items_for(&entries, rate_limit.as_ref())
})
.await
{
*cache.lock().unwrap_or_else(PoisonError::into_inner) = Some(items);
}
tokio::select! {
() = loop_token.cancelled() => break,
() = tokio::time::sleep(interval) => {}
}
}
});
*guard = Some(RefreshTask { token, handle });
}
pub fn start_pr_poller(&self) {
self.start_pr_poller_with(
pr_poll_interval(),
pr_debounce_interval(),
crate::pr_status::resolve_gh_binary(),
);
}
fn start_pr_poller_with(&self, base: Duration, debounce: Duration, gh_bin: PathBuf) {
if tokio::runtime::Handle::try_current().is_err() {
tracing::debug!("no tokio runtime; worktrees PR poller not started");
return;
}
let mut guard = self.poller.lock().unwrap_or_else(PoisonError::into_inner);
if guard.is_some() {
return;
}
let token = CancellationToken::new();
let loop_token = token.clone();
let registry = self.registry.clone();
let tree_cache = self.tree_cache.clone();
let pr_cache = self.pr_cache.clone();
let rate_limit_cache = self.rate_limit_cache.clone();
let pr_cache_path = self
.pr_cache_path
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone();
let warm_start = self
.pr_warm_start
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
let mut changes = self.registry.subscribe_changes();
let handle = tokio::spawn(async move {
let mut backoff = base;
let (mut watched, mut last_poll): (Option<Vec<PrWatch>>, Option<Instant>) =
match warm_start {
Some(ws) => {
let elapsed = (Utc::now() - ws.polled_at)
.to_std()
.unwrap_or(Duration::ZERO);
(Some(ws.watched), Instant::now().checked_sub(elapsed))
}
None => (None, None),
};
let mut moved_at: Option<Instant> = None;
'poll: loop {
tokio::select! {
() = loop_token.cancelled() => break,
() = tokio::time::sleep(base) => {}
result = changes.changed() => {
if result.is_err() {
break;
}
let overall_deadline = Instant::now() + debounce.saturating_mul(4);
loop {
tokio::select! {
() = loop_token.cancelled() => break 'poll,
() = tokio::time::sleep(debounce) => break,
r = changes.changed() => {
if r.is_err() {
break 'poll;
}
if Instant::now() >= overall_deadline {
break;
}
}
}
}
}
}
let snapshot = tree_cache.snapshot().await;
let watch = pr_watch_from_snapshot(&snapshot);
if watch.is_empty() {
backoff = base;
last_poll = None;
moved_at = None;
continue;
}
let grew = pr_watch_grew(watched.as_deref().unwrap_or(&[]), &watch);
let keep: HashSet<PrTarget> = watch.iter().map(|w| w.target.clone()).collect();
pr_cache.retain_targets(&keep);
let rate_limit = rate_limit_cache.get();
let over_budget = rate_limit.is_some_and(|s| s.over_warn());
let effective_backoff = budget_throttled_delay(backoff, rate_limit.as_ref());
let trigger = grew && !over_budget;
if !pr_should_fetch(trigger, last_poll.map(|at| at.elapsed()), effective_backoff) {
if !grew {
watched = Some(watch);
}
continue;
}
if grew {
backoff = base;
moved_at = Some(Instant::now());
}
let targets: Vec<PrTarget> = watch.iter().map(|w| w.target.clone()).collect();
let bin = gh_bin.clone();
let resolved = tokio::task::spawn_blocking(move || {
crate::pr_status::resolve_with_budget(&bin, &targets)
})
.await
.unwrap_or_else(|err| Err(anyhow!("blocking poll task failed: {err}")));
let (pending, resolved_ok) = match resolved {
Ok((resolutions, budget)) => {
if let Some(b) = budget {
tracing::debug!(
"PR poll cost {} point(s); graphql {}/{} used, {} remaining",
b.cost,
b.used,
b.limit,
b.remaining
);
rate_limit_cache.observe_graphql(RateLimitResource::new(
b.used,
b.limit,
b.remaining,
b.reset,
));
}
if pr_cache.replace(resolutions) {
registry.bump();
}
(pr_cache.any_pending(), true)
}
Err(err) => {
tracing::debug!("PR badge poll failed: {err:#}");
(false, false)
}
};
last_poll = Some(Instant::now());
watched = Some(watch);
let since_moved = moved_at.map(|at| at.elapsed());
backoff = next_pr_poll_delay(backoff, base, pending, since_moved);
if resolved_ok {
if let Some(path) = &pr_cache_path {
persist_pr_cache(
path,
&pr_cache,
watched.as_deref().unwrap_or(&[]),
Utc::now(),
);
}
}
}
});
*guard = Some(PollerTask { token, handle });
}
pub fn start_rate_limit_poller(&self) {
self.start_rate_limit_poller_with(
rate_limit_poll_interval(),
crate::pr_status::resolve_gh_binary(),
);
}
fn start_rate_limit_poller_with(&self, interval: Duration, gh_bin: PathBuf) {
if tokio::runtime::Handle::try_current().is_err() {
tracing::debug!("no tokio runtime; worktrees rate-limit poller not started");
return;
}
let mut guard = self
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner);
if guard.is_some() {
return;
}
let token = CancellationToken::new();
let loop_token = token.clone();
let cache = self.rate_limit_cache.clone();
let registry = self.registry.clone();
let handle = tokio::spawn(async move {
let mut prev: Option<RateLimitSnapshot> = None;
loop {
if !registry.list().is_empty() || !registry.polling_snapshot().is_empty() {
let bin = gh_bin.clone();
let resolved =
tokio::task::spawn_blocking(move || resolve_rate_limit_with(&bin))
.await
.unwrap_or_else(|err| {
Err(anyhow!("blocking rate-limit poll task failed: {err}"))
});
match resolved {
Ok(snap) => {
if rate_limit_crossed_warn(prev.as_ref(), &snap) {
let summary = snap.summary_line();
tracing::warn!(
"GitHub API rate limit high: {summary} (querying \
/rate_limit is free; the daemon's gh usage is not)"
);
}
cache.replace(snap);
prev = Some(snap);
}
Err(err) => tracing::debug!("GitHub rate-limit poll failed: {err:#}"),
}
}
tokio::select! {
() = loop_token.cancelled() => break,
() = tokio::time::sleep(interval) => {}
}
}
});
*guard = Some(PollerTask { token, handle });
}
async fn close(&self, req: CloseRequest) -> Result<Value> {
let entries = self.registry.list();
let scan_path = req.path.clone();
let open_windows =
tokio::task::spawn_blocking(move || windows_with_path(&entries, &scan_path))
.await
.unwrap_or_default();
let open = !open_windows.is_empty();
let window_key = open_windows.first().map(|(k, _)| k.clone());
let window_folder_count = open_windows.first().map_or(0, |(_, c)| *c);
if req.remove && !req.confirmed {
let path = req.path.clone();
let git = tokio::task::spawn_blocking(move || git_safety(&path))
.await
.map_err(|e| anyhow!("safety check task panicked: {e}"))??;
return Ok(serde_json::to_value(SafetyReport {
removable: git.removable,
is_main: git.is_main,
open,
window_key,
window_folder_count,
risks: git.risks,
info: git.info,
})
.unwrap_or_else(|_| json!({})));
}
let others: Vec<String> = open_windows
.iter()
.map(|(k, _)| k.clone())
.filter(|k| req.requester_key.as_deref() != Some(k))
.collect();
for key in &others {
self.registry.mark_close_pending(key);
}
if !others.is_empty() {
await_windows_closed(
&self.registry,
&req.path,
req.requester_key.as_deref(),
CLOSE_WAIT_TIMEOUT,
CLOSE_WAIT_POLL,
)
.await?;
}
if req.remove {
let path = req.path.clone();
let _guard = self.prune_lock.lock().await;
tokio::task::spawn_blocking(move || remove_worktree(&path))
.await
.map_err(|e| anyhow!("worktree removal task panicked: {e}"))??;
Ok(json!({ "removed": true }))
} else {
Ok(json!({ "closed": true }))
}
}
}
impl Default for WorktreesService {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl DaemonService for WorktreesService {
fn name(&self) -> &'static str {
SERVICE_NAME
}
async fn handle(&self, op: &str, payload: Value) -> Result<Value> {
match op {
"register" => {
let req: RegisterRequest =
serde_json::from_value(payload).context("invalid `register` payload")?;
if req.key.trim().is_empty() {
bail!("`register` requires a non-empty `key`");
}
self.registry.register(req);
Ok(json!({ "ok": true }))
}
"heartbeat" => {
let key = require_str(&payload, "key", "heartbeat")?;
let known = self.registry.heartbeat(key);
let mut reply = json!({ "known": known });
if self.registry.take_close_pending(key) {
reply["close"] = Value::Bool(true);
}
Ok(reply)
}
"unregister" => {
let key = require_str(&payload, "key", "unregister")?;
Ok(json!({ "removed": self.registry.unregister(key) }))
}
"list" => Ok(json!({ "windows": enriched_windows(self.registry.list()).await })),
"tree" => {
Ok(tree_snapshot(&self.registry, self.pr_cache.clone()).await)
}
"ahead-behind" => {
let paths = payload
.get("paths")
.and_then(Value::as_array)
.map(|arr| {
arr.iter()
.filter_map(Value::as_str)
.map(PathBuf::from)
.collect::<Vec<_>>()
})
.unwrap_or_default();
Ok(json!({ "results": ahead_behind_results(paths).await }))
}
"set-show-closed" => {
let show_closed = payload
.get("show_closed")
.and_then(Value::as_bool)
.ok_or_else(|| anyhow!("`set-show-closed` requires a boolean `show_closed`"))?;
self.registry.set_show_closed(show_closed);
Ok(json!({ "ok": true }))
}
"set-polling" => {
let owner = require_str(&payload, "owner", "set-polling")?;
let name = require_str(&payload, "name", "set-polling")?;
let enabled = payload
.get("enabled")
.and_then(Value::as_bool)
.ok_or_else(|| anyhow!("`set-polling` requires a boolean `enabled`"))?;
if owner.trim().is_empty() || name.trim().is_empty() {
bail!("`set-polling` requires a non-empty `owner` and `name`");
}
if self.registry.set_polling(owner, name, enabled) {
self.persist_polling_prefs();
}
Ok(json!({ "ok": true }))
}
"open-prs" => {
let owner = require_str(&payload, "owner", "open-prs")?;
let name = require_str(&payload, "name", "open-prs")?;
if owner.trim().is_empty() || name.trim().is_empty() {
bail!("`open-prs` requires a non-empty `owner` and `name`");
}
Ok(json!({ "pull_requests": self.open_prs(owner, name).await? }))
}
"open" => {
let path = require_str(&payload, "path", "open")?;
focus_window(Path::new(path))?;
Ok(json!({ "ok": true }))
}
"close" => {
let req: CloseRequest =
serde_json::from_value(payload).context("invalid `close` payload")?;
self.close(req).await
}
other => bail!("unknown worktrees op: {other}"),
}
}
fn subscribe(&self, op: &str, _payload: &Value) -> Option<Box<dyn ServiceStream>> {
if op != "subscribe" {
return None;
}
Some(Box::new(WorktreesStream {
cache: self.tree_cache.clone(),
changes: self.registry.subscribe_changes(),
}))
}
fn menu(&self) -> MenuSnapshot {
let cached = self
.menu_cache
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone();
let items = cached.unwrap_or_else(|| {
menu_items_for(&self.registry.list(), self.rate_limit_cache.get().as_ref())
});
MenuSnapshot {
title: SUBMENU_TITLE.to_string(),
items,
}
}
async fn menu_action(&self, action_id: &str) -> Result<()> {
if let Some(key) = action_id.strip_prefix("focus:") {
let folder = self
.registry
.first_folder(key)
.ok_or_else(|| anyhow!("no open window with key {key} (it may have closed)"))?;
focus_window(&folder)?;
return Ok(());
}
bail!("unknown worktrees menu action: {action_id}")
}
async fn status(&self) -> ServiceStatus {
let entries = self.registry.list();
let repos: BTreeSet<&str> = entries.iter().filter_map(|e| e.repo.as_deref()).collect();
let summary = format!("{} window(s) across {} repo(s)", entries.len(), repos.len());
let windows = enriched_windows(entries).await;
ServiceStatus {
name: SERVICE_NAME.to_string(),
healthy: true,
summary,
detail: json!({ "windows": windows }),
}
}
async fn shutdown(&self) {
let task = self
.refresh
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(task) = task {
task.token.cancel();
let _ = task.handle.await;
}
let poller = self
.poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(poller) = poller {
poller.token.cancel();
let _ = poller.handle.await;
}
let rate_limit_poller = self
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take();
if let Some(poller) = rate_limit_poller {
poller.token.cancel();
let _ = poller.handle.await;
}
}
}
fn require_str<'a>(payload: &'a Value, field: &str, op: &str) -> Result<&'a str> {
payload
.get(field)
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("`{op}` requires `{field}`"))
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)]
struct GitStatus {
#[serde(skip_serializing_if = "Option::is_none")]
branch: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
head_sha: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_sha: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
ahead: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
behind: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
main_repo: Option<String>,
#[serde(skip_serializing_if = "is_false")]
is_worktree: bool,
}
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_false(b: &bool) -> bool {
!*b
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PollingLease {
repo: String,
expires_at: DateTime<Utc>,
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PollingPrefs {
#[serde(default)]
enabled: Vec<PollingLease>,
}
fn write_polling_prefs(path: &Path, prefs: &PollingPrefs) -> Result<()> {
if let Some(parent) = path.parent() {
crate::daemon::paths::ensure_dir_0700(parent)?;
}
let json = serde_json::to_vec_pretty(prefs).context("failed to serialize polling prefs")?;
crate::daemon::paths::write_file_0600(path, &json)
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PersistedBadge {
number: u64,
is_draft: bool,
checks: PrCheckState,
url: String,
head_oid: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
enum PersistedResolution {
Pr(PersistedBadge),
NoPr,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PersistedEntry {
target: PrTarget,
resolution: PersistedResolution,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PersistedWatch {
target: PrTarget,
#[serde(default, skip_serializing_if = "Option::is_none")]
upstream_sha: Option<String>,
}
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct PrCachePrefs {
#[serde(default)]
entries: Vec<PersistedEntry>,
#[serde(default)]
watched: Vec<PersistedWatch>,
#[serde(default, skip_serializing_if = "Option::is_none")]
polled_at: Option<DateTime<Utc>>,
}
impl PersistedResolution {
fn from_resolution(r: &PrResolution) -> Self {
match r {
PrResolution::Pr(b) => Self::Pr(PersistedBadge {
number: b.number,
is_draft: b.is_draft,
checks: b.checks,
url: b.url.clone(),
head_oid: b.head_oid.clone(),
}),
PrResolution::NoPr => Self::NoPr,
}
}
fn into_resolution(self) -> PrResolution {
match self {
Self::Pr(b) => PrResolution::Pr(PrBadge {
number: b.number,
is_draft: b.is_draft,
checks: b.checks,
url: b.url,
head_oid: b.head_oid,
}),
Self::NoPr => PrResolution::NoPr,
}
}
}
fn pr_cache_prefs_from(
entries: Vec<(PrTarget, PrResolution)>,
watched: &[PrWatch],
polled_at: DateTime<Utc>,
) -> PrCachePrefs {
let mut entries: Vec<PersistedEntry> = entries
.into_iter()
.map(|(target, resolution)| PersistedEntry {
target,
resolution: PersistedResolution::from_resolution(&resolution),
})
.collect();
entries.sort_by(|a, b| a.target.cmp(&b.target));
let mut watched: Vec<PersistedWatch> = watched
.iter()
.map(|w| PersistedWatch {
target: w.target.clone(),
upstream_sha: w.upstream_sha.clone(),
})
.collect();
watched.sort_by(|a, b| a.target.cmp(&b.target));
PrCachePrefs {
entries,
watched,
polled_at: Some(polled_at),
}
}
fn write_pr_cache(path: &Path, prefs: &PrCachePrefs) -> Result<()> {
if let Some(parent) = path.parent() {
crate::daemon::paths::ensure_dir_0700(parent)?;
}
let json = serde_json::to_vec_pretty(prefs).context("failed to serialize PR cache")?;
crate::daemon::paths::write_file_0600(path, &json)
}
fn persist_pr_cache(
path: &Path,
pr_cache: &PrStatusCache,
watched: &[PrWatch],
polled_at: DateTime<Utc>,
) {
let prefs = pr_cache_prefs_from(pr_cache.entries(), watched, polled_at);
if let Err(err) = write_pr_cache(path, &prefs) {
let at = path.display();
tracing::warn!("could not persist worktrees PR cache to {at}: {err:#}");
}
}
#[derive(Debug, Clone)]
struct PrWarmStart {
watched: Vec<PrWatch>,
polled_at: DateTime<Utc>,
}
#[derive(Debug, Clone)]
struct OpenPrEntry {
at: Instant,
prs: Vec<Value>,
}
#[derive(Debug)]
struct OpenPrCache {
entries: Mutex<HashMap<String, OpenPrEntry>>,
ttl: Duration,
}
impl OpenPrCache {
fn new(ttl: Duration) -> Self {
Self {
entries: Mutex::new(HashMap::new()),
ttl,
}
}
fn fresh(&self, key: &str) -> Option<Vec<Value>> {
self.entries
.lock()
.unwrap_or_else(PoisonError::into_inner)
.get(key)
.filter(|e| e.at.elapsed() < self.ttl)
.map(|e| e.prs.clone())
}
fn store(&self, key: String, prs: Vec<Value>) {
self.entries
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(
key,
OpenPrEntry {
at: Instant::now(),
prs,
},
);
}
}
fn open_pr_list(bin: &Path, slug: &str) -> Result<Vec<Value>> {
let output = crate::github_metrics::run_gh(
bin,
[
"pr",
"list",
"--repo",
slug,
"--state",
"open",
"--json",
OPEN_PR_JSON_FIELDS,
"--limit",
OPEN_PR_LIST_LIMIT,
],
"pr list",
None,
)
.with_context(|| {
format!(
"failed to run {} (is the GitHub CLI installed?)",
bin.display()
)
})?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!("gh pr list failed: {}", stderr.trim());
}
match serde_json::from_slice(&output.stdout).context("gh pr list returned invalid JSON")? {
Value::Array(arr) => Ok(arr),
_ => bail!("gh pr list did not return a JSON array"),
}
}
fn git_status(folder: &Path) -> GitStatus {
git_status_impl(folder, true)
}
fn git_status_cheap(folder: &Path) -> GitStatus {
git_status_impl(folder, false)
}
fn git_status_impl(folder: &Path, with_ahead_behind: bool) -> GitStatus {
let Ok(repo) = Repository::discover(folder) else {
return GitStatus::default();
};
let base = GitStatus {
main_repo: main_repo_name(repo.commondir()),
is_worktree: repo.is_worktree(),
..GitStatus::default()
};
let Ok(head) = repo.head() else {
return base;
};
let base = GitStatus {
head_sha: head.target().map(|oid| oid.to_string()),
..base
};
let Some(name) = head
.shorthand()
.ok()
.filter(|_| head.is_branch())
.map(str::to_string)
else {
return base;
};
let branch = git2::Branch::wrap(head);
let upstream_sha = upstream_target(&branch);
let (ahead, behind) = if with_ahead_behind {
match upstream_ahead_behind(&repo, &branch) {
Some((ahead, behind)) => (Some(ahead), Some(behind)),
None => (None, None),
}
} else {
(None, None)
};
GitStatus {
branch: Some(name),
upstream_sha,
ahead,
behind,
..base
}
}
fn upstream_target(branch: &git2::Branch<'_>) -> Option<String> {
Some(branch.upstream().ok()?.get().target()?.to_string())
}
fn folder_ahead_behind(folder: &Path) -> Option<(usize, usize)> {
let repo = Repository::discover(folder).ok()?;
let head = repo.head().ok()?;
if !head.is_branch() {
return None;
}
let branch = git2::Branch::wrap(head);
upstream_ahead_behind(&repo, &branch)
}
fn main_repo_name(commondir: &Path) -> Option<String> {
let file_name = commondir.file_name()?.to_string_lossy().into_owned();
if file_name == ".git" {
commondir
.parent()
.and_then(Path::file_name)
.map(|n| n.to_string_lossy().into_owned())
} else {
Some(
file_name
.strip_suffix(".git")
.unwrap_or(&file_name)
.to_string(),
)
}
}
fn upstream_ahead_behind(repo: &Repository, branch: &git2::Branch<'_>) -> Option<(usize, usize)> {
let upstream = branch.upstream().ok()?;
let local_oid = branch.get().target()?;
let upstream_oid = upstream.get().target()?;
repo.graph_ahead_behind(local_oid, upstream_oid).ok()
}
#[derive(Serialize)]
struct EnrichedEntry<'a> {
#[serde(flatten)]
entry: &'a WindowEntry,
#[serde(flatten)]
git: GitStatus,
}
fn enriched_entry(entry: &WindowEntry) -> Value {
let git = entry
.folders
.first()
.map(|folder| git_status(folder))
.unwrap_or_default();
serde_json::to_value(EnrichedEntry { entry, git }).unwrap_or_else(|_| json!({}))
}
async fn enriched_windows(entries: Vec<WindowEntry>) -> Vec<Value> {
tokio::task::spawn_blocking(move || entries.iter().map(enriched_entry).collect())
.await
.unwrap_or_default()
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct GithubIdentity {
owner: String,
name: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct TreeWorktree {
path: String,
#[serde(skip_serializing_if = "Option::is_none")]
branch: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
head_sha: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
upstream_sha: Option<String>,
is_main: bool,
open: bool,
#[serde(skip_serializing_if = "Option::is_none")]
window_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pr: Option<PrBadge>,
#[serde(skip_serializing_if = "is_false")]
pr_none: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct TreeRepo {
main_repo: String,
#[serde(skip_serializing_if = "Option::is_none")]
github: Option<GithubIdentity>,
root: String,
#[serde(skip_serializing_if = "is_false")]
polling_enabled: bool,
worktrees: Vec<TreeWorktree>,
}
fn github_identity(url: &str) -> Option<GithubIdentity> {
let url = url.trim();
let rest = [
"https://github.com/",
"http://github.com/",
"ssh://git@github.com/",
"git://github.com/",
"git@github.com:",
]
.iter()
.find_map(|prefix| url.strip_prefix(prefix))?;
let rest = rest.strip_suffix(".git").unwrap_or(rest);
let rest = rest.trim_end_matches('/');
let mut parts = rest.splitn(2, '/');
let owner = parts.next()?.trim();
let name = parts.next()?.trim();
if owner.is_empty() || name.is_empty() || name.contains('/') {
return None;
}
Some(GithubIdentity {
owner: owner.to_string(),
name: name.to_string(),
})
}
fn remote_github_identity(repo: &Repository) -> Option<GithubIdentity> {
if let Ok(origin) = repo.find_remote("origin") {
if let Some(id) = origin.url().ok().and_then(github_identity) {
return Some(id);
}
}
let names = repo.remotes().ok();
names
.iter()
.flat_map(|arr| arr.iter())
.flatten()
.flatten()
.filter_map(|name| repo.find_remote(name).ok())
.find_map(|remote| remote.url().ok().and_then(github_identity))
}
fn canonical(path: &Path) -> PathBuf {
std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
}
fn open_window_index(entries: &[WindowEntry]) -> HashMap<PathBuf, String> {
let mut index = HashMap::new();
for entry in entries {
for folder in &entry.folders {
index
.entry(canonical(folder))
.or_insert_with(|| entry.key.clone());
}
}
index
}
fn worktree_entry(
path: &Path,
is_main: bool,
open_index: &HashMap<PathBuf, String>,
) -> TreeWorktree {
let status = git_status_cheap(path);
let window_key = open_index.get(&canonical(path)).cloned();
TreeWorktree {
path: path.display().to_string(),
branch: status.branch,
head_sha: status.head_sha,
upstream_sha: status.upstream_sha,
is_main,
open: window_key.is_some(),
window_key,
pr: None,
pr_none: false,
}
}
fn stamp_polling(repos: &mut [TreeRepo], enabled: &HashSet<String>) {
for repo in repos {
if let Some(github) = &repo.github {
repo.polling_enabled = enabled.contains(&format!("{}/{}", github.owner, github.name));
}
}
}
fn fold_pr_badges(repos: &mut [TreeRepo], pr_cache: &PrStatusCache) {
for repo in repos {
if !repo.polling_enabled {
continue;
}
let Some(github) = repo.github.clone() else {
continue;
};
for worktree in &mut repo.worktrees {
let Some(branch) = &worktree.branch else {
continue;
};
match pr_cache.get(&github.owner, &github.name, branch) {
Some(PrResolution::Pr(mut badge)) => {
if badge.is_stale_for(worktree.head_sha.as_deref()) {
badge.checks = PrCheckState::Pending;
}
worktree.pr = Some(badge);
}
Some(PrResolution::NoPr) => worktree.pr_none = true,
None => {}
}
}
}
}
fn repo_tree(discovered: &Repository, open_index: &HashMap<PathBuf, String>) -> Option<TreeRepo> {
let commondir = canonical(discovered.commondir());
let main_root = commondir.parent()?.to_path_buf();
let main_repo = Repository::open(&main_root).ok()?;
let mut worktrees = vec![worktree_entry(&main_root, true, open_index)];
let names = main_repo.worktrees().ok();
let mut linked: Vec<PathBuf> = names
.iter()
.flat_map(|arr| arr.iter())
.flatten() .flatten() .filter_map(|name| main_repo.find_worktree(name).ok())
.map(|wt| wt.path().to_path_buf())
.collect();
linked.sort();
worktrees.extend(
linked
.iter()
.map(|path| worktree_entry(path, false, open_index)),
);
Some(TreeRepo {
main_repo: main_repo_name(&commondir)?,
github: remote_github_identity(&main_repo),
root: main_root.display().to_string(),
polling_enabled: false,
worktrees,
})
}
fn build_tree(folders: Vec<PathBuf>, windows: Vec<WindowEntry>) -> Vec<TreeRepo> {
let open_index = open_window_index(&windows);
let mut repos: BTreeMap<PathBuf, TreeRepo> = BTreeMap::new();
for folder in &folders {
let Ok(repo) = Repository::discover(folder) else {
continue;
};
let key = canonical(repo.commondir());
if repos.contains_key(&key) {
continue;
}
if let Some(tree) = repo_tree(&repo, &open_index) {
repos.insert(key, tree);
}
}
repos.into_values().collect()
}
async fn tree_repos(
folders: Vec<PathBuf>,
windows: Vec<WindowEntry>,
pr_cache: Arc<PrStatusCache>,
enabled_polling: HashSet<String>,
) -> Vec<Value> {
tokio::task::spawn_blocking(move || {
let mut repos = build_tree(folders, windows);
stamp_polling(&mut repos, &enabled_polling);
fold_pr_badges(&mut repos, &pr_cache);
repos
.iter()
.map(|repo| serde_json::to_value(repo).unwrap_or_else(|_| json!({})))
.collect()
})
.await
.unwrap_or_default()
}
async fn ahead_behind_results(paths: Vec<PathBuf>) -> Value {
tokio::task::spawn_blocking(move || {
let mut results = serde_json::Map::new();
for path in paths {
if let Some((ahead, behind)) = folder_ahead_behind(&path) {
results.insert(
path.display().to_string(),
json!({ "ahead": ahead, "behind": behind }),
);
}
}
Value::Object(results)
})
.await
.unwrap_or_else(|_| json!({}))
}
struct WorktreesStream {
cache: Arc<TreeSnapshotCache>,
changes: watch::Receiver<u64>,
}
#[async_trait]
impl ServiceStream for WorktreesStream {
async fn changed(&mut self) {
if self.changes.changed().await.is_err() {
std::future::pending::<()>().await;
}
}
async fn snapshot(&self) -> Value {
self.cache.snapshot().await
}
}
struct TreeSnapshotCache {
registry: Arc<WorktreesRegistry>,
pr_cache: Arc<PrStatusCache>,
ttl: Duration,
state: AsyncMutex<Option<CachedTree>>,
computes: AtomicU64,
}
struct CachedTree {
generation: u64,
computed_at: Instant,
value: Arc<Value>,
}
impl TreeSnapshotCache {
fn new(registry: Arc<WorktreesRegistry>, pr_cache: Arc<PrStatusCache>) -> Self {
Self::with_ttl(registry, pr_cache, crate::daemon::server::stream_tick())
}
fn with_ttl(
registry: Arc<WorktreesRegistry>,
pr_cache: Arc<PrStatusCache>,
ttl: Duration,
) -> Self {
Self {
registry,
pr_cache,
ttl,
state: AsyncMutex::new(None),
computes: AtomicU64::new(0),
}
}
async fn snapshot(&self) -> Value {
let mut state = self.state.lock().await;
let generation = self.registry.change_generation();
let fresh = state.as_ref().and_then(|cached| {
(cached.generation == generation && cached.computed_at.elapsed() < self.ttl)
.then(|| Arc::clone(&cached.value))
});
let value = if let Some(value) = fresh {
value
} else {
let value = Arc::new(tree_snapshot(&self.registry, self.pr_cache.clone()).await);
self.computes.fetch_add(1, Ordering::Relaxed);
*state = Some(CachedTree {
generation,
computed_at: Instant::now(),
value: Arc::clone(&value),
});
value
};
drop(state);
(*value).clone()
}
#[cfg(test)]
fn compute_count(&self) -> u64 {
self.computes.load(Ordering::Relaxed)
}
}
async fn tree_snapshot(registry: &WorktreesRegistry, pr_cache: Arc<PrStatusCache>) -> Value {
let folders = registry.open_folders();
let windows = registry.list();
let show_closed = registry.show_closed();
let enabled_polling = registry.enabled_polling_repos();
json!({
"repos": tree_repos(folders, windows, pr_cache, enabled_polling).await,
"show_closed": show_closed,
})
}
fn display_name(entry: &WindowEntry) -> String {
if let Some(repo) = &entry.repo {
return repo.clone();
}
if let Some(folder) = entry.folders.first() {
return folder.file_name().map_or_else(
|| folder.display().to_string(),
|n| n.to_string_lossy().into_owned(),
);
}
"(no folder)".to_string()
}
const REPO_SEP: char = '·';
const WORKTREE_SEP: char = '⑂';
fn menu_items_for(
entries: &[WindowEntry],
rate_limit: Option<&RateLimitSnapshot>,
) -> Vec<MenuItem> {
let mut items = Vec::new();
if let Some(label) = rate_limit.map(RateLimitSnapshot::tray_label) {
if !label.is_empty() {
items.push(MenuItem::Label(label));
items.push(MenuItem::Separator);
}
}
if entries.is_empty() {
items.push(MenuItem::Label("No open windows".to_string()));
} else {
items.extend(window_menu_items(entries));
}
items
}
fn window_menu_items(entries: &[WindowEntry]) -> Vec<MenuItem> {
entries
.iter()
.map(|entry| {
let label = window_label(entry);
if entry.folders.is_empty() {
MenuItem::Label(label)
} else {
MenuItem::Action(MenuAction {
id: format!("focus:{}", entry.key),
label,
enabled: true,
})
}
})
.collect()
}
fn window_label(entry: &WindowEntry) -> String {
let status = entry
.folders
.first()
.map(|folder| git_status(folder))
.unwrap_or_default();
let name = status
.main_repo
.clone()
.unwrap_or_else(|| display_name(entry));
if let Some(branch) = &status.branch {
let sep = if status.is_worktree {
WORKTREE_SEP
} else {
REPO_SEP
};
return match sync_indicator(status.ahead, status.behind) {
Some(sync) => format!("{name} {sep} {branch} {sync}"),
None => format!("{name} {sep} {branch}"),
};
}
match &entry.title {
Some(title) if title != &name => format!("{name} {REPO_SEP} {title}"),
_ => name,
}
}
fn sync_indicator(ahead: Option<usize>, behind: Option<usize>) -> Option<String> {
match (ahead, behind) {
(Some(ahead), Some(behind)) => Some(format!("(+{ahead} -{behind})")),
_ => None,
}
}
const CODE_BINARY_CANDIDATES: &[&str] = &[
"/usr/local/bin/code",
"/opt/homebrew/bin/code",
"/Applications/Visual Studio Code.app/Contents/Resources/app/bin/code",
"/usr/bin/code",
];
pub(crate) fn focus_window(folder: &Path) -> Result<()> {
focus_window_with(&resolve_code_binary(), folder)
}
fn focus_window_with(program: &Path, folder: &Path) -> Result<()> {
if !folder.is_absolute() {
bail!(
"refusing to focus a non-absolute folder path: {}",
folder.display()
);
}
if !folder.is_dir() {
bail!("worktree folder no longer exists: {}", folder.display());
}
let child = Command::new(program)
.arg(folder)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.with_context(|| {
format!(
"failed to launch `{}` to focus {}",
program.display(),
folder.display()
)
})?;
std::thread::spawn(move || {
let mut child = child;
let _ = child.wait();
});
Ok(())
}
fn resolve_code_binary() -> PathBuf {
resolve_code_binary_from(std::env::var_os(VSCODE_BIN_ENV), CODE_BINARY_CANDIDATES)
}
fn resolve_code_binary_from(
env_override: Option<std::ffi::OsString>,
candidates: &[&str],
) -> PathBuf {
if let Some(path) = env_override {
return PathBuf::from(path);
}
for candidate in candidates {
let path = Path::new(candidate);
if path.exists() {
return path.to_path_buf();
}
}
PathBuf::from("code")
}
#[derive(Debug, Clone, Deserialize)]
struct CloseRequest {
path: PathBuf,
#[serde(default)]
requester_key: Option<String>,
#[serde(default)]
remove: bool,
#[serde(default)]
confirmed: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct Note {
kind: String,
detail: String,
}
impl Note {
fn new(kind: &str, detail: impl Into<String>) -> Self {
Self {
kind: kind.to_string(),
detail: detail.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct SafetyReport {
removable: bool,
is_main: bool,
open: bool,
#[serde(skip_serializing_if = "Option::is_none")]
window_key: Option<String>,
window_folder_count: usize,
risks: Vec<Note>,
info: Vec<Note>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct GitSafety {
is_main: bool,
removable: bool,
risks: Vec<Note>,
info: Vec<Note>,
}
fn windows_with_path(entries: &[WindowEntry], path: &Path) -> Vec<(String, usize)> {
let target = canonical(path);
entries
.iter()
.filter(|e| e.folders.iter().any(|f| canonical(f) == target))
.map(|e| (e.key.clone(), e.folders.len()))
.collect()
}
const CLOSE_WAIT_TIMEOUT: Duration = Duration::from_secs(20);
const CLOSE_WAIT_POLL: Duration = Duration::from_millis(250);
async fn await_windows_closed(
registry: &WorktreesRegistry,
path: &Path,
requester: Option<&str>,
timeout: Duration,
poll: Duration,
) -> Result<()> {
let deadline = std::time::Instant::now() + timeout;
loop {
let entries = registry.list();
let path = path.to_path_buf();
let requester = requester.map(str::to_string);
let remaining: Vec<String> = tokio::task::spawn_blocking(move || {
windows_with_path(&entries, &path)
.into_iter()
.map(|(k, _)| k)
.filter(|k| requester.as_deref() != Some(k))
.collect()
})
.await
.unwrap_or_default();
if remaining.is_empty() {
return Ok(());
}
if std::time::Instant::now() >= deadline {
bail!("window(s) did not close in time: {}", remaining.join(", "));
}
tokio::time::sleep(poll).await;
}
}
fn git_safety(path: &Path) -> Result<GitSafety> {
if !path.exists() {
return Ok(GitSafety {
is_main: false,
removable: true,
risks: vec![],
info: vec![Note::new("already-removed", "worktree no longer exists")],
});
}
let repo = Repository::open(path)
.with_context(|| format!("not a git worktree: {}", path.display()))?;
if !repo.is_worktree() {
return Ok(GitSafety {
is_main: true,
removable: false,
risks: vec![],
info: vec![Note::new(
"main-working-tree",
"the repository's main working tree is never deleted",
)],
});
}
let mut risks = Vec::new();
let mut info = Vec::new();
let (dirty, untracked) = count_dirty_untracked(&repo);
if dirty > 0 {
risks.push(Note::new(
"dirty",
format!("{dirty} modified tracked file(s) would be lost"),
));
}
if untracked > 0 {
risks.push(Note::new(
"untracked",
format!("{untracked} untracked file(s) would be lost"),
));
}
let state = repo.state();
if state != RepositoryState::Clean {
risks.push(Note::new(
"in-progress",
format!("an in-progress {state:?} operation would be lost"),
));
}
if repo.head_detached().unwrap_or(false) {
let lost = unreachable_commit_count(&repo).unwrap_or(0);
if lost > 0 {
risks.push(Note::new(
"unreachable-commits",
format!("{lost} commit(s) on a detached HEAD will be permanently lost"),
));
}
}
if let Some(ahead) = current_branch_ahead(&repo) {
if ahead > 0 {
info.push(Note::new(
"unpushed",
format!("{ahead} unpushed commit(s) on the branch (kept — the branch survives)"),
));
}
}
Ok(GitSafety {
is_main: false,
removable: true,
risks,
info,
})
}
fn count_dirty_untracked(repo: &Repository) -> (usize, usize) {
let mut opts = StatusOptions::new();
opts.include_untracked(true)
.recurse_untracked_dirs(true)
.include_ignored(false)
.exclude_submodules(true);
let Ok(statuses) = repo.statuses(Some(&mut opts)) else {
return (0, 0);
};
let tracked = Status::INDEX_NEW
| Status::INDEX_MODIFIED
| Status::INDEX_DELETED
| Status::INDEX_RENAMED
| Status::INDEX_TYPECHANGE
| Status::WT_MODIFIED
| Status::WT_DELETED
| Status::WT_TYPECHANGE
| Status::WT_RENAMED
| Status::CONFLICTED;
let mut dirty = 0;
let mut untracked = 0;
for entry in statuses.iter() {
let s = entry.status();
if s.contains(Status::WT_NEW) {
untracked += 1;
}
if s.intersects(tracked) {
dirty += 1;
}
}
(dirty, untracked)
}
fn unreachable_commit_count(repo: &Repository) -> Option<usize> {
let head_oid = repo.head().ok()?.target()?;
let mut walk = repo.revwalk().ok()?;
walk.push(head_oid).ok()?;
for reference in repo.references().ok()? {
let Ok(reference) = reference else { continue };
if matches!(reference.name(), Ok("HEAD")) {
continue;
}
if let Some(oid) = reference.target() {
let _ = walk.hide(oid);
}
}
Some(walk.flatten().count())
}
fn current_branch_ahead(repo: &Repository) -> Option<usize> {
let head = repo.head().ok()?;
if !head.is_branch() {
return None;
}
let branch = git2::Branch::wrap(head);
upstream_ahead_behind(repo, &branch).map(|(ahead, _behind)| ahead)
}
fn worktree_name_for_path(main_repo: &Repository, target: &Path) -> Result<String> {
let names = main_repo.worktrees()?;
names
.iter()
.flatten() .flatten() .find(|name| {
main_repo
.find_worktree(name)
.is_ok_and(|wt| canonical(wt.path()) == target)
})
.map(str::to_string)
.ok_or_else(|| {
anyhow!(
"worktree {} is not registered in {}",
target.display(),
main_repo.path().display()
)
})
}
const WORKTREE_RMDIR_BACKOFF: &[Duration] = &[
Duration::from_millis(250),
Duration::from_millis(500),
Duration::from_secs(1),
Duration::from_secs(1),
];
fn is_transient_rmdir_error(e: &std::io::Error) -> bool {
matches!(
e.raw_os_error(),
Some(nix::libc::ENOTEMPTY | nix::libc::EEXIST | nix::libc::EBUSY)
)
}
fn remove_dir_all_retrying(dir: &Path) -> Result<()> {
remove_dir_all_retrying_with(dir, WORKTREE_RMDIR_BACKOFF, || std::fs::remove_dir_all(dir))
}
fn remove_dir_all_retrying_with(
dir: &Path,
backoff: &[Duration],
mut remove: impl FnMut() -> std::io::Result<()>,
) -> Result<()> {
let mut backoff = backoff.iter();
loop {
match remove() {
Ok(()) => return Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => {
if is_transient_rmdir_error(&e) {
if let Some(delay) = backoff.next() {
std::thread::sleep(*delay);
continue;
}
}
return Err(e).with_context(|| {
format!("failed to remove worktree directory {}", dir.display())
});
}
}
}
}
fn is_orphaned_worktree(path: &Path) -> bool {
let Ok(contents) = std::fs::read_to_string(path.join(".git")) else {
return false;
};
let Some(admin) = contents.strip_prefix("gitdir:").map(str::trim) else {
return false;
};
let admin = Path::new(admin);
admin.components().any(|c| c.as_os_str() == "worktrees") && !admin.exists()
}
fn remove_worktree(path: &Path) -> Result<()> {
if !path.exists() {
return Ok(());
}
let repo = match Repository::open(path) {
Ok(repo) => repo,
Err(_) if is_orphaned_worktree(path) => return remove_dir_all_retrying(path),
Err(e) => return Err(e).context(format!("not a git worktree: {}", path.display())),
};
if !repo.is_worktree() {
bail!(
"refusing to delete the main working tree: {}",
path.display()
);
}
let commondir = canonical(repo.commondir());
let main_root = commondir
.parent()
.ok_or_else(|| anyhow!("no repository root for {}", path.display()))?
.to_path_buf();
drop(repo);
let main_repo = Repository::open(&main_root)
.with_context(|| format!("failed to open repository at {}", main_root.display()))?;
let name = worktree_name_for_path(&main_repo, &canonical(path))?;
let worktree = main_repo.find_worktree(&name)?;
if let WorktreeLockStatus::Locked(reason) = worktree.is_locked()? {
let because = reason.map(|r| format!(" ({r})")).unwrap_or_default();
bail!("worktree is locked{because}; unlock it first (git worktree unlock)");
}
remove_dir_all_retrying(path)?;
let mut opts = git2::WorktreePruneOptions::new();
opts.valid(true).working_tree(false);
worktree
.prune(Some(&mut opts))
.with_context(|| format!("failed to prune worktree metadata for {}", path.display()))?;
Ok(())
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use crate::test_support::shim::{shim_lock, write_exec_script};
use std::sync::MutexGuard;
fn register_payload(key: &str, repo: Option<&str>, folder: &str) -> Value {
json!({
"key": key,
"folders": [folder],
"repo": repo,
"title": format!("{key}-title"),
"pid": 1234,
})
}
fn windows_of(payload: &Value) -> &Vec<Value> {
payload
.get("windows")
.and_then(Value::as_array)
.expect("windows array")
}
#[tokio::test]
async fn name_and_unknown_op() {
let svc = WorktreesService::new();
assert_eq!(svc.name(), "worktrees");
assert!(svc.handle("frobnicate", Value::Null).await.is_err());
}
#[tokio::test]
async fn handle_routes_ops_and_shapes_payloads() {
let svc = WorktreesService::new();
let payload = svc.handle("list", Value::Null).await.unwrap();
assert_eq!(payload, json!({ "windows": [] }));
let reply = svc
.handle("register", register_payload("w1", Some("repo-a"), "/tmp/a"))
.await
.unwrap();
assert_eq!(reply, json!({ "ok": true }));
let windows = windows_of(&svc.handle("list", Value::Null).await.unwrap()).clone();
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].get("key").and_then(Value::as_str), Some("w1"));
assert!(windows[0].get("last_seen").is_some());
let known = svc
.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap();
assert_eq!(known, json!({ "known": true }));
let unknown = svc
.handle("heartbeat", json!({ "key": "nope" }))
.await
.unwrap();
assert_eq!(unknown, json!({ "known": false }));
let gone = svc
.handle("unregister", json!({ "key": "w1" }))
.await
.unwrap();
assert_eq!(gone, json!({ "removed": true }));
let again = svc
.handle("unregister", json!({ "key": "w1" }))
.await
.unwrap();
assert_eq!(again, json!({ "removed": false }));
}
#[tokio::test]
async fn handle_rejects_missing_or_empty_key() {
let svc = WorktreesService::new();
assert!(svc.handle("register", json!({})).await.is_err());
assert!(svc
.handle("register", json!({ "key": " " }))
.await
.is_err());
assert!(svc.handle("heartbeat", json!({})).await.is_err());
assert!(svc.handle("unregister", json!({})).await.is_err());
}
#[test]
fn display_name_prefers_repo_then_folder_basename() {
let base = WindowEntry {
key: "k".to_string(),
folders: vec![PathBuf::from("/home/me/project")],
repo: Some("my-repo".to_string()),
title: None,
pid: None,
last_seen: Utc::now(),
};
assert_eq!(display_name(&base), "my-repo");
let no_repo = WindowEntry {
repo: None,
..base.clone()
};
assert_eq!(display_name(&no_repo), "project");
let nothing = WindowEntry {
repo: None,
folders: vec![],
..base.clone()
};
assert_eq!(display_name(¬hing), "(no folder)");
let rootish = WindowEntry {
repo: None,
folders: vec![PathBuf::from("/")],
..base
};
assert_eq!(display_name(&rootish), "/");
}
#[test]
fn window_menu_items_merge_stats_and_focus_into_one_clickable_line() {
let now = Utc::now();
let entries = vec![
WindowEntry {
key: "k2".to_string(),
folders: vec![],
repo: Some("solo".to_string()),
title: Some("solo".to_string()),
pid: None,
last_seen: now,
},
WindowEntry {
key: "k1".to_string(),
folders: vec![PathBuf::from("/tmp/a")],
repo: Some("repo".to_string()),
title: Some("a branch".to_string()),
pid: None,
last_seen: now,
},
];
let items = window_menu_items(&entries);
assert_eq!(items.len(), 2);
assert!(!items.iter().any(|i| matches!(i, MenuItem::Separator)));
let action = items
.iter()
.find_map(|i| match i {
MenuItem::Action(a) => Some(a),
_ => None,
})
.expect("a focus action");
assert_eq!(action.id, "focus:k1");
assert_eq!(action.label, "repo · a branch");
let labels: Vec<&str> = items
.iter()
.filter_map(|i| match i {
MenuItem::Label(t) => Some(t.as_str()),
_ => None,
})
.collect();
assert_eq!(labels, vec!["solo"]);
}
#[tokio::test]
async fn menu_and_status_shapes() {
let svc = WorktreesService::new();
let menu = svc.menu();
assert_eq!(menu.title, "Worktrees");
assert!(matches!(
menu.items.first(),
Some(MenuItem::Label(text)) if text == "No open windows"
));
let status = svc.status().await;
assert_eq!(status.name, "worktrees");
assert!(status.healthy);
assert_eq!(status.summary, "0 window(s) across 0 repo(s)");
svc.handle("register", register_payload("w1", Some("repo-a"), "/tmp/a"))
.await
.unwrap();
svc.handle("register", register_payload("w2", Some("repo-a"), "/tmp/b"))
.await
.unwrap();
svc.handle(
"register",
json!({ "key": "w3", "repo": "repo-a", "folders": [] }),
)
.await
.unwrap();
let status = svc.status().await;
assert_eq!(status.summary, "3 window(s) across 1 repo(s)");
let menu = svc.menu();
assert_eq!(menu.items.len(), 3);
assert!(!menu.items.iter().any(|i| matches!(i, MenuItem::Separator)));
let action_ids: Vec<&str> = menu
.items
.iter()
.filter_map(|i| match i {
MenuItem::Action(a) => Some(a.id.as_str()),
_ => None,
})
.collect();
assert!(action_ids.contains(&"focus:w1"));
assert!(action_ids.contains(&"focus:w2"));
assert!(!action_ids.contains(&"focus:w3"));
}
#[test]
fn start_menu_refresh_is_a_noop_outside_a_runtime() {
let svc = WorktreesService::new();
svc.start_menu_refresh();
assert!(svc.refresh.lock().unwrap().is_none());
}
#[tokio::test]
async fn start_menu_refresh_populates_cache_and_shutdown_stops_it() {
let svc = WorktreesService::new();
svc.handle("register", register_payload("w1", Some("repo-a"), "/tmp/a"))
.await
.unwrap();
assert!(svc.menu_cache.lock().unwrap().is_none());
svc.start_menu_refresh();
svc.start_menu_refresh();
let mut filled = false;
for _ in 0..100 {
if svc.menu_cache.lock().unwrap().is_some() {
filled = true;
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
assert!(filled, "background refresh should populate the menu cache");
let menu = svc.menu();
assert_eq!(menu.title, "Worktrees");
assert!(menu
.items
.iter()
.any(|i| matches!(i, MenuItem::Action(a) if a.id == "focus:w1")));
svc.shutdown().await;
assert!(svc.refresh.lock().unwrap().is_none());
}
#[tokio::test]
async fn default_constructs_an_empty_service() {
let svc = WorktreesService::default();
let payload = svc.handle("list", Value::Null).await.unwrap();
assert_eq!(payload, json!({ "windows": [] }));
}
#[tokio::test]
async fn subscribe_streams_only_for_the_subscribe_op() {
let svc = WorktreesService::new();
assert!(svc.subscribe("subscribe", &Value::Null).is_some());
assert!(svc.subscribe("list", &Value::Null).is_none());
assert!(svc.subscribe("register", &Value::Null).is_none());
assert!(svc.subscribe("bogus", &Value::Null).is_none());
}
#[tokio::test]
async fn subscribe_snapshot_matches_the_tree_op() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let svc = WorktreesService::new();
let stream = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
assert_eq!(
stream.snapshot().await,
json!({ "repos": [], "show_closed": true })
);
svc.handle(
"register",
json!({ "key": "w1", "folders": [dir.path()], "repo": "r" }),
)
.await
.unwrap();
let snap = stream.snapshot().await;
let tree = svc.handle("tree", Value::Null).await.unwrap();
assert_eq!(snap, tree);
let repos = snap["repos"].as_array().expect("repos array");
assert_eq!(repos.len(), 1);
assert_eq!(repos[0]["worktrees"][0]["branch"], json!("main"));
}
#[tokio::test]
async fn subscribe_changed_wakes_on_register() {
let svc = WorktreesService::new();
let mut stream = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
tokio::select! {
() = stream.changed() => panic!("changed resolved with no registry change"),
() = tokio::time::sleep(Duration::from_millis(50)) => {}
}
svc.handle("register", register_payload("w1", Some("r"), "/tmp/a"))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), stream.changed())
.await
.expect("changed should resolve after a register");
}
#[tokio::test]
async fn tree_cache_coalesces_reads_within_ttl_and_generation() {
let reg = Arc::new(WorktreesRegistry::new());
let cache = TreeSnapshotCache::with_ttl(
reg,
Arc::new(PrStatusCache::new()),
Duration::from_secs(60),
);
let first = cache.snapshot().await;
assert_eq!(cache.compute_count(), 1);
let second = cache.snapshot().await;
assert_eq!(
cache.compute_count(),
1,
"an unchanged read must not rebuild"
);
assert_eq!(first, second);
}
#[tokio::test]
async fn tree_cache_single_flights_a_read_burst() {
let reg = Arc::new(WorktreesRegistry::new());
let cache = Arc::new(TreeSnapshotCache::with_ttl(
reg,
Arc::new(PrStatusCache::new()),
Duration::from_secs(60),
));
let mut handles = Vec::new();
for _ in 0..16 {
let cache = cache.clone();
handles.push(tokio::spawn(async move { cache.snapshot().await }));
}
let mut results = Vec::new();
for handle in handles {
results.push(handle.await.unwrap());
}
assert_eq!(
cache.compute_count(),
1,
"a concurrent read burst must build the tree once"
);
assert!(
results.windows(2).all(|w| w[0] == w[1]),
"every reader must observe the identical snapshot"
);
}
#[tokio::test]
async fn tree_cache_rebuilds_on_registry_change() {
let reg = Arc::new(WorktreesRegistry::new());
let cache = TreeSnapshotCache::with_ttl(
reg.clone(),
Arc::new(PrStatusCache::new()),
Duration::from_secs(60),
);
cache.snapshot().await;
assert_eq!(cache.compute_count(), 1);
assert!(reg.set_show_closed(false));
cache.snapshot().await;
assert_eq!(
cache.compute_count(),
2,
"a generation bump must force a rebuild"
);
}
#[tokio::test]
async fn tree_cache_rebuilds_after_ttl_expiry() {
let reg = Arc::new(WorktreesRegistry::new());
let cache =
TreeSnapshotCache::with_ttl(reg, Arc::new(PrStatusCache::new()), Duration::ZERO);
cache.snapshot().await;
cache.snapshot().await;
assert_eq!(
cache.compute_count(),
2,
"an expired TTL must force a rebuild each read"
);
}
#[tokio::test]
async fn subscribe_streams_share_one_build_per_generation() {
let svc = WorktreesService::new();
let s1 = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
let s2 = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
let a = s1.snapshot().await;
let b = s2.snapshot().await;
assert_eq!(a, b);
assert_eq!(
svc.tree_cache.compute_count(),
1,
"N streams on one generation must share a single build"
);
}
#[tokio::test]
async fn set_show_closed_toggles_the_snapshot_field() {
let svc = WorktreesService::new();
assert_eq!(
svc.handle("tree", Value::Null).await.unwrap()["show_closed"],
json!(true)
);
let reply = svc
.handle("set-show-closed", json!({ "show_closed": false }))
.await
.unwrap();
assert_eq!(reply, json!({ "ok": true }));
assert_eq!(
svc.handle("tree", Value::Null).await.unwrap()["show_closed"],
json!(false)
);
}
#[tokio::test]
async fn set_show_closed_rejects_a_non_boolean_payload() {
let svc = WorktreesService::new();
assert!(svc.handle("set-show-closed", json!({})).await.is_err());
assert!(svc
.handle("set-show-closed", json!({ "show_closed": "yes" }))
.await
.is_err());
}
#[tokio::test]
async fn set_show_closed_wakes_the_subscription() {
let svc = WorktreesService::new();
let mut stream = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
svc.handle("set-show-closed", json!({ "show_closed": false }))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), stream.changed())
.await
.expect("changed should resolve after a toggle flip");
assert_eq!(stream.snapshot().await["show_closed"], json!(false));
}
#[tokio::test]
async fn set_polling_toggles_the_snapshot_field_for_a_repo() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
let repo = repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0].clone();
assert!(
repo.get("polling_enabled").is_none(),
"default off omits the flag: {repo:?}"
);
let reply = svc
.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": true }),
)
.await
.unwrap();
assert_eq!(reply, json!({ "ok": true }));
let repo = repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0].clone();
assert_eq!(repo["polling_enabled"], json!(true));
svc.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": false }),
)
.await
.unwrap();
let repo = repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0].clone();
assert!(repo.get("polling_enabled").is_none());
}
#[tokio::test]
async fn set_polling_rejects_missing_or_empty_fields() {
let svc = WorktreesService::new();
assert!(svc
.handle("set-polling", json!({ "owner": "o", "name": "n" }))
.await
.is_err());
assert!(svc
.handle("set-polling", json!({ "enabled": true }))
.await
.is_err());
assert!(svc
.handle(
"set-polling",
json!({ "owner": " ", "name": "n", "enabled": true })
)
.await
.is_err());
}
#[tokio::test]
async fn set_polling_wakes_the_subscription() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
let mut stream = svc
.subscribe("subscribe", &Value::Null)
.expect("subscribe stream");
svc.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": true }),
)
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), stream.changed())
.await
.expect("changed should resolve after enabling a repo");
let repo = repos_of(&stream.snapshot().await)[0].clone();
assert_eq!(repo["polling_enabled"], json!(true));
}
#[tokio::test]
async fn disabling_a_repo_drops_its_pr_badges_immediately() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap().to_string();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(pending_badge(7, &head)),
);
svc.pr_cache.replace(badges);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr"]["number"], json!(7));
svc.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": false }),
)
.await
.unwrap();
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(
wt.get("pr").is_none(),
"a disabled repo must carry no badge: {wt:?}"
);
}
#[tokio::test]
async fn an_expired_lease_drops_the_flag_and_badges() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap().to_string();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(pending_badge(7, &head)),
);
svc.pr_cache.replace(badges);
let snap = svc.handle("tree", Value::Null).await.unwrap();
assert_eq!(repos_of(&snap)[0]["polling_enabled"], json!(true));
assert_eq!(repos_of(&snap)[0]["worktrees"][0]["pr"]["number"], json!(7));
assert_eq!(pr_targets_from_snapshot(&snap).len(), 1);
svc.registry.set_polling_expiry(
"rust-works",
"omni-dev",
Utc::now() - chrono::Duration::minutes(1),
);
let snap = svc.handle("tree", Value::Null).await.unwrap();
let repo0 = &repos_of(&snap)[0];
assert!(
repo0.get("polling_enabled").is_none(),
"expired lease drops the flag: {repo0:?}"
);
assert!(
repo0["worktrees"][0].get("pr").is_none(),
"expired lease drops the badge"
);
assert!(
pr_targets_from_snapshot(&snap).is_empty(),
"the poller no longer watches an expired repo"
);
}
#[tokio::test]
async fn polling_prefs_persist_across_reloads_with_0600() {
let dir = tempfile::tempdir().unwrap();
let prefs = dir.path().join("worktrees-polling.json");
let svc = WorktreesService::new();
svc.load_polling_prefs(prefs.clone());
assert!(!svc.registry.is_polling_enabled("rust-works", "omni-dev"));
svc.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": true }),
)
.await
.unwrap();
assert!(prefs.exists());
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
assert_eq!(
std::fs::metadata(&prefs).unwrap().permissions().mode() & 0o777,
0o600
);
}
let svc2 = WorktreesService::new();
svc2.load_polling_prefs(prefs.clone());
assert!(svc2.registry.is_polling_enabled("rust-works", "omni-dev"));
svc2.handle(
"set-polling",
json!({ "owner": "rust-works", "name": "omni-dev", "enabled": false }),
)
.await
.unwrap();
let svc3 = WorktreesService::new();
svc3.load_polling_prefs(prefs);
assert!(!svc3.registry.is_polling_enabled("rust-works", "omni-dev"));
}
#[test]
fn load_polling_prefs_tolerates_a_corrupt_or_unreadable_file() {
let dir = tempfile::tempdir().unwrap();
let corrupt = dir.path().join("worktrees-polling.json");
std::fs::write(&corrupt, b"{ not valid json ]").unwrap();
let svc = WorktreesService::new();
svc.load_polling_prefs(corrupt);
assert!(svc.registry.enabled_polling_repos().is_empty());
let as_dir = dir.path().join("is-a-directory");
std::fs::create_dir(&as_dir).unwrap();
let svc2 = WorktreesService::new();
svc2.load_polling_prefs(as_dir);
assert!(svc2.registry.enabled_polling_repos().is_empty());
}
#[tokio::test]
async fn pr_poller_asks_nothing_for_a_registered_but_not_enabled_repo() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let marker = bin_dir.path().join("spawned");
let fake = bin_dir.path().join("fake-gh");
std::fs::write(
&fake,
format!("#!/bin/sh\ntouch '{}'\necho '{{}}'\n", marker.display()),
)
.unwrap();
let mut perms = std::fs::metadata(&fake).unwrap().permissions();
std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o755);
std::fs::set_permissions(&fake, perms).unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.start_pr_poller_with(Duration::from_millis(20), Duration::from_millis(10), fake);
tokio::time::sleep(Duration::from_millis(200)).await;
svc.shutdown().await;
assert!(
!marker.exists(),
"a registered-but-not-enabled repo must drive zero gh"
);
}
#[tokio::test]
async fn menu_action_rejects_unknown_and_missing_window() {
let svc = WorktreesService::new();
assert!(svc.menu_action("bogus").await.is_err());
assert!(svc.menu_action("focus:nope").await.is_err());
svc.shutdown().await;
}
struct VscodeBinGuard(Option<std::ffi::OsString>);
impl Drop for VscodeBinGuard {
fn drop(&mut self) {
match self.0.take() {
Some(v) => std::env::set_var(VSCODE_BIN_ENV, v),
None => std::env::remove_var(VSCODE_BIN_ENV),
}
}
}
#[tokio::test]
async fn menu_action_focus_resolves_folder_and_spawns() {
let dir = tempfile::tempdir().unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w1", "folders": [dir.path()], "repo": "r" }),
)
.await
.unwrap();
let _g = VscodeBinGuard(std::env::var_os(VSCODE_BIN_ENV));
std::env::set_var(VSCODE_BIN_ENV, "/bin/sh");
svc.menu_action("focus:w1").await.unwrap();
}
#[tokio::test]
async fn open_rejects_missing_relative_or_nonexistent_path() {
let svc = WorktreesService::new();
assert!(svc.handle("open", json!({})).await.is_err());
assert!(svc.handle("open", json!({ "path": 42 })).await.is_err());
assert!(svc
.handle("open", json!({ "path": "relative/dir" }))
.await
.is_err());
assert!(svc
.handle("open", json!({ "path": "-flag" }))
.await
.is_err());
assert!(svc
.handle("open", json!({ "path": "/no/such/abs/dir/xyzzy" }))
.await
.is_err());
svc.shutdown().await;
}
#[tokio::test]
async fn open_focuses_an_existing_absolute_dir() {
let dir = tempfile::tempdir().unwrap();
let svc = WorktreesService::new();
let _g = VscodeBinGuard(std::env::var_os(VSCODE_BIN_ENV));
std::env::set_var(VSCODE_BIN_ENV, "/bin/sh");
let reply = svc
.handle("open", json!({ "path": dir.path() }))
.await
.unwrap();
assert_eq!(reply, json!({ "ok": true }));
svc.shutdown().await;
}
#[test]
fn focus_window_with_validates_folder_then_spawns() {
let dir = tempfile::tempdir().unwrap();
assert!(focus_window_with(Path::new("/bin/sh"), Path::new("relative/dir")).is_err());
assert!(
focus_window_with(Path::new("/bin/sh"), Path::new("/no/such/abs/dir/xyzzy")).is_err()
);
focus_window_with(Path::new("/bin/sh"), dir.path()).unwrap();
assert!(focus_window_with(Path::new("/no/such/launcher/xyzzy"), dir.path()).is_err());
}
#[test]
fn resolve_code_binary_from_prefers_env_then_candidate_then_fallback() {
assert_eq!(
resolve_code_binary_from(Some("/custom/code".into()), &["/usr/bin/code"]),
PathBuf::from("/custom/code")
);
let existing = tempfile::NamedTempFile::new().unwrap();
let existing_path = existing.path().to_str().unwrap();
assert_eq!(
resolve_code_binary_from(None, &["/no/such/candidate/xyzzy", existing_path]),
PathBuf::from(existing_path)
);
assert_eq!(
resolve_code_binary_from(None, &["/no/such/candidate/xyzzy"]),
PathBuf::from("code")
);
let _ = resolve_code_binary();
}
fn init_repo(dir: &Path) -> Repository {
let repo = Repository::init(dir).unwrap();
let mut cfg = repo.config().unwrap();
cfg.set_str("user.name", "Test").unwrap();
cfg.set_str("user.email", "test@example.com").unwrap();
repo
}
fn empty_commit(
repo: &Repository,
refname: Option<&str>,
parents: &[&git2::Commit<'_>],
msg: &str,
) -> git2::Oid {
let sig = git2::Signature::now("Test", "test@example.com").unwrap();
let tree = repo
.find_tree(repo.treebuilder(None).unwrap().write().unwrap())
.unwrap();
repo.commit(refname, &sig, &sig, msg, &tree, parents)
.unwrap()
}
fn commit_file(
repo: &Repository,
refname: &str,
name: &str,
content: &[u8],
msg: &str,
) -> git2::Oid {
let sig = git2::Signature::now("Test", "test@example.com").unwrap();
let blob = repo.blob(content).unwrap();
let mut builder = repo.treebuilder(None).unwrap();
builder.insert(name, blob, 0o100_644).unwrap();
let tree = repo.find_tree(builder.write().unwrap()).unwrap();
let parent = repo
.refname_to_id(refname)
.ok()
.and_then(|oid| repo.find_commit(oid).ok());
let parents: Vec<&git2::Commit<'_>> = parent.iter().collect();
repo.commit(Some(refname), &sig, &sig, msg, &tree, &parents)
.unwrap()
}
fn diverging_repo(dir: &Path) -> Repository {
let repo = init_repo(dir);
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
let a_commit = repo.find_commit(a).unwrap();
let c = empty_commit(&repo, None, &[&a_commit], "C");
repo.reference("refs/remotes/origin/main", c, true, "origin main")
.unwrap();
empty_commit(&repo, Some("refs/heads/main"), &[&a_commit], "B");
drop(a_commit);
repo.set_head("refs/heads/main").unwrap();
let mut cfg = repo.config().unwrap();
cfg.set_str("remote.origin.url", "https://example.invalid/x.git")
.unwrap();
cfg.set_str("remote.origin.fetch", "+refs/heads/*:refs/remotes/origin/*")
.unwrap();
cfg.set_str("branch.main.remote", "origin").unwrap();
cfg.set_str("branch.main.merge", "refs/heads/main").unwrap();
repo
}
#[test]
fn git_status_reads_branch_and_ahead_behind() {
let dir = tempfile::tempdir().unwrap();
let _repo = diverging_repo(dir.path());
let status = git_status(dir.path());
assert_eq!(status.branch.as_deref(), Some("main"));
assert_eq!(status.ahead, Some(1));
assert_eq!(status.behind, Some(1));
assert_eq!(
status.main_repo.as_deref(),
dir.path().file_name().and_then(|n| n.to_str())
);
assert!(!status.is_worktree);
}
#[test]
fn git_status_empty_repo_is_unborn() {
let dir = tempfile::tempdir().unwrap();
init_repo(dir.path());
let status = git_status(dir.path());
assert_eq!(status.branch, None);
assert_eq!(status.head_sha, None);
assert_eq!(status.ahead, None);
assert_eq!(status.behind, None);
assert_eq!(
status.main_repo.as_deref(),
dir.path().file_name().and_then(|n| n.to_str())
);
assert!(!status.is_worktree);
}
#[test]
fn git_status_no_upstream_reports_branch_only() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let status = git_status(dir.path());
assert_eq!(status.branch.as_deref(), Some("main"));
assert_eq!(status.ahead, None);
assert_eq!(status.behind, None);
assert_eq!(status.upstream_sha, None);
}
#[test]
fn git_status_non_repo_is_empty_detached_reports_repo_without_branch() {
let plain = tempfile::tempdir().unwrap();
assert_eq!(git_status(plain.path()), GitStatus::default());
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head_detached(a).unwrap();
let status = git_status(dir.path());
assert_eq!(status.branch, None);
assert_eq!(status.head_sha.as_deref(), Some(a.to_string().as_str()));
assert_eq!(status.ahead, None);
assert_eq!(status.behind, None);
assert_eq!(status.upstream_sha, None);
assert_eq!(
status.main_repo.as_deref(),
dir.path().file_name().and_then(|n| n.to_str())
);
assert!(!status.is_worktree);
}
#[test]
fn git_status_cheap_reads_branch_but_skips_the_divergence_walk() {
let dir = tempfile::tempdir().unwrap();
let repo = diverging_repo(dir.path());
let status = git_status_cheap(dir.path());
assert_eq!(status.branch.as_deref(), Some("main"));
assert_eq!(status.ahead, None);
assert_eq!(status.behind, None);
assert_eq!(
status.main_repo.as_deref(),
dir.path().file_name().and_then(|n| n.to_str())
);
let head = repo.head().unwrap().target().unwrap();
assert_eq!(status.head_sha.as_deref(), Some(head.to_string().as_str()));
}
#[test]
fn git_status_head_sha_tracks_new_commits() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let before = git_status_cheap(dir.path());
assert_eq!(before.head_sha.as_deref(), Some(a.to_string().as_str()));
let head = repo.find_commit(a).unwrap();
let b = empty_commit(&repo, Some("refs/heads/main"), &[&head], "B");
let after = git_status_cheap(dir.path());
assert_eq!(after.head_sha.as_deref(), Some(b.to_string().as_str()));
assert_ne!(before.head_sha, after.head_sha);
assert_eq!(before.branch, after.branch);
}
fn simulate_push(repo: &Repository, oid: git2::Oid) {
repo.reference("refs/remotes/origin/main", oid, true, "push")
.unwrap();
}
#[test]
fn git_status_upstream_sha_tracks_a_push() {
let dir = tempfile::tempdir().unwrap();
let repo = diverging_repo(dir.path());
let before = git_status(dir.path());
assert_eq!(before.ahead, Some(1));
assert_eq!(before.behind, Some(1));
let head = repo.head().unwrap().target().unwrap();
simulate_push(&repo, head);
let after = git_status(dir.path());
assert_eq!(
after.upstream_sha.as_deref(),
Some(head.to_string().as_str())
);
assert_ne!(before.upstream_sha, after.upstream_sha);
assert_eq!(after.ahead, Some(0));
assert_eq!(after.behind, Some(0));
assert_eq!(before.branch, after.branch);
assert_eq!(before.head_sha, after.head_sha);
}
#[test]
fn git_status_cheap_reports_upstream_sha() {
let dir = tempfile::tempdir().unwrap();
let repo = diverging_repo(dir.path());
let status = git_status_cheap(dir.path());
let upstream = repo
.find_branch("origin/main", git2::BranchType::Remote)
.unwrap()
.get()
.target()
.unwrap();
assert_eq!(
status.upstream_sha.as_deref(),
Some(upstream.to_string().as_str())
);
assert_eq!(status.ahead, None);
assert_eq!(status.behind, None);
}
#[test]
fn folder_ahead_behind_computes_divergence_and_degrades() {
let dir = tempfile::tempdir().unwrap();
let _repo = diverging_repo(dir.path());
assert_eq!(folder_ahead_behind(dir.path()), Some((1, 1)));
let no_up = tempfile::tempdir().unwrap();
let repo = init_repo(no_up.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
assert_eq!(folder_ahead_behind(no_up.path()), None);
let detached = tempfile::tempdir().unwrap();
let drepo = init_repo(detached.path());
let a = empty_commit(&drepo, Some("refs/heads/main"), &[], "A");
drepo.set_head_detached(a).unwrap();
assert_eq!(folder_ahead_behind(detached.path()), None);
let plain = tempfile::tempdir().unwrap();
assert_eq!(folder_ahead_behind(plain.path()), None);
}
#[tokio::test]
async fn ahead_behind_op_returns_divergence_keyed_by_path_and_omits_no_upstream() {
let diverging = tempfile::tempdir().unwrap();
let _d = diverging_repo(diverging.path());
let no_up = tempfile::tempdir().unwrap();
let repo = init_repo(no_up.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let svc = WorktreesService::new();
let diverging_path = diverging.path().display().to_string();
let no_up_path = no_up.path().display().to_string();
let reply = svc
.handle(
"ahead-behind",
json!({ "paths": [&diverging_path, &no_up_path] }),
)
.await
.unwrap();
let results = reply.get("results").unwrap();
let d = results.get(diverging_path.as_str()).unwrap();
assert_eq!(d.get("ahead").and_then(Value::as_u64), Some(1));
assert_eq!(d.get("behind").and_then(Value::as_u64), Some(1));
assert!(results.get(no_up_path.as_str()).is_none(), "{results:?}");
let empty = svc.handle("ahead-behind", json!({})).await.unwrap();
assert_eq!(empty.get("results"), Some(&json!({})));
}
#[tokio::test]
async fn tree_snapshot_omits_ahead_behind_for_a_diverging_worktree() {
let dir = tempfile::tempdir().unwrap();
let _repo = diverging_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "x" }),
)
.await
.unwrap();
let repos = repos_of(&svc.handle("tree", Value::Null).await.unwrap());
let worktrees = repos[0].get("worktrees").and_then(Value::as_array).unwrap();
let main_wt = &worktrees[0];
assert_eq!(main_wt.get("branch").and_then(Value::as_str), Some("main"));
assert!(main_wt.get("ahead").is_none(), "{main_wt:?}");
assert!(main_wt.get("behind").is_none(), "{main_wt:?}");
}
#[tokio::test]
async fn tree_snapshot_carries_head_sha_so_a_commit_is_a_real_delta() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "x" }),
)
.await
.unwrap();
let before = svc.handle("tree", Value::Null).await.unwrap();
let wt = &repos_of(&before)[0]["worktrees"][0];
assert_eq!(
wt.get("head_sha").and_then(Value::as_str),
Some(a.to_string().as_str())
);
let head = repo.find_commit(a).unwrap();
let b = empty_commit(&repo, Some("refs/heads/main"), &[&head], "B");
let after = svc.handle("tree", Value::Null).await.unwrap();
assert_eq!(
repos_of(&after)[0]["worktrees"][0]
.get("head_sha")
.and_then(Value::as_str),
Some(b.to_string().as_str())
);
assert_ne!(before, after, "a commit must be a visible snapshot delta");
}
#[tokio::test]
async fn tree_snapshot_omits_head_sha_for_an_unborn_repo() {
let dir = tempfile::tempdir().unwrap();
init_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "x" }),
)
.await
.unwrap();
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(wt.get("head_sha").is_none(), "{wt:?}");
}
#[tokio::test]
async fn tree_snapshot_carries_upstream_sha_so_a_push_is_a_real_delta() {
let dir = tempfile::tempdir().unwrap();
let repo = diverging_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "x" }),
)
.await
.unwrap();
let before = svc.handle("tree", Value::Null).await.unwrap();
let head = repo.head().unwrap().target().unwrap();
assert_ne!(
repos_of(&before)[0]["worktrees"][0]
.get("upstream_sha")
.and_then(Value::as_str),
Some(head.to_string().as_str()),
"the fixture must start un-pushed for this to prove anything"
);
simulate_push(&repo, head);
let after = svc.handle("tree", Value::Null).await.unwrap();
let wt = &repos_of(&after)[0]["worktrees"][0];
assert_eq!(
wt.get("upstream_sha").and_then(Value::as_str),
Some(head.to_string().as_str())
);
assert_eq!(
wt.get("head_sha").and_then(Value::as_str),
repos_of(&before)[0]["worktrees"][0]
.get("head_sha")
.and_then(Value::as_str)
);
assert_ne!(before, after, "a push must be a visible snapshot delta");
}
#[tokio::test]
async fn tree_snapshot_omits_upstream_sha_without_an_upstream() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "x" }),
)
.await
.unwrap();
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(wt.get("upstream_sha").is_none(), "{wt:?}");
assert!(wt.get("head_sha").is_some(), "{wt:?}");
}
fn fake_gh(dir: &Path, stdout: &str) -> (PathBuf, MutexGuard<'static, ()>) {
let guard = shim_lock();
let path = dir.join("fake-gh");
write_exec_script(&path, &format!("#!/bin/sh\ncat <<'JSON'\n{stdout}\nJSON\n"));
(path, guard)
}
fn counting_fake_gh(dir: &Path, stdout: &str) -> (PathBuf, MutexGuard<'static, ()>, PathBuf) {
let guard = shim_lock();
let path = dir.join("fake-gh");
let counter = dir.join("gh-calls");
write_exec_script(
&path,
&format!(
"#!/bin/sh\nprintf x >> {counter:?}\ncat <<'JSON'\n{stdout}\nJSON\n",
counter = counter.display()
),
);
(path, guard, counter)
}
fn gh_spawn_count(counter: &Path) -> usize {
std::fs::read(counter).map_or(0, |b| b.len())
}
fn counted_gh_records(log: &Path) -> usize {
std::fs::read_to_string(log)
.unwrap_or_default()
.lines()
.filter(|l| l.contains(r#""kind":"gh""#) && l.contains(r#""exit_code":0"#))
.count()
}
fn github_repo(dir: &Path) -> Repository {
github_repo_with_remote(dir, "git@github.com:rust-works/omni-dev.git")
}
fn github_repo_with_remote(dir: &Path, url: &str) -> Repository {
let repo = init_repo(dir);
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
repo.remote("origin", url).unwrap();
repo
}
fn pending_badge(number: u64, head_oid: &str) -> PrBadge {
PrBadge {
number,
is_draft: false,
checks: PrCheckState::Pending,
url: "u".into(),
head_oid: head_oid.to_string(),
}
}
fn pr(badge: PrBadge) -> PrResolution {
PrResolution::Pr(badge)
}
#[test]
fn pr_targets_from_snapshot_reads_github_branches_and_dedupes() {
let snapshot = json!({"repos":[
{
"main_repo":"omni-dev",
"github":{"owner":"rust-works","name":"omni-dev"},
"root":"/r",
"polling_enabled":true,
"worktrees":[
{"path":"/r","branch":"main","is_main":true,"open":true},
{"path":"/w1","branch":"main","is_main":false,"open":true},
{"path":"/w2","branch":"feature","is_main":false,"open":true},
{"path":"/w3","is_main":false,"open":true}
]
},
{
"main_repo":"local","root":"/l",
"worktrees":[{"path":"/l","branch":"main","is_main":true,"open":true}]
}
]});
let targets = pr_targets_from_snapshot(&snapshot);
assert_eq!(
targets,
vec![
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "feature".into()
},
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into()
},
]
);
}
#[test]
fn pr_targets_from_snapshot_is_empty_without_repos() {
assert!(pr_targets_from_snapshot(&json!({"repos":[]})).is_empty());
assert!(pr_targets_from_snapshot(&json!({})).is_empty());
}
#[test]
fn pr_targets_from_snapshot_skips_a_malformed_github_identity() {
for github in [
json!({}),
json!({"owner": "o"}),
json!({"owner": 1, "name": 2}),
] {
let snapshot = json!({"repos":[{
"main_repo":"r","github":github,"root":"/r","polling_enabled":true,
"worktrees":[{"path":"/r","branch":"main","is_main":true,"open":true}]
}]});
assert!(
pr_targets_from_snapshot(&snapshot).is_empty(),
"{snapshot:?}"
);
}
}
#[test]
fn pr_watch_from_snapshot_skips_a_not_polled_repo() {
for repo in [
json!({
"main_repo":"omni-dev","github":{"owner":"o","name":"n"},"root":"/r",
"worktrees":[{"path":"/r","branch":"main","is_main":true,"open":true}]
}),
json!({
"main_repo":"omni-dev","github":{"owner":"o","name":"n"},"root":"/r",
"polling_enabled":false,
"worktrees":[{"path":"/r","branch":"main","is_main":true,"open":true}]
}),
] {
let snapshot = json!({ "repos": [repo] });
assert!(
pr_targets_from_snapshot(&snapshot).is_empty(),
"not-polled repo must yield no targets: {snapshot:?}"
);
}
let enabled = json!({"repos":[{
"main_repo":"omni-dev","github":{"owner":"o","name":"n"},"root":"/r",
"polling_enabled":true,
"worktrees":[{"path":"/r","branch":"main","is_main":true,"open":true}]
}]});
assert_eq!(pr_targets_from_snapshot(&enabled).len(), 1);
}
#[test]
fn pr_should_fetch_when_the_watch_grew_or_the_backoff_elapsed() {
let backoff = Duration::from_secs(600);
assert!(pr_should_fetch(false, None, backoff));
assert!(!pr_should_fetch(
false,
Some(Duration::from_secs(1)),
backoff
));
assert!(pr_should_fetch(false, Some(backoff), backoff));
assert!(pr_should_fetch(false, Some(backoff * 2), backoff));
assert!(pr_should_fetch(true, Some(Duration::ZERO), backoff));
assert!(pr_should_fetch(
true,
Some(Duration::from_millis(1)),
backoff
));
}
#[test]
fn next_pr_poll_delay_escalates_within_pending_and_backs_off_when_terminal() {
let base = Duration::from_secs(10);
let fresh = Some(Duration::ZERO);
let stale = Some(PENDING_FAST_WINDOW);
assert_eq!(next_pr_poll_delay(base, base, true, fresh), base);
assert_eq!(
next_pr_poll_delay(PENDING_MAX_INTERVAL, base, true, fresh),
base
);
assert_eq!(next_pr_poll_delay(base, base, true, stale), base * 2);
assert_eq!(
next_pr_poll_delay(PENDING_MAX_INTERVAL, base, true, stale),
PENDING_MAX_INTERVAL
);
assert_eq!(next_pr_poll_delay(base, base, true, None), base * 2);
assert_eq!(next_pr_poll_delay(base, base, false, fresh), base * 2);
assert_eq!(next_pr_poll_delay(base * 2, base, false, fresh), base * 4);
assert_eq!(
next_pr_poll_delay(MAX_PR_POLL_INTERVAL, base, false, fresh),
MAX_PR_POLL_INTERVAL
);
assert_eq!(
next_pr_poll_delay(Duration::MAX, base, false, None),
MAX_PR_POLL_INTERVAL
);
}
fn watch(branch: &str, upstream: Option<&str>) -> PrWatch {
PrWatch {
target: PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: branch.into(),
},
upstream_sha: upstream.map(str::to_string),
}
}
#[test]
fn pr_watch_grew_fires_on_additions_and_pushes_but_never_on_removals() {
let a = watch("a", Some("111"));
let b = watch("b", Some("222"));
let ab = [a.clone(), b.clone()];
let just_a = std::slice::from_ref(&a);
let just_b = std::slice::from_ref(&b);
assert!(!pr_watch_grew(&ab, &ab));
assert!(pr_watch_grew(just_a, &ab));
assert!(!pr_watch_grew(&ab, just_a));
assert!(pr_watch_grew(&[], just_a));
let a_pushed = [watch("a", Some("999"))];
assert!(pr_watch_grew(just_a, &a_pushed));
assert!(pr_watch_grew(just_a, just_b));
}
#[test]
fn budget_throttled_delay_holds_the_floor_only_when_over_warn() {
let base = Duration::from_secs(10);
let over = RateLimitSnapshot {
graphql: Some(rl_resource(90)),
core: Some(rl_resource(3)),
search: None,
};
let under = RateLimitSnapshot {
graphql: Some(rl_resource(50)),
core: Some(rl_resource(3)),
search: None,
};
assert_eq!(budget_throttled_delay(base, None), base);
assert_eq!(budget_throttled_delay(base, Some(&under)), base);
assert_eq!(
budget_throttled_delay(base, Some(&over)),
BUDGET_THROTTLE_INTERVAL
);
let long = BUDGET_THROTTLE_INTERVAL * 2;
assert_eq!(budget_throttled_delay(long, Some(&over)), long);
}
#[test]
fn pr_cache_prefs_round_trips_through_json_including_head_oid() {
let target = PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
};
let badge = PrResolution::Pr(PrBadge {
number: 1337,
is_draft: true,
checks: PrCheckState::Pending,
url: "http://x/1337".into(),
head_oid: "deadbeef".into(),
});
let watched = vec![watch("main", Some("abc"))];
let polled_at = DateTime::parse_from_rfc3339("2026-07-21T00:00:00Z")
.unwrap()
.with_timezone(&Utc);
let prefs = pr_cache_prefs_from(vec![(target, badge.clone())], &watched, polled_at);
let json = serde_json::to_vec(&prefs).unwrap();
let back: PrCachePrefs = serde_json::from_slice(&json).unwrap();
assert_eq!(back, prefs);
assert_eq!(back.polled_at, Some(polled_at));
assert_eq!(back.watched[0].upstream_sha.as_deref(), Some("abc"));
assert_eq!(back.entries[0].resolution.clone().into_resolution(), badge);
}
#[test]
fn pr_cache_prefs_round_trip_an_explicit_no_pr_verdict() {
let target = PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "feature".into(),
};
let prefs = pr_cache_prefs_from(vec![(target, PrResolution::NoPr)], &[], Utc::now());
let json = serde_json::to_vec(&prefs).unwrap();
let back: PrCachePrefs = serde_json::from_slice(&json).unwrap();
assert_eq!(back.entries[0].resolution, PersistedResolution::NoPr);
assert_eq!(
back.entries[0].resolution.clone().into_resolution(),
PrResolution::NoPr
);
}
#[test]
fn load_pr_cache_without_polled_at_restores_badges_but_no_warm_start() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("pr-cache.json");
let target = PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
};
let mut prefs = pr_cache_prefs_from(
vec![(target, PrResolution::Pr(pending_badge(7, "abc")))],
&[watch("main", None)],
Utc::now(),
);
prefs.polled_at = None;
write_pr_cache(&path, &prefs).unwrap();
let svc = WorktreesService::new();
svc.load_pr_cache(path);
assert!(
svc.pr_cache.get("rust-works", "omni-dev", "main").is_some(),
"the badge itself must still restore"
);
assert!(
svc.pr_warm_start
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none(),
"no poll time means a cold start, not a trusted warm one"
);
}
fn warn_subscriber() -> tracing::subscriber::DefaultGuard {
tracing::subscriber::set_default(
tracing_subscriber::fmt()
.with_max_level(tracing::Level::WARN)
.with_writer(std::io::sink)
.finish(),
)
}
#[test]
fn load_pr_cache_tolerates_a_corrupt_or_unreadable_file() {
let _trace = warn_subscriber();
let dir = tempfile::tempdir().unwrap();
let corrupt = dir.path().join("pr-cache.json");
std::fs::write(&corrupt, b"not json").unwrap();
let svc = WorktreesService::new();
svc.load_pr_cache(corrupt.clone());
assert!(svc.pr_cache.entries().is_empty());
assert_eq!(
svc.pr_cache_path
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_deref(),
Some(corrupt.as_path()),
"the path must be stored even when the load fails, so persistence recovers"
);
let svc = WorktreesService::new();
svc.load_pr_cache(dir.path().to_path_buf());
assert!(svc.pr_cache.entries().is_empty());
}
#[test]
fn persist_pr_cache_swallows_a_write_failure() {
let _trace = warn_subscriber();
let dir = tempfile::tempdir().unwrap();
let blocker = dir.path().join("blocker");
std::fs::write(&blocker, b"").unwrap();
let path = blocker.join("pr-cache.json");
persist_pr_cache(&path, &PrStatusCache::new(), &[], Utc::now());
assert!(!path.exists());
persist_pr_cache(Path::new("/"), &PrStatusCache::new(), &[], Utc::now());
}
#[tokio::test]
async fn tree_snapshot_folds_cached_pr_badges_onto_matching_branches() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap().to_string();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(wt.get("pr").is_none(), "{wt:?}");
assert!(wt.get("pr_none").is_none(), "{wt:?}");
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(pending_badge(1337, &head)),
);
assert!(svc.pr_cache.replace(badges));
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr"]["number"], json!(1337));
assert_eq!(wt["pr"]["checks"], json!("pending"));
assert_eq!(wt["pr"]["isDraft"], json!(false));
assert!(wt.get("pr_none").is_none(), "{wt:?}");
}
#[tokio::test]
async fn tree_snapshot_omits_a_badge_for_a_detached_worktree() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap();
repo.set_head_detached(head).unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(pending_badge(1, &head.to_string())),
);
svc.pr_cache.replace(badges);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(wt.get("branch").is_none(), "{wt:?}");
assert_eq!(
wt.get("head_sha").and_then(Value::as_str),
Some(head.to_string().as_str())
);
assert!(wt.get("pr").is_none(), "{wt:?}");
assert!(wt.get("pr_none").is_none(), "{wt:?}");
}
#[tokio::test]
async fn tree_snapshot_omits_a_badge_for_an_unmatched_branch() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "other".into(),
},
pr(pending_badge(1, "irrelevant")),
);
svc.pr_cache.replace(badges);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert!(wt.get("pr").is_none(), "{wt:?}");
assert!(wt.get("pr_none").is_none(), "{wt:?}");
}
#[tokio::test]
async fn tree_snapshot_reports_an_explicit_negative_for_a_branch_with_no_pr() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut resolutions = HashMap::new();
resolutions.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
PrResolution::NoPr,
);
assert!(svc.pr_cache.replace(resolutions));
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr_none"], json!(true));
assert!(wt.get("pr").is_none(), "{wt:?}");
}
#[tokio::test]
async fn a_commit_does_not_drop_a_negative_resolution() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let first = repo.head().unwrap().target().unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut resolutions = HashMap::new();
resolutions.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
PrResolution::NoPr,
);
svc.pr_cache.replace(resolutions);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr_none"], json!(true));
let head = repo.find_commit(first).unwrap();
empty_commit(&repo, Some("refs/heads/main"), &[&head], "B");
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(
wt["pr_none"],
json!(true),
"a local commit must not drop the negative"
);
}
#[tokio::test]
async fn pr_poller_asks_nothing_while_no_window_is_registered() {
let bin_dir = tempfile::tempdir().unwrap();
let marker = bin_dir.path().join("spawned");
let fake = bin_dir.path().join("fake-gh");
std::fs::write(
&fake,
format!("#!/bin/sh\ntouch '{}'\necho '{{}}'\n", marker.display()),
)
.unwrap();
let mut perms = std::fs::metadata(&fake).unwrap().permissions();
std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o755);
std::fs::set_permissions(&fake, perms).unwrap();
let svc = WorktreesService::new();
svc.start_pr_poller_with(Duration::from_millis(20), Duration::from_millis(10), fake);
tokio::time::sleep(Duration::from_millis(200)).await;
svc.shutdown().await;
assert!(
!marker.exists(),
"the poller must not spawn gh with no windows registered"
);
}
#[tokio::test]
async fn pr_poller_survives_a_failing_gh_and_keeps_the_last_good_badges() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap().to_string();
let bin_dir = tempfile::tempdir().unwrap();
let fake = bin_dir.path().join("fake-gh");
std::fs::write(
&fake,
"#!/bin/sh\necho 'gh: not authenticated' >&2\nexit 1\n",
)
.unwrap();
let mut perms = std::fs::metadata(&fake).unwrap().permissions();
std::os::unix::fs::PermissionsExt::set_mode(&mut perms, 0o755);
std::fs::set_permissions(&fake, perms).unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut seeded = HashMap::new();
seeded.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(pending_badge(7, &head)),
);
svc.pr_cache.replace(seeded);
svc.start_pr_poller_with(Duration::from_millis(20), Duration::from_millis(10), fake);
tokio::time::sleep(Duration::from_millis(200)).await;
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr"]["number"], json!(7));
assert!(wt.get("pr_none").is_none(), "{wt:?}");
svc.shutdown().await;
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_wakes_when_the_first_window_opens_after_an_idle_start() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"a","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"IN_PROGRESS","conclusion":null}
]}}},
"associatedPullRequests":{"nodes":[{"number":99,"isDraft":false,"url":"u"}]}
}}}}"#,
);
let svc = WorktreesService::new();
svc.start_pr_poller_with(Duration::from_millis(50), Duration::from_millis(10), fake);
tokio::time::sleep(Duration::from_millis(150)).await;
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let badge = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Some(PrResolution::Pr(badge)) =
svc.pr_cache.get("rust-works", "omni-dev", "main")
{
return badge;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("a window opening must wake the poller out of its idle backoff");
assert_eq!(badge.number, 99);
svc.shutdown().await;
}
#[tokio::test]
async fn a_commit_invalidates_the_previous_verdict_without_a_poll() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let first = repo.head().unwrap().target().unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
let mut badges = HashMap::new();
badges.insert(
PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
},
pr(PrBadge {
number: 1337,
is_draft: false,
checks: PrCheckState::Success,
url: "u".into(),
head_oid: first.to_string(),
}),
);
svc.pr_cache.replace(badges);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(
wt["pr"]["checks"],
json!("success"),
"green for its own commit"
);
let head = repo.find_commit(first).unwrap();
empty_commit(&repo, Some("refs/heads/main"), &[&head], "B");
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(
wt["pr"]["checks"],
json!("pending"),
"the previous commit's ✓ must not stand after a new commit"
);
assert_eq!(wt["pr"]["number"], json!(1337));
}
#[test]
fn is_stale_for_compares_the_commit_the_verdict_describes() {
let badge = pending_badge(1, "aaa");
assert!(!badge.is_stale_for(Some("aaa")));
assert!(badge.is_stale_for(Some("bbb")));
assert!(!badge.is_stale_for(None));
}
#[test]
fn pr_watch_ignores_the_head_so_a_local_commit_asks_nothing() {
let snap = |sha: &str| {
json!({"repos":[{
"main_repo":"omni-dev",
"github":{"owner":"rust-works","name":"omni-dev"},
"root":"/r",
"polling_enabled":true,
"worktrees":[{"path":"/r","branch":"main","head_sha":sha,"is_main":true,"open":true}]
}]})
};
let before = pr_watch_from_snapshot(&snap("aaa"));
let after = pr_watch_from_snapshot(&snap("bbb"));
assert_eq!(before.len(), 1);
assert_eq!(before[0].target, after[0].target);
assert_eq!(before, after);
assert!(!pr_watch_grew(&before, &after));
}
#[test]
fn pr_watch_tracks_the_upstream_so_a_push_is_visible_to_the_poller() {
let snap = |upstream: &str| {
json!({"repos":[{
"main_repo":"omni-dev",
"github":{"owner":"rust-works","name":"omni-dev"},
"root":"/r",
"polling_enabled":true,
"worktrees":[{"path":"/r","branch":"main","head_sha":"aaa",
"upstream_sha":upstream,"is_main":true,"open":true}]
}]})
};
let before = pr_watch_from_snapshot(&snap("aaa"));
let after = pr_watch_from_snapshot(&snap("bbb"));
assert_eq!(before.len(), 1);
assert_eq!(before[0].target, after[0].target);
assert_ne!(before, after);
assert!(pr_watch_grew(&before, &after));
assert_eq!(before, pr_watch_from_snapshot(&snap("aaa")));
assert!(!pr_watch_grew(
&before,
&pr_watch_from_snapshot(&snap("aaa"))
));
}
#[test]
fn pr_watch_omits_an_absent_upstream_rather_than_erroring() {
let snap = json!({"repos":[{
"main_repo":"omni-dev",
"github":{"owner":"rust-works","name":"omni-dev"},
"root":"/r",
"polling_enabled":true,
"worktrees":[{"path":"/r","branch":"main","head_sha":"aaa","is_main":true,"open":true}]
}]});
let watch = pr_watch_from_snapshot(&snap);
assert_eq!(watch.len(), 1);
assert_eq!(watch[0].upstream_sha, None);
}
#[test]
fn start_pr_poller_is_a_noop_outside_a_runtime() {
let svc = WorktreesService::new();
svc.start_pr_poller();
assert!(svc
.poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none());
}
#[tokio::test]
async fn start_pr_poller_is_idempotent_and_shutdown_stops_it() {
let svc = WorktreesService::new();
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
PathBuf::from("/bin/true"),
);
let token = svc
.poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.map(|t| t.token.clone())
.expect("poller started");
token.cancel();
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
PathBuf::from("/bin/true"),
);
assert!(svc
.poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.is_some_and(|t| t.token.is_cancelled()));
svc.shutdown().await;
assert!(svc
.poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none());
}
fn rl_resource(used: u64) -> RateLimitResource {
RateLimitResource {
used,
limit: 100,
remaining: 100 - used,
percent: used as f64,
reset: 0,
}
}
#[test]
fn rate_limit_crossed_warn_fires_only_on_the_rising_edge() {
let snap = |graphql: u64, core: u64| RateLimitSnapshot {
graphql: Some(rl_resource(graphql)),
core: Some(rl_resource(core)),
search: None,
};
assert!(rate_limit_crossed_warn(None, &snap(85, 3)));
assert!(!rate_limit_crossed_warn(None, &snap(50, 3)));
assert!(rate_limit_crossed_warn(Some(&snap(70, 3)), &snap(85, 3)));
assert!(!rate_limit_crossed_warn(Some(&snap(85, 3)), &snap(90, 3)));
assert!(!rate_limit_crossed_warn(Some(&snap(85, 3)), &snap(50, 3)));
assert!(rate_limit_crossed_warn(Some(&snap(85, 3)), &snap(50, 90)));
}
#[test]
fn start_rate_limit_poller_is_a_noop_outside_a_runtime() {
let svc = WorktreesService::new();
svc.start_rate_limit_poller();
assert!(svc
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none());
}
#[tokio::test]
async fn start_rate_limit_poller_is_idempotent_and_shutdown_stops_it() {
let svc = WorktreesService::new();
svc.start_rate_limit_poller_with(Duration::from_millis(50), PathBuf::from("/bin/true"));
let token = svc
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.map(|t| t.token.clone())
.expect("poller started");
token.cancel();
svc.start_rate_limit_poller_with(Duration::from_millis(50), PathBuf::from("/bin/true"));
assert!(svc
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.is_some_and(|t| t.token.is_cancelled()));
svc.shutdown().await;
assert!(svc
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none());
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn rate_limit_poller_resolves_via_gh_populates_the_cache_and_stops_on_shutdown() {
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"resources":{
"graphql":{"limit":5000,"used":4100,"remaining":900,"reset":1700000000},
"core":{"limit":5000,"used":27,"remaining":4973,"reset":1700000000}
}}"#,
);
let svc = WorktreesService::new();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_rate_limit_poller_with(Duration::from_millis(50), fake.clone());
let snap = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Some(snap) = svc.rate_limit_cache.get() {
return snap;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("poller should populate the cache through the fake gh");
assert_eq!(snap.graphql.unwrap().used, 4100);
assert_eq!(snap.core.unwrap().used, 27);
assert!(svc.rate_limit_cache().get().is_some());
svc.shutdown().await;
assert!(svc
.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_none());
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn rate_limit_poller_stays_idle_with_nothing_registered() {
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(
bin_dir.path(),
r#"{"resources":{"graphql":{"limit":5000,"used":1,"remaining":4999,"reset":1}}}"#,
);
let svc = WorktreesService::new();
svc.start_rate_limit_poller_with(Duration::from_millis(20), fake);
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(
gh_spawn_count(&counter),
0,
"idle daemon must not poll (#1389, fix 8b)"
);
assert!(svc.rate_limit_cache.get().is_none());
svc.registry.set_polling("rust-works", "omni-dev", true);
tokio::time::timeout(Duration::from_secs(30), async {
loop {
if svc.rate_limit_cache.get().is_some() {
return;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("an active lease should resume polling");
assert!(gh_spawn_count(&counter) >= 1);
svc.shutdown().await;
}
#[tokio::test]
async fn rate_limit_poller_survives_a_failing_gh() {
let svc = WorktreesService::new();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_rate_limit_poller_with(
Duration::from_millis(20),
PathBuf::from("/no/such/gh/xyzzy"),
);
tokio::time::sleep(Duration::from_millis(150)).await;
assert!(svc.rate_limit_cache.get().is_none());
assert!(
svc.rate_limit_poller
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_some(),
"the loop must survive a failing gh, not panic out"
);
svc.shutdown().await;
}
#[test]
fn menu_prepends_the_rate_limit_line_only_when_the_cache_is_populated() {
let svc = WorktreesService::new();
let items = svc.menu().items;
assert!(
!items
.iter()
.any(|i| matches!(i, MenuItem::Label(l) if l.contains("github:"))),
"no github line before the first poll"
);
svc.rate_limit_cache.replace(RateLimitSnapshot {
graphql: Some(rl_resource(82)),
core: Some(rl_resource(3)),
search: None,
});
let items = svc.menu().items;
assert!(
matches!(items.first(), Some(MenuItem::Label(l)) if l.starts_with("github: graphql 82%")),
"expected the github line first, got {items:?}"
);
assert!(
matches!(items.get(1), Some(MenuItem::Separator)),
"expected a separator after the github line"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_resolves_via_gh_populates_the_cache_and_stops_on_shutdown() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"IN_PROGRESS","conclusion":null}
]}}},
"associatedPullRequests":{"nodes":[{"number":1337,"isDraft":false,"url":"http://x/1337"}]}
}}}}"#,
);
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
fake.clone(),
);
let badge = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Some(PrResolution::Pr(badge)) =
svc.pr_cache.get("rust-works", "omni-dev", "main")
{
return badge;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("poller should resolve a badge through the fake gh");
assert_eq!(badge.number, 1337);
assert_eq!(badge.checks, crate::pr_status::PrCheckState::Pending);
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr"]["number"], json!(1337));
svc.shutdown().await;
let generation = svc.registry.change_generation();
tokio::time::sleep(Duration::from_millis(120)).await;
assert_eq!(
svc.registry.change_generation(),
generation,
"no bumps after shutdown"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poll_folds_its_graphql_budget_into_the_rate_limit_cache() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{
"rateLimit":{"limit":5000,"cost":1,"remaining":4877,"used":123,
"resetAt":"2026-07-21T16:00:00Z"},
"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"IN_PROGRESS","conclusion":null}
]}}},
"associatedPullRequests":{"nodes":[{"number":1337,"isDraft":false,"url":"http://x/1337"}]}
}}
}}"#,
);
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
fake.clone(),
);
let graphql = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Some(g) = svc.rate_limit_cache.get().and_then(|s| s.graphql) {
return g;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("the PR poll should fold its budget into the cache");
assert_eq!(graphql.used, 123);
assert_eq!(graphql.limit, 5000);
assert_eq!(graphql.remaining, 4877);
svc.shutdown().await;
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poll_counts_every_gh_call_exactly_once() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"IN_PROGRESS","conclusion":null}
]}}},
"associatedPullRequests":{"nodes":[{"number":1337,"isDraft":false,"url":"http://x/1337"}]}
}}}}"#,
);
let log = bin_dir.path().join("log.jsonl");
std::env::set_var("OMNI_DEV_LOG_FILE", &log);
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(Duration::from_millis(30), Duration::from_millis(10), fake);
tokio::time::timeout(Duration::from_secs(30), async {
loop {
if svc.pr_cache.get("rust-works", "omni-dev", "main").is_some() {
return;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("poller should fetch through the fake gh");
svc.shutdown().await;
let spawns = gh_spawn_count(&counter);
let counted = counted_gh_records(&log);
std::env::remove_var("OMNI_DEV_LOG_FILE");
assert!(
spawns >= 1,
"the poll should have spent at least one gh call"
);
assert_eq!(
counted, spawns,
"#1387: every gh call ({spawns}) must be counted exactly once, got {counted}"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poll_debounces_a_registration_storm_into_one_fetch() {
let dir_a = tempfile::tempdir().unwrap();
let dir_b = tempfile::tempdir().unwrap();
github_repo(dir_a.path()); github_repo_with_remote(dir_b.path(), "git@github.com:rust-works/other-repo.git"); let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(
bin_dir.path(),
r#"{"data":{
"r0":{"b0":{"target":{"oid":"a","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"COMPLETED","conclusion":"SUCCESS"}]}}},
"associatedPullRequests":{"nodes":[{"number":1,"isDraft":false,"url":"http://x/1"}]}}},
"r1":{"b0":{"target":{"oid":"b","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"COMPLETED","conclusion":"SUCCESS"}]}}},
"associatedPullRequests":{"nodes":[{"number":2,"isDraft":false,"url":"http://x/2"}]}}}
}}"#,
);
let svc = WorktreesService::new();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.registry.set_polling("rust-works", "other-repo", true);
svc.start_pr_poller_with(Duration::from_secs(30), Duration::from_millis(200), fake);
svc.handle(
"register",
json!({ "key": "a", "folders": [dir_a.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
svc.handle(
"register",
json!({ "key": "b", "folders": [dir_b.path()], "repo": "other-repo" }),
)
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(30), async {
loop {
let a = svc.pr_cache.get("rust-works", "omni-dev", "main").is_some();
let b = svc
.pr_cache
.get("rust-works", "other-repo", "main")
.is_some();
if a && b {
return;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("the debounced fetch should resolve both repos");
svc.shutdown().await;
assert_eq!(
gh_spawn_count(&counter),
1,
"the registration storm must collapse into exactly one fetch (#1389, fix 2)"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poll_debounce_deadline_bounds_a_steady_drip_of_changes() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(bin_dir.path(), "{}");
let svc = WorktreesService::new();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(Duration::from_secs(30), Duration::from_millis(50), fake);
let register = json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" });
svc.handle("register", register.clone()).await.unwrap();
for _ in 0..24 {
tokio::time::sleep(Duration::from_millis(25)).await;
svc.handle("register", register.clone()).await.unwrap();
}
let spawned_mid_drip = gh_spawn_count(&counter);
svc.shutdown().await;
assert!(
spawned_mid_drip >= 1,
"the deadline must force a fetch while the drip is still running (#1389, fix 2)"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_skips_the_immediate_fetch_when_the_warm_cache_is_fresh() {
let dir = tempfile::tempdir().unwrap();
let repo = github_repo(dir.path());
let head = repo.head().unwrap().target().unwrap().to_string();
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(bin_dir.path(), "{}");
let cache_path = bin_dir.path().join("pr-cache.json");
let target = PrTarget {
owner: "rust-works".into(),
name: "omni-dev".into(),
branch: "main".into(),
};
let prefs = pr_cache_prefs_from(
vec![(target, PrResolution::Pr(pending_badge(1337, &head)))],
&[watch("main", None)],
Utc::now(),
);
write_pr_cache(&cache_path, &prefs).unwrap();
let svc = WorktreesService::new();
svc.load_pr_cache(cache_path);
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.start_pr_poller_with(Duration::from_secs(30), Duration::from_millis(10), fake);
let number = tokio::time::timeout(Duration::from_secs(30), async {
loop {
let tree = svc.handle("tree", Value::Null).await.unwrap();
if let Some(n) = repos_of(&tree)
.first()
.and_then(|r| r["worktrees"][0]["pr"]["number"].as_u64())
{
return n;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("the restored badge should render from the warm cache");
assert_eq!(number, 1337);
tokio::time::sleep(Duration::from_millis(300)).await;
svc.shutdown().await;
assert_eq!(
gh_spawn_count(&counter),
0,
"a fresh warm cache must skip the immediate re-poll (#1389, fix 4)"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_persists_fresh_verdicts_for_the_next_warm_start() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"COMPLETED","conclusion":"SUCCESS"}]}}},
"associatedPullRequests":{"nodes":[{"number":41,"isDraft":false,"url":"u"}]}
}}}}"#,
);
let svc = WorktreesService::new();
let cache_path = bin_dir.path().join("runtime").join("pr-cache.json");
svc.load_pr_cache(cache_path.clone());
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.start_pr_poller_with(Duration::from_millis(50), Duration::from_millis(10), fake);
let prefs = tokio::time::timeout(Duration::from_secs(30), async {
loop {
if let Ok(bytes) = std::fs::read(&cache_path) {
if let Ok(prefs) = serde_json::from_slice::<PrCachePrefs>(&bytes) {
if !prefs.entries.is_empty() {
return prefs;
}
}
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("a successful resolve should persist the cache file");
svc.shutdown().await;
assert_eq!(prefs.entries[0].target.branch, "main");
assert!(
matches!(&prefs.entries[0].resolution, PersistedResolution::Pr(b) if b.number == 41),
"{:?}",
prefs.entries[0].resolution
);
assert_eq!(
prefs.watched,
vec![PersistedWatch {
target: prefs.entries[0].target.clone(),
upstream_sha: None
}]
);
assert!(
prefs.polled_at.is_some(),
"the poll time is what ages the next warm start"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn open_prs_op_serves_from_gh_then_dedupes_within_the_ttl() {
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(
bin_dir.path(),
r#"[{"number":42,"title":"T","url":"http://x/42","headRefName":"feat",
"baseRefName":"main","isDraft":false,"state":"OPEN","author":{"login":"me"}}]"#,
);
let svc = WorktreesService::new();
let prs = svc
.open_prs_with("rust-works", "omni-dev", fake.clone())
.await
.expect("gh pr list should resolve");
assert_eq!(prs.len(), 1);
assert_eq!(prs[0]["number"], json!(42));
assert_eq!(prs[0]["url"], json!("http://x/42"));
assert_eq!(gh_spawn_count(&counter), 1, "first call spends one gh");
let again = svc
.open_prs_with("rust-works", "omni-dev", fake.clone())
.await
.expect("cache hit should resolve");
assert_eq!(again, prs);
assert_eq!(
gh_spawn_count(&counter),
1,
"the second lookup must dedupe to the cached result, not a new gh (#1389, fix 7)"
);
let reply = svc
.handle(
"open-prs",
json!({ "owner": "rust-works", "name": "omni-dev" }),
)
.await
.expect("open-prs op should route");
assert_eq!(reply["pull_requests"][0]["number"], json!(42));
assert!(svc
.handle("open-prs", json!({ "owner": " ", "name": "x" }))
.await
.is_err());
}
#[test]
fn open_pr_list_surfaces_a_missing_binary_a_failed_run_and_bad_json() {
let err = open_pr_list(Path::new("/nonexistent/gh"), "rust-works/omni-dev").unwrap_err();
assert!(
err.to_string().contains("is the GitHub CLI installed"),
"{err:#}"
);
let bin_dir = tempfile::tempdir().unwrap();
let _guard = shim_lock();
let failing = bin_dir.path().join("fake-gh-fails");
write_exec_script(&failing, "#!/bin/sh\necho 'boom' >&2\nexit 1\n");
let err = open_pr_list(&failing, "rust-works/omni-dev").unwrap_err();
assert!(err.to_string().contains("gh pr list failed"), "{err:#}");
assert!(err.to_string().contains("boom"), "{err:#}");
let object = bin_dir.path().join("fake-gh-object");
write_exec_script(&object, "#!/bin/sh\necho '{}'\n");
let err = open_pr_list(&object, "rust-works/omni-dev").unwrap_err();
assert!(
err.to_string().contains("did not return a JSON array"),
"{err:#}"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_throttles_when_the_budget_is_over_warn() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim, counter) = counting_fake_gh(bin_dir.path(), "{}");
let svc = WorktreesService::new();
*svc.pr_warm_start
.lock()
.unwrap_or_else(PoisonError::into_inner) = Some(PrWarmStart {
watched: vec![],
polled_at: Utc::now(),
});
svc.rate_limit_cache.replace(RateLimitSnapshot {
graphql: Some(rl_resource(90)),
core: Some(rl_resource(3)),
search: None,
});
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.start_pr_poller_with(Duration::from_secs(30), Duration::from_millis(10), fake);
tokio::time::sleep(Duration::from_millis(300)).await;
svc.shutdown().await;
assert_eq!(
gh_spawn_count(&counter),
0,
"over WARN_PERCENT the poller must not fetch a grown watch (#1389, fix 6)"
);
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_bumps_only_when_a_verdict_actually_moves() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":{"contexts":{"nodes":[
{"__typename":"CheckRun","status":"IN_PROGRESS","conclusion":null}
]}}},
"associatedPullRequests":{"nodes":[{"number":1,"isDraft":false,"url":"u"}]}
}}}}"#,
);
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
fake.clone(),
);
tokio::time::timeout(Duration::from_secs(30), async {
loop {
if svc.pr_cache.get("rust-works", "omni-dev", "main").is_some() {
return;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("poller should resolve a badge through the fake gh");
let settled = svc.registry.change_generation();
tokio::time::sleep(Duration::from_millis(150)).await;
assert_eq!(
svc.registry.change_generation(),
settled,
"an unchanged poll must not bump the change-notify"
);
svc.shutdown().await;
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn pr_poller_resolves_a_negative_through_gh_and_bumps_once() {
let dir = tempfile::tempdir().unwrap();
github_repo(dir.path());
let bin_dir = tempfile::tempdir().unwrap();
let (fake, _shim) = fake_gh(
bin_dir.path(),
r#"{"data":{"r0":{"b0":{
"target":{"oid":"abc","statusCheckRollup":null},
"associatedPullRequests":{"nodes":[]}
}}}}"#,
);
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w", "folders": [dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.registry.set_polling("rust-works", "omni-dev", true);
svc.start_pr_poller_with(
Duration::from_millis(50),
Duration::from_millis(10),
fake.clone(),
);
tokio::time::timeout(Duration::from_secs(30), async {
loop {
if svc.pr_cache.get("rust-works", "omni-dev", "main") == Some(PrResolution::NoPr) {
return;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.expect("poller should resolve the negative through the fake gh");
let wt = &repos_of(&svc.handle("tree", Value::Null).await.unwrap())[0]["worktrees"][0];
assert_eq!(wt["pr_none"], json!(true));
assert!(wt.get("pr").is_none(), "{wt:?}");
let settled = svc.registry.change_generation();
tokio::time::sleep(Duration::from_millis(150)).await;
assert_eq!(
svc.registry.change_generation(),
settled,
"an unchanged negative must not bump the change-notify"
);
svc.shutdown().await;
}
#[test]
fn sync_indicator_formats_only_with_upstream() {
assert_eq!(sync_indicator(Some(2), Some(1)).as_deref(), Some("(+2 -1)"));
assert_eq!(sync_indicator(Some(0), Some(0)).as_deref(), Some("(+0 -0)"));
assert_eq!(sync_indicator(None, None), None);
assert_eq!(sync_indicator(Some(1), None), None);
}
#[tokio::test]
async fn list_enriches_entries_with_git_status() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w1", "folders": [dir.path()], "repo": "r" }),
)
.await
.unwrap();
let payload = svc.handle("list", Value::Null).await.unwrap();
let windows = windows_of(&payload);
assert_eq!(windows.len(), 1);
assert_eq!(
windows[0].get("branch").and_then(Value::as_str),
Some("main")
);
assert!(windows[0].get("ahead").is_none());
assert!(windows[0].get("behind").is_none());
assert_eq!(
windows[0].get("main_repo").and_then(Value::as_str),
dir.path().file_name().and_then(|n| n.to_str())
);
let plain = tempfile::tempdir().unwrap();
svc.handle(
"register",
json!({ "key": "w2", "folders": [plain.path()], "repo": "plain" }),
)
.await
.unwrap();
let windows = windows_of(&svc.handle("list", Value::Null).await.unwrap()).clone();
let w2 = windows
.iter()
.find(|w| w.get("key").and_then(Value::as_str) == Some("w2"))
.unwrap();
assert!(w2.get("branch").is_none());
assert!(w2.get("main_repo").is_none());
}
#[test]
fn window_label_prefers_git_branch_over_title() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let repo_name = dir.path().file_name().unwrap().to_str().unwrap();
let entry = WindowEntry {
key: "k".to_string(),
folders: vec![dir.path().to_path_buf()],
repo: Some("companion-repo".to_string()),
title: Some("ignored title".to_string()),
pid: None,
last_seen: Utc::now(),
};
assert_eq!(window_label(&entry), format!("{repo_name} · main"));
}
#[tokio::test]
async fn list_includes_ahead_behind_for_tracking_branch() {
let dir = tempfile::tempdir().unwrap();
let _repo = diverging_repo(dir.path());
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w1", "folders": [dir.path()], "repo": "r" }),
)
.await
.unwrap();
let payload = svc.handle("list", Value::Null).await.unwrap();
let windows = windows_of(&payload);
assert_eq!(
windows[0].get("branch").and_then(Value::as_str),
Some("main")
);
assert_eq!(windows[0].get("ahead").and_then(Value::as_u64), Some(1));
assert_eq!(windows[0].get("behind").and_then(Value::as_u64), Some(1));
}
#[test]
fn window_label_includes_sync_for_tracking_branch() {
let dir = tempfile::tempdir().unwrap();
let _repo = diverging_repo(dir.path());
let repo_name = dir.path().file_name().unwrap().to_str().unwrap();
let entry = WindowEntry {
key: "k".to_string(),
folders: vec![dir.path().to_path_buf()],
repo: Some("companion-repo".to_string()),
title: None,
pid: None,
last_seen: Utc::now(),
};
assert_eq!(window_label(&entry), format!("{repo_name} · main (+1 -1)"));
}
fn add_worktree(repo: &Repository, base: git2::Oid, wt_path: &Path, branch: &str) {
let commit = repo.find_commit(base).unwrap();
repo.branch(branch, &commit, false).unwrap();
let reference = repo
.find_reference(&format!("refs/heads/{branch}"))
.unwrap();
let mut opts = git2::WorktreeAddOptions::new();
opts.reference(Some(&reference));
repo.worktree(branch, wt_path, Some(&opts)).unwrap();
}
#[test]
fn git_status_marks_linked_worktree_and_names_parent_repo() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
let status = git_status(&wt_path);
assert!(status.is_worktree);
assert_eq!(status.branch.as_deref(), Some("feature"));
assert_eq!(
status.main_repo.as_deref(),
main_dir.path().file_name().and_then(|n| n.to_str())
);
let main_status = git_status(main_dir.path());
assert!(!main_status.is_worktree);
assert_eq!(main_status.main_repo, status.main_repo);
}
#[test]
fn window_label_marks_worktree_with_fork_glyph() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
let repo_name = main_dir.path().file_name().unwrap().to_str().unwrap();
let entry = WindowEntry {
key: "k".to_string(),
folders: vec![wt_path],
repo: Some("feature-wt".to_string()),
title: None,
pid: None,
last_seen: Utc::now(),
};
assert_eq!(window_label(&entry), format!("{repo_name} ⑂ feature"));
}
#[test]
fn main_repo_name_derives_from_common_dir() {
assert_eq!(
main_repo_name(Path::new("/home/me/omni-dev/.git")).as_deref(),
Some("omni-dev")
);
assert_eq!(
main_repo_name(Path::new("/home/me/omni-dev/.git/")).as_deref(),
Some("omni-dev")
);
assert_eq!(
main_repo_name(Path::new("/srv/git/omni-dev.git")).as_deref(),
Some("omni-dev")
);
assert_eq!(main_repo_name(Path::new("/.git")), None);
}
fn repos_of(payload: &Value) -> Vec<Value> {
payload
.get("repos")
.and_then(Value::as_array)
.expect("repos array")
.clone()
}
fn github(owner: &str, name: &str) -> Option<GithubIdentity> {
Some(GithubIdentity {
owner: owner.to_string(),
name: name.to_string(),
})
}
#[test]
fn github_identity_parses_supported_forms() {
assert_eq!(
github_identity("https://github.com/rust-works/omni-dev.git"),
github("rust-works", "omni-dev")
);
assert_eq!(
github_identity("https://github.com/rust-works/omni-dev"),
github("rust-works", "omni-dev")
);
assert_eq!(github_identity("http://github.com/o/r"), github("o", "r"));
assert_eq!(
github_identity("git@github.com:rust-works/omni-dev.git"),
github("rust-works", "omni-dev")
);
assert_eq!(
github_identity("ssh://git@github.com/o/r.git"),
github("o", "r")
);
assert_eq!(github_identity("git://github.com/o/r"), github("o", "r"));
assert_eq!(
github_identity(" https://github.com/o/r/ "),
github("o", "r")
);
}
#[test]
fn github_identity_rejects_non_github_and_malformed() {
assert_eq!(github_identity("https://gitlab.com/o/r.git"), None);
assert_eq!(github_identity("git@example.com:o/r.git"), None);
assert_eq!(github_identity("https://github.com/onlyowner"), None);
assert_eq!(github_identity("https://github.com/o/r/extra"), None);
assert_eq!(github_identity("https://github.com/"), None);
assert_eq!(github_identity("not a url"), None);
}
#[test]
fn remote_github_identity_reads_origin_then_falls_back() {
let dir = tempfile::tempdir().unwrap();
let repo = init_repo(dir.path());
assert_eq!(remote_github_identity(&repo), None);
repo.remote("origin", "https://gitlab.com/o/r.git").unwrap();
assert_eq!(remote_github_identity(&repo), None);
repo.remote_set_url("origin", "git@github.com:rust-works/omni-dev.git")
.unwrap();
assert_eq!(
remote_github_identity(&repo),
github("rust-works", "omni-dev")
);
repo.remote_set_url("origin", "https://gitlab.com/o/r.git")
.unwrap();
repo.remote("upstream", "https://github.com/other/proj.git")
.unwrap();
assert_eq!(remote_github_identity(&repo), github("other", "proj"));
}
#[tokio::test]
async fn tree_is_empty_with_no_windows_and_skips_non_repos() {
let svc = WorktreesService::new();
assert_eq!(
svc.handle("tree", Value::Null).await.unwrap(),
json!({ "repos": [], "show_closed": true })
);
let plain = tempfile::tempdir().unwrap();
svc.handle(
"register",
json!({ "key": "w1", "folders": [plain.path()], "repo": "plain" }),
)
.await
.unwrap();
assert!(repos_of(&svc.handle("tree", Value::Null).await.unwrap()).is_empty());
}
#[tokio::test]
async fn tree_enumerates_main_and_linked_with_open_join_and_github() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
repo.remote("origin", "git@github.com:rust-works/omni-dev.git")
.unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "wm", "folders": [main_dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
svc.handle(
"register",
json!({ "key": "wf", "folders": [wt_path], "repo": "feature-wt" }),
)
.await
.unwrap();
let repos = repos_of(&svc.handle("tree", Value::Null).await.unwrap());
assert_eq!(
repos.len(),
1,
"two worktrees of one repo dedupe: {repos:?}"
);
let repo0 = &repos[0];
assert_eq!(
repo0.get("main_repo").and_then(Value::as_str),
main_dir.path().file_name().and_then(|n| n.to_str())
);
assert_eq!(
repo0.pointer("/github/owner").and_then(Value::as_str),
Some("rust-works")
);
assert_eq!(
repo0.pointer("/github/name").and_then(Value::as_str),
Some("omni-dev")
);
assert!(repo0.get("root").and_then(Value::as_str).is_some());
let worktrees = repo0.get("worktrees").and_then(Value::as_array).unwrap();
assert_eq!(worktrees.len(), 2);
let main_wt = &worktrees[0];
assert_eq!(main_wt.get("is_main").and_then(Value::as_bool), Some(true));
assert_eq!(main_wt.get("open").and_then(Value::as_bool), Some(true));
assert_eq!(
main_wt.get("window_key").and_then(Value::as_str),
Some("wm")
);
assert_eq!(main_wt.get("branch").and_then(Value::as_str), Some("main"));
let linked = &worktrees[1];
assert_eq!(linked.get("is_main").and_then(Value::as_bool), Some(false));
assert_eq!(linked.get("open").and_then(Value::as_bool), Some(true));
assert_eq!(linked.get("window_key").and_then(Value::as_str), Some("wf"));
assert_eq!(
linked.get("branch").and_then(Value::as_str),
Some("feature")
);
}
#[tokio::test]
async fn tree_marks_unopened_linked_worktree_closed_and_omits_github() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/main"), &[], "A");
repo.set_head("refs/heads/main").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "wm", "folders": [main_dir.path()], "repo": "omni-dev" }),
)
.await
.unwrap();
let repos = repos_of(&svc.handle("tree", Value::Null).await.unwrap());
assert_eq!(repos.len(), 1);
assert!(repos[0].get("github").is_none(), "no remote → no github");
let worktrees = repos[0].get("worktrees").and_then(Value::as_array).unwrap();
let linked = worktrees
.iter()
.find(|w| w.get("is_main").and_then(Value::as_bool) == Some(false))
.expect("the linked worktree");
assert_eq!(linked.get("open").and_then(Value::as_bool), Some(false));
assert!(linked.get("window_key").is_none());
}
fn repo_with_linked_worktree() -> (tempfile::TempDir, tempfile::TempDir, PathBuf) {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/trunk"), &[], "A");
repo.set_head("refs/heads/trunk").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
(main_dir, wt_parent, wt_path)
}
fn repo_with_two_linked_worktrees() -> (tempfile::TempDir, tempfile::TempDir, PathBuf, PathBuf)
{
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/trunk"), &[], "A");
repo.set_head("refs/heads/trunk").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let first = wt_parent.path().join("first-wt");
let second = wt_parent.path().join("second-wt");
add_worktree(&repo, a, &first, "first");
add_worktree(&repo, a, &second, "second");
(main_dir, wt_parent, first, second)
}
#[tokio::test]
async fn close_removes_two_linked_worktrees_of_one_repo_concurrently() {
let (main_dir, _wtp, first, second) = repo_with_two_linked_worktrees();
let svc = Arc::new(WorktreesService::new());
let close = |path: PathBuf| {
let svc = svc.clone();
async move {
svc.handle(
"close",
json!({ "path": path, "remove": true, "confirmed": true }),
)
.await
}
};
let (a, b) = tokio::join!(close(first.clone()), close(second.clone()));
assert_eq!(a.unwrap(), json!({ "removed": true }));
assert_eq!(b.unwrap(), json!({ "removed": true }));
assert!(!first.exists());
assert!(!second.exists());
let repo = Repository::open(main_dir.path()).unwrap();
assert!(repo.worktrees().unwrap().is_empty());
}
#[tokio::test]
async fn concurrent_closes_overlap_their_heartbeat_waits() {
let (_main, _wtp, first, second) = repo_with_two_linked_worktrees();
let svc = Arc::new(WorktreesService::new());
for (key, path) in [("w2", &first), ("w3", &second)] {
svc.handle("register", json!({ "key": key, "folders": [path] }))
.await
.unwrap();
}
let spawn_close = |path: PathBuf| {
let svc = svc.clone();
tokio::spawn(async move {
svc.handle(
"close",
json!({
"path": path,
"remove": true,
"confirmed": true,
"requester_key": "w1",
}),
)
.await
})
};
let a = spawn_close(first.clone());
let b = spawn_close(second.clone());
for key in ["w2", "w3"] {
let mut saw_close = false;
for _ in 0..400 {
let hb = svc
.handle("heartbeat", json!({ "key": key }))
.await
.unwrap();
if hb.get("close").and_then(Value::as_bool) == Some(true) {
saw_close = true;
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert!(saw_close, "{key} should have been told to close while the other target's close was still waiting");
}
assert!(
!a.is_finished() && !b.is_finished(),
"neither close can have finished: both windows are still registered"
);
for key in ["w2", "w3"] {
svc.handle("unregister", json!({ "key": key }))
.await
.unwrap();
}
assert_eq!(a.await.unwrap().unwrap(), json!({ "removed": true }));
assert_eq!(b.await.unwrap().unwrap(), json!({ "removed": true }));
assert!(!first.exists());
assert!(!second.exists());
}
#[tokio::test]
async fn close_safety_check_reports_clean_linked_as_removable_with_no_risks() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
assert_eq!(report.get("removable").and_then(Value::as_bool), Some(true));
assert_eq!(report.get("is_main").and_then(Value::as_bool), Some(false));
assert_eq!(report.get("open").and_then(Value::as_bool), Some(false));
assert!(report
.get("risks")
.and_then(Value::as_array)
.unwrap()
.is_empty());
assert!(wt_path.exists());
}
#[tokio::test]
async fn close_removes_a_clean_linked_worktree() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
let reply = svc
.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap();
assert_eq!(reply, json!({ "removed": true }));
assert!(
!wt_path.exists(),
"the worktree directory should be deleted"
);
}
#[test]
fn remove_worktree_deletes_the_directory_and_prunes_the_admin_metadata() {
let (main, _wtp, wt_path) = repo_with_linked_worktree();
let admin = main.path().join(".git").join("worktrees").join("feature");
assert!(admin.exists(), "admin metadata should exist before removal");
remove_worktree(&wt_path).unwrap();
assert!(!wt_path.exists(), "the working directory should be gone");
assert!(!admin.exists(), "the admin metadata should be pruned");
let main_repo = Repository::open(main.path()).unwrap();
assert_eq!(
main_repo.worktrees().unwrap().len(),
0,
"git should no longer track the worktree"
);
}
#[test]
fn remove_worktree_recovers_a_half_removed_orphan() {
let (main, _wtp, wt_path) = repo_with_linked_worktree();
let admin = main.path().join(".git").join("worktrees").join("feature");
std::fs::remove_dir_all(&admin).unwrap();
assert!(wt_path.join(".git").is_file(), "dangling gitlink remains");
assert!(
Repository::open(&wt_path).is_err(),
"the orphan should not open as a repo"
);
remove_worktree(&wt_path).unwrap();
assert!(
!wt_path.exists(),
"the leftover directory should be removed"
);
}
#[test]
fn is_orphaned_worktree_only_matches_a_dangling_linked_gitlink() {
let (main, _wtp, wt_path) = repo_with_linked_worktree();
assert!(!is_orphaned_worktree(&wt_path));
assert!(!is_orphaned_worktree(main.path()));
std::fs::remove_dir_all(main.path().join(".git").join("worktrees").join("feature"))
.unwrap();
assert!(is_orphaned_worktree(&wt_path));
}
#[test]
fn remove_dir_all_retrying_is_idempotent_on_a_missing_directory() {
let tmp = tempfile::tempdir().unwrap();
let missing = tmp.path().join("gone");
assert!(remove_dir_all_retrying(&missing).is_ok());
}
#[test]
fn is_transient_rmdir_error_matches_only_the_repopulated_directory_race() {
use std::io::Error;
for errno in [nix::libc::ENOTEMPTY, nix::libc::EEXIST, nix::libc::EBUSY] {
assert!(
is_transient_rmdir_error(&Error::from_raw_os_error(errno)),
"errno {errno} is the concurrent-writer race and must be retried"
);
}
for errno in [
nix::libc::EACCES,
nix::libc::EPERM,
nix::libc::EROFS,
nix::libc::ENOTDIR,
] {
assert!(
!is_transient_rmdir_error(&Error::from_raw_os_error(errno)),
"errno {errno} is permanent and must not be retried"
);
}
assert!(!is_transient_rmdir_error(&Error::other("synthetic")));
}
#[test]
fn remove_dir_all_retrying_surfaces_a_non_transient_error_without_retrying() {
let tmp = tempfile::tempdir().unwrap();
let file = tmp.path().join("not-a-directory");
std::fs::write(&file, b"x").unwrap();
let mut attempts = 0;
let err = remove_dir_all_retrying_with(&file, WORKTREE_RMDIR_BACKOFF, || {
attempts += 1;
std::fs::remove_dir_all(&file)
})
.unwrap_err();
assert_eq!(attempts, 1, "a permanent error must not be retried");
assert!(
err.to_string()
.contains("failed to remove worktree directory"),
"unexpected error: {err:#}"
);
assert!(err.source().is_some(), "the io::Error cause is preserved");
assert!(file.exists());
}
#[test]
fn remove_dir_all_retrying_gives_up_after_the_backoff_is_exhausted() {
let tmp = tempfile::tempdir().unwrap();
let mut attempts = 0;
let backoff = [Duration::ZERO, Duration::ZERO];
let err = remove_dir_all_retrying_with(tmp.path(), &backoff, || {
attempts += 1;
Err(std::io::Error::from_raw_os_error(nix::libc::ENOTEMPTY))
})
.unwrap_err();
assert_eq!(attempts, backoff.len() + 1);
assert!(
err.to_string()
.contains("failed to remove worktree directory"),
"unexpected error: {err:#}"
);
}
#[test]
fn remove_dir_all_retrying_succeeds_once_the_writer_quiesces() {
let tmp = tempfile::tempdir().unwrap();
let mut attempts = 0;
let result = remove_dir_all_retrying_with(tmp.path(), WORKTREE_RMDIR_BACKOFF, || {
attempts += 1;
if attempts < 3 {
Err(std::io::Error::from_raw_os_error(nix::libc::ENOTEMPTY))
} else {
Ok(())
}
});
assert!(result.is_ok(), "{result:?}");
assert_eq!(attempts, 3);
}
#[test]
fn is_orphaned_worktree_ignores_a_git_file_that_is_not_a_gitlink() {
let tmp = tempfile::tempdir().unwrap();
std::fs::write(tmp.path().join(".git"), b"not a gitlink\n").unwrap();
assert!(!is_orphaned_worktree(tmp.path()));
}
#[test]
fn remove_worktree_rejects_a_path_that_is_not_a_worktree() {
let tmp = tempfile::tempdir().unwrap();
let plain = tmp.path().join("plain");
std::fs::create_dir(&plain).unwrap();
let err = remove_worktree(&plain).unwrap_err();
assert!(
err.to_string().contains("not a git worktree"),
"unexpected error: {err:#}"
);
assert!(plain.exists(), "a non-worktree path must be left alone");
}
#[test]
fn remove_worktree_succeeds_while_a_concurrent_writer_winds_down() {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let target = wt_path.join("target");
std::fs::create_dir_all(&target).unwrap();
let stop = Arc::new(AtomicBool::new(false));
let writer_stop = Arc::clone(&stop);
let writer_dir = target;
let writer = std::thread::spawn(move || {
let mut n = 0u64;
let deadline = std::time::Instant::now() + Duration::from_millis(400);
while !writer_stop.load(Ordering::Relaxed) && std::time::Instant::now() < deadline {
let nested = writer_dir.join("nested");
let _ = std::fs::create_dir_all(&nested);
let _ = std::fs::write(nested.join(format!("artifact-{n}.tmp")), b"x");
n += 1;
}
});
let result = remove_worktree(&wt_path);
stop.store(true, Ordering::Relaxed);
writer.join().unwrap();
assert!(
result.is_ok(),
"removal should retry past the writer: {result:?}"
);
assert!(!wt_path.exists(), "the worktree directory should be gone");
}
#[tokio::test]
async fn close_safety_check_flags_untracked_and_does_not_remove_without_confirmation() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
std::fs::write(wt_path.join("scratch.txt"), b"work in progress").unwrap();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let risks = report.get("risks").and_then(Value::as_array).unwrap();
assert!(
risks
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("untracked")),
"expected an untracked risk: {report}"
);
assert_eq!(report.get("removable").and_then(Value::as_bool), Some(true));
assert!(wt_path.exists());
}
#[tokio::test]
async fn close_confirmed_removes_a_dirty_worktree() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
std::fs::write(wt_path.join("scratch.txt"), b"discard me").unwrap();
let svc = WorktreesService::new();
let reply = svc
.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap();
assert_eq!(reply, json!({ "removed": true }));
assert!(!wt_path.exists());
}
#[tokio::test]
async fn close_refuses_to_remove_the_main_working_tree() {
let (main, _wtp, _wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": main.path(), "remove": true }))
.await
.unwrap();
assert_eq!(report.get("is_main").and_then(Value::as_bool), Some(true));
assert_eq!(
report.get("removable").and_then(Value::as_bool),
Some(false)
);
assert!(svc
.handle(
"close",
json!({ "path": main.path(), "remove": true, "confirmed": true }),
)
.await
.is_err());
assert!(main.path().exists());
}
#[tokio::test]
async fn close_removes_a_linked_worktree_on_the_default_branch_and_keeps_the_branch() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/trunk"), &[], "A");
repo.set_head("refs/heads/trunk").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("main-wt");
add_worktree(&repo, a, &wt_path, "main");
let svc = WorktreesService::new();
let reply = svc
.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap();
assert_eq!(reply, json!({ "removed": true }));
assert!(!wt_path.exists());
assert!(
repo.find_branch("main", git2::BranchType::Local).is_ok(),
"the default branch must survive worktree removal"
);
}
#[tokio::test]
async fn close_is_idempotent_when_the_worktree_is_already_gone() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
svc.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap();
let reply = svc
.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap();
assert_eq!(reply, json!({ "removed": true }));
}
#[tokio::test]
async fn close_safety_check_detects_detached_head_unreachable_commits() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let wt_repo = Repository::open(&wt_path).unwrap();
let parent_oid = wt_repo.head().unwrap().target().unwrap();
let parent = wt_repo.find_commit(parent_oid).unwrap();
let orphan = empty_commit(&wt_repo, None, &[&parent], "orphan");
wt_repo.set_head_detached(orphan).unwrap();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let risks = report.get("risks").and_then(Value::as_array).unwrap();
assert!(
risks
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("unreachable-commits")),
"expected an unreachable-commits risk: {report}"
);
}
#[tokio::test]
async fn close_self_close_removes_when_the_requester_owns_the_target() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w1", "folders": [wt_path], "repo": "feature-wt" }),
)
.await
.unwrap();
let reply = svc
.handle(
"close",
json!({
"path": wt_path,
"remove": true,
"confirmed": true,
"requester_key": "w1",
}),
)
.await
.unwrap();
assert_eq!(reply, json!({ "removed": true }));
assert!(!wt_path.exists());
}
#[tokio::test]
async fn close_safety_check_surfaces_the_owning_window() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w2", "folders": [&wt_path, "/tmp/other"], "repo": "feature-wt" }),
)
.await
.unwrap();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
assert_eq!(report.get("open").and_then(Value::as_bool), Some(true));
assert_eq!(report.get("window_key").and_then(Value::as_str), Some("w2"));
assert_eq!(
report.get("window_folder_count").and_then(Value::as_u64),
Some(2)
);
}
#[tokio::test]
async fn heartbeat_op_surfaces_a_pending_close_directive_once() {
let svc = WorktreesService::new();
svc.handle("register", register_payload("w1", Some("r"), "/tmp/a"))
.await
.unwrap();
assert_eq!(
svc.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap(),
json!({ "known": true })
);
svc.registry.mark_close_pending("w1");
assert_eq!(
svc.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap(),
json!({ "known": true, "close": true })
);
assert_eq!(
svc.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap(),
json!({ "known": true })
);
}
#[tokio::test]
async fn close_signals_a_cross_window_target_then_removes_after_it_closes() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = Arc::new(WorktreesService::new());
svc.handle(
"register",
json!({ "key": "w2", "folders": [&wt_path], "repo": "feature-wt" }),
)
.await
.unwrap();
let svc2 = svc.clone();
let path = wt_path.clone();
let close = tokio::spawn(async move {
svc2.handle(
"close",
json!({
"path": path,
"remove": true,
"confirmed": true,
"requester_key": "w1",
}),
)
.await
});
let mut saw_close = false;
for _ in 0..200 {
let hb = svc
.handle("heartbeat", json!({ "key": "w2" }))
.await
.unwrap();
if hb.get("close").and_then(Value::as_bool) == Some(true) {
saw_close = true;
svc.handle("unregister", json!({ "key": "w2" }))
.await
.unwrap();
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert!(saw_close, "w2 should have been told to close");
let reply = close.await.unwrap().unwrap();
assert_eq!(reply, json!({ "removed": true }));
assert!(!wt_path.exists());
}
#[tokio::test]
async fn await_windows_closed_times_out_when_a_window_never_closes() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
svc.handle(
"register",
json!({ "key": "w2", "folders": [&wt_path], "repo": "feature-wt" }),
)
.await
.unwrap();
let err = await_windows_closed(
&svc.registry,
&wt_path,
Some("w1"),
Duration::from_millis(150),
Duration::from_millis(25),
)
.await
.unwrap_err();
assert!(
err.to_string().contains("w2"),
"error names the window: {err}"
);
await_windows_closed(
&svc.registry,
&wt_path,
Some("w2"),
Duration::from_millis(150),
Duration::from_millis(25),
)
.await
.unwrap();
}
#[tokio::test]
async fn close_window_without_remove_replies_closed_and_never_deletes() {
let (main, _wtp, _wt_path) = repo_with_linked_worktree();
let svc = WorktreesService::new();
let reply = svc
.handle("close", json!({ "path": main.path(), "remove": false }))
.await
.unwrap();
assert_eq!(reply, json!({ "closed": true }));
assert!(main.path().exists());
}
#[tokio::test]
async fn close_safety_check_flags_modified_tracked_files() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = commit_file(&repo, "refs/heads/trunk", "tracked.txt", b"original\n", "A");
repo.set_head("refs/heads/trunk").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
std::fs::write(wt_path.join("tracked.txt"), b"uncommitted change\n").unwrap();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let risks = report.get("risks").and_then(Value::as_array).unwrap();
assert!(
risks
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("dirty")),
"expected a dirty risk: {report}"
);
}
#[tokio::test]
async fn close_safety_check_flags_an_in_progress_operation() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let wt_repo = Repository::open(&wt_path).unwrap();
let head = wt_repo.head().unwrap().target().unwrap();
std::fs::write(wt_repo.path().join("MERGE_HEAD"), format!("{head}\n")).unwrap();
assert_ne!(wt_repo.state(), RepositoryState::Clean);
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let risks = report.get("risks").and_then(Value::as_array).unwrap();
assert!(
risks
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("in-progress")),
"expected an in-progress risk: {report}"
);
}
#[tokio::test]
async fn close_safety_check_reports_unpushed_commits_as_info_not_a_risk() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = empty_commit(&repo, Some("refs/heads/trunk"), &[], "A");
repo.set_head("refs/heads/trunk").unwrap();
let a_commit = repo.find_commit(a).unwrap();
repo.branch("feature", &a_commit, false).unwrap();
repo.reference("refs/remotes/origin/feature", a, true, "origin feature")
.unwrap();
empty_commit(&repo, Some("refs/heads/feature"), &[&a_commit], "B");
drop(a_commit);
let mut cfg = repo.config().unwrap();
cfg.set_str("remote.origin.url", "https://example.invalid/x.git")
.unwrap();
cfg.set_str("remote.origin.fetch", "+refs/heads/*:refs/remotes/origin/*")
.unwrap();
cfg.set_str("branch.feature.remote", "origin").unwrap();
cfg.set_str("branch.feature.merge", "refs/heads/feature")
.unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
let reference = repo.find_reference("refs/heads/feature").unwrap();
let mut opts = git2::WorktreeAddOptions::new();
opts.reference(Some(&reference));
repo.worktree("feature", &wt_path, Some(&opts)).unwrap();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let info = report.get("info").and_then(Value::as_array).unwrap();
assert!(
info.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("unpushed")),
"expected an unpushed info note: {report}"
);
assert!(
report
.get("risks")
.and_then(Value::as_array)
.unwrap()
.is_empty(),
"unpushed commits alone must not block: {report}"
);
assert_eq!(report.get("removable").and_then(Value::as_bool), Some(true));
}
#[tokio::test]
async fn close_safety_check_ignores_gitignored_files() {
let main_dir = tempfile::tempdir().unwrap();
let repo = init_repo(main_dir.path());
let a = commit_file(&repo, "refs/heads/trunk", ".gitignore", b"build/\n", "A");
repo.set_head("refs/heads/trunk").unwrap();
let wt_parent = tempfile::tempdir().unwrap();
let wt_path = wt_parent.path().join("feature-wt");
add_worktree(&repo, a, &wt_path, "feature");
std::fs::create_dir(wt_path.join("build")).unwrap();
std::fs::write(wt_path.join("build/artifact.o"), b"junk").unwrap();
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
assert!(
report
.get("risks")
.and_then(Value::as_array)
.unwrap()
.is_empty(),
"a gitignored file must not create a risk: {report}"
);
assert_eq!(report.get("removable").and_then(Value::as_bool), Some(true));
}
#[tokio::test]
async fn close_safety_check_treats_a_missing_path_as_already_removed() {
let svc = WorktreesService::new();
let report = svc
.handle(
"close",
json!({ "path": "/no/such/worktree/xyzzy", "remove": true }),
)
.await
.unwrap();
assert_eq!(report.get("removable").and_then(Value::as_bool), Some(true));
assert_eq!(report.get("is_main").and_then(Value::as_bool), Some(false));
assert!(report
.get("risks")
.and_then(Value::as_array)
.unwrap()
.is_empty());
let info = report.get("info").and_then(Value::as_array).unwrap();
assert!(info
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("already-removed")));
}
#[tokio::test]
async fn close_refuses_a_locked_worktree() {
let (main, _wtp, wt_path) = repo_with_linked_worktree();
let main_repo = Repository::open(main.path()).unwrap();
main_repo
.find_worktree("feature")
.unwrap()
.lock(Some("under test"))
.unwrap();
let svc = WorktreesService::new();
let err = svc
.handle(
"close",
json!({ "path": wt_path, "remove": true, "confirmed": true }),
)
.await
.unwrap_err();
assert!(
err.to_string().contains("locked"),
"expected a locked error: {err}"
);
assert!(wt_path.exists(), "a locked worktree must not be removed");
}
#[tokio::test]
async fn close_safety_check_does_not_flag_a_detached_head_reachable_from_a_branch() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let wt_repo = Repository::open(&wt_path).unwrap();
let tip = wt_repo.head().unwrap().target().unwrap();
wt_repo.set_head_detached(tip).unwrap();
assert!(wt_repo.head_detached().unwrap());
let svc = WorktreesService::new();
let report = svc
.handle("close", json!({ "path": wt_path, "remove": true }))
.await
.unwrap();
let risks = report.get("risks").and_then(Value::as_array).unwrap();
assert!(
!risks
.iter()
.any(|r| r.get("kind").and_then(Value::as_str) == Some("unreachable-commits")),
"a detached HEAD reachable from a branch must not be flagged: {report}"
);
}
#[test]
fn worktree_name_for_path_resolves_a_real_worktree_and_errors_otherwise() {
let (main, _wtp, wt_path) = repo_with_linked_worktree();
let main_repo = Repository::open(main.path()).unwrap();
assert_eq!(
worktree_name_for_path(&main_repo, &canonical(&wt_path)).unwrap(),
"feature"
);
let err =
worktree_name_for_path(&main_repo, Path::new("/no/such/worktree/xyzzy")).unwrap_err();
assert!(
err.to_string().contains("not registered"),
"expected a not-registered error: {err}"
);
}
#[test]
fn count_dirty_untracked_degrades_to_zero_on_an_unreadable_index() {
let (_main, _wtp, wt_path) = repo_with_linked_worktree();
let repo = Repository::open(&wt_path).unwrap();
std::fs::write(repo.path().join("index"), b"not a valid git index").unwrap();
assert!(
repo.statuses(Some(&mut StatusOptions::new())).is_err(),
"a corrupt index should make statuses() fail"
);
assert_eq!(count_dirty_untracked(&repo), (0, 0));
}
}