"""Code-domain staging index — ``muse code add`` persistence layer. The staging index lives at ``.muse/code/stage.msgpack``. 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. The file is internal to the VCS and never shown directly to the user, so it uses msgpack (binary, compact) rather than JSON. A one-time migration from the legacy ``stage.json`` format is performed transparently on the first ``muse code add`` after upgrading. Format (msgpack-encoded dict):: { "version": 2, "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) """ from __future__ import annotations import datetime import logging import os import pathlib import tempfile from typing import Literal, TypedDict import msgpack from muse.core.store import MsgpackDict, MsgpackValue, read_msgpack_file 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 (msgpack-encoded on disk).""" version: int entries: StagedFileMap _STAGE_VERSION = 2 def stage_path(root: pathlib.Path) -> pathlib.Path: """Return the absolute path to ``.muse/code/stage.msgpack``.""" return root / ".muse" / "code" / "stage.msgpack" def _legacy_json_path(root: pathlib.Path) -> pathlib.Path: """Return the old JSON stage path (used only during migration).""" return root / ".muse" / "code" / "stage.json" 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). Raises a warning and clears the stage if the file is corrupt — this surfaces the data loss explicitly rather than silently proceeding with an empty stage (which would look like "nothing staged"). Performs a one-time transparent migration from the legacy JSON format. """ path = stage_path(root) if path.exists(): try: _raw = read_msgpack_file(path, strict_map_key=False) if not isinstance(_raw, dict): return {} raw: MsgpackDict = _raw raw_entries_obj = raw.get("entries") raw_entries: MsgpackDict = ( 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", "") 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, ) path.unlink(missing_ok=True) return {} # One-time migration from the legacy JSON format. legacy = _legacy_json_path(root) if legacy.exists(): try: import json data: MsgpackDict = json.loads(legacy.read_text(encoding="utf-8")) legacy_entries_obj = data.get("entries") legacy_entries: MsgpackDict = ( legacy_entries_obj if isinstance(legacy_entries_obj, dict) else {} ) migrated: StagedFileMap = {} for k, v in legacy_entries.items(): if not isinstance(v, dict): continue raw_mode = v.get("mode") if raw_mode == "A": leg_mode: Literal["A", "M", "D"] = "A" elif raw_mode == "D": leg_mode = "D" else: leg_mode = "M" obj_id = v.get("object_id", "") staged = v.get("staged_at", "") migrated[k] = StagedEntry( object_id=str(obj_id) if isinstance(obj_id, str) else "", mode=leg_mode, staged_at=str(staged) if isinstance(staged, str) else "", ) # Persist as msgpack and remove the old JSON file. write_stage(root, migrated) legacy.unlink(missing_ok=True) logger.debug("✅ Migrated stage index from JSON to msgpack.") return migrated except Exception as exc: logger.warning("⚠️ Legacy stage.json is corrupt (%s) — clearing.", exc) legacy.unlink(missing_ok=True) return {} return {} def write_stage(root: pathlib.Path, entries: StagedFileMap) -> None: """Persist *entries* to ``.muse/code/stage.msgpack``. Creates the ``.muse/code/`` directory if it does not exist. Writing an empty dict clears the stage file (equivalent to calling :func:`clear_stage`). 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) _legacy_json_path(root).unlink(missing_ok=True) return path.parent.mkdir(parents=True, exist_ok=True) payload = StageIndex( version=_STAGE_VERSION, entries=entries, ) fd, tmp_str = tempfile.mkstemp( dir=path.parent, prefix=".stage-tmp-", suffix=".msgpack" ) tmp = pathlib.Path(tmp_str) try: os.close(fd) tmp.write_bytes(msgpack.packb(payload, use_bin_type=True)) tmp.replace(path) except Exception: tmp.unlink(missing_ok=True) raise # Remove any leftover legacy JSON on every successful write. _legacy_json_path(root).unlink(missing_ok=True) def clear_stage(root: pathlib.Path) -> None: """Remove the stage index, resetting to full-snapshot mode.""" stage_path(root).unlink(missing_ok=True) _legacy_json_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=datetime.datetime.now(datetime.timezone.utc).isoformat(), )