use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use anyhow::{anyhow, bail, Context, Result};
use async_trait::async_trait;
use git2::Repository;
use serde::Serialize;
use serde_json::{json, Value};
use crate::daemon::service::{DaemonService, MenuAction, MenuItem, MenuSnapshot, ServiceStatus};
use crate::worktrees::{RegisterRequest, WindowEntry, WorktreesRegistry};
pub const SERVICE_NAME: &str = "worktrees";
const VSCODE_BIN_ENV: &str = "OMNI_DEV_VSCODE_BIN";
pub struct WorktreesService {
registry: WorktreesRegistry,
}
impl WorktreesService {
#[must_use]
pub fn new() -> Self {
Self {
registry: WorktreesRegistry::new(),
}
}
}
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_key(&payload, "heartbeat")?;
Ok(json!({ "known": self.registry.heartbeat(key) }))
}
"unregister" => {
let key = require_key(&payload, "unregister")?;
Ok(json!({ "removed": self.registry.unregister(key) }))
}
"list" => Ok(json!({ "windows": enriched_windows(self.registry.list()).await })),
other => bail!("unknown worktrees op: {other}"),
}
}
fn menu(&self) -> MenuSnapshot {
let entries = self.registry.list();
let items = if entries.is_empty() {
vec![MenuItem::Label("No open windows".to_string())]
} else {
window_menu_items(&entries)
};
MenuSnapshot {
title: "Worktrees".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) {
}
}
fn require_key<'a>(payload: &'a Value, op: &str) -> Result<&'a str> {
payload
.get("key")
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("`{op}` requires `key`"))
}
#[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")]
ahead: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
behind: Option<usize>,
}
fn git_status(folder: &Path) -> GitStatus {
let Ok(repo) = Repository::discover(folder) else {
return GitStatus::default();
};
let Ok(head) = repo.head() else {
return GitStatus::default();
};
let Some(name) = head
.shorthand()
.ok()
.filter(|_| head.is_branch())
.map(str::to_string)
else {
return GitStatus::default();
};
let branch = git2::Branch::wrap(head);
let (ahead, behind) = match upstream_ahead_behind(&repo, &branch) {
Some((ahead, behind)) => (Some(ahead), Some(behind)),
None => (None, None),
};
GitStatus {
branch: Some(name),
ahead,
behind,
}
}
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()
}
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()
}
fn window_menu_items(entries: &[WindowEntry]) -> Vec<MenuItem> {
let mut items = Vec::new();
for entry in entries {
items.push(MenuItem::Label(window_label(entry)));
}
items.push(MenuItem::Separator);
for entry in entries {
if !entry.folders.is_empty() {
items.push(MenuItem::Action(MenuAction {
id: format!("focus:{}", entry.key),
label: format!("Focus {}", display_name(entry)),
enabled: true,
}));
}
}
items
}
fn window_label(entry: &WindowEntry) -> String {
let name = display_name(entry);
let status = entry
.folders
.first()
.map(|folder| git_status(folder))
.unwrap_or_default();
if let Some(branch) = &status.branch {
return match sync_indicator(status.ahead, status.behind) {
Some(sync) => format!("{name} · {branch} {sync}"),
None => format!("{name} · {branch}"),
};
}
match &entry.title {
Some(title) if title != &name => format!("{name} · {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",
];
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")
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use chrono::Utc;
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_label_omits_redundant_title_and_skips_folderless_actions() {
let now = Utc::now();
let entries = vec![
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,
},
WindowEntry {
key: "k2".to_string(),
folders: vec![],
repo: Some("solo".to_string()),
title: Some("solo".to_string()),
pid: None,
last_seen: now,
},
];
let items = window_menu_items(&entries);
let labels: Vec<&str> = items
.iter()
.filter_map(|i| match i {
MenuItem::Label(t) => Some(t.as_str()),
_ => None,
})
.collect();
assert!(labels.contains(&"repo · a branch"));
assert!(labels.contains(&"solo"));
let action_ids: Vec<&str> = items
.iter()
.filter_map(|i| match i {
MenuItem::Action(a) => Some(a.id.as_str()),
_ => None,
})
.collect();
assert_eq!(action_ids, vec!["focus:k1"]);
}
#[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();
let status = svc.status().await;
assert_eq!(status.summary, "2 window(s) across 1 repo(s)");
let menu = svc.menu();
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"));
}
#[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 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();
}
#[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 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));
}
#[test]
fn git_status_empty_repo_is_unborn() {
let dir = tempfile::tempdir().unwrap();
init_repo(dir.path());
assert_eq!(git_status(dir.path()), GitStatus::default());
}
#[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);
}
#[test]
fn git_status_non_repo_and_detached_are_empty() {
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();
assert_eq!(git_status(dir.path()), GitStatus::default());
}
#[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());
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());
}
#[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 entry = WindowEntry {
key: "k".to_string(),
folders: vec![dir.path().to_path_buf()],
repo: Some("my-repo".to_string()),
title: Some("ignored title".to_string()),
pid: None,
last_seen: Utc::now(),
};
assert_eq!(window_label(&entry), "my-repo · 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 entry = WindowEntry {
key: "k".to_string(),
folders: vec![dir.path().to_path_buf()],
repo: Some("my-repo".to_string()),
title: None,
pid: None,
last_seen: Utc::now(),
};
assert_eq!(window_label(&entry), "my-repo · main (+1 -1)");
}
}