gabriel / muse public
stage.py python
223 lines 7.2 KB
Raw
sha256:a154bc65916614c833d5a40a10d81ba3eae0d0495b0afddd34dc34f18d5e91b8 fix: test suite alignment and typing audit — zero violations Sonnet 4.6 minor ⚠ breaking 50 days ago
1 """Code-domain staging index — ``muse code add`` persistence layer.
2
3 The staging index lives at ``.muse/code/stage.msgpack``. It records which
4 files the user has explicitly staged for the next ``muse commit``, along
5 with the content-addressed object ID of each staged version.
6
7 The file is internal to the VCS and never shown directly to the user, so it
8 uses msgpack (binary, compact) rather than JSON. A one-time migration from
9 the legacy ``stage.json`` format is performed transparently on the first
10 ``muse code add`` after upgrading.
11
12 Format (msgpack-encoded dict)::
13
14 {
15 "version": 2,
16 "entries": {
17 "src/auth.py": {
18 "object_id": "<sha256>",
19 "mode": "M",
20 "staged_at": "2026-03-21T14:32:00+00:00"
21 }
22 }
23 }
24
25 Modes:
26
27 - ``"A"`` — added (file is new; not in the previous commit)
28 - ``"M"`` — modified (file exists in the previous commit)
29 - ``"D"`` — deleted (file will be removed from the next commit)
30 """
31
32 from __future__ import annotations
33
34 import datetime
35 import logging
36 import os
37 import pathlib
38 import tempfile
39 from typing import Literal, TypedDict
40
41 import msgpack
42
43 from muse.core.store import MsgpackDict, MsgpackValue, read_msgpack_file
44
45 type StagedFileMap = dict[str, "StagedEntry"]
46
47 logger = logging.getLogger(__name__)
48
49
50 class StagedEntry(TypedDict):
51 """One file's staging record."""
52
53 object_id: str
54 mode: Literal["A", "M", "D"]
55 staged_at: str # ISO-8601
56
57
58 class StageIndex(TypedDict):
59 """Full contents of the stage index (msgpack-encoded on disk)."""
60
61 version: int
62 entries: StagedFileMap
63
64
65 _STAGE_VERSION = 2
66
67
68 def stage_path(root: pathlib.Path) -> pathlib.Path:
69 """Return the absolute path to ``.muse/code/stage.msgpack``."""
70 return root / ".muse" / "code" / "stage.msgpack"
71
72
73 def _legacy_json_path(root: pathlib.Path) -> pathlib.Path:
74 """Return the old JSON stage path (used only during migration)."""
75 return root / ".muse" / "code" / "stage.json"
76
77
78 def read_stage(root: pathlib.Path) -> StagedFileMap:
79 """Read the stage index from disk.
80
81 Returns an empty dict when no stage exists (fresh repo or after a commit).
82
83 Raises a warning and clears the stage if the file is corrupt — this
84 surfaces the data loss explicitly rather than silently proceeding with
85 an empty stage (which would look like "nothing staged").
86
87 Performs a one-time transparent migration from the legacy JSON format.
88 """
89 path = stage_path(root)
90
91 if path.exists():
92 try:
93 _raw = read_msgpack_file(path, strict_map_key=False)
94 if not isinstance(_raw, dict):
95 return {}
96 raw: MsgpackDict = _raw
97 raw_entries_obj = raw.get("entries")
98 raw_entries: MsgpackDict = (
99 raw_entries_obj if isinstance(raw_entries_obj, dict) else {}
100 )
101 entries: StagedFileMap = {}
102 for k, v in raw_entries.items():
103 if not isinstance(v, dict):
104 continue
105 raw_mode = v.get("mode")
106 if raw_mode == "A":
107 mode: Literal["A", "M", "D"] = "A"
108 elif raw_mode == "D":
109 mode = "D"
110 else:
111 mode = "M"
112 obj_id = v.get("object_id", "")
113 staged = v.get("staged_at", "")
114 entries[str(k)] = StagedEntry(
115 object_id=str(obj_id) if isinstance(obj_id, str) else "",
116 mode=mode,
117 staged_at=str(staged) if isinstance(staged, str) else "",
118 )
119 return entries
120 except Exception as exc:
121 logger.warning(
122 "⚠️ Stage index at %s is corrupt (%s) — clearing stage to prevent "
123 "phantom staged entries. Run 'muse code add' again to re-stage.",
124 path,
125 exc,
126 )
127 path.unlink(missing_ok=True)
128 return {}
129
130 # One-time migration from the legacy JSON format.
131 legacy = _legacy_json_path(root)
132 if legacy.exists():
133 try:
134 import json
135 data: MsgpackDict = json.loads(legacy.read_text(encoding="utf-8"))
136 legacy_entries_obj = data.get("entries")
137 legacy_entries: MsgpackDict = (
138 legacy_entries_obj if isinstance(legacy_entries_obj, dict) else {}
139 )
140 migrated: StagedFileMap = {}
141 for k, v in legacy_entries.items():
142 if not isinstance(v, dict):
143 continue
144 raw_mode = v.get("mode")
145 if raw_mode == "A":
146 leg_mode: Literal["A", "M", "D"] = "A"
147 elif raw_mode == "D":
148 leg_mode = "D"
149 else:
150 leg_mode = "M"
151 obj_id = v.get("object_id", "")
152 staged = v.get("staged_at", "")
153 migrated[k] = StagedEntry(
154 object_id=str(obj_id) if isinstance(obj_id, str) else "",
155 mode=leg_mode,
156 staged_at=str(staged) if isinstance(staged, str) else "",
157 )
158 # Persist as msgpack and remove the old JSON file.
159 write_stage(root, migrated)
160 legacy.unlink(missing_ok=True)
161 logger.debug("✅ Migrated stage index from JSON to msgpack.")
162 return migrated
163 except Exception as exc:
164 logger.warning("⚠️ Legacy stage.json is corrupt (%s) — clearing.", exc)
165 legacy.unlink(missing_ok=True)
166 return {}
167
168 return {}
169
170
171 def write_stage(root: pathlib.Path, entries: StagedFileMap) -> None:
172 """Persist *entries* to ``.muse/code/stage.msgpack``.
173
174 Creates the ``.muse/code/`` directory if it does not exist. Writing
175 an empty dict clears the stage file (equivalent to calling
176 :func:`clear_stage`).
177
178 Writes are atomic (temp file + ``os.replace``) so a process crash
179 mid-write never leaves a corrupt stage file.
180 """
181 path = stage_path(root)
182 if not entries:
183 path.unlink(missing_ok=True)
184 _legacy_json_path(root).unlink(missing_ok=True)
185 return
186
187 path.parent.mkdir(parents=True, exist_ok=True)
188 payload = StageIndex(
189 version=_STAGE_VERSION,
190 entries=entries,
191 )
192 fd, tmp_str = tempfile.mkstemp(
193 dir=path.parent, prefix=".stage-tmp-", suffix=".msgpack"
194 )
195 tmp = pathlib.Path(tmp_str)
196 try:
197 os.close(fd)
198 tmp.write_bytes(msgpack.packb(payload, use_bin_type=True))
199 tmp.replace(path)
200 except Exception:
201 tmp.unlink(missing_ok=True)
202 raise
203
204 # Remove any leftover legacy JSON on every successful write.
205 _legacy_json_path(root).unlink(missing_ok=True)
206
207
208 def clear_stage(root: pathlib.Path) -> None:
209 """Remove the stage index, resetting to full-snapshot mode."""
210 stage_path(root).unlink(missing_ok=True)
211 _legacy_json_path(root).unlink(missing_ok=True)
212
213
214 def make_entry(
215 object_id: str,
216 mode: Literal["A", "M", "D"],
217 ) -> StagedEntry:
218 """Build a :class:`StagedEntry` with the current UTC timestamp."""
219 return StagedEntry(
220 object_id=object_id,
221 mode=mode,
222 staged_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
223 )
File History 3 commits
sha256:a154bc65916614c833d5a40a10d81ba3eae0d0495b0afddd34dc34f18d5e91b8 fix: test suite alignment and typing audit — zero violations Sonnet 4.6 minor 50 days ago
sha256:3f46367650ccd121654f3bbe06ed3471a9007c3229fe9556d1069d64b6a2550a refactor: directories are proper content-addressed objects … Sonnet 4.6 patch 50 days ago
sha256:8c872e4dffa2db45a9629956256fa1c99a3d2ff33b80c055252e58d94a0e8d1b feat: staged directory renames shown as renamed in muse status Sonnet 4.6 minor 50 days ago