"""Code-domain staging index — ``muse code add`` persistence layer. The staging index lives at ``.muse/code/stage.json``. It records which files the user has explicitly staged for the next ``muse commit``, along with the content-addressed object ID of each staged version. Format (UTF-8 JSON):: { "version": 3, "entries": { "src/auth.py": { "object_id": "", "mode": "M", "staged_at": "2026-03-21T14:32:00+00:00" } } } Modes: - ``"A"`` — added (file is new; not in the previous commit) - ``"M"`` — modified (file exists in the previous commit) - ``"D"`` — deleted (file will be removed from the next commit) """ import json as _json import logging import os import pathlib import tempfile from collections.abc import Mapping from typing import Literal, TypedDict from muse.core.types import blob_id, now_utc_iso from muse.core.paths import code_stage_path as _code_stage_path # Canonical object ID for an empty directory — sha256 of zero bytes. # Every tracked empty directory uses this object ID in the stage index, # making directories proper content-addressed objects (all share one object). EMPTY_DIR_OID: str = blob_id(b"") # Legacy sentinel written by older Muse versions — migrated on read. _DIR_SENTINEL_LEGACY = "dir:" type StagedFileMap = dict[str, "StagedEntry"] logger = logging.getLogger(__name__) class StagedEntry(TypedDict): """One file's staging record.""" object_id: str mode: Literal["A", "M", "D"] staged_at: str # ISO-8601 class StageIndex(TypedDict): """Full contents of the stage index (JSON on disk).""" version: int entries: StagedFileMap dir_renames: dict[str, str] # old_dir_path → new_dir_path (no trailing slash) _STAGE_VERSION = 3 def stage_path(root: pathlib.Path) -> pathlib.Path: """Return the absolute path to ``.muse/code/stage.json``.""" return _code_stage_path(root) def read_stage(root: pathlib.Path) -> StagedFileMap: """Read the stage index from disk. Returns an empty dict when no stage exists (fresh repo or after a commit). Returns an empty dict and logs a warning when the file is corrupt or contains old binary msgpack content — this is a safe reset; the user re-stages. """ path = stage_path(root) if not path.exists(): return {} try: raw = path.read_bytes() # Msgpack first byte is always 0x80–0xFF (fixmap, fixarray, fixstr, etc.) # for any dict/list/string. If the file starts with such a byte, it's # old binary msgpack — treat as stale and return empty. if raw and raw[0] > 0x7F: logger.warning( "⚠️ Stage index at %s contains old binary format — " "clearing stage. Run 'muse code add' again to re-stage.", path, ) path.unlink(missing_ok=True) return {} data = _json.loads(raw.decode("utf-8")) if not isinstance(data, dict): return {} raw_entries_obj = data.get("entries") raw_entries: Mapping[str, object] = raw_entries_obj if isinstance(raw_entries_obj, dict) else {} entries: StagedFileMap = {} for k, v in raw_entries.items(): if not isinstance(v, dict): continue raw_mode = v.get("mode") if raw_mode == "A": mode: Literal["A", "M", "D"] = "A" elif raw_mode == "D": mode = "D" else: mode = "M" obj_id = v.get("object_id", "") if obj_id == _DIR_SENTINEL_LEGACY: obj_id = EMPTY_DIR_OID # migrate legacy sentinel to real content hash staged = v.get("staged_at", "") entries[str(k)] = StagedEntry( object_id=str(obj_id) if isinstance(obj_id, str) else "", mode=mode, staged_at=str(staged) if isinstance(staged, str) else "", ) return entries except Exception as exc: logger.warning( "⚠️ Stage index at %s is corrupt (%s) — clearing stage to prevent " "phantom staged entries. Run 'muse code add' again to re-stage.", path, exc, ) return {} def read_stage_dir_renames(root: pathlib.Path) -> Mapping[str, str]: """Return the directory rename map stored in the stage index. Returns ``{}`` when no stage exists or when the file contains no renames. Keys and values are repo-relative paths without trailing slash. """ path = stage_path(root) if not path.exists(): return {} try: data = _json.loads(path.read_bytes().decode("utf-8")) if not isinstance(data, dict): return {} renames = data.get("dir_renames", {}) if not isinstance(renames, dict): return {} return {str(k): str(v) for k, v in renames.items()} except Exception: return {} def write_stage( root: pathlib.Path, entries: StagedFileMap, dir_renames: dict[str, str] | None = None, ) -> None: """Persist *entries* (and optional *dir_renames*) to ``.muse/code/stage.json``. Creates the ``.muse/code/`` directory if it does not exist. Writing an empty dict clears the stage file (equivalent to calling :func:`clear_stage`). ``dir_renames`` records explicit directory rename pairs produced by ``muse mv`` on tracked directories (old_path → new_path, no trailing slash). When ``None`` the existing on-disk renames are preserved; pass ``{}`` to clear them. Writes are atomic (temp file + ``os.replace``) so a process crash mid-write never leaves a corrupt stage file. """ path = stage_path(root) if not entries: path.unlink(missing_ok=True) return # Preserve existing dir_renames when the caller doesn't supply new ones. if dir_renames is None: dir_renames = read_stage_dir_renames(root) path.parent.mkdir(parents=True, exist_ok=True) payload = StageIndex( version=_STAGE_VERSION, entries=entries, dir_renames=dir_renames, ) encoded = _json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") fd, tmp_str = tempfile.mkstemp( dir=path.parent, prefix=".stage-tmp-", suffix=".json" ) tmp = pathlib.Path(tmp_str) try: with os.fdopen(fd, "wb") as fh: fh.write(encoded) fh.flush() os.fsync(fh.fileno()) tmp.replace(path) except Exception: tmp.unlink(missing_ok=True) raise def clear_stage(root: pathlib.Path) -> None: """Remove the stage index, resetting to full-snapshot mode.""" stage_path(root).unlink(missing_ok=True) def make_entry( object_id: str, mode: Literal["A", "M", "D"], ) -> StagedEntry: """Build a :class:`StagedEntry` with the current UTC timestamp.""" return StagedEntry( object_id=object_id, mode=mode, staged_at=now_utc_iso(), )