aaronrene / gabriel-muse public
rebase.py python
409 lines 13.8 KB
Raw
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago
1 """Rebase engine for ``muse rebase``.
2
3 A Muse rebase replays a sequence of commits onto a new base. Because commits
4 are content-addressed, replaying a commit produces a *new* commit with a new
5 ID — the original commits are untouched in the store.
6
7 Algorithm
8 ---------
9 Given::
10
11 A ─── B ─── C ─── D (current branch HEAD = D)
12 \\
13 E ─── F (upstream = F)
14
15 After ``muse rebase F`` (or ``muse rebase --onto F A`` where A is the merge
16 base)::
17
18 E ─── F ─── B' ─── C' ─── D' (current branch HEAD = D')
19
20 Each replayed commit ``X'`` is produced by:
21
22 1. Taking the delta between ``X`` and its parent ``X-1`` (what changed).
23 2. Applying that delta on top of the current tip via the domain plugin's
24 three-way merge (same logic as cherry-pick).
25 3. Writing a new ``CommitRecord`` with the new parent pointer.
26
27 State
28 -----
29 When a conflict occurs mid-replay, the rebase pauses and writes
30 ``.muse/REBASE_STATE.json``. The user resolves the conflict and runs
31 ``muse rebase --continue`` to resume, or ``muse rebase --abort`` to undo.
32
33 Squash mode
34 -----------
35 When ``squash=True``, all commits are replayed without writing intermediate
36 commits — only the final merged state is committed. Squash mode does not
37 support ``--continue`` because no incremental state is tracked.
38
39 Security model
40 --------------
41 - ``REBASE_STATE.json`` is checked for symlinks before any read or write to
42 prevent path-traversal via a planted symlink.
43 - File size is capped at ``_MAX_STATE_BYTES`` before parsing to prevent OOM
44 via a crafted or corrupt state file.
45 """
46
47 from __future__ import annotations
48
49 import datetime
50 import json
51 import logging
52 import pathlib
53 from typing import TypedDict
54
55 from muse.core.snapshot import compute_commit_id, compute_snapshot_id, directories_from_manifest
56 from muse.core.store import (
57 CommitRecord,
58 SnapshotRecord,
59 read_commit,
60 read_snapshot,
61 write_branch_ref,
62 write_commit,
63 write_snapshot,
64 write_text_atomic,
65 )
66 from muse.core.validation import validate_branch_name
67 from muse.core.workdir import apply_manifest
68 from muse.domain import MergeResult, MuseDomainPlugin, SnapshotManifest
69 from muse.core._types import Manifest
70
71 logger = logging.getLogger(__name__)
72
73 _REBASE_STATE_FILE = ".muse/REBASE_STATE.json"
74
75 # 4 MiB — a state file with 10 000 commit IDs at ~64 bytes each is ~640 KiB.
76 _MAX_STATE_BYTES = 4 * 1024 * 1024
77
78
79 # ---------------------------------------------------------------------------
80 # State TypedDict
81 # ---------------------------------------------------------------------------
82
83
84 class RebaseState(TypedDict):
85 """Serialisable state for an in-progress rebase session."""
86
87 original_branch: str
88 original_head: str
89 onto: str
90 remaining: list[str]
91 completed: list[str]
92 squash: bool
93
94
95 class RebaseProgress(TypedDict):
96 """Snapshot of rebase progress for ``--status`` output.
97
98 Attributes:
99 active: Whether a rebase is in progress.
100 original_branch: Branch being rebased (empty string if not active).
101 original_head: HEAD before the rebase started.
102 onto: Target commit ID.
103 total: Total commits to replay.
104 done: Commits already replayed.
105 remaining: Commits yet to replay.
106 squash: Whether this is a squash rebase.
107 """
108
109 active: bool
110 original_branch: str
111 original_head: str
112 onto: str
113 total: int
114 done: int
115 remaining: int
116 squash: bool
117
118
119 # ---------------------------------------------------------------------------
120 # State file I/O
121 # ---------------------------------------------------------------------------
122
123
124 def load_rebase_state(root: pathlib.Path) -> RebaseState | None:
125 """Return the current rebase state, or ``None`` if none is active.
126
127 Security guards applied before any read:
128
129 - Symlink check: a symlink at the state path could redirect reads to
130 sensitive files outside the repo or writes to unintended locations.
131 - Size cap (``_MAX_STATE_BYTES``): a tampered or corrupt state file cannot
132 exhaust memory.
133 """
134 path = root / _REBASE_STATE_FILE
135 if not path.exists():
136 return None
137 if path.is_symlink():
138 logger.warning(
139 "⚠️ REBASE_STATE.json is a symlink — ignoring to prevent path traversal"
140 )
141 return None
142 try:
143 size = path.stat().st_size
144 if size > _MAX_STATE_BYTES:
145 logger.warning(
146 "⚠️ REBASE_STATE.json is %.1f MiB — exceeds cap of %d MiB; ignoring",
147 size / (1024 * 1024),
148 _MAX_STATE_BYTES // (1024 * 1024),
149 )
150 return None
151 raw = json.loads(path.read_text(encoding="utf-8"))
152 except (json.JSONDecodeError, OSError):
153 return None
154 if not isinstance(raw, dict):
155 return None
156 remaining = raw.get("remaining")
157 completed = raw.get("completed")
158 if not isinstance(remaining, list) or not isinstance(completed, list):
159 return None
160 return RebaseState(
161 original_branch=str(raw.get("original_branch", "")),
162 original_head=str(raw.get("original_head", "")),
163 onto=str(raw.get("onto", "")),
164 remaining=[str(x) for x in remaining if isinstance(x, str)],
165 completed=[str(x) for x in completed if isinstance(x, str)],
166 squash=bool(raw.get("squash", False)),
167 )
168
169
170 def save_rebase_state(root: pathlib.Path, state: RebaseState) -> None:
171 """Write rebase state to ``.muse/REBASE_STATE.json`` atomically.
172
173 Raises ``OSError`` if the target path is a symlink (would redirect the
174 write to an unintended location).
175 """
176 path = root / _REBASE_STATE_FILE
177 if path.exists() and path.is_symlink():
178 raise OSError(
179 f"Refusing to write rebase state — {path} is a symlink"
180 )
181 write_text_atomic(path, json.dumps(dict(state), indent=2))
182
183
184 def clear_rebase_state(root: pathlib.Path) -> None:
185 """Remove ``.muse/REBASE_STATE.json``.
186
187 No-op if the state file does not exist. Refuses to unlink a symlink
188 to prevent unintentional deletion of a file outside the repo.
189 """
190 path = root / _REBASE_STATE_FILE
191 if not path.exists():
192 return
193 if path.is_symlink():
194 logger.warning(
195 "⚠️ REBASE_STATE.json is a symlink — refusing to unlink"
196 )
197 return
198 path.unlink()
199 logger.debug("✅ Cleared REBASE_STATE.json")
200
201
202 def get_rebase_progress(root: pathlib.Path) -> RebaseProgress:
203 """Return a :class:`RebaseProgress` describing the current rebase state.
204
205 Always returns a valid ``RebaseProgress`` — ``active=False`` when no
206 rebase is in progress.
207
208 Used by ``muse rebase --status`` to give agents and humans a structured
209 view of an ongoing rebase without re-reading the raw state file.
210 """
211 state = load_rebase_state(root)
212 if state is None:
213 return RebaseProgress(
214 active=False,
215 original_branch="",
216 original_head="",
217 onto="",
218 total=0,
219 done=0,
220 remaining=0,
221 squash=False,
222 )
223 total = len(state["completed"]) + len(state["remaining"])
224 return RebaseProgress(
225 active=True,
226 original_branch=state["original_branch"],
227 original_head=state["original_head"],
228 onto=state["onto"],
229 total=total,
230 done=len(state["completed"]),
231 remaining=len(state["remaining"]),
232 squash=state["squash"],
233 )
234
235
236 # ---------------------------------------------------------------------------
237 # Commit collection
238 # ---------------------------------------------------------------------------
239
240
241 def collect_commits_to_replay(
242 root: pathlib.Path,
243 stop_at: str,
244 tip: str,
245 max_commits: int = 10_000,
246 ) -> list[CommitRecord]:
247 """Return commits from *tip* back to (but not including) *stop_at*.
248
249 The result is in chronological order (oldest first) so the replay loop
250 can iterate forward.
251
252 Only the first-parent chain is walked. Merge commits are replayed as
253 a single commit (their second-parent history is not re-played).
254
255 Args:
256 root: Repository root.
257 stop_at: Commit ID to stop at (exclusive — the merge base).
258 tip: Starting commit ID (the current branch HEAD).
259 max_commits: Safety cap on the number of commits returned. Prevents
260 unbounded traversal on very long histories.
261
262 Returns:
263 List of ``CommitRecord`` objects, oldest first (ready to replay).
264 """
265 commits: list[CommitRecord] = []
266 seen: set[str] = set()
267 current: str | None = tip
268
269 while current and current not in seen and len(commits) < max_commits:
270 seen.add(current)
271 if current == stop_at:
272 break
273 commit = read_commit(root, current)
274 if commit is None:
275 break
276 commits.append(commit)
277 current = commit.parent_commit_id
278
279 # Reverse so oldest is first.
280 commits.reverse()
281 return commits
282
283
284 # ---------------------------------------------------------------------------
285 # Single-commit replay
286 # ---------------------------------------------------------------------------
287
288
289 def replay_one(
290 root: pathlib.Path,
291 commit: CommitRecord,
292 parent_id: str,
293 plugin: MuseDomainPlugin,
294 domain: str,
295 repo_id: str,
296 branch: str,
297 ) -> CommitRecord | list[str]:
298 """Replay *commit* on top of *parent_id* using the domain plugin.
299
300 Performs a three-way merge where:
301
302 - ``base`` = commit's original parent snapshot (what existed before)
303 - ``ours`` = the current rebased tip snapshot (what we've built so far)
304 - ``theirs`` = commit's snapshot (what we want to apply)
305
306 When the merge is clean, writes the new commit and snapshot and returns
307 the new ``CommitRecord``. When conflicts exist, returns the list of
308 conflicting paths — the caller is responsible for writing
309 ``MERGE_STATE.json`` and stopping the rebase.
310
311 Args:
312 root: Repository root.
313 commit: The original commit being replayed.
314 parent_id: The new parent commit ID (last replayed commit or onto base).
315 plugin: The active domain plugin instance.
316 domain: Domain name string.
317 repo_id: Repository UUID.
318 branch: Current branch name.
319
320 Returns:
321 New ``CommitRecord`` on clean merge; ``list[str]`` of conflict paths
322 on conflict.
323
324 Raises:
325 TypeError: If *plugin* is not a ``MuseDomainPlugin``.
326 """
327 if not isinstance(plugin, MuseDomainPlugin):
328 raise TypeError(
329 f"replay_one: plugin {type(plugin).__name__!r} is not a MuseDomainPlugin"
330 )
331
332 # Resolve original parent snapshot (the "base" for the merge).
333 base_manifest: Manifest = {}
334 if commit.parent_commit_id:
335 parent_commit = read_commit(root, commit.parent_commit_id)
336 if parent_commit:
337 parent_snap = read_snapshot(root, parent_commit.snapshot_id)
338 if parent_snap:
339 base_manifest = parent_snap.manifest
340
341 # "Theirs" = the original commit's snapshot.
342 theirs_snap = read_snapshot(root, commit.snapshot_id)
343 if theirs_snap is None:
344 # A missing snapshot means we cannot reconstruct the commit's content.
345 # Falling back to {} would silently delete all files the commit added or
346 # modified — producing a wrong rebased history. Raise so the caller
347 # surfaces a clear error rather than silently corrupting the rebase.
348 raise ValueError(
349 f"rebase: snapshot {commit.snapshot_id[:8]} for commit "
350 f"{commit.commit_id[:8]} ({commit.message!r}) is missing or corrupt — "
351 "cannot replay this commit. Run `muse verify-pack` to audit the store."
352 )
353 theirs_manifest = theirs_snap.manifest
354
355 # "Ours" = the current rebased tip.
356 ours_manifest: Manifest = {}
357 if parent_id:
358 parent_rec = read_commit(root, parent_id)
359 if parent_rec:
360 ours_snap = read_snapshot(root, parent_rec.snapshot_id)
361 if ours_snap:
362 ours_manifest = ours_snap.manifest
363
364 base_snap_obj = SnapshotManifest(files=base_manifest, domain=domain, directories=directories_from_manifest(base_manifest))
365 ours_snap_obj = SnapshotManifest(files=ours_manifest, domain=domain, directories=directories_from_manifest(ours_manifest))
366 theirs_snap_obj = SnapshotManifest(files=theirs_manifest, domain=domain, directories=directories_from_manifest(theirs_manifest))
367
368 result: MergeResult = plugin.merge(
369 base_snap_obj, ours_snap_obj, theirs_snap_obj, repo_root=root
370 )
371
372 if not result.is_clean:
373 return result.conflicts
374
375 merged_manifest = result.merged["files"]
376
377 # Apply the merged state to the working tree.
378 apply_manifest(root, merged_manifest)
379
380 merged_dirs = directories_from_manifest(merged_manifest)
381 snapshot_id = compute_snapshot_id(merged_manifest, merged_dirs)
382 committed_at = datetime.datetime.now(datetime.timezone.utc)
383 new_commit_id = compute_commit_id(
384 parent_ids=[parent_id] if parent_id else [],
385 snapshot_id=snapshot_id,
386 message=commit.message,
387 committed_at_iso=committed_at.isoformat(),
388 )
389
390 write_snapshot(root, SnapshotRecord(snapshot_id=snapshot_id, manifest=merged_manifest, directories=merged_dirs))
391 new_commit = CommitRecord(
392 commit_id=new_commit_id,
393 repo_id=repo_id,
394 branch=branch,
395 snapshot_id=snapshot_id,
396 message=commit.message,
397 committed_at=committed_at,
398 parent_commit_id=parent_id if parent_id else None,
399 author=commit.author,
400 agent_id=commit.agent_id,
401 model_id=commit.model_id,
402 )
403 write_commit(root, new_commit)
404 return new_commit
405
406
407 def _write_branch_ref(root: pathlib.Path, branch: str, commit_id: str) -> None:
408 """Write *commit_id* to the branch ref file atomically and durably."""
409 write_branch_ref(root, branch, commit_id)
File History 1 commit
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 12 days ago