aaronrene / gabriel-muse public
worktree.py python
514 lines 17.5 KB
Raw
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago
1 """Worktree management — multiple simultaneous branch checkouts.
2
3 A *worktree* is a second (or third, …) checked-out working tree linked to
4 the same ``.muse/`` repository. Each worktree has its own branch and its own
5 working directory, so multiple agents — or multiple human engineers — can
6 work on different branches simultaneously without interfering with each other.
7
8 Layout
9 ------
10 Each linked worktree lives in a sibling directory of the repository root::
11
12 myproject/ ← main worktree (holds .muse/)
13 .muse/
14 worktrees/
15 <name>.json ← metadata for each linked worktree
16 <name>.HEAD ← HEAD ref for the linked worktree
17
18 myproject-<name>/ ← linked worktree directory
19 ← worktree working directory
20
21 The shared ``.muse/`` directory is the single source of truth for commits,
22 snapshots, objects, and branch refs. Each worktree has its own HEAD file
23 stored inside the main ``.muse/worktrees/<name>.HEAD``.
24
25 Security model
26 --------------
27 - **Name validation**: Worktree names pass through ``validate_branch_name``
28 (no path separators, no null bytes, no control characters). This ensures
29 the derived meta path and HEAD path cannot escape ``.muse/worktrees/``.
30 - **Symlink guard on meta files**: ``_load_meta`` rejects a symlink at the
31 meta file path before any read.
32 - **Size cap on meta files**: ``_load_meta`` refuses files larger than
33 ``_MAX_META_BYTES`` to guard against memory exhaustion.
34 - **Path safety on delete**: ``remove_worktree`` and ``prune_worktrees``
35 call ``_safe_delete_path`` which refuses to ``rmtree`` a path that is a
36 symlink itself or that resolves inside the shared ``.muse/`` store.
37
38 Agent concurrency
39 -----------------
40 Multiple agents can operate on separate worktrees simultaneously. Each
41 worktree's HEAD is independent; commits from one worktree appear immediately
42 in all others (they share the object store).
43 """
44
45 from __future__ import annotations
46
47 import json
48 import logging
49 import pathlib
50 import shutil
51 from dataclasses import dataclass
52 from typing import TypedDict
53
54 from muse.core.object_store import restore_object
55 from muse.core.store import (
56 get_head_commit_id,
57 read_commit,
58 read_current_branch,
59 read_snapshot,
60 write_text_atomic,
61 )
62 from muse.core.validation import contain_path, validate_branch_name
63
64 logger = logging.getLogger(__name__)
65
66 # Directory inside .muse/ that holds per-worktree metadata and HEAD files.
67 _WORKTREES_META_DIR = ".muse/worktrees"
68
69 # Guard against tampered or pathologically large meta files.
70 _MAX_META_BYTES: int = 4 * 1024 # 4 KiB — more than enough for any meta record
71
72
73 # ---------------------------------------------------------------------------
74 # Types
75 # ---------------------------------------------------------------------------
76
77
78 class WorktreeRecord(TypedDict):
79 """Persisted metadata for a linked worktree."""
80
81 name: str
82 branch: str
83 path: str # absolute path to the worktree directory
84
85
86 @dataclass
87 class WorktreeInfo:
88 """Runtime information about a worktree."""
89
90 name: str
91 branch: str
92 path: pathlib.Path
93 head_commit: str | None
94 is_main: bool = False
95
96
97 class WorktreeStatusResult(TypedDict):
98 """Machine-readable status of a single worktree."""
99
100 name: str
101 branch: str
102 path: str
103 head_commit: str | None
104 present: bool
105 is_main: bool
106
107
108 # ---------------------------------------------------------------------------
109 # Paths
110 # ---------------------------------------------------------------------------
111
112
113 def _worktrees_dir(repo_root: pathlib.Path) -> pathlib.Path:
114 return repo_root / _WORKTREES_META_DIR
115
116
117 def _worktree_meta_path(repo_root: pathlib.Path, name: str) -> pathlib.Path:
118 return _worktrees_dir(repo_root) / f"{name}.json"
119
120
121 def _worktree_head_path(repo_root: pathlib.Path, name: str) -> pathlib.Path:
122 return _worktrees_dir(repo_root) / f"{name}.HEAD"
123
124
125 def _worktree_dir(repo_root: pathlib.Path, name: str) -> pathlib.Path:
126 """Return the default path of the linked worktree directory (sibling of repo_root)."""
127 parent = repo_root.parent
128 repo_name = repo_root.name
129 return parent / f"{repo_name}-{name}"
130
131
132 # ---------------------------------------------------------------------------
133 # Internal helpers
134 # ---------------------------------------------------------------------------
135
136
137 def _load_meta(repo_root: pathlib.Path, name: str) -> WorktreeRecord | None:
138 """Read and validate the metadata file for *name*.
139
140 Safety guards applied before any read:
141
142 - **Symlink check**: a symlink at the meta path could redirect writes to
143 arbitrary locations or reads to sensitive files.
144 - **Size cap** (``_MAX_META_BYTES``): a tampered or corrupt meta file
145 cannot be used to exhaust memory.
146 """
147 meta_path = _worktree_meta_path(repo_root, name)
148 if not meta_path.exists():
149 return None
150 if meta_path.is_symlink():
151 logger.warning("⚠️ Worktree meta file for %r is a symlink — ignoring", name)
152 return None
153 try:
154 if meta_path.stat().st_size > _MAX_META_BYTES:
155 logger.warning(
156 "⚠️ Worktree meta file for %r exceeds size cap (%d bytes) — ignoring",
157 name,
158 _MAX_META_BYTES,
159 )
160 return None
161 raw = json.loads(meta_path.read_text(encoding="utf-8"))
162 return WorktreeRecord(
163 name=str(raw["name"]),
164 branch=str(raw["branch"]),
165 path=str(raw["path"]),
166 )
167 except (KeyError, ValueError, OSError) as exc:
168 logger.warning("⚠️ Could not read worktree metadata for %r: %s", name, exc)
169 return None
170
171
172 def _save_meta(repo_root: pathlib.Path, record: WorktreeRecord) -> None:
173 """Write *record* to the metadata file atomically."""
174 meta_path = _worktree_meta_path(repo_root, record["name"])
175 write_text_atomic(meta_path, json.dumps(record, indent=2))
176
177
178 def _write_worktree_pointer(wt_dir: pathlib.Path, repo_root: pathlib.Path) -> None:
179 """Write a ``.muse`` pointer file in *wt_dir* pointing to *repo_root*'s store.
180
181 The file contains a single line::
182
183 musestore: /absolute/path/to/main/.muse
184
185 This mirrors git's ``.git`` worktree file, enabling ``find_repo_root``
186 to resolve the shared object store from any worktree directory.
187
188 No-ops when ``wt_dir`` is the main repo itself or when ``.muse`` already
189 exists as a directory (a legacy worktree with its own store).
190 """
191 pointer_path = wt_dir / ".muse"
192 # Never clobber a real .muse/ store directory.
193 if pointer_path.is_dir():
194 logger.debug("Skipping pointer write — %s is already a store directory", pointer_path)
195 return
196 store_path = (repo_root / ".muse").resolve()
197 # Don't write a self-referential pointer (wt_dir IS the main repo).
198 if pointer_path.resolve() == store_path:
199 logger.debug("Skipping pointer write — wt_dir is the main repo root")
200 return
201 write_text_atomic(pointer_path, f"musestore: {store_path}\n")
202 logger.debug("Wrote worktree pointer %s → %s", pointer_path, store_path)
203
204
205 def _safe_delete_path(repo_root: pathlib.Path, path: pathlib.Path) -> bool:
206 """Delete *path* and its contents, with safety guards.
207
208 Refuses deletion when:
209
210 - *path* is a symlink (could target an unrelated directory).
211 - *path* resolves to be inside the shared ``.muse/`` store — deleting it
212 would corrupt the repository.
213 - *path* does not exist (no-op, returns True so callers can proceed).
214
215 Returns:
216 ``True`` if the directory was deleted or did not exist.
217 ``False`` if the deletion was refused for safety reasons.
218 """
219 if not path.exists():
220 return True
221 if path.is_symlink():
222 logger.warning(
223 "⚠️ Refusing to delete worktree path %s — it is a symlink", path
224 )
225 return False
226 muse_dir = (repo_root / ".muse").resolve()
227 try:
228 resolved = path.resolve()
229 except OSError:
230 logger.warning("⚠️ Could not resolve worktree path %s", path)
231 return False
232 try:
233 resolved.relative_to(muse_dir)
234 # Path is inside .muse/ — refuse.
235 logger.warning(
236 "⚠️ Refusing to delete worktree path %s — it resolves inside .muse/", path
237 )
238 return False
239 except ValueError:
240 pass # Not inside .muse/ — safe to delete.
241 shutil.rmtree(path)
242 return True
243
244
245 def _read_main_branch(repo_root: pathlib.Path) -> str:
246 return read_current_branch(repo_root)
247
248
249 # ---------------------------------------------------------------------------
250 # Public API
251 # ---------------------------------------------------------------------------
252
253
254 def add_worktree(
255 repo_root: pathlib.Path,
256 name: str,
257 branch: str,
258 path: pathlib.Path | None = None,
259 ) -> pathlib.Path:
260 """Create and populate a new linked worktree.
261
262 Args:
263 repo_root: Main repository root (where ``.muse/`` lives).
264 name: Short identifier for the worktree (validated like a branch
265 name — no path separators, no null bytes).
266 branch: Branch to check out in the new worktree.
267 path: Explicit filesystem path for the worktree directory. When
268 ``None`` (default) the standard sibling layout is used:
269 ``<repo_root.parent>/<repo_root.name>-<name>``.
270
271 Returns:
272 The path to the newly created worktree directory.
273
274 Raises:
275 ValueError: If the name is invalid, the worktree already exists, or
276 the branch does not exist.
277 """
278 validate_branch_name(name)
279
280 wt_dir = path if path is not None else _worktree_dir(repo_root, name)
281 meta_path = _worktree_meta_path(repo_root, name)
282
283 if meta_path.exists():
284 raise ValueError(f"Worktree '{name}' already exists.")
285 if wt_dir.exists():
286 raise ValueError(f"Directory '{wt_dir}' already exists.")
287
288 # Verify the branch exists.
289 branch_ref = repo_root / ".muse" / "refs" / "heads" / branch
290 if not branch_ref.exists():
291 raise ValueError(f"Branch '{branch}' does not exist.")
292
293 # Create the worktree directory.
294 wt_dir.mkdir(parents=True)
295
296 # Write .muse pointer file so `muse` commands work inside the worktree.
297 _write_worktree_pointer(wt_dir, repo_root)
298
299 # Write the worktree HEAD file.
300 head_path = _worktree_head_path(repo_root, name)
301 write_text_atomic(head_path, f"refs/heads/{branch}\n")
302
303 # Populate the worktree from the branch's latest snapshot.
304 commit_id = get_head_commit_id(repo_root, branch)
305 if commit_id:
306 commit = read_commit(repo_root, commit_id)
307 if commit:
308 snap = read_snapshot(repo_root, commit.snapshot_id)
309 if snap:
310 for rel_path, object_id in snap.manifest.items():
311 try:
312 dest = contain_path(wt_dir, rel_path)
313 except ValueError as exc:
314 logger.warning("⚠️ Skipping unsafe path %r: %s", rel_path, exc)
315 continue
316 restore_object(repo_root, object_id, dest)
317
318 # Persist metadata.
319 record: WorktreeRecord = {
320 "name": name,
321 "branch": branch,
322 "path": str(wt_dir),
323 }
324 _save_meta(repo_root, record)
325
326 logger.info("✅ Worktree '%s' created at %s (branch: %s)", name, wt_dir, branch)
327 return wt_dir
328
329
330 def list_worktrees(repo_root: pathlib.Path) -> list[WorktreeInfo]:
331 """Return all worktrees (main + linked), sorted by name.
332
333 The main worktree is always first; linked worktrees follow in
334 lexicographic order of name.
335 """
336 results: list[WorktreeInfo] = []
337
338 # Main worktree.
339 main_branch = _read_main_branch(repo_root)
340 main_head = get_head_commit_id(repo_root, main_branch)
341 results.append(WorktreeInfo(
342 name="(main)",
343 branch=main_branch,
344 path=repo_root,
345 head_commit=main_head,
346 is_main=True,
347 ))
348
349 wt_dir = _worktrees_dir(repo_root)
350 if not wt_dir.exists():
351 return results
352
353 for meta_file in sorted(wt_dir.glob("*.json")):
354 name = meta_file.stem
355 record = _load_meta(repo_root, name)
356 if record is None:
357 continue
358 wt_path = pathlib.Path(record["path"])
359 branch = record["branch"]
360 head_path = _worktree_head_path(repo_root, name)
361 commit_id = get_head_commit_id(repo_root, branch) if head_path.exists() else None
362 results.append(WorktreeInfo(
363 name=name,
364 branch=branch,
365 path=wt_path,
366 head_commit=commit_id,
367 ))
368 return results
369
370
371 def remove_worktree(repo_root: pathlib.Path, name: str, force: bool = False) -> None:
372 """Remove a linked worktree.
373
374 The branch itself is not deleted — only the worktree directory and its
375 metadata are removed. Commits already made in the worktree remain in the
376 shared object store.
377
378 Args:
379 repo_root: Main repository root.
380 name: Name of the worktree to remove.
381 force: Accepted for interface compatibility. Currently has no
382 effect since Muse does not track working-tree dirtiness
383 per-worktree.
384
385 Raises:
386 ValueError: If the worktree does not exist or its metadata is corrupt.
387 """
388 validate_branch_name(name)
389
390 meta_path = _worktree_meta_path(repo_root, name)
391 if not meta_path.exists():
392 raise ValueError(f"Worktree '{name}' does not exist.")
393
394 record = _load_meta(repo_root, name)
395 if record is None:
396 raise ValueError(f"Could not read metadata for worktree '{name}'.")
397
398 wt_path = pathlib.Path(record["path"])
399 if not _safe_delete_path(repo_root, wt_path):
400 raise ValueError(
401 f"Refusing to delete worktree path '{wt_path}' — "
402 "it is a symlink or resolves inside .muse/."
403 )
404
405 meta_path.unlink(missing_ok=True)
406 head_path = _worktree_head_path(repo_root, name)
407 head_path.unlink(missing_ok=True)
408 # Belt-and-suspenders: remove the pointer file if it still exists
409 # (the directory deletion above should have removed it already).
410 pointer_path = wt_path / ".muse"
411 pointer_path.unlink(missing_ok=True)
412
413 logger.info("Worktree '%s' removed.", name)
414
415
416 def prune_worktrees(repo_root: pathlib.Path, *, dry_run: bool = False) -> list[str]:
417 """Remove metadata for worktrees whose directories no longer exist.
418
419 Args:
420 repo_root: Main repository root.
421 dry_run: When ``True``, report what would be pruned without
422 making any filesystem changes.
423
424 Returns:
425 Names of pruned (or would-be-pruned, when *dry_run* is ``True``)
426 worktrees.
427 """
428 pruned: list[str] = []
429 wt_dir = _worktrees_dir(repo_root)
430 if not wt_dir.exists():
431 return pruned
432 for meta_file in list(wt_dir.glob("*.json")):
433 name = meta_file.stem
434 record = _load_meta(repo_root, name)
435 if record is None:
436 if not dry_run:
437 meta_file.unlink(missing_ok=True)
438 pruned.append(name)
439 continue
440 wt_path = pathlib.Path(record["path"])
441 if not wt_path.exists():
442 if not dry_run:
443 meta_file.unlink(missing_ok=True)
444 _worktree_head_path(repo_root, name).unlink(missing_ok=True)
445 pruned.append(name)
446 return pruned
447
448
449 def get_worktree_status(repo_root: pathlib.Path, name: str) -> WorktreeStatusResult:
450 """Return the status of a single named worktree.
451
452 Covers both linked worktrees (by name) and the implicit main worktree
453 when *name* is ``"(main)"`` or ``"main"``.
454
455 Returns:
456 A ``WorktreeStatusResult`` with present/absent flag, current branch,
457 and HEAD commit — ready for JSON serialisation by the CLI.
458
459 Raises:
460 ValueError: If no worktree with *name* exists.
461 """
462 # Main worktree shortcut.
463 if name in {"(main)", "main"}:
464 branch = _read_main_branch(repo_root)
465 head = get_head_commit_id(repo_root, branch)
466 return WorktreeStatusResult(
467 name="(main)",
468 branch=branch,
469 path=str(repo_root),
470 head_commit=head,
471 present=repo_root.exists(),
472 is_main=True,
473 )
474
475 validate_branch_name(name)
476 meta_path = _worktree_meta_path(repo_root, name)
477 if not meta_path.exists():
478 raise ValueError(f"Worktree '{name}' does not exist.")
479
480 record = _load_meta(repo_root, name)
481 if record is None:
482 raise ValueError(f"Could not read metadata for worktree '{name}'.")
483
484 wt_path = pathlib.Path(record["path"])
485 branch = record["branch"]
486 head_path = _worktree_head_path(repo_root, name)
487 head = get_head_commit_id(repo_root, branch) if head_path.exists() else None
488 return WorktreeStatusResult(
489 name=name,
490 branch=branch,
491 path=str(wt_path),
492 head_commit=head,
493 present=wt_path.exists() and not wt_path.is_symlink(),
494 is_main=False,
495 )
496
497
498 def repair_worktree_pointers(repo_root: pathlib.Path) -> list[str]:
499 """Write missing ``.muse`` pointer files for all registered worktrees.
500
501 Idempotent — safe to run multiple times. Returns a list of worktree
502 names that were repaired (pointer file written or overwritten).
503 """
504 repaired: list[str] = []
505 for wt in list_worktrees(repo_root):
506 if wt.is_main:
507 continue # Main worktree has a real .muse/ store — never needs a pointer.
508 wt_path = wt.path
509 if not wt_path.is_dir():
510 continue
511 _write_worktree_pointer(wt_path, repo_root)
512 repaired.append(wt.name)
513 logger.info("✅ Repaired worktree pointer: %s", wt_path / ".muse")
514 return repaired
File History 1 commit
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago