muse_restore.py
python
| 1 | """Muse Restore Service — restore specific files from a commit or index. |
| 2 | |
| 3 | ``muse restore`` is surgical: restore one instrument track from a specific |
| 4 | commit while keeping everything else at HEAD. Critical for music production |
| 5 | where you want "the bass from take 3, everything else from take 7." |
| 6 | |
| 7 | Restore modes |
| 8 | ------------- |
| 9 | **worktree** (default, ``--worktree``) |
| 10 | Copy the file content recorded in the *source* snapshot directly into |
| 11 | ``muse-work/``. Branch pointer and index are not changed. This is the |
| 12 | primary use case: "put the bass from take 3 back into my working tree." |
| 13 | |
| 14 | **staged** (``--staged``) |
| 15 | In a full VCS with an explicit staging area this would reset the index |
| 16 | entry for the path from the source snapshot without touching ``muse-work/``. |
| 17 | In the current Muse model (no separate staging area) ``--staged`` is |
| 18 | documented for forward-compatibility and behaves identically to |
| 19 | ``--worktree``: it restores the file in ``muse-work/`` from the source |
| 20 | snapshot. When a staging index is added this module will be updated. |
| 21 | |
| 22 | **source** (``--source <commit>``) |
| 23 | Selects which snapshot to extract the file from. Defaults to ``HEAD`` |
| 24 | when omitted. |
| 25 | |
| 26 | Object store contract |
| 27 | --------------------- |
| 28 | Restore reads objects from ``.muse/objects/`` exactly like ``muse reset |
| 29 | --hard``. If an object is missing, :class:`MissingObjectError` is raised |
| 30 | and ``muse-work/`` is left unchanged (the restore is atomic per path). |
| 31 | |
| 32 | This module is a pure service layer — no Typer, no CLI, no StateStore. |
| 33 | Import boundary: may import muse_cli.{db,models,object_store}, muse_reset |
| 34 | (for MissingObjectError and resolve_ref), but NOT executor, |
| 35 | maestro_handlers, mcp, or StateStore. |
| 36 | """ |
| 37 | from __future__ import annotations |
| 38 | |
| 39 | import logging |
| 40 | import pathlib |
| 41 | from dataclasses import dataclass, field |
| 42 | |
| 43 | from sqlalchemy.ext.asyncio import AsyncSession |
| 44 | |
| 45 | from maestro.muse_cli.db import get_commit_snapshot_manifest |
| 46 | from maestro.muse_cli.object_store import has_object, restore_object |
| 47 | from maestro.services.muse_reset import ( |
| 48 | MissingObjectError, |
| 49 | resolve_ref, |
| 50 | ) |
| 51 | |
| 52 | logger = logging.getLogger(__name__) |
| 53 | |
| 54 | # --------------------------------------------------------------------------- |
| 55 | # Public types |
| 56 | # --------------------------------------------------------------------------- |
| 57 | |
| 58 | |
| 59 | @dataclass(frozen=True) |
| 60 | class RestoreResult: |
| 61 | """Outcome of a completed ``muse restore`` operation. |
| 62 | |
| 63 | Attributes: |
| 64 | source_commit_id: Full SHA of the commit the files were extracted from. |
| 65 | paths_restored: Relative paths (within ``muse-work/``) that were |
| 66 | written to disk. |
| 67 | staged: Whether ``--staged`` mode was active. |
| 68 | """ |
| 69 | |
| 70 | source_commit_id: str |
| 71 | paths_restored: list[str] = field(default_factory=list) |
| 72 | staged: bool = False |
| 73 | |
| 74 | |
| 75 | class PathNotInSnapshotError(Exception): |
| 76 | """Raised when a requested path is absent from the source snapshot. |
| 77 | |
| 78 | Attributes: |
| 79 | rel_path: The path that was not found. |
| 80 | source_commit_id: The commit that was searched. |
| 81 | """ |
| 82 | |
| 83 | def __init__(self, rel_path: str, source_commit_id: str) -> None: |
| 84 | super().__init__( |
| 85 | f"Path {rel_path!r} not found in snapshot of commit " |
| 86 | f"{source_commit_id[:8]}. " |
| 87 | "Use 'muse log' to list commits and 'muse show <commit>' to inspect " |
| 88 | "the snapshot manifest." |
| 89 | ) |
| 90 | self.rel_path = rel_path |
| 91 | self.source_commit_id = source_commit_id |
| 92 | |
| 93 | |
| 94 | # --------------------------------------------------------------------------- |
| 95 | # Core restore logic |
| 96 | # --------------------------------------------------------------------------- |
| 97 | |
| 98 | |
| 99 | async def perform_restore( |
| 100 | *, |
| 101 | root: pathlib.Path, |
| 102 | session: AsyncSession, |
| 103 | paths: list[str], |
| 104 | source_ref: str | None, |
| 105 | staged: bool, |
| 106 | ) -> RestoreResult: |
| 107 | """Restore specific files from a source commit into ``muse-work/``. |
| 108 | |
| 109 | Resolves the source commit reference, validates that every requested path |
| 110 | exists in the snapshot manifest, verifies all required objects are present |
| 111 | in the object store (fail-fast before touching ``muse-work/``), then copies |
| 112 | each file atomically. |
| 113 | |
| 114 | The branch pointer is never modified — only ``muse-work/`` files are written. |
| 115 | |
| 116 | Args: |
| 117 | root: Muse repository root (directory containing ``.muse/``). |
| 118 | session: Open async DB session. |
| 119 | paths: Relative paths within ``muse-work/`` to restore. Must be |
| 120 | non-empty. Paths may be given as ``muse-work/bass/bassline.mid`` |
| 121 | or bare ``bass/bassline.mid`` — the ``muse-work/`` prefix is |
| 122 | stripped if present. |
| 123 | source_ref: Commit reference to restore from (``HEAD``, ``HEAD~N``, SHA, |
| 124 | or ``None`` for HEAD). |
| 125 | staged: When ``True``, ``--staged`` mode is active. In the current |
| 126 | Muse model (no separate staging area) this is semantically |
| 127 | equivalent to ``--worktree``. |
| 128 | |
| 129 | Returns: |
| 130 | :class:`RestoreResult` describing the completed operation. |
| 131 | |
| 132 | Raises: |
| 133 | typer.Exit: On user-facing errors (ref not found, no commits). |
| 134 | PathNotInSnapshotError: When any requested path is absent from the source. |
| 135 | MissingObjectError: When a required blob is absent from the object store. |
| 136 | """ |
| 137 | import json |
| 138 | |
| 139 | import typer |
| 140 | |
| 141 | from maestro.muse_cli.errors import ExitCode |
| 142 | |
| 143 | muse_dir = root / ".muse" |
| 144 | |
| 145 | # ── Repo identity ──────────────────────────────────────────────────── |
| 146 | repo_data: dict[str, str] = json.loads((muse_dir / "repo.json").read_text()) |
| 147 | repo_id = repo_data["repo_id"] |
| 148 | |
| 149 | # ── Current branch ─────────────────────────────────────────────────── |
| 150 | head_ref = (muse_dir / "HEAD").read_text().strip() # "refs/heads/main" |
| 151 | branch = head_ref.rsplit("/", 1)[-1] # "main" |
| 152 | ref_path = muse_dir / pathlib.Path(head_ref) |
| 153 | |
| 154 | if not ref_path.exists() or not ref_path.read_text().strip(): |
| 155 | typer.echo("❌ Current branch has no commits. Nothing to restore.") |
| 156 | raise typer.Exit(code=ExitCode.USER_ERROR) |
| 157 | |
| 158 | # ── Resolve source commit ───────────────────────────────────────────── |
| 159 | effective_ref = source_ref if source_ref is not None else "HEAD" |
| 160 | source_commit = await resolve_ref(session, repo_id, branch, effective_ref) |
| 161 | if source_commit is None: |
| 162 | typer.echo(f"❌ Could not resolve source ref: {effective_ref!r}") |
| 163 | raise typer.Exit(code=ExitCode.USER_ERROR) |
| 164 | |
| 165 | source_commit_id = source_commit.commit_id |
| 166 | |
| 167 | # ── Load snapshot manifest ──────────────────────────────────────────── |
| 168 | manifest = await get_commit_snapshot_manifest(session, source_commit_id) |
| 169 | if manifest is None: |
| 170 | typer.echo( |
| 171 | f"❌ Could not load snapshot for commit {source_commit_id[:8]}. " |
| 172 | "Database may be corrupt." |
| 173 | ) |
| 174 | raise typer.Exit(code=ExitCode.INTERNAL_ERROR) |
| 175 | |
| 176 | # ── Normalise paths — strip leading "muse-work/" prefix if present ──── |
| 177 | normalised: list[str] = [] |
| 178 | for p in paths: |
| 179 | stripped = p.removeprefix("muse-work/") |
| 180 | normalised.append(stripped) |
| 181 | |
| 182 | # ── Validate: every path must be in the manifest ───────────────────── |
| 183 | for rel_path in normalised: |
| 184 | if rel_path not in manifest: |
| 185 | raise PathNotInSnapshotError(rel_path, source_commit_id) |
| 186 | |
| 187 | # ── Validate: every object must exist (fail-fast before touching disk) ─ |
| 188 | for rel_path in normalised: |
| 189 | object_id = manifest[rel_path] |
| 190 | if not has_object(root, object_id): |
| 191 | raise MissingObjectError(object_id, rel_path) |
| 192 | |
| 193 | # ── Restore files into muse-work/ ───────────────────────────────────── |
| 194 | workdir = root / "muse-work" |
| 195 | workdir.mkdir(parents=True, exist_ok=True) |
| 196 | |
| 197 | restored: list[str] = [] |
| 198 | for rel_path in normalised: |
| 199 | object_id = manifest[rel_path] |
| 200 | dest = workdir / rel_path |
| 201 | restore_object(root, object_id, dest) |
| 202 | restored.append(rel_path) |
| 203 | logger.debug( |
| 204 | "✅ Restored %s from object %s (commit %s)", |
| 205 | rel_path, |
| 206 | object_id[:8], |
| 207 | source_commit_id[:8], |
| 208 | ) |
| 209 | |
| 210 | mode_label = "--staged" if staged else "--worktree" |
| 211 | logger.info( |
| 212 | "✅ muse restore %s: %d file(s) from commit %s", |
| 213 | mode_label, |
| 214 | len(restored), |
| 215 | source_commit_id[:8], |
| 216 | ) |
| 217 | |
| 218 | return RestoreResult( |
| 219 | source_commit_id=source_commit_id, |
| 220 | paths_restored=restored, |
| 221 | staged=staged, |
| 222 | ) |