cgcardona / muse public
muse_stash.py python
574 lines 20.0 KB
Raw
12901c5a Initial extraction from tellurstori/maestro cgcardona <gabriel@tellurstori.com> 4d ago
1 """Muse Stash Service — temporarily shelve uncommitted muse-work/ changes.
2
3 Stash lets a producer save in-progress work without committing, switch
4 context (e.g. fix the intro for a client call), then restore the shelved
5 state with ``muse stash pop``.
6
7 Design
8 ------
9 - Stash entries live purely on the filesystem in ``.muse/stash/``.
10 - Each entry is a JSON file named ``stash-<epoch_ns>.json``.
11 - The stack is ordered by creation time: index 0 is most recent.
12 - File content is preserved by writing blobs to the existing
13 ``.muse/objects/<oid[:2]>/<oid[2:]>`` content-addressed store
14 (same layout used by ``muse commit`` and ``muse reset --hard``).
15 - Restoring HEAD after push reads manifests from the DB (same as hard reset).
16 If the branch has no commits, muse-work/ is simply cleared.
17 - On apply/pop, files are copied from the object store back into muse-work/.
18 Files whose objects are absent from the store are reported as missing.
19 - Track/section scoping on push limits which files are saved and what is
20 restored afterward (only scoped paths are erased; others stay untouched).
21
22 Path-scoped stash (``--track`` / ``--section``)
23 ------------------------------------------------
24 When a scope is supplied, only files under ``tracks/<track>/`` or
25 ``sections/<section>/`` are saved to the stash. After saving the scope,
26 the HEAD snapshot is restored only for those paths (other working-tree
27 files are left untouched). Applying a scoped stash similarly only writes
28 paths that match the original scope.
29
30 Boundary rules:
31 - No Typer imports.
32 - No StateStore, EntityRegistry, or get_or_create_store.
33 - May import muse_cli.{db,models,object_store,snapshot}.
34 - Filesystem stash store is independent of the Postgres schema.
35 """
36 from __future__ import annotations
37
38 import datetime
39 import json
40 import logging
41 import pathlib
42 import shutil
43 import uuid
44 from dataclasses import dataclass, field
45 from typing import Optional
46
47 from sqlalchemy.ext.asyncio import AsyncSession
48
49 from maestro.muse_cli.snapshot import hash_file
50
51 logger = logging.getLogger(__name__)
52
53 # ---------------------------------------------------------------------------
54 # Constants
55 # ---------------------------------------------------------------------------
56
57 _STASH_DIR = "stash"
58
59
60 # ---------------------------------------------------------------------------
61 # Types
62 # ---------------------------------------------------------------------------
63
64
65 @dataclass(frozen=True)
66 class StashEntry:
67 """A single stash entry persisted in ``.muse/stash/``.
68
69 Attributes:
70 stash_id: Unique filename stem (``stash-<epoch_ns>-<uuid8>``).
71 index: Position in the stack (0 = most recent).
72 branch: Branch name at the time of stash.
73 message: Human-readable label (``On <branch>: <text>``).
74 created_at: ISO-8601 timestamp.
75 manifest: ``{rel_path: sha256_object_id}`` of stashed files.
76 track: Optional track scope used during push.
77 section: Optional section scope used during push.
78 """
79
80 stash_id: str
81 index: int
82 branch: str
83 message: str
84 created_at: str
85 manifest: dict[str, str]
86 track: Optional[str]
87 section: Optional[str]
88
89
90 @dataclass(frozen=True)
91 class StashPushResult:
92 """Outcome of ``muse stash push``.
93
94 Attributes:
95 stash_ref: Human label (``stash@{0}``).
96 message: The label stored in the entry.
97 branch: Branch at the time of push.
98 files_stashed: Number of files saved into the stash.
99 head_restored: Whether HEAD snapshot was restored to muse-work/.
100 missing_head: Paths that could not be restored from the object store
101 (object bytes not present; stash push succeeded but
102 HEAD restoration is incomplete).
103 """
104
105 stash_ref: str
106 message: str
107 branch: str
108 files_stashed: int
109 head_restored: bool
110 missing_head: tuple[str, ...]
111
112
113 @dataclass(frozen=True)
114 class StashApplyResult:
115 """Outcome of ``muse stash apply`` or ``muse stash pop``.
116
117 Attributes:
118 stash_ref: Human label of the entry that was applied.
119 message: The entry's label.
120 files_applied: Number of files written to muse-work/.
121 missing: Paths whose object bytes were absent from the store.
122 dropped: True when the entry was removed (pop); False for apply.
123 """
124
125 stash_ref: str
126 message: str
127 files_applied: int
128 missing: tuple[str, ...]
129 dropped: bool
130
131
132 # ---------------------------------------------------------------------------
133 # Filesystem helpers
134 # ---------------------------------------------------------------------------
135
136
137 def _stash_dir(root: pathlib.Path) -> pathlib.Path:
138 """Return the stash directory path (does not create it)."""
139 return root / ".muse" / _STASH_DIR
140
141
142 def _entry_path(root: pathlib.Path, stash_id: str) -> pathlib.Path:
143 """Return the filesystem path of a single stash entry JSON file."""
144 return _stash_dir(root) / f"{stash_id}.json"
145
146
147 def _object_path(root: pathlib.Path, object_id: str) -> pathlib.Path:
148 """Return the sharded object store path for *object_id*.
149
150 Matches the layout used by ``muse commit`` and ``muse reset --hard``:
151 ``.muse/objects/<oid[:2]>/<oid[2:]>``.
152 """
153 return root / ".muse" / "objects" / object_id[:2] / object_id[2:]
154
155
156 def _read_entry(entry_file: pathlib.Path, index: int) -> StashEntry:
157 """Deserialize a stash entry JSON file."""
158 raw: dict[str, object] = json.loads(entry_file.read_text())
159 raw_manifest = raw.get("manifest", {})
160 manifest: dict[str, str] = (
161 {str(k): str(v) for k, v in raw_manifest.items()}
162 if isinstance(raw_manifest, dict)
163 else {}
164 )
165 return StashEntry(
166 stash_id=str(raw["stash_id"]),
167 index=index,
168 branch=str(raw["branch"]),
169 message=str(raw["message"]),
170 created_at=str(raw["created_at"]),
171 manifest=manifest,
172 track=str(raw["track"]) if raw.get("track") else None,
173 section=str(raw["section"]) if raw.get("section") else None,
174 )
175
176
177 def _write_entry(root: pathlib.Path, entry_data: dict[str, object]) -> str:
178 """Serialize and write a stash entry. Returns the stash_id."""
179 stash_dir = _stash_dir(root)
180 stash_dir.mkdir(parents=True, exist_ok=True)
181 stash_id = str(entry_data["stash_id"])
182 (_stash_dir(root) / f"{stash_id}.json").write_text(
183 json.dumps(entry_data, indent=2)
184 )
185 return stash_id
186
187
188 def _sorted_entries(root: pathlib.Path) -> list[pathlib.Path]:
189 """Return stash JSON files sorted newest-first (index 0 = most recent)."""
190 stash_dir = _stash_dir(root)
191 if not stash_dir.exists():
192 return []
193 files = sorted(stash_dir.glob("stash-*.json"), reverse=True)
194 return files
195
196
197 # ---------------------------------------------------------------------------
198 # Path-scope filter (mirrors muse_revert._filter_paths)
199 # ---------------------------------------------------------------------------
200
201
202 def _filter_paths(
203 manifest: dict[str, str],
204 track: Optional[str],
205 section: Optional[str],
206 ) -> set[str]:
207 """Return paths in *manifest* matching the scope (all paths when no scope)."""
208 if not track and not section:
209 return set(manifest.keys())
210
211 matched: set[str] = set()
212 for path in manifest:
213 if track and path.startswith(f"tracks/{track}/"):
214 matched.add(path)
215 if section and path.startswith(f"sections/{section}/"):
216 matched.add(path)
217 return matched
218
219
220 # ---------------------------------------------------------------------------
221 # Object store writes
222 # ---------------------------------------------------------------------------
223
224
225 def _store_files(
226 root: pathlib.Path,
227 workdir: pathlib.Path,
228 paths: set[str],
229 ) -> dict[str, str]:
230 """Copy *paths* from *workdir* into the object store.
231
232 Returns a manifest ``{rel_path: object_id}`` for the stored files.
233 Objects already in the store are not overwritten (content-addressed).
234 """
235 manifest: dict[str, str] = {}
236 for rel_path in sorted(paths):
237 abs_path = workdir / rel_path
238 if not abs_path.exists():
239 logger.warning("⚠️ Stash: skipping missing file %s", rel_path)
240 continue
241 oid = hash_file(abs_path)
242 dest = _object_path(root, oid)
243 if not dest.exists():
244 dest.parent.mkdir(parents=True, exist_ok=True)
245 shutil.copy2(abs_path, dest)
246 logger.debug("✅ Stash stored object %s ← %s", oid[:8], rel_path)
247 manifest[rel_path] = oid
248 return manifest
249
250
251 # ---------------------------------------------------------------------------
252 # HEAD snapshot restoration helper
253 # ---------------------------------------------------------------------------
254
255
256 def _restore_from_manifest(
257 root: pathlib.Path,
258 workdir: pathlib.Path,
259 target_manifest: dict[str, str],
260 scope_paths: Optional[set[str]],
261 ) -> tuple[int, list[str]]:
262 """Write *target_manifest* files from the object store into *workdir*.
263
264 When *scope_paths* is given, only paths in that set are touched; other
265 working-tree files are left as-is.
266
267 Returns:
268 ``(files_written, missing_paths)`` where *missing_paths* lists files
269 that could not be restored because their objects are absent.
270 """
271 workdir.mkdir(parents=True, exist_ok=True)
272
273 if scope_paths is not None:
274 paths_to_restore = {p: oid for p, oid in target_manifest.items() if p in scope_paths}
275 else:
276 paths_to_restore = dict(target_manifest)
277
278 missing: list[str] = []
279 written = 0
280
281 for rel_path, oid in sorted(paths_to_restore.items()):
282 obj_file = _object_path(root, oid)
283 if not obj_file.exists():
284 missing.append(rel_path)
285 logger.warning(
286 "⚠️ Stash: object %s missing from store for %s", oid[:8], rel_path
287 )
288 continue
289 dest = workdir / rel_path
290 dest.parent.mkdir(parents=True, exist_ok=True)
291 shutil.copy2(obj_file, dest)
292 written += 1
293
294 # When scope_paths is None (full restore), delete files not in the manifest.
295 if scope_paths is None:
296 target_abs = {workdir / p for p in target_manifest}
297 for f in list(workdir.rglob("*")):
298 if f.is_file() and not f.name.startswith(".") and f not in target_abs:
299 f.unlink(missing_ok=True)
300
301 return written, missing
302
303
304 # ---------------------------------------------------------------------------
305 # Public API — pure filesystem operations
306 # ---------------------------------------------------------------------------
307
308
309 def list_stash(root: pathlib.Path) -> list[StashEntry]:
310 """Return all stash entries, newest first (index 0 = most recent).
311
312 Returns an empty list when the stash is empty or ``.muse/stash/``
313 does not exist.
314
315 Args:
316 root: Muse repository root.
317 """
318 files = _sorted_entries(root)
319 return [_read_entry(f, i) for i, f in enumerate(files)]
320
321
322 def drop_stash(root: pathlib.Path, index: int) -> StashEntry:
323 """Remove the stash entry at *index* from the stack.
324
325 Args:
326 root: Muse repository root.
327 index: 0-based stash index (0 = most recent).
328
329 Returns:
330 The dropped :class:`StashEntry`.
331
332 Raises:
333 IndexError: When *index* is out of range.
334 """
335 files = _sorted_entries(root)
336 if index < 0 or index >= len(files):
337 raise IndexError(
338 f"stash@{{{index}}} does not exist (stack has {len(files)} entries)"
339 )
340 entry = _read_entry(files[index], index)
341 files[index].unlink()
342 logger.info("✅ muse stash drop stash@{%d}: %s", index, entry.message)
343 return entry
344
345
346 def clear_stash(root: pathlib.Path) -> int:
347 """Remove all stash entries.
348
349 Args:
350 root: Muse repository root.
351
352 Returns:
353 Number of entries removed.
354 """
355 files = _sorted_entries(root)
356 for f in files:
357 f.unlink()
358 count = len(files)
359 if count:
360 logger.info("✅ muse stash clear: removed %d entries", count)
361 return count
362
363
364 # ---------------------------------------------------------------------------
365 # Push — save working state and restore HEAD
366 # ---------------------------------------------------------------------------
367
368
369 def push_stash(
370 root: pathlib.Path,
371 *,
372 message: Optional[str] = None,
373 track: Optional[str] = None,
374 section: Optional[str] = None,
375 head_manifest: Optional[dict[str, str]] = None,
376 ) -> StashPushResult:
377 """Save muse-work/ changes to the stash and restore HEAD snapshot.
378
379 Algorithm:
380 1. Walk muse-work/ and identify paths to stash (all, or scoped by
381 ``track``/``section``).
382 2. Copy each file into the object store (sharded ``.muse/objects/``).
383 3. Build a stash entry JSON and write it to ``.muse/stash/``.
384 4. Restore HEAD snapshot to muse-work/:
385 - Full push: restore full HEAD manifest (deletes extra files).
386 - Scoped push: restore HEAD for scoped paths only; other files
387 remain untouched.
388
389 Args:
390 root: Muse repository root.
391 message: Optional label; defaults to ``On <branch>: stash``.
392 track: Optional track name scope (e.g. ``"drums"``).
393 section: Optional section name scope (e.g. ``"chorus"``).
394 head_manifest: Snapshot manifest for the current HEAD commit, used
395 to restore muse-work/ after stashing. When ``None``
396 the branch has no commits and muse-work/ is cleared
397 (full push) or left untouched (scoped push).
398
399 Returns:
400 :class:`StashPushResult` describing what was saved and restored.
401 """
402 muse_dir = root / ".muse"
403 workdir = root / "muse-work"
404
405 # ── Resolve current branch ───────────────────────────────────────────
406 head_text = (muse_dir / "HEAD").read_text().strip()
407 branch = head_text.rsplit("/", 1)[-1]
408
409 # ── Build working-tree manifest ──────────────────────────────────────
410 from maestro.muse_cli.snapshot import walk_workdir
411
412 current_manifest: dict[str, str] = {}
413 if workdir.exists():
414 current_manifest = walk_workdir(workdir)
415
416 # ── Identify paths to stash ──────────────────────────────────────────
417 scope_paths = _filter_paths(current_manifest, track, section)
418 if not scope_paths:
419 # Nothing to stash (either empty workdir or scope matched nothing).
420 return StashPushResult(
421 stash_ref="",
422 message="",
423 branch=branch,
424 files_stashed=0,
425 head_restored=False,
426 missing_head=(),
427 )
428
429 # ── Copy files into object store ─────────────────────────────────────
430 stash_manifest = _store_files(root, workdir, scope_paths)
431
432 # ── Build and write stash entry ──────────────────────────────────────
433 created_at = datetime.datetime.now(datetime.timezone.utc).isoformat()
434 stash_id = f"stash-{created_at.replace(':', '').replace('+', 'Z').replace('.', '')}-{uuid.uuid4().hex[:8]}"
435 label = message or f"On {branch}: stash"
436 entry_data: dict[str, object] = {
437 "stash_id": stash_id,
438 "branch": branch,
439 "message": label,
440 "created_at": created_at,
441 "track": track,
442 "section": section,
443 "manifest": stash_manifest,
444 }
445 _write_entry(root, entry_data)
446
447 # ── Compute stash@{0} index (newly pushed is always newest) ──────────
448 files = _sorted_entries(root)
449 stash_ref = "stash@{0}"
450
451 # ── Restore HEAD snapshot to muse-work/ ──────────────────────────────
452 missing_head: list[str] = []
453 head_restored = False
454
455 if head_manifest is not None:
456 # Determine scope for restore
457 restore_scope: Optional[set[str]] = None
458 if track or section:
459 # Only restore paths that match the scope (others untouched)
460 restore_scope = _filter_paths(head_manifest, track, section)
461 # Also remove scoped paths that are NOT in head_manifest
462 # (new files in workdir under the scope must be deleted)
463 paths_to_clear = scope_paths - set(head_manifest.keys())
464 for rel_path in paths_to_clear:
465 abs_path = workdir / rel_path
466 if abs_path.exists():
467 abs_path.unlink(missing_ok=True)
468 else:
469 restore_scope = None # full restore
470
471 _, missing_head = _restore_from_manifest(
472 root, workdir, head_manifest, restore_scope
473 )
474 head_restored = True
475 else:
476 # No HEAD commit: full push clears muse-work/, scoped push does nothing.
477 if not track and not section:
478 for f in list(workdir.rglob("*")):
479 if f.is_file() and not f.name.startswith("."):
480 f.unlink(missing_ok=True)
481
482 logger.info(
483 "✅ muse stash push: %s (%d files, branch=%r)",
484 stash_ref,
485 len(stash_manifest),
486 branch,
487 )
488
489 return StashPushResult(
490 stash_ref=stash_ref,
491 message=label,
492 branch=branch,
493 files_stashed=len(stash_manifest),
494 head_restored=head_restored,
495 missing_head=tuple(sorted(missing_head)),
496 )
497
498
499 # ---------------------------------------------------------------------------
500 # Apply / Pop — restore stash to working tree
501 # ---------------------------------------------------------------------------
502
503
504 def apply_stash(
505 root: pathlib.Path,
506 index: int = 0,
507 *,
508 drop: bool = False,
509 ) -> StashApplyResult:
510 """Restore a stash entry to muse-work/.
511
512 Algorithm:
513 1. Resolve *index* → ``StashEntry``.
514 2. For each path in the entry's manifest, copy the object from the
515 store back into muse-work/ (overwriting any conflicting file).
516 3. If *drop* is True, remove the stash entry (this is ``pop`` semantics).
517
518 Conflict strategy: last-write-wins. Files in muse-work/ that are NOT
519 in the stash manifest are left untouched; only stash paths are written.
520
521 Args:
522 root: Muse repository root.
523 index: 0-based stash index (0 = most recent).
524 drop: Remove the entry after applying (True → pop, False → apply).
525
526 Returns:
527 :class:`StashApplyResult` describing what was applied.
528
529 Raises:
530 IndexError: When *index* is out of range.
531 """
532 files = _sorted_entries(root)
533 if index < 0 or index >= len(files):
534 raise IndexError(
535 f"stash@{{{index}}} does not exist (stack has {len(files)} entries)"
536 )
537
538 entry = _read_entry(files[index], index)
539 stash_ref = f"stash@{{{index}}}"
540 workdir = root / "muse-work"
541 workdir.mkdir(parents=True, exist_ok=True)
542
543 # ── Restore files from object store ──────────────────────────────────
544 missing: list[str] = []
545 written = 0
546
547 for rel_path, oid in sorted(entry.manifest.items()):
548 obj_file = _object_path(root, oid)
549 if not obj_file.exists():
550 missing.append(rel_path)
551 logger.warning(
552 "⚠️ Stash apply: object %s missing for %s", oid[:8], rel_path
553 )
554 continue
555 dest = workdir / rel_path
556 dest.parent.mkdir(parents=True, exist_ok=True)
557 shutil.copy2(obj_file, dest)
558 written += 1
559 logger.debug("✅ Stash apply: restored %s from object %s", rel_path, oid[:8])
560
561 # ── Drop entry if pop semantics ───────────────────────────────────────
562 if drop:
563 files[index].unlink()
564 logger.info("✅ muse stash pop: applied and dropped %s", stash_ref)
565 else:
566 logger.info("✅ muse stash apply: applied %s (kept)", stash_ref)
567
568 return StashApplyResult(
569 stash_ref=stash_ref,
570 message=entry.message,
571 files_applied=written,
572 missing=tuple(sorted(missing)),
573 dropped=drop,
574 )