aaronrene / gabriel-muse public
worktree.py python
664 lines 24.5 KB
Raw
sha256:0e3697f29ae6c62f39818a19861d4db6276465de16481fadf866eb93cbf4aae4 test: validate staging push and MP workflow for aaronrene Human 17 days ago
1 """``muse worktree`` — manage multiple simultaneous branch checkouts.
2
3 Worktrees let you work on multiple branches at once without stashing or
4 switching — each worktree is an independent working directory, but they all
5 share the same ``.muse/`` object store.
6
7 This is especially powerful for agents: one agent per worktree, each
8 autonomously developing a feature on its own branch, with zero interference.
9 Use ``--json`` on any subcommand for machine-readable output.
10
11 Subcommands::
12
13 muse worktree add <name> <branch> [--path PATH] [-b NEW_BRANCH] [--json]
14 muse worktree list [--json]
15 muse worktree prune [--dry-run] [--json]
16 muse worktree remove <name> [--force] [--json]
17 muse worktree repair [--json]
18 muse worktree status <name> [--json]
19
20 Layout::
21
22 myproject/ ← main worktree (holds .muse/)
23 myproject-feat-audio/ ← linked worktree for feat/audio
24
25 Exit codes::
26
27 0 — success
28 1 — user error (invalid name/branch, worktree not found, path conflict)
29 2 — internal error
30 """
31
32 from __future__ import annotations
33
34 import argparse
35 import json
36 import logging
37 import pathlib
38 import sys
39 from typing import TypedDict
40
41 from muse.core.errors import ExitCode
42 from muse.core.repo import require_repo
43 from muse.core.store import get_head_commit_id, write_text_atomic
44 from muse.core.validation import sanitize_display, validate_branch_name
45 from muse.core.worktree import (
46 WorktreeInfo,
47 WorktreeStatusResult,
48 add_worktree,
49 get_worktree_status,
50 list_worktrees,
51 prune_worktrees,
52 remove_worktree,
53 repair_worktree_pointers,
54 )
55
56 logger = logging.getLogger(__name__)
57
58
59 # ---------------------------------------------------------------------------
60 # Typed JSON schemas
61 # ---------------------------------------------------------------------------
62
63
64 class _WorktreeAddJson(TypedDict):
65 """JSON output of ``muse worktree add --json``."""
66
67 name: str
68 branch: str
69 path: str
70 head_commit: str | None
71
72
73 class _WorktreeListEntryJson(TypedDict):
74 """One entry in the JSON array from ``muse worktree list --json``."""
75
76 name: str
77 branch: str
78 path: str
79 head_commit: str | None
80 is_main: bool
81
82
83 class _WorktreeRemoveJson(TypedDict):
84 """JSON output of ``muse worktree remove --json``."""
85
86 name: str
87 status: str
88
89
90 class _WorktreePruneJson(TypedDict):
91 """JSON output of ``muse worktree prune --json``."""
92
93 pruned: list[str]
94 count: int
95 dry_run: bool
96
97
98 # ---------------------------------------------------------------------------
99 # Internal helpers
100 # ---------------------------------------------------------------------------
101
102
103 def _fmt_info(wt: WorktreeInfo) -> str:
104 """Format a worktree row for human-readable ``list`` output."""
105 prefix = "* " if wt.is_main else " "
106 head = wt.head_commit[:12] if wt.head_commit else "(no commits)"
107 return (
108 f"{prefix}{sanitize_display(wt.name):<24} "
109 f"{sanitize_display(wt.branch):<30} "
110 f"{head} "
111 f"{sanitize_display(str(wt.path))}"
112 )
113
114
115 def _info_to_json(wt: WorktreeInfo) -> _WorktreeListEntryJson:
116 return _WorktreeListEntryJson(
117 name=sanitize_display(wt.name),
118 branch=sanitize_display(wt.branch),
119 path=str(wt.path),
120 head_commit=wt.head_commit,
121 is_main=wt.is_main,
122 )
123
124
125 # ---------------------------------------------------------------------------
126 # Subcommand handlers
127 # ---------------------------------------------------------------------------
128
129
130 def run_worktree_add(args: argparse.Namespace) -> None:
131 """Create a new linked worktree checked out at *branch*.
132
133 The new worktree is placed at ``<repo-parent>/<repo-name>-<name>`` by
134 default, or at ``--path PATH`` when specified. Its working directory is
135 pre-populated from the branch's latest snapshot.
136
137 When ``-b`` / ``--create-branch`` is given, a new branch is created from
138 BRANCH (used as the start point) and the worktree checks out that new
139 branch. This mirrors ``git worktree add -b <new-branch> <path> <base>``.
140
141 JSON schema::
142
143 {
144 "name": "<worktree-name>",
145 "branch": "<branch>",
146 "path": "<absolute-path>",
147 "head_commit": "<sha256> | null"
148 }
149
150 Exit codes::
151
152 0 — success
153 1 — invalid branch name, branch does not exist, worktree already
154 exists, directory already exists, or path validation error
155 2 — not inside a Muse repository
156 """
157 name: str = args.name
158 branch: str = args.branch
159 create_branch: str | None = getattr(args, "create_branch", None)
160 worktree_path_arg: str | None = getattr(args, "worktree_path", None)
161 explicit_path: pathlib.Path | None = (
162 pathlib.Path(worktree_path_arg).resolve() if worktree_path_arg else None
163 )
164 output_json: bool = args.output_json
165
166 root = require_repo()
167
168 # When -b is given, create the new branch from BRANCH (start point) first.
169 checkout_branch = branch
170 if create_branch is not None:
171 try:
172 validate_branch_name(create_branch)
173 except ValueError as exc:
174 print(f"❌ Invalid branch name: {sanitize_display(str(exc))}", file=sys.stderr)
175 raise SystemExit(ExitCode.USER_ERROR)
176
177 new_ref = root / ".muse" / "refs" / "heads" / create_branch
178 if new_ref.exists():
179 print(
180 f"❌ Branch '{sanitize_display(create_branch)}' already exists.",
181 file=sys.stderr,
182 )
183 raise SystemExit(ExitCode.USER_ERROR)
184
185 start_commit = get_head_commit_id(root, branch)
186 if start_commit is None:
187 print(
188 f"❌ Start-point branch '{sanitize_display(branch)}' has no commits.",
189 file=sys.stderr,
190 )
191 raise SystemExit(ExitCode.USER_ERROR)
192
193 new_ref.parent.mkdir(parents=True, exist_ok=True)
194 write_text_atomic(new_ref, f"{start_commit}\n")
195 checkout_branch = create_branch
196
197 try:
198 wt_path = add_worktree(root, name, checkout_branch, path=explicit_path)
199 except ValueError as exc:
200 print(f"❌ {exc}", file=sys.stderr)
201 raise SystemExit(ExitCode.USER_ERROR)
202
203 if output_json:
204 print(json.dumps(_WorktreeAddJson(
205 name=name,
206 branch=checkout_branch,
207 path=str(wt_path),
208 head_commit=get_head_commit_id(root, checkout_branch),
209 )))
210 else:
211 print(f"✅ Worktree '{sanitize_display(name)}' created at {sanitize_display(str(wt_path))}")
212 print(f" Branch: {sanitize_display(checkout_branch)}")
213
214
215 def run_worktree_list(args: argparse.Namespace) -> None:
216 """List all worktrees (main + linked).
217
218 The main worktree is always first; linked worktrees follow in
219 lexicographic order of name. ``name`` and ``branch`` in both text and
220 JSON output are sanitized — ANSI control sequences are stripped.
221
222 JSON schema (array element)::
223
224 {
225 "name": "<worktree-name>",
226 "branch": "<branch>",
227 "path": "<absolute-path>",
228 "head_commit": "<sha256> | null",
229 "is_main": true | false
230 }
231
232 Exit codes::
233
234 0 — success (even when no linked worktrees exist)
235 2 — not inside a Muse repository
236 """
237 output_json: bool = args.output_json
238
239 root = require_repo()
240 worktrees = list_worktrees(root)
241
242 if output_json:
243 print(json.dumps([_info_to_json(wt) for wt in worktrees]))
244 return
245
246 if not worktrees:
247 print("No worktrees.")
248 return
249 header = f"{' name':<26} {'branch':<30} {'HEAD':12} path"
250 print(header)
251 print("-" * len(header))
252 for wt in worktrees:
253 print(_fmt_info(wt))
254
255
256 def run_worktree_status(args: argparse.Namespace) -> None:
257 """Show the status of a single worktree.
258
259 Reports the worktree's branch, HEAD commit, and whether the working
260 directory is present on disk. Useful for agents checking whether a
261 peer agent's worktree is still active.
262
263 ``name`` and ``branch`` in both text and JSON output are sanitized —
264 ANSI control sequences are stripped before display or serialisation.
265
266 JSON schema::
267
268 {
269 "name": "<worktree-name>",
270 "branch": "<branch>",
271 "path": "<absolute-path>",
272 "head_commit": "<sha256> | null",
273 "present": true | false,
274 "is_main": true | false
275 }
276
277 Exit codes::
278
279 0 — success
280 1 — worktree does not exist, or name is invalid
281 2 — not inside a Muse repository
282 """
283 name: str = args.name
284 output_json: bool = args.output_json
285
286 root = require_repo()
287 try:
288 status = get_worktree_status(root, name)
289 except ValueError as exc:
290 print(f"❌ {exc}", file=sys.stderr)
291 raise SystemExit(ExitCode.USER_ERROR)
292
293 if output_json:
294 print(json.dumps(WorktreeStatusResult(
295 name=sanitize_display(status["name"]),
296 branch=sanitize_display(status["branch"]),
297 path=sanitize_display(status["path"]),
298 head_commit=status["head_commit"],
299 present=status["present"],
300 is_main=status["is_main"],
301 )))
302 return
303
304 present_flag = "✅ present" if status["present"] else "❌ missing"
305 head = status["head_commit"][:12] if status["head_commit"] else "(no commits)"
306 main_flag = " [main]" if status["is_main"] else ""
307 print(f"worktree: {sanitize_display(status['name'])}{main_flag}")
308 print(f"branch: {sanitize_display(status['branch'])}")
309 print(f"HEAD: {head}")
310 print(f"path: {sanitize_display(status['path'])}")
311 print(f"status: {present_flag}")
312
313
314 def run_worktree_remove(args: argparse.Namespace) -> None:
315 """Remove a linked worktree and its working directory.
316
317 The branch itself is not deleted — only the worktree directory and its
318 metadata are removed. Commits already pushed from the worktree remain
319 in the shared object store.
320
321 The main worktree cannot be removed. Removal is refused if the worktree
322 path is a symlink or resolves inside ``.muse/`` (safety guard).
323
324 JSON schema::
325
326 {
327 "name": "<worktree-name>",
328 "status": "removed"
329 }
330
331 Exit codes::
332
333 0 — success
334 1 — worktree does not exist, name is invalid, or removal refused
335 (symlink / .muse overlap)
336 2 — not inside a Muse repository
337 """
338 name: str = args.name
339 force: bool = args.force
340 output_json: bool = args.output_json
341
342 root = require_repo()
343 try:
344 remove_worktree(root, name, force=force)
345 except ValueError as exc:
346 print(f"❌ {exc}", file=sys.stderr)
347 raise SystemExit(ExitCode.USER_ERROR)
348
349 if output_json:
350 print(json.dumps(_WorktreeRemoveJson(name=sanitize_display(name), status="removed")))
351 else:
352 print(f"✅ Worktree '{sanitize_display(name)}' removed.")
353
354
355 def run_worktree_prune(args: argparse.Namespace) -> None:
356 """Remove metadata entries for worktrees whose directories no longer exist.
357
358 A worktree is considered stale when its registered path is absent from
359 the filesystem. Prune deletes the metadata file and HEAD pointer for
360 each stale entry; it does not touch any branch refs.
361
362 Pass ``--dry-run`` to preview which entries would be pruned without
363 making any filesystem changes. All names in both text and JSON output
364 are sanitized — ANSI control sequences are stripped.
365
366 JSON schema::
367
368 {
369 "pruned": ["<name>", ...],
370 "count": <int>,
371 "dry_run": true | false
372 }
373
374 Exit codes::
375
376 0 — success (even when nothing was pruned)
377 2 — not inside a Muse repository
378 """
379 output_json: bool = args.output_json
380 dry_run: bool = args.dry_run
381
382 root = require_repo()
383 pruned = prune_worktrees(root, dry_run=dry_run)
384
385 if output_json:
386 print(json.dumps(_WorktreePruneJson(
387 pruned=[sanitize_display(n) for n in pruned],
388 count=len(pruned),
389 dry_run=dry_run,
390 )))
391 return
392
393 if not pruned:
394 print("Nothing to prune.")
395 return
396 prefix = "[dry-run] would prune" if dry_run else "pruned"
397 for name in pruned:
398 print(f" {prefix}: {sanitize_display(name)}")
399 action = "Would prune" if dry_run else "Pruned"
400 print(f"{action} {len(pruned)} stale worktree(s).")
401
402
403 def run_worktree_repair(args: argparse.Namespace) -> None:
404 """Write missing .muse pointer files for all registered linked worktrees.
405
406 Idempotent — safe to run multiple times. Re-writes the pointer even
407 when it already exists. Skip worktrees whose directories are absent
408 (they should be pruned instead).
409
410 Useful to retroactively fix worktrees created before pointer-file support
411 was added, or to repair a pointer that was accidentally deleted.
412
413 All names in both text and JSON output are sanitized — ANSI control
414 sequences are stripped.
415
416 JSON schema::
417
418 {
419 "repaired": ["<worktree-name>", ...]
420 }
421
422 Exit codes::
423
424 0 — success (even when nothing needed repair)
425 2 — not inside a Muse repository
426 """
427 output_json: bool = args.output_json
428 repo_root = require_repo()
429 repaired = repair_worktree_pointers(repo_root)
430 if output_json:
431 print(json.dumps({"repaired": [sanitize_display(n) for n in repaired]}))
432 else:
433 if repaired:
434 for name in repaired:
435 print(f" repaired {sanitize_display(name)}")
436 print(f"\n✅ {len(repaired)} worktree(s) repaired.")
437 else:
438 print("✅ All worktrees already have pointer files.")
439
440
441 # ---------------------------------------------------------------------------
442 # Registration
443 # ---------------------------------------------------------------------------
444
445
446 def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None:
447 """Register the ``worktree`` subcommand."""
448 parser = subparsers.add_parser(
449 "worktree",
450 help="Manage multiple simultaneous branch checkouts.",
451 description=__doc__,
452 formatter_class=argparse.RawDescriptionHelpFormatter,
453 )
454 subs = parser.add_subparsers(dest="subcommand", metavar="SUBCOMMAND")
455 subs.required = True
456
457 # ── add ───────────────────────────────────────────────────────────────────
458 add_p = subs.add_parser(
459 "add",
460 help="Create a new linked worktree checked out at a branch.",
461 description=(
462 "Create a new linked worktree checked out at BRANCH.\n\n"
463 "The worktree is placed at <repo-parent>/<repo-name>-NAME by default,\n"
464 "or at --path PATH when specified.\n\n"
465 "Branch creation\n"
466 "---------------\n"
467 " -b NEW_BRANCH Create NEW_BRANCH from BRANCH (start point) and\n"
468 " check it out in the new worktree.\n\n"
469 "Agent quickstart\n"
470 "----------------\n"
471 " muse worktree add feat-x feat/x --json\n"
472 " muse worktree add my-wt main -b feat/new --json\n"
473 " muse worktree add custom main --path /tmp/my-wt --json\n\n"
474 "JSON output schema\n"
475 "------------------\n"
476 ' {"name": "<name>", "branch": "<branch>",\n'
477 ' "path": "<absolute-path>", "head_commit": "<sha256> | null"}\n\n'
478 "Exit codes\n"
479 "----------\n"
480 " 0 — success\n"
481 " 1 — invalid name/branch, branch missing, worktree/dir already exists\n"
482 " 2 — not inside a Muse repository\n"
483 ),
484 formatter_class=argparse.RawDescriptionHelpFormatter,
485 )
486 add_p.add_argument("name", metavar="NAME", help="Worktree name.")
487 add_p.add_argument("branch", metavar="BRANCH", help="Branch to check out.")
488 add_p.add_argument(
489 "--path", metavar="PATH", dest="worktree_path", default=None,
490 help=(
491 "Explicit filesystem path for the worktree directory. "
492 "When omitted, placed at <repo-parent>/<repo-name>-<name>."
493 ),
494 )
495 add_p.add_argument(
496 "-b", "--create-branch", metavar="NEW_BRANCH", dest="create_branch", default=None,
497 help=(
498 "Create NEW_BRANCH from BRANCH (start point) and check out the "
499 "new branch in the worktree."
500 ),
501 )
502 add_p.add_argument(
503 "--json", "-j", action="store_true", dest="output_json", default=False,
504 help="Emit machine-readable JSON to stdout.",
505 )
506 add_p.set_defaults(func=run_worktree_add)
507
508 # ── list ──────────────────────────────────────────────────────────────────
509 list_p = subs.add_parser(
510 "list",
511 help="List all worktrees (main + linked).",
512 description=(
513 "List all worktrees: the main worktree and every linked worktree.\n\n"
514 "The main worktree is always first; linked worktrees follow in\n"
515 "lexicographic order of name. Output is sanitized — ANSI control\n"
516 "sequences are stripped from name, branch, and path.\n\n"
517 "Agent quickstart\n"
518 "----------------\n"
519 " muse worktree list --json\n\n"
520 "JSON output schema (array element)\n"
521 "----------------------------------\n"
522 ' {"name": "<name>", "branch": "<branch>", "path": "<absolute-path>",\n'
523 ' "head_commit": "<sha256> | null", "is_main": true|false}\n\n'
524 "Exit codes\n"
525 "----------\n"
526 " 0 — success (even when no linked worktrees exist)\n"
527 " 2 — not inside a Muse repository\n"
528 ),
529 formatter_class=argparse.RawDescriptionHelpFormatter,
530 )
531 list_p.add_argument(
532 "--json", "-j", action="store_true", dest="output_json", default=False,
533 help="Emit machine-readable JSON array to stdout.",
534 )
535 list_p.set_defaults(func=run_worktree_list)
536
537 # ── prune ─────────────────────────────────────────────────────────────────
538 prune_p = subs.add_parser(
539 "prune",
540 help="Remove metadata entries for missing worktrees.",
541 description=(
542 "Remove metadata for worktrees whose directories no longer exist.\n\n"
543 "A worktree is stale when its registered path is absent from disk.\n"
544 "Prune removes the metadata file and HEAD pointer for each stale\n"
545 "entry; branch refs are never touched.\n\n"
546 "Use --dry-run to preview without making any changes.\n\n"
547 "Agent quickstart\n"
548 "----------------\n"
549 " muse worktree prune --json\n"
550 " muse worktree prune --dry-run --json\n\n"
551 "JSON output schema\n"
552 "------------------\n"
553 ' {"pruned": ["<name>", ...], "count": <int>, "dry_run": true|false}\n\n'
554 "Exit codes\n"
555 "----------\n"
556 " 0 — success (even when nothing was pruned)\n"
557 " 2 — not inside a Muse repository\n"
558 ),
559 formatter_class=argparse.RawDescriptionHelpFormatter,
560 )
561 prune_p.add_argument(
562 "--dry-run", action="store_true", dest="dry_run", default=False,
563 help="Preview what would be pruned without making changes.",
564 )
565 prune_p.add_argument(
566 "--json", "-j", action="store_true", dest="output_json", default=False,
567 help="Emit machine-readable JSON to stdout.",
568 )
569 prune_p.set_defaults(func=run_worktree_prune)
570
571 # ── remove ────────────────────────────────────────────────────────────────
572 remove_p = subs.add_parser(
573 "remove",
574 help="Remove a linked worktree and its working directory.",
575 description=(
576 "Remove a linked worktree: deletes the working directory and its\n"
577 "metadata. The branch is NOT deleted. Already-pushed commits\n"
578 "remain in the shared object store.\n\n"
579 "Safety guards refuse removal when the worktree path is a symlink\n"
580 "or resolves inside .muse/ (would corrupt the repository).\n\n"
581 "Agent quickstart\n"
582 "----------------\n"
583 " muse worktree remove feat-x --json\n\n"
584 "JSON output schema\n"
585 "------------------\n"
586 ' {"name": "<name>", "status": "removed"}\n\n'
587 "Exit codes\n"
588 "----------\n"
589 " 0 — success\n"
590 " 1 — worktree missing, invalid name, or removal refused\n"
591 " 2 — not inside a Muse repository\n"
592 ),
593 formatter_class=argparse.RawDescriptionHelpFormatter,
594 )
595 remove_p.add_argument("name", metavar="NAME", help="Worktree name to remove.")
596 remove_p.add_argument(
597 "--force", action="store_true",
598 help="Force removal (accepted for interface compatibility).",
599 )
600 remove_p.add_argument(
601 "--json", "-j", action="store_true", dest="output_json", default=False,
602 help="Emit machine-readable JSON to stdout.",
603 )
604 remove_p.set_defaults(func=run_worktree_remove)
605
606 # ── repair ────────────────────────────────────────────────────────────────
607 repair_p = subs.add_parser(
608 "repair",
609 help="Write missing .muse pointer files into existing linked worktrees.",
610 description=(
611 "Write (or re-write) the .muse pointer file in every registered\n"
612 "linked worktree whose directory is present on disk.\n\n"
613 "Idempotent — safe to run multiple times. Worktrees whose\n"
614 "directories are absent are skipped (use 'prune' for those).\n\n"
615 "Agent quickstart\n"
616 "----------------\n"
617 " muse worktree repair --json\n\n"
618 "JSON output schema\n"
619 "------------------\n"
620 ' {"repaired": ["<name>", ...]}\n\n'
621 "Exit codes\n"
622 "----------\n"
623 " 0 — success (even when nothing needed repair)\n"
624 " 2 — not inside a Muse repository\n"
625 ),
626 formatter_class=argparse.RawDescriptionHelpFormatter,
627 )
628 repair_p.add_argument(
629 "--json", "-j", action="store_true", dest="output_json", default=False,
630 help="Emit machine-readable JSON to stdout.",
631 )
632 repair_p.set_defaults(func=run_worktree_repair)
633
634 # ── status ────────────────────────────────────────────────────────────────
635 status_p = subs.add_parser(
636 "status",
637 help="Show the status of a single worktree.",
638 description=(
639 "Show branch, HEAD commit, path, and presence of a single worktree.\n\n"
640 "Pass 'main' or '(main)' to query the main worktree. Output fields\n"
641 "are sanitized — ANSI control sequences are stripped.\n\n"
642 "Agent quickstart\n"
643 "----------------\n"
644 " muse worktree status feat-x --json\n"
645 " muse worktree status main --json\n\n"
646 "JSON output schema\n"
647 "------------------\n"
648 ' {"name": "<name>", "branch": "<branch>", "path": "<absolute-path>",\n'
649 ' "head_commit": "<sha256> | null", "present": true|false,\n'
650 ' "is_main": true|false}\n\n'
651 "Exit codes\n"
652 "----------\n"
653 " 0 — success\n"
654 " 1 — worktree does not exist, or name is invalid\n"
655 " 2 — not inside a Muse repository\n"
656 ),
657 formatter_class=argparse.RawDescriptionHelpFormatter,
658 )
659 status_p.add_argument("name", metavar="NAME", help="Worktree name (or 'main' / '(main)').")
660 status_p.add_argument(
661 "--json", "-j", action="store_true", dest="output_json", default=False,
662 help="Emit machine-readable JSON to stdout.",
663 )
664 status_p.set_defaults(func=run_worktree_status)
File History 1 commit
sha256:0e3697f29ae6c62f39818a19861d4db6276465de16481fadf866eb93cbf4aae4 test: validate staging push and MP workflow for aaronrene Human 17 days ago