cgcardona / muse public
worktree.py python
526 lines 17.4 KB
Raw
12901c5a Initial extraction from tellurstori/maestro cgcardona <gabriel@tellurstori.com> 4d ago
1 """muse worktree — manage local Muse worktrees from the CLI.
2
3 Muse worktrees let a producer work on two arrangements simultaneously — one
4 worktree for "radio edit" mixing, another for "extended club version" — without
5 switching branches back and forth.
6
7 Architecture
8 ------------
9 Main repo (.muse/ directory):
10 .muse/worktrees/<slug>/path — absolute path to the linked worktree directory
11 .muse/worktrees/<slug>/branch — branch name checked out there
12 .muse/objects/ — shared content-addressed object store
13
14 Linked worktree directory:
15 .muse — plain text file: "gitdir: <abs-path-to-main-.muse>"
16 muse-work/ — per-worktree working files (independent)
17
18 The same-branch exclusivity constraint mirrors git: a branch can only be
19 checked out in one worktree at a time. Attempting to add a second worktree
20 on an already-checked-out branch exits with code 1.
21
22 Subcommands
23 -----------
24 muse worktree add <path> [branch] — create a linked worktree
25 muse worktree remove <path> — delete a linked worktree
26 muse worktree list — show all worktrees
27 muse worktree prune — remove stale registrations
28 """
29 from __future__ import annotations
30
31 import dataclasses
32 import logging
33 import pathlib
34 import re
35
36 import typer
37
38 from maestro.muse_cli._repo import require_repo
39 from maestro.muse_cli.errors import ExitCode
40
41 logger = logging.getLogger(__name__)
42
43 app = typer.Typer(invoke_without_command=True, help="Manage local Muse worktrees.")
44
45
46 # ---------------------------------------------------------------------------
47 # Result types — consumed by agents and tests
48 # ---------------------------------------------------------------------------
49
50
51 @dataclasses.dataclass(frozen=True)
52 class WorktreeInfo:
53 """A single worktree entry returned by ``list_worktrees``.
54
55 ``is_main`` is True for the primary repo directory (which owns the
56 ``.muse/`` directory tree), False for linked worktrees.
57
58 ``slug`` is the sanitized key used inside ``.muse/worktrees/``. It is
59 the empty string for the main worktree.
60
61 ``branch`` is the symbolic name of the checked-out branch (e.g.
62 ``"main"`` or ``"feature/club-mix"``). When HEAD is detached (no branch
63 ref), the value is ``"(detached)"``.
64
65 ``head_commit`` is the current HEAD commit SHA for the worktree (empty
66 string when the branch has no commits yet).
67 """
68
69 path: pathlib.Path
70 branch: str
71 head_commit: str
72 is_main: bool
73 slug: str
74
75
76 # ---------------------------------------------------------------------------
77 # Slug helpers
78 # ---------------------------------------------------------------------------
79
80
81 def _slugify(path: pathlib.Path) -> str:
82 """Derive a filesystem-safe registration key from a worktree path.
83
84 Converts the absolute path to a short ASCII slug by replacing every
85 non-alphanumeric run with a single hyphen and stripping leading/trailing
86 hyphens. Collision is extremely unlikely given unique absolute paths.
87 """
88 return re.sub(r"[^a-zA-Z0-9]+", "-", str(path.resolve())).strip("-")[:64]
89
90
91 def _worktrees_dir(muse_dir: pathlib.Path) -> pathlib.Path:
92 """Return ``<main-repo>/.muse/worktrees/``, creating it if absent."""
93 wt_dir = muse_dir / "worktrees"
94 wt_dir.mkdir(parents=True, exist_ok=True)
95 return wt_dir
96
97
98 # ---------------------------------------------------------------------------
99 # Core — testable, no Typer coupling
100 # ---------------------------------------------------------------------------
101
102
103 def _read_head_branch(muse_dir: pathlib.Path) -> str:
104 """Return the current branch name from a ``.muse/HEAD`` file.
105
106 Returns ``"(detached)"`` when HEAD is a bare commit rather than a
107 symbolic ref.
108 """
109 head_path = muse_dir / "HEAD"
110 if not head_path.exists():
111 return "(detached)"
112 text = head_path.read_text().strip()
113 if text.startswith("refs/heads/"):
114 return text[len("refs/heads/"):]
115 return "(detached)"
116
117
118 def _read_head_commit(muse_dir: pathlib.Path) -> str:
119 """Return the commit SHA that HEAD resolves to, or empty string."""
120 head_path = muse_dir / "HEAD"
121 if not head_path.exists():
122 return ""
123 text = head_path.read_text().strip()
124 if text.startswith("refs/heads/"):
125 ref_path = muse_dir / text
126 if ref_path.exists():
127 return ref_path.read_text().strip()
128 return ""
129 # Detached HEAD — text is the commit SHA directly.
130 return text
131
132
133 def _all_checked_out_branches(root: pathlib.Path) -> list[str]:
134 """Return all branch names currently checked out across all worktrees.
135
136 Includes the main worktree plus every registered linked worktree. Used
137 to enforce the single-checkout-per-branch constraint.
138 """
139 muse_dir = root / ".muse"
140 branches: list[str] = []
141
142 # Main worktree.
143 main_branch = _read_head_branch(muse_dir)
144 if main_branch != "(detached)":
145 branches.append(main_branch)
146
147 # Linked worktrees.
148 wt_dir = muse_dir / "worktrees"
149 if not wt_dir.is_dir():
150 return branches
151 for entry in wt_dir.iterdir():
152 if not entry.is_dir():
153 continue
154 branch_file = entry / "branch"
155 if branch_file.exists():
156 b = branch_file.read_text().strip()
157 if b:
158 branches.append(b)
159
160 return branches
161
162
163 def list_worktrees(root: pathlib.Path) -> list[WorktreeInfo]:
164 """Return all worktrees: main first, then linked (in registration order).
165
166 The main worktree is always first. Linked worktrees are included even
167 when their target directory no longer exists (``path.exists()`` may be
168 False for stale entries — callers should check before accessing files).
169
170 Args:
171 root: Repository root (parent of ``.muse/``).
172
173 Returns:
174 List of :class:`WorktreeInfo` objects.
175 """
176 muse_dir = root / ".muse"
177 result: list[WorktreeInfo] = []
178
179 # Main worktree.
180 result.append(
181 WorktreeInfo(
182 path=root,
183 branch=_read_head_branch(muse_dir),
184 head_commit=_read_head_commit(muse_dir),
185 is_main=True,
186 slug="",
187 )
188 )
189
190 # Linked worktrees.
191 wt_dir = muse_dir / "worktrees"
192 if not wt_dir.is_dir():
193 return result
194
195 for entry in sorted(wt_dir.iterdir()):
196 if not entry.is_dir():
197 continue
198 path_file = entry / "path"
199 branch_file = entry / "branch"
200 if not path_file.exists():
201 continue
202 linked_path = pathlib.Path(path_file.read_text().strip())
203 branch = branch_file.read_text().strip() if branch_file.exists() else "(detached)"
204
205 # Read HEAD commit from the linked worktree's own HEAD file if present.
206 linked_muse_dir = linked_path / ".muse"
207 head_commit = ""
208 if linked_muse_dir.is_dir():
209 head_commit = _read_head_commit(linked_muse_dir)
210
211 result.append(
212 WorktreeInfo(
213 path=linked_path,
214 branch=branch,
215 head_commit=head_commit,
216 is_main=False,
217 slug=entry.name,
218 )
219 )
220
221 return result
222
223
224 def add_worktree(
225 *,
226 root: pathlib.Path,
227 link_path: pathlib.Path,
228 branch: str,
229 ) -> WorktreeInfo:
230 """Create a new linked worktree at ``link_path`` checked out to ``branch``.
231
232 Enforces same-branch exclusivity: if ``branch`` is already checked out
233 in any worktree (main or linked), raises ``typer.Exit(1)``.
234
235 Creates the branch ref in the main repo if it does not yet exist,
236 mirroring ``git worktree add --orphan``-like behaviour so that producers
237 can start a completely fresh arrangement on a new branch.
238
239 Layout of the new directory::
240
241 <link_path>/
242 .muse (plain-text file: "gitdir: <main-repo>/.muse")
243 muse-work/ (empty working directory)
244
245 Registration in main repo::
246
247 .muse/worktrees/<slug>/path → abs path to link_path
248 .muse/worktrees/<slug>/branch → branch name
249
250 Args:
251 root: Repository root (parent of ``.muse/``).
252 link_path: Absolute path for the new linked worktree directory.
253 branch: Branch name to check out. Created from HEAD if absent.
254
255 Returns:
256 :class:`WorktreeInfo` describing the newly created worktree.
257
258 Raises:
259 typer.Exit(1): branch already checked out, path already exists, or
260 link_path is inside the main repo's ``.muse/`` tree.
261 """
262 muse_dir = root / ".muse"
263 link_path = link_path.resolve()
264
265 # Guard: link_path must not be the main repo itself or inside .muse/.
266 if link_path == root or link_path.is_relative_to(muse_dir):
267 typer.echo("❌ Worktree path must be outside the main repository.")
268 raise typer.Exit(code=ExitCode.USER_ERROR)
269
270 # Guard: path must not already exist.
271 if link_path.exists():
272 typer.echo(f"❌ Path already exists: {link_path}\n"
273 " Choose a different location for the linked worktree.")
274 raise typer.Exit(code=ExitCode.USER_ERROR)
275
276 # Guard: branch exclusivity.
277 checked_out = _all_checked_out_branches(root)
278 if branch in checked_out:
279 typer.echo(
280 f"❌ Branch '{branch}' is already checked out in another worktree.\n"
281 " A branch can only be active in one worktree at a time."
282 )
283 raise typer.Exit(code=ExitCode.USER_ERROR)
284
285 # Ensure the branch ref exists in the main repo (create from HEAD if not).
286 ref_path = muse_dir / "refs" / "heads" / branch
287 if not ref_path.exists():
288 head_commit = _read_head_commit(muse_dir)
289 ref_path.parent.mkdir(parents=True, exist_ok=True)
290 ref_path.write_text(head_commit)
291 logger.info("✅ Created branch ref %r (from HEAD=%s)", branch, head_commit[:8] or "(empty)")
292
293 # Create the linked worktree directory structure.
294 link_path.mkdir(parents=True)
295 (link_path / "muse-work").mkdir()
296
297 # Write the .muse gitdir file so the linked worktree points back to main.
298 slug = _slugify(link_path)
299 registration_dir = _worktrees_dir(muse_dir) / slug
300 registration_dir.mkdir(parents=True, exist_ok=True)
301
302 gitdir_content = f"gitdir: {muse_dir}\n"
303 (link_path / ".muse").write_text(gitdir_content)
304
305 # Write HEAD and branch ref inside a minimal .muse/ inside the linked dir.
306 # Wait — per design, .muse is a *file* pointing at the main repo.
307 # The linked worktree's HEAD lives in the registration entry.
308 # For list_worktrees to read the HEAD commit we also write it to the
309 # registration dir so no filesystem traversal to the linked path is needed.
310 head_commit = ref_path.read_text().strip()
311 (registration_dir / "path").write_text(str(link_path))
312 (registration_dir / "branch").write_text(branch)
313
314 logger.info("✅ muse worktree add %s (branch=%r, slug=%r)", link_path, branch, slug)
315
316 return WorktreeInfo(
317 path=link_path,
318 branch=branch,
319 head_commit=head_commit,
320 is_main=False,
321 slug=slug,
322 )
323
324
325 def remove_worktree(*, root: pathlib.Path, link_path: pathlib.Path) -> None:
326 """Remove a linked worktree: delete its directory and de-register it.
327
328 The main worktree cannot be removed. Only the linked worktree's own
329 directory and its registration entry are removed — the shared objects
330 store and branch refs in the main repo are left intact, so the branch
331 remains accessible from the main worktree.
332
333 Args:
334 root: Repository root (parent of ``.muse/``).
335 link_path: Path to the linked worktree to remove.
336
337 Raises:
338 typer.Exit(1): path is not a registered linked worktree or is the
339 main repo root.
340 """
341 muse_dir = root / ".muse"
342 link_path = link_path.resolve()
343
344 if link_path == root:
345 typer.echo("❌ Cannot remove the main worktree.")
346 raise typer.Exit(code=ExitCode.USER_ERROR)
347
348 # Find the registration entry for this path.
349 wt_dir = muse_dir / "worktrees"
350 registration_dir: pathlib.Path | None = None
351 if wt_dir.is_dir():
352 for entry in wt_dir.iterdir():
353 path_file = entry / "path"
354 if path_file.exists() and pathlib.Path(path_file.read_text().strip()) == link_path:
355 registration_dir = entry
356 break
357
358 if registration_dir is None:
359 typer.echo(
360 f"❌ '{link_path}' is not a registered linked worktree.\n"
361 " Run 'muse worktree list' to see registered worktrees."
362 )
363 raise typer.Exit(code=ExitCode.USER_ERROR)
364
365 # Remove the linked worktree directory.
366 if link_path.exists():
367 _remove_tree(link_path)
368
369 # Remove the registration entry.
370 _remove_tree(registration_dir)
371
372 logger.info("✅ muse worktree remove %s", link_path)
373
374
375 def prune_worktrees(*, root: pathlib.Path) -> list[str]:
376 """Remove stale worktree registrations (directory no longer exists).
377
378 Scans ``.muse/worktrees/`` for entries whose target path is absent and
379 removes those registration directories. This is a local-only operation
380 that does not touch branch refs or the objects store.
381
382 Args:
383 root: Repository root (parent of ``.muse/``).
384
385 Returns:
386 List of absolute path strings that were pruned.
387 """
388 muse_dir = root / ".muse"
389 wt_dir = muse_dir / "worktrees"
390 pruned: list[str] = []
391
392 if not wt_dir.is_dir():
393 return pruned
394
395 for entry in list(wt_dir.iterdir()):
396 if not entry.is_dir():
397 continue
398 path_file = entry / "path"
399 if not path_file.exists():
400 _remove_tree(entry)
401 pruned.append(str(entry))
402 continue
403 target = pathlib.Path(path_file.read_text().strip())
404 if not target.exists():
405 pruned.append(str(target))
406 _remove_tree(entry)
407 logger.info("⚠️ muse worktree prune: removed stale entry %s → %s", entry.name, target)
408
409 return pruned
410
411
412 # ---------------------------------------------------------------------------
413 # Internal helpers
414 # ---------------------------------------------------------------------------
415
416
417 def _remove_tree(path: pathlib.Path) -> None:
418 """Recursively remove a directory tree or a single file.
419
420 This is a pure Python implementation to avoid shelling out. We intentionally
421 do not use ``shutil.rmtree`` to remain dependency-free (shutil is stdlib but
422 the explicit implementation is clearer in tests and error traces).
423 """
424 if path.is_file() or path.is_symlink():
425 path.unlink()
426 return
427 for child in path.iterdir():
428 _remove_tree(child)
429 path.rmdir()
430
431
432 # ---------------------------------------------------------------------------
433 # Typer commands
434 # ---------------------------------------------------------------------------
435
436
437 @app.callback(invoke_without_command=True)
438 def worktree_callback(ctx: typer.Context) -> None:
439 """Manage local Muse worktrees.
440
441 Run ``muse worktree --help`` to see available subcommands.
442 """
443 if ctx.invoked_subcommand is None:
444 typer.echo(ctx.get_help())
445 raise typer.Exit(code=ExitCode.SUCCESS)
446
447
448 @app.command("add")
449 def worktree_add(
450 path: str = typer.Argument(..., help="Directory to create the linked worktree in."),
451 branch: str = typer.Argument(
452 ...,
453 help="Branch name to check out. Created from HEAD if it does not exist.",
454 ),
455 ) -> None:
456 """Create a new linked worktree at PATH checked out to BRANCH.
457
458 The new directory is created at PATH with an independent ``muse-work/``
459 working directory. The shared ``.muse/objects/`` store remains in the
460 main repository.
461
462 Example::
463
464 muse worktree add ../club-mix feature/extended
465 """
466 root = require_repo()
467 link_path = pathlib.Path(path).expanduser().resolve()
468 info = add_worktree(root=root, link_path=link_path, branch=branch)
469 typer.echo(f"✅ Linked worktree '{info.branch}' created at {info.path}")
470
471
472 @app.command("remove")
473 def worktree_remove(
474 path: str = typer.Argument(..., help="Path of the linked worktree to remove."),
475 ) -> None:
476 """Remove a linked worktree and de-register it.
477
478 The branch ref and shared objects store are preserved. Only the linked
479 worktree directory and its registration entry are deleted.
480
481 Example::
482
483 muse worktree remove ../club-mix
484 """
485 root = require_repo()
486 link_path = pathlib.Path(path).expanduser().resolve()
487 remove_worktree(root=root, link_path=link_path)
488 typer.echo(f"✅ Worktree at {link_path} removed.")
489
490
491 @app.command("list")
492 def worktree_list() -> None:
493 """List all worktrees: main and linked, with path, branch, and HEAD.
494
495 Example output::
496
497 /path/to/project [main] branch: main head: a1b2c3d4
498 /path/to/club-mix branch: feature/club head: a1b2c3d4
499 """
500 root = require_repo()
501 worktrees = list_worktrees(root)
502 for wt in worktrees:
503 tag = "[main]" if wt.is_main else " "
504 head = wt.head_commit[:8] if wt.head_commit else "(no commits)"
505 typer.echo(f"{wt.path!s:<50} {tag} branch: {wt.branch:<30} head: {head}")
506
507
508 @app.command("prune")
509 def worktree_prune() -> None:
510 """Remove stale worktree registrations (directory no longer exists).
511
512 Scans ``.muse/worktrees/`` for entries whose target directory is absent
513 and removes them. Safe to run any time — no data loss risk.
514
515 Example::
516
517 muse worktree prune
518 """
519 root = require_repo()
520 pruned = prune_worktrees(root=root)
521 if pruned:
522 for p in pruned:
523 typer.echo(f"⚠️ Pruned stale worktree: {p}")
524 typer.echo(f"✅ Pruned {len(pruned)} stale worktree registration(s).")
525 else:
526 typer.echo("✅ No stale worktrees found.")