Skip to main content

omni_dev/cli/
worktrees.rs

1//! `omni-dev worktrees` — a thin client for the daemon's cross-window worktree
2//! registry.
3//!
4//! Lifecycle stays on `omni-dev daemon` (`start`/`stop`/`status`/`restart`);
5//! this command only sends the `worktrees` service's read ops (`list` for the
6//! open windows, `tree` for the repos-and-all-their-worktrees view) over the
7//! daemon's Unix control socket. The companion VS Code extension is what *feeds*
8//! the registry (`register`/`heartbeat`/`unregister`), talking to the same
9//! socket directly from each window.
10
11use std::path::{Path, PathBuf};
12
13use anyhow::{bail, Result};
14use chrono::Utc;
15use clap::{Parser, Subcommand};
16use serde_json::{json, Value};
17
18use crate::cli::format::TableOrJson;
19use crate::daemon::client::DaemonClient;
20use crate::daemon::protocol::{DaemonEnvelope, DaemonReply};
21use crate::daemon::server;
22
23/// The `worktrees` service routing key on the daemon control socket.
24const SERVICE: &str = "worktrees";
25
26/// Worktrees: see the repos/worktrees open across every VS Code window, kept
27/// live by the daemon.
28#[derive(Parser)]
29pub struct WorktreesCommand {
30    /// The worktrees subcommand to execute.
31    #[command(subcommand)]
32    pub command: WorktreesSubcommands,
33}
34
35/// Worktrees subcommands.
36#[derive(Subcommand)]
37pub enum WorktreesSubcommands {
38    /// List the repos/worktrees currently open across all windows.
39    List(ListCommand),
40    /// Show every repository and all its worktrees, grouped by repository.
41    Tree(TreeCommand),
42}
43
44impl WorktreesCommand {
45    /// Executes the worktrees command.
46    pub async fn execute(self) -> Result<()> {
47        match self.command {
48            WorktreesSubcommands::List(cmd) => cmd.execute().await,
49            WorktreesSubcommands::Tree(cmd) => cmd.execute().await,
50        }
51    }
52}
53
54/// Lists the live cross-window set of open worktrees/repos.
55#[derive(Parser)]
56pub struct ListCommand {
57    /// Control-socket path. Defaults to the per-user runtime location.
58    #[arg(long, value_name = "PATH")]
59    pub socket: Option<PathBuf>,
60    /// Output format.
61    #[arg(short = 'o', long, value_enum, default_value_t = TableOrJson::Table)]
62    pub output: TableOrJson,
63    /// Deprecated: use `-o`/`--output json` instead.
64    #[arg(long, hide = true)]
65    pub json: bool,
66}
67
68impl ListCommand {
69    /// Executes the list command.
70    pub async fn execute(mut self) -> Result<()> {
71        if self.json {
72            eprintln!("warning: --json is deprecated; use -o/--output json instead");
73            self.output = TableOrJson::Json;
74        }
75        let socket = server::resolve_socket(self.socket)?;
76        let result = call(&socket, "list", Value::Null).await?;
77        match self.output {
78            TableOrJson::Json => println!("{}", serde_json::to_string_pretty(&result)?),
79            TableOrJson::Table => println!("{}", render_windows(&result)),
80        }
81        Ok(())
82    }
83}
84
85/// Shows every repository and all of its worktrees (open or not), grouped by
86/// repository — the daemon's `tree` op, which derives the repos from the open
87/// windows and enumerates each repo's worktrees.
88#[derive(Parser)]
89pub struct TreeCommand {
90    /// Control-socket path. Defaults to the per-user runtime location.
91    #[arg(long, value_name = "PATH")]
92    pub socket: Option<PathBuf>,
93    /// Output format.
94    #[arg(short = 'o', long, value_enum, default_value_t = TableOrJson::Table)]
95    pub output: TableOrJson,
96}
97
98impl TreeCommand {
99    /// Executes the tree command.
100    pub async fn execute(self) -> Result<()> {
101        let socket = server::resolve_socket(self.socket)?;
102        let mut result = call(&socket, "tree", Value::Null).await?;
103        // Ahead/behind is no longer part of the (cheap) streamed `tree` snapshot
104        // (#1306); fetch it on demand for the worktrees we are about to render and
105        // fold it back in, so `worktrees tree` shows the same `+ahead -behind` sync
106        // state as before. Best-effort: an older daemon without the `ahead-behind`
107        // op just renders `-`.
108        enrich_ahead_behind(&socket, &mut result).await;
109        match self.output {
110            TableOrJson::Json => println!("{}", serde_json::to_string_pretty(&result)?),
111            TableOrJson::Table => println!("{}", render_tree(&result)),
112        }
113        Ok(())
114    }
115}
116
117/// Fetches ahead/behind on demand for every worktree in a `tree` reply and folds
118/// the counts back into each worktree object, so `worktrees tree` renders the same
119/// `+ahead -behind` sync state the cheap snapshot no longer carries (#1306). A
120/// best-effort enrichment: if there are no worktrees, the daemon lacks the
121/// `ahead-behind` op (older daemon), or the call fails, `result` is left as-is and
122/// the tree still renders — just with `-` for sync.
123async fn enrich_ahead_behind(socket: &Path, result: &mut Value) {
124    let paths = worktree_paths(result);
125    if paths.is_empty() {
126        return;
127    }
128    let Ok(reply) = call(socket, "ahead-behind", json!({ "paths": paths })).await else {
129        return;
130    };
131    if let Some(results) = reply.get("results").and_then(Value::as_object) {
132        merge_ahead_behind(result, results);
133    }
134}
135
136/// Every worktree path in a `tree` reply, in render order — the batch the
137/// on-demand `ahead-behind` op is asked about.
138fn worktree_paths(result: &Value) -> Vec<String> {
139    let mut paths = Vec::new();
140    for repo in result
141        .get("repos")
142        .and_then(Value::as_array)
143        .map(Vec::as_slice)
144        .unwrap_or_default()
145    {
146        for worktree in repo
147            .get("worktrees")
148            .and_then(Value::as_array)
149            .map(Vec::as_slice)
150            .unwrap_or_default()
151        {
152            if let Some(path) = worktree.get("path").and_then(Value::as_str) {
153                paths.push(path.to_string());
154            }
155        }
156    }
157    paths
158}
159
160/// Folds `{ ahead, behind }` counts (keyed by worktree path) from an `ahead-behind`
161/// reply back into a `tree` reply's worktree objects. A worktree whose path is
162/// absent from `results` (no upstream) is left untouched. Pure, so the merge is
163/// unit-testable without a socket.
164fn merge_ahead_behind(result: &mut Value, results: &serde_json::Map<String, Value>) {
165    for repo in result
166        .get_mut("repos")
167        .and_then(Value::as_array_mut)
168        .into_iter()
169        .flatten()
170    {
171        for worktree in repo
172            .get_mut("worktrees")
173            .and_then(Value::as_array_mut)
174            .into_iter()
175            .flatten()
176        {
177            // Take the worktree object up front so the insert reuses this handle
178            // rather than a second, always-succeeding `as_object_mut` (a non-object
179            // element in the array is skipped here).
180            let Some(obj) = worktree.as_object_mut() else {
181                continue;
182            };
183            let Some(path) = obj.get("path").and_then(Value::as_str).map(str::to_string) else {
184                continue;
185            };
186            let Some(counts) = results.get(&path) else {
187                continue;
188            };
189            // Fold both counts in together, or neither — a malformed entry missing
190            // a side is left as no-sync rather than half-applied.
191            if let (Some(ahead), Some(behind)) =
192                (counts.get("ahead").cloned(), counts.get("behind").cloned())
193            {
194                obj.insert("ahead".to_string(), ahead);
195                obj.insert("behind".to_string(), behind);
196            }
197        }
198    }
199}
200
201/// Sends one `worktrees` service op over the control socket, returning its
202/// payload or turning an `ok: false` reply into an error.
203async fn call(socket: &Path, op: &str, payload: Value) -> Result<Value> {
204    let reply = DaemonClient::new(socket)
205        .request(DaemonEnvelope::service(SERVICE, op, payload))
206        .await?;
207    reply_payload(reply)
208}
209
210/// Unwraps a daemon reply into its payload, turning an `ok: false` reply into an
211/// error. Pure (no socket), so both mappings are unit-testable.
212fn reply_payload(reply: DaemonReply) -> Result<Value> {
213    if reply.ok {
214        Ok(reply.payload)
215    } else {
216        bail!(
217            "daemon returned an error: {}",
218            reply.error.as_deref().unwrap_or("unknown error")
219        )
220    }
221}
222
223/// Renders a `list` reply as a human-readable table: a header and one row per
224/// open window (repo, the daemon-computed branch and its ahead/behind sync
225/// state, the primary folder, and how long ago it was last seen). Returns a
226/// placeholder line when nothing is open.
227fn render_windows(result: &Value) -> String {
228    let windows = result
229        .get("windows")
230        .and_then(Value::as_array)
231        .map(Vec::as_slice)
232        .unwrap_or_default();
233    if windows.is_empty() {
234        return "No open windows.".to_string();
235    }
236    let mut out = format!(
237        "{:<22} {:<24} {:<9} {:<40} {:>5}",
238        "REPO", "BRANCH", "SYNC", "FOLDER", "AGE"
239    );
240    for window in windows {
241        let repo = sanitize(repo_name(window));
242        let branch = sanitize(window.get("branch").and_then(Value::as_str).unwrap_or("-"));
243        let sync = sync_summary(window);
244        let folder_disp = folder_summary(window);
245        let age = age_secs(window.get("last_seen").and_then(Value::as_str));
246        out.push_str(&format!(
247            "\n{repo:<22} {branch:<24} {sync:<9} {folder_disp:<40} {age:>4}s"
248        ));
249    }
250    out
251}
252
253/// Renders a `tree` reply as a repo-grouped view: a header line per repository
254/// (its name, GitHub `owner/name` when present, and root path), then one indented
255/// row per worktree — a `*` marks the main working tree, followed by the branch,
256/// its `+ahead -behind` sync state, an `open` flag when a live window has it open,
257/// and the worktree path. Returns a placeholder when no repository is open.
258fn render_tree(result: &Value) -> String {
259    let repos = result
260        .get("repos")
261        .and_then(Value::as_array)
262        .map(Vec::as_slice)
263        .unwrap_or_default();
264    if repos.is_empty() {
265        return "No repositories open.".to_string();
266    }
267    let mut out = String::new();
268    for (i, repo) in repos.iter().enumerate() {
269        // A blank line separates repositories (but not before the first): the
270        // previous worktree row has no trailing newline, so two are needed.
271        if i > 0 {
272            out.push_str("\n\n");
273        }
274        out.push_str(&repo_header(repo));
275        for worktree in repo
276            .get("worktrees")
277            .and_then(Value::as_array)
278            .map(Vec::as_slice)
279            .unwrap_or_default()
280        {
281            out.push('\n');
282            out.push_str(&worktree_row(worktree));
283        }
284    }
285    out
286}
287
288/// The header line for one repo in the tree view: `<name>  (github: owner/name)
289/// <root>`, with the GitHub clause omitted for a non-GitHub repo.
290fn repo_header(repo: &Value) -> String {
291    let name = sanitize(repo.get("main_repo").and_then(Value::as_str).unwrap_or("-"));
292    let root = sanitize(repo.get("root").and_then(Value::as_str).unwrap_or(""));
293    match github_summary(repo) {
294        Some(github) => format!("{name}  ({github})  {root}"),
295        None => format!("{name}  {root}"),
296    }
297}
298
299/// A `github: owner/name` summary for a repo, or `None` when it has no GitHub
300/// identity (a non-GitHub or remote-less repo).
301fn github_summary(repo: &Value) -> Option<String> {
302    let owner = repo.pointer("/github/owner").and_then(Value::as_str)?;
303    let name = repo.pointer("/github/name").and_then(Value::as_str)?;
304    Some(format!("github: {}/{}", sanitize(owner), sanitize(name)))
305}
306
307/// One indented worktree row: a `*` for the main working tree, the branch, the
308/// `+ahead -behind` sync state, an `open` flag when a window has it open, and the
309/// worktree path.
310fn worktree_row(worktree: &Value) -> String {
311    let marker = if worktree.get("is_main").and_then(Value::as_bool) == Some(true) {
312        '*'
313    } else {
314        ' '
315    };
316    let branch = sanitize(
317        worktree
318            .get("branch")
319            .and_then(Value::as_str)
320            .unwrap_or("-"),
321    );
322    let sync = sync_summary(worktree);
323    let open = if worktree.get("open").and_then(Value::as_bool) == Some(true) {
324        "open"
325    } else {
326        ""
327    };
328    let path = sanitize(worktree.get("path").and_then(Value::as_str).unwrap_or(""));
329    format!("  {marker} {branch:<24} {sync:<9} {open:<5} {path}")
330}
331
332/// The repo name to show for a window: the daemon-computed `main_repo` (which
333/// names the *parent* repository of a linked worktree, not its worktree-folder
334/// basename) when present, else the companion-reported `repo`, else `-`.
335fn repo_name(window: &Value) -> &str {
336    window
337        .get("main_repo")
338        .and_then(Value::as_str)
339        .or_else(|| window.get("repo").and_then(Value::as_str))
340        .unwrap_or("-")
341}
342
343/// A compact `+ahead -behind` divergence indicator for a window, or `-` when
344/// the branch tracks no upstream (or there is no branch at all). The counts are
345/// daemon-computed integers, so no sanitizing is needed.
346fn sync_summary(window: &Value) -> String {
347    let ahead = window.get("ahead").and_then(Value::as_u64);
348    let behind = window.get("behind").and_then(Value::as_u64);
349    match (ahead, behind) {
350        (Some(ahead), Some(behind)) => format!("+{ahead} -{behind}"),
351        _ => "-".to_string(),
352    }
353}
354
355/// The primary folder of a window, with a `(+N)` suffix when it has more than
356/// one workspace folder.
357fn folder_summary(window: &Value) -> String {
358    let folders = window
359        .get("folders")
360        .and_then(Value::as_array)
361        .map(Vec::as_slice)
362        .unwrap_or_default();
363    let first = sanitize(folders.first().and_then(Value::as_str).unwrap_or(""));
364    let extra = folders.len().saturating_sub(1);
365    if extra > 0 {
366        format!("{first} (+{extra})")
367    } else {
368        first
369    }
370}
371
372/// Strips control characters (C0, DEL, C1) from an untrusted registry string so
373/// a malicious `register` payload cannot inject terminal escape sequences into
374/// the rendered table (#1137). The `--json` path stays verbatim.
375fn sanitize(s: &str) -> String {
376    s.chars().filter(|c| !c.is_control()).collect()
377}
378
379/// Seconds elapsed since an RFC 3339 timestamp (0 if absent/unparseable).
380fn age_secs(ts: Option<&str>) -> i64 {
381    ts.and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
382        .map_or(0, |t| {
383            (Utc::now() - t.with_timezone(&Utc)).num_seconds().max(0)
384        })
385}
386
387#[cfg(test)]
388#[allow(clippy::unwrap_used, clippy::expect_used)]
389mod tests {
390    use super::*;
391    use serde_json::json;
392
393    /// Mirrors the `omni-dev worktrees` argv surface for parse tests.
394    #[derive(Parser)]
395    struct Wrapper {
396        #[command(subcommand)]
397        cmd: WorktreesSubcommands,
398    }
399
400    fn parse(args: &[&str]) -> WorktreesSubcommands {
401        let mut full = vec!["omni-dev"];
402        full.extend_from_slice(args);
403        Wrapper::try_parse_from(full).unwrap().cmd
404    }
405
406    #[test]
407    fn list_parses_flags_and_defaults() {
408        // Routing: `worktrees list` maps to the List variant.
409        assert!(matches!(parse(&["list"]), WorktreesSubcommands::List(_)));
410        // Flags, via the leaf parser (clap treats argv[0] as the command name).
411        let cmd = ListCommand::try_parse_from(["list"]).unwrap();
412        assert_eq!(cmd.output, TableOrJson::Table);
413        assert!(!cmd.json);
414        assert!(cmd.socket.is_none());
415
416        let cmd =
417            ListCommand::try_parse_from(["list", "-o", "json", "--socket", "/tmp/d.sock"]).unwrap();
418        assert_eq!(cmd.output, TableOrJson::Json);
419        assert_eq!(cmd.socket.as_deref(), Some(Path::new("/tmp/d.sock")));
420    }
421
422    #[test]
423    fn list_deprecated_json_flag_still_parses() {
424        // `--json` is captured separately; `execute` folds it into `output`.
425        let cmd = ListCommand::try_parse_from(["list", "--json"]).unwrap();
426        assert!(cmd.json);
427        assert_eq!(cmd.output, TableOrJson::Table);
428    }
429
430    #[test]
431    fn tree_parses_flags_and_defaults() {
432        // Routing: `worktrees tree` maps to the Tree variant.
433        assert!(matches!(parse(&["tree"]), WorktreesSubcommands::Tree(_)));
434        let cmd = TreeCommand::try_parse_from(["tree"]).unwrap();
435        assert_eq!(cmd.output, TableOrJson::Table);
436        assert!(cmd.socket.is_none());
437
438        let cmd =
439            TreeCommand::try_parse_from(["tree", "-o", "json", "--socket", "/tmp/d.sock"]).unwrap();
440        assert_eq!(cmd.output, TableOrJson::Json);
441        assert_eq!(cmd.socket.as_deref(), Some(Path::new("/tmp/d.sock")));
442    }
443
444    #[test]
445    fn render_windows_handles_empty_replies() {
446        assert_eq!(
447            render_windows(&json!({ "windows": [] })),
448            "No open windows."
449        );
450        assert_eq!(render_windows(&json!({})), "No open windows.");
451    }
452
453    #[test]
454    fn render_windows_renders_rows() {
455        let result = json!({ "windows": [{
456            "key": "w1",
457            "repo": "omni-dev",
458            "branch": "issue-1011",
459            "ahead": 2,
460            "behind": 1,
461            "folders": ["/home/me/omni-dev", "/home/me/docs"],
462            "last_seen": "2000-01-01T00:00:00Z",
463        }]});
464        let table = render_windows(&result);
465        assert!(table.contains("omni-dev"), "{table}");
466        // The computed branch and its sync state both render.
467        assert!(table.contains("issue-1011"), "{table}");
468        assert!(table.contains("+2 -1"), "{table}");
469        // Primary folder plus a (+1) for the second workspace folder.
470        assert!(table.contains("/home/me/omni-dev (+1)"), "{table}");
471        // A header line plus exactly one data row.
472        assert_eq!(table.lines().count(), 2, "{table}");
473    }
474
475    #[test]
476    fn render_windows_prefers_main_repo_over_companion_repo() {
477        // A linked worktree: the companion reports the worktree-folder basename,
478        // but the daemon-computed `main_repo` names the parent repo, and that is
479        // what the REPO column shows.
480        let result = json!({ "windows": [{
481            "key": "w1",
482            "repo": "issue-1250",
483            "main_repo": "omni-dev",
484            "branch": "issue-1250",
485            "folders": ["/home/me/worktrees/issue-1250"],
486            "last_seen": "2000-01-01T00:00:00Z",
487        }]});
488        let table = render_windows(&result);
489        assert!(table.contains("omni-dev"), "{table}");
490        // The misleading worktree-folder basename does not appear in REPO (it is
491        // still visible in the FOLDER column path).
492        let data_row = table.lines().nth(1).unwrap();
493        assert!(data_row.starts_with("omni-dev"), "{data_row}");
494    }
495
496    #[test]
497    fn repo_name_falls_back_to_companion_repo_then_dash() {
498        assert_eq!(
499            repo_name(&json!({ "main_repo": "omni-dev", "repo": "wt" })),
500            "omni-dev"
501        );
502        assert_eq!(repo_name(&json!({ "repo": "wt" })), "wt");
503        assert_eq!(repo_name(&json!({})), "-");
504    }
505
506    #[test]
507    fn render_windows_strips_control_bytes() {
508        // C0 (ESC, CR, BEL), DEL, and C1 (CSI) bytes in every string-valued
509        // field must not reach the terminal (#1137).
510        let result = json!({ "windows": [{
511            "key": "w1",
512            "repo": "evil\x1b[31mrepo",
513            "branch": "br\ranch\x07\u{9b}2J",
514            "folders": ["/tmp/a\x1b]0;owned\x07\u{7f}", "/tmp/b"],
515            "last_seen": "2000-01-01T00:00:00Z",
516        }]});
517        let table = render_windows(&result);
518        assert!(
519            !table.contains(|c: char| c.is_control() && c != '\n'),
520            "{table:?}"
521        );
522        // Visible text survives with only the control bytes removed.
523        assert!(table.contains("evil[31mrepo"), "{table:?}");
524        assert!(table.contains("branch2J"), "{table:?}");
525        assert!(table.contains("/tmp/a]0;owned (+1)"), "{table:?}");
526        // Embedded CR/LF cannot forge extra rows: header plus one data row.
527        assert_eq!(table.lines().count(), 2, "{table:?}");
528    }
529
530    #[test]
531    fn sync_summary_formats_or_dashes() {
532        assert_eq!(sync_summary(&json!({ "ahead": 2, "behind": 1 })), "+2 -1");
533        assert_eq!(sync_summary(&json!({ "ahead": 0, "behind": 0 })), "+0 -0");
534        // Branch present but no upstream, or nothing at all → a dash.
535        assert_eq!(sync_summary(&json!({ "branch": "main" })), "-");
536        assert_eq!(sync_summary(&json!({})), "-");
537    }
538
539    #[test]
540    fn folder_summary_strips_control_bytes() {
541        assert_eq!(
542            folder_summary(&json!({ "folders": ["/a\x1b[2J/b"] })),
543            "/a[2J/b"
544        );
545    }
546
547    #[test]
548    fn folder_summary_counts_extra_folders() {
549        assert_eq!(folder_summary(&json!({ "folders": [] })), "");
550        assert_eq!(folder_summary(&json!({ "folders": ["/a"] })), "/a");
551        assert_eq!(
552            folder_summary(&json!({ "folders": ["/a", "/b", "/c"] })),
553            "/a (+2)"
554        );
555    }
556
557    #[test]
558    fn age_secs_handles_absent_and_unparseable_and_past() {
559        assert_eq!(age_secs(None), 0);
560        assert_eq!(age_secs(Some("not-a-timestamp")), 0);
561        assert!(age_secs(Some("2000-01-01T00:00:00Z")) > 0);
562    }
563
564    #[test]
565    fn render_tree_handles_empty_replies() {
566        assert_eq!(
567            render_tree(&json!({ "repos": [] })),
568            "No repositories open."
569        );
570        assert_eq!(render_tree(&json!({})), "No repositories open.");
571    }
572
573    #[test]
574    fn worktree_paths_collects_every_worktree_in_render_order() {
575        let result = json!({ "repos": [
576            // The middle worktree has no `path` and is skipped, not collected.
577            { "worktrees": [ { "path": "/a" }, { "branch": "detached" }, { "path": "/b" } ] },
578            { "worktrees": [ { "path": "/c" } ] },
579        ]});
580        assert_eq!(worktree_paths(&result), vec!["/a", "/b", "/c"]);
581        // No repos / no worktrees → an empty batch (nothing to fetch).
582        assert!(worktree_paths(&json!({})).is_empty());
583        assert!(worktree_paths(&json!({ "repos": [{ "worktrees": [] }] })).is_empty());
584    }
585
586    #[test]
587    fn merge_ahead_behind_folds_counts_by_path_and_leaves_others() {
588        // The on-demand `ahead-behind` op reports one worktree diverging and omits
589        // the other (no upstream). The merge folds the counts onto the matching
590        // path and leaves the untracked worktree without sync fields.
591        let mut result = json!({ "repos": [{ "worktrees": [
592            { "path": "/a", "branch": "main" },
593            { "path": "/b", "branch": "feature" },
594        ]}]});
595        let results = json!({ "/a": { "ahead": 2, "behind": 1 } });
596        merge_ahead_behind(&mut result, results.as_object().unwrap());
597
598        let worktrees = result.pointer("/repos/0/worktrees").unwrap();
599        let a = &worktrees[0];
600        assert_eq!(a.get("ahead").and_then(Value::as_u64), Some(2));
601        assert_eq!(a.get("behind").and_then(Value::as_u64), Some(1));
602        // And it renders exactly as an eager snapshot would have.
603        assert_eq!(sync_summary(a), "+2 -1");
604        let b = &worktrees[1];
605        assert!(b.get("ahead").is_none(), "{b:?}");
606        assert!(b.get("behind").is_none(), "{b:?}");
607        assert_eq!(sync_summary(b), "-");
608    }
609
610    #[test]
611    fn merge_ahead_behind_skips_malformed_worktrees_and_counts() {
612        // Every defensive guard, on malformed input that never comes from a real
613        // daemon: a non-object array element, a worktree with no `path`, and a
614        // results entry missing a side. None panics; none is half-applied.
615        let mut result = json!({ "repos": [{ "worktrees": [
616            "not-an-object",                       // non-object element → skipped
617            { "branch": "detached" },              // object, but no path → skipped
618            { "path": "/a", "branch": "main" },    // matched, but counts malformed
619        ]}]});
620        let results = json!({ "/a": { "ahead": 2 } }); // missing `behind`
621        merge_ahead_behind(&mut result, results.as_object().unwrap());
622
623        let worktrees = result.pointer("/repos/0/worktrees").unwrap();
624        // Non-object element is untouched.
625        assert_eq!(worktrees[0], json!("not-an-object"));
626        // Pathless worktree: no sync fields inserted.
627        assert!(worktrees[1].get("ahead").is_none(), "{:?}", worktrees[1]);
628        // Malformed counts: neither side folded in (both-or-nothing).
629        assert!(worktrees[2].get("ahead").is_none(), "{:?}", worktrees[2]);
630        assert!(worktrees[2].get("behind").is_none(), "{:?}", worktrees[2]);
631    }
632
633    #[tokio::test]
634    async fn enrich_ahead_behind_is_a_noop_when_there_are_no_worktrees() {
635        // No worktrees → no batch to fetch → early return before any socket call,
636        // so even a nonexistent socket leaves the tree untouched.
637        let mut result = json!({ "repos": [] });
638        let before = result.clone();
639        enrich_ahead_behind(Path::new("/nonexistent/omni-dev-ab.sock"), &mut result).await;
640        assert_eq!(result, before);
641    }
642
643    #[tokio::test]
644    async fn enrich_ahead_behind_leaves_the_tree_when_the_daemon_is_unreachable() {
645        // A real worktree but no daemon at the socket → the call fails and the tree
646        // is returned as-is (rendered with `-` for sync), never erroring.
647        let mut result =
648            json!({ "repos": [{ "worktrees": [{ "path": "/x", "branch": "main" }] }] });
649        enrich_ahead_behind(Path::new("/nonexistent/omni-dev-ab.sock"), &mut result).await;
650        let wt = result.pointer("/repos/0/worktrees/0").unwrap();
651        assert!(wt.get("ahead").is_none(), "{wt:?}");
652        assert!(wt.get("behind").is_none(), "{wt:?}");
653    }
654
655    /// Spawns a minimal fake daemon on a short-path Unix socket that answers the
656    /// one `ahead-behind` request with `reply` (the daemon's NDJSON reply shape).
657    /// Returns the temp dir (kept alive for the socket's lifetime), the socket
658    /// path, and the server task.
659    fn fake_daemon_reply(
660        reply: Value,
661    ) -> (tempfile::TempDir, PathBuf, tokio::task::JoinHandle<()>) {
662        use futures::{SinkExt, StreamExt};
663        use tokio::net::UnixListener;
664        use tokio_util::codec::{Framed, LinesCodec};
665
666        // A short base path keeps the socket under the 104-byte `sockaddr_un` limit.
667        let dir = tempfile::tempdir_in("/tmp").unwrap();
668        let sock = dir.path().join("d.sock");
669        let listener = UnixListener::bind(&sock).unwrap();
670        let server = tokio::spawn(async move {
671            let (stream, _) = listener.accept().await.unwrap();
672            let mut framed = Framed::new(stream, LinesCodec::new());
673            let _req = framed.next().await.unwrap().unwrap();
674            framed
675                .send(serde_json::to_string(&reply).unwrap())
676                .await
677                .unwrap();
678        });
679        (dir, sock, server)
680    }
681
682    #[tokio::test]
683    async fn enrich_ahead_behind_folds_counts_from_a_live_socket() {
684        let (_dir, sock, server) = fake_daemon_reply(
685            json!({ "ok": true, "payload": { "results": { "/x": { "ahead": 3, "behind": 4 } } } }),
686        );
687        let mut result =
688            json!({ "repos": [{ "worktrees": [{ "path": "/x", "branch": "main" }] }] });
689        enrich_ahead_behind(&sock, &mut result).await;
690        server.await.unwrap();
691
692        let wt = result.pointer("/repos/0/worktrees/0").unwrap();
693        assert_eq!(wt.get("ahead").and_then(Value::as_u64), Some(3));
694        assert_eq!(wt.get("behind").and_then(Value::as_u64), Some(4));
695    }
696
697    #[tokio::test]
698    async fn enrich_ahead_behind_ignores_a_reply_without_results() {
699        // An `ok` reply carrying no `results` object (an older/oddly-shaped daemon)
700        // leaves the tree unchanged rather than erroring.
701        let (_dir, sock, server) = fake_daemon_reply(json!({ "ok": true, "payload": {} }));
702        let mut result =
703            json!({ "repos": [{ "worktrees": [{ "path": "/x", "branch": "main" }] }] });
704        enrich_ahead_behind(&sock, &mut result).await;
705        server.await.unwrap();
706
707        let wt = result.pointer("/repos/0/worktrees/0").unwrap();
708        assert!(wt.get("ahead").is_none(), "{wt:?}");
709        assert!(wt.get("behind").is_none(), "{wt:?}");
710    }
711
712    #[test]
713    fn render_tree_groups_repos_and_worktrees() {
714        let result = json!({ "repos": [{
715            "main_repo": "omni-dev",
716            "github": { "owner": "rust-works", "name": "omni-dev" },
717            "root": "/home/me/omni-dev",
718            "worktrees": [
719                { "path": "/home/me/omni-dev", "branch": "main", "ahead": 2, "behind": 0,
720                  "is_main": true, "open": true, "window_key": "w1" },
721                { "path": "/home/me/wt/issue-1300", "branch": "issue-1300", "ahead": 1, "behind": 3,
722                  "is_main": false, "open": false },
723            ],
724        }]});
725        let out = render_tree(&result);
726        // Repo header carries the GitHub identity and root.
727        let header = out.lines().next().unwrap();
728        assert!(header.contains("omni-dev"), "{out}");
729        assert!(header.contains("github: rust-works/omni-dev"), "{out}");
730        assert!(header.contains("/home/me/omni-dev"), "{out}");
731        // The main working tree is marked with `*`, its sync, and `open`.
732        assert!(
733            out.lines()
734                .any(|l| l.contains("* main") && l.contains("+2 -0") && l.contains("open")),
735            "{out}"
736        );
737        // The linked worktree is unmarked and not flagged open.
738        let linked = out
739            .lines()
740            .find(|l| l.contains("issue-1300"))
741            .unwrap_or_default();
742        assert!(!linked.contains('*'), "{linked}");
743        assert!(!linked.contains("open"), "{linked}");
744        assert!(linked.contains("+1 -3"), "{linked}");
745        // Header + two worktree rows.
746        assert_eq!(out.lines().count(), 3, "{out}");
747    }
748
749    #[test]
750    fn render_tree_separates_multiple_repos_with_blank_line() {
751        let result = json!({ "repos": [
752            {
753                "main_repo": "alpha",
754                "root": "/r/alpha",
755                "worktrees": [
756                    { "path": "/r/alpha", "branch": "main", "is_main": true, "open": false },
757                ],
758            },
759            {
760                "main_repo": "beta",
761                "root": "/r/beta",
762                "worktrees": [
763                    { "path": "/r/beta", "branch": "main", "is_main": true, "open": false },
764                ],
765            },
766        ]});
767        let out = render_tree(&result);
768        // Two headers, two worktree rows, and one blank separator between repos.
769        assert!(
770            out.contains("\n\nbeta"),
771            "repos not blank-separated: {out:?}"
772        );
773        let alpha = out.find("alpha").unwrap();
774        let beta = out.find("beta").unwrap();
775        assert!(alpha < beta, "repo order not preserved: {out}");
776        assert_eq!(out.lines().count(), 5, "{out:?}");
777    }
778
779    #[test]
780    fn render_tree_omits_github_for_non_github_repo() {
781        let result = json!({ "repos": [{
782            "main_repo": "internal",
783            "root": "/srv/internal",
784            "worktrees": [
785                { "path": "/srv/internal", "branch": "main", "is_main": true, "open": false },
786            ],
787        }]});
788        let out = render_tree(&result);
789        assert!(!out.contains("github:"), "{out}");
790        assert!(out.lines().next().unwrap().contains("internal"), "{out}");
791    }
792
793    #[test]
794    fn render_tree_strips_control_bytes() {
795        // Control bytes in the repo name, github identity, branch, and path must
796        // not reach the terminal (#1137), matching the `list` renderer.
797        let result = json!({ "repos": [{
798            "main_repo": "evil\x1b[31mrepo",
799            "github": { "owner": "ow\x07ner", "name": "na\u{9b}2Jme" },
800            "root": "/tmp/r\x1b]0;x\x07oot",
801            "worktrees": [
802                { "path": "/tmp/w\rt", "branch": "br\x1b[2Janch", "is_main": true, "open": true },
803            ],
804        }]});
805        let out = render_tree(&result);
806        assert!(
807            !out.contains(|c: char| c.is_control() && c != '\n'),
808            "{out:?}"
809        );
810        // Embedded CR/LF cannot forge extra lines: header plus one worktree row.
811        assert_eq!(out.lines().count(), 2, "{out:?}");
812    }
813
814    #[test]
815    fn github_summary_needs_both_owner_and_name() {
816        assert_eq!(
817            github_summary(&json!({ "github": { "owner": "o", "name": "n" } })).as_deref(),
818            Some("github: o/n")
819        );
820        assert_eq!(github_summary(&json!({ "github": { "owner": "o" } })), None);
821        assert_eq!(github_summary(&json!({})), None);
822    }
823
824    #[test]
825    fn reply_payload_unwraps_ok_and_maps_errors() {
826        // ok → payload.
827        assert_eq!(
828            reply_payload(DaemonReply::ok(json!({ "a": 1 }))).unwrap(),
829            json!({ "a": 1 })
830        );
831        // ok: false with a message → that message.
832        let err = reply_payload(DaemonReply::err("boom")).unwrap_err();
833        assert!(err.to_string().contains("boom"), "{err}");
834        // ok: false with no message → the "unknown error" fallback.
835        let err = reply_payload(DaemonReply {
836            ok: false,
837            payload: Value::Null,
838            error: None,
839        })
840        .unwrap_err();
841        assert!(err.to_string().contains("unknown error"), "{err}");
842    }
843}