use std::collections::{BTreeSet, HashMap};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::{Mutex, MutexGuard, PoisonError};
use std::time::Duration;
use anyhow::{anyhow, bail, Context, Result};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use crate::daemon::service::{DaemonService, MenuAction, MenuItem, MenuSnapshot, ServiceStatus};
pub const SERVICE_NAME: &str = "worktrees";
const DEFAULT_TTL: Duration = Duration::from_secs(30);
const VSCODE_BIN_ENV: &str = "OMNI_DEV_VSCODE_BIN";
#[derive(Debug, Clone, Deserialize)]
struct RegisterRequest {
key: String,
#[serde(default)]
folders: Vec<PathBuf>,
#[serde(default)]
repo: Option<String>,
#[serde(default)]
title: Option<String>,
#[serde(default)]
pid: Option<u32>,
}
#[derive(Debug, Clone, Serialize)]
struct WindowEntry {
key: String,
folders: Vec<PathBuf>,
#[serde(skip_serializing_if = "Option::is_none")]
repo: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pid: Option<u32>,
last_seen: DateTime<Utc>,
}
pub struct WorktreesService {
windows: Mutex<HashMap<String, WindowEntry>>,
ttl: Duration,
}
impl WorktreesService {
#[must_use]
pub fn new() -> Self {
Self {
windows: Mutex::new(HashMap::new()),
ttl: DEFAULT_TTL,
}
}
fn lock(&self) -> MutexGuard<'_, HashMap<String, WindowEntry>> {
self.windows.lock().unwrap_or_else(PoisonError::into_inner)
}
fn live_entries(&self, now: DateTime<Utc>) -> Vec<WindowEntry> {
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
sorted_entries(&windows)
}
}
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> {
let now = Utc::now();
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`");
}
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
windows.insert(
req.key.clone(),
WindowEntry {
key: req.key,
folders: req.folders,
repo: req.repo,
title: req.title,
pid: req.pid,
last_seen: now,
},
);
Ok(json!({ "ok": true }))
}
"heartbeat" => {
let key = require_key(&payload, "heartbeat")?;
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
let known = match windows.get_mut(key) {
Some(entry) => {
entry.last_seen = now;
true
}
None => false,
};
Ok(json!({ "known": known }))
}
"unregister" => {
let key = require_key(&payload, "unregister")?;
let mut windows = self.lock();
let removed = windows.remove(key).is_some();
reap(&mut windows, self.ttl, now);
Ok(json!({ "removed": removed }))
}
"list" => Ok(json!({ "windows": self.live_entries(now) })),
other => bail!("unknown worktrees op: {other}"),
}
}
fn menu(&self) -> MenuSnapshot {
let entries = self.live_entries(Utc::now());
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 = {
let windows = self.lock();
windows.get(key).and_then(|e| e.folders.first().cloned())
};
let folder = folder
.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.live_entries(Utc::now());
let repos: BTreeSet<&str> = entries.iter().filter_map(|e| e.repo.as_deref()).collect();
ServiceStatus {
name: SERVICE_NAME.to_string(),
healthy: true,
summary: format!("{} window(s) across {} repo(s)", entries.len(), repos.len()),
detail: json!({ "windows": entries }),
}
}
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`"))
}
fn reap(windows: &mut HashMap<String, WindowEntry>, ttl: Duration, now: DateTime<Utc>) {
let max_age = ttl.as_secs() as i64;
windows.retain(|_, e| (now - e.last_seen).num_seconds() <= max_age);
}
fn sorted_entries(windows: &HashMap<String, WindowEntry>) -> Vec<WindowEntry> {
let mut entries: Vec<WindowEntry> = windows.values().cloned().collect();
entries.sort_by(|a, b| a.repo.cmp(&b.repo).then_with(|| a.key.cmp(&b.key)));
entries
}
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 {
let name = display_name(entry);
let label = match &entry.title {
Some(title) if title != &name => format!("{name} · {title}"),
_ => name,
};
items.push(MenuItem::Label(label));
}
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
}
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::*;
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 list_is_empty_initially() {
let svc = WorktreesService::new();
let payload = svc.handle("list", Value::Null).await.unwrap();
assert_eq!(payload, json!({ "windows": [] }));
}
#[tokio::test]
async fn register_then_list_round_trips() {
let svc = WorktreesService::new();
let reply = svc
.handle("register", register_payload("w1", Some("repo-a"), "/tmp/a"))
.await
.unwrap();
assert_eq!(reply, json!({ "ok": true }));
let payload = svc.handle("list", Value::Null).await.unwrap();
let windows = windows_of(&payload);
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].get("key").and_then(Value::as_str), Some("w1"));
assert_eq!(
windows[0].get("repo").and_then(Value::as_str),
Some("repo-a")
);
assert!(windows[0].get("last_seen").is_some());
}
#[tokio::test]
async fn register_is_idempotent_upsert() {
let svc = WorktreesService::new();
svc.handle("register", register_payload("w1", Some("repo-a"), "/tmp/a"))
.await
.unwrap();
svc.handle("register", register_payload("w1", Some("repo-b"), "/tmp/b"))
.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("repo").and_then(Value::as_str),
Some("repo-b")
);
}
#[tokio::test]
async fn register_requires_key() {
let svc = WorktreesService::new();
assert!(svc.handle("register", json!({})).await.is_err());
assert!(svc
.handle("register", json!({ "key": " " }))
.await
.is_err());
}
#[tokio::test]
async fn heartbeat_reports_known_and_unknown() {
let svc = WorktreesService::new();
let unknown = svc
.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap();
assert_eq!(unknown, json!({ "known": false }));
svc.handle("register", register_payload("w1", None, "/tmp/a"))
.await
.unwrap();
let known = svc
.handle("heartbeat", json!({ "key": "w1" }))
.await
.unwrap();
assert_eq!(known, json!({ "known": true }));
assert!(svc.handle("heartbeat", json!({})).await.is_err());
}
#[tokio::test]
async fn unregister_removes() {
let svc = WorktreesService::new();
svc.handle("register", register_payload("w1", None, "/tmp/a"))
.await
.unwrap();
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 }));
assert!(svc.handle("unregister", json!({})).await.is_err());
}
#[test]
fn reap_evicts_only_stale_entries() {
let now = Utc::now();
let mut windows = HashMap::new();
windows.insert(
"fresh".to_string(),
WindowEntry {
key: "fresh".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(5),
},
);
windows.insert(
"stale".to_string(),
WindowEntry {
key: "stale".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(120),
},
);
reap(&mut windows, DEFAULT_TTL, now);
assert!(windows.contains_key("fresh"));
assert!(!windows.contains_key("stale"));
}
#[test]
fn sorted_entries_orders_by_repo_then_key() {
let now = Utc::now();
let mut windows = HashMap::new();
for (key, repo) in [("z", "repo-a"), ("a", "repo-b"), ("m", "repo-a")] {
windows.insert(
key.to_string(),
WindowEntry {
key: key.to_string(),
folders: vec![],
repo: Some(repo.to_string()),
title: None,
pid: None,
last_seen: now,
},
);
}
let entries = sorted_entries(&windows);
let ordered: Vec<(&str, &str)> = entries
.iter()
.map(|e| (e.key.as_str(), e.repo.as_deref().unwrap()))
.collect();
assert_eq!(
ordered,
vec![("m", "repo-a"), ("z", "repo-a"), ("a", "repo-b")]
);
}
#[test]
fn 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 default_constructs_an_empty_service() {
let svc = WorktreesService::default();
assert!(svc.lock().is_empty());
}
#[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 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();
}
}