969 lines
41 KiB
Python
969 lines
41 KiB
Python
"""GA 世界线(Rewind / checkpoint)统一后端 —— **UI 无关,单文件,tui/gui 共用**。
|
||
|
||
三块合一(本文件不含任何 Textual/前端交互逻辑,只产数据与做存储):
|
||
1. **持久化**(`RewindStore`):content-addressable blob + checkpoint 树 + 全持久化。
|
||
- 统一 **checkpoint 树**;**content-addressable blob**(内容 sha256 去重);
|
||
- 只追 `file_write`/`file_patch`(路径由调用方在 `tool_before` 钩子喂入);
|
||
- **global baseline**(每文件「首次改前」)解决「回退到该文件被追踪之前」。
|
||
2. **视图模型 + 压缩**(`CheckpointTree`/`tree_from_store`/`CompressedTree`):把真实
|
||
store 投影成只读树并强制折叠线性段,供任何前端渲染。
|
||
3. **导航数学 + 恢复编排**(`next_same_depth`/`nearest_depth_node`/`restore_plan`):
|
||
纯计算;`restore_plan` 算出回退后的 history/文件/prefill 并落地(改文件+移 HEAD+
|
||
重写投影),前端只管把结果刷到自己的界面。
|
||
|
||
参考 Claude Code file-history 思路,按 GA 改造。纯 Python、可离线测试;线程安全交调用方
|
||
(agent 跑在子线程,约定仅 agent 空闲时 restore)。详见需求文档(ga_rewind_tui_requirements)。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import difflib
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import shutil
|
||
import time
|
||
from datetime import datetime
|
||
import zlib
|
||
from dataclasses import dataclass, field
|
||
from typing import Dict, List, Optional
|
||
|
||
from rich.cells import cell_len, set_cell_size
|
||
|
||
# 该值为 None 表示「该版本文件不存在」（对应 CC 的 backupFileName: null）。
|
||
_ABSENT = None
|
||
|
||
|
||
class RewindStore:
|
||
"""一个 session 的 checkpoint 树 + blob 库。
|
||
|
||
用法（典型时序）：
|
||
store = RewindStore(root, cwd)
|
||
# 改前（tool_before 钩子，file_write/file_patch 执行前）：
|
||
store.track_pre_edit(abs_path)
|
||
# ...工具实际写盘...
|
||
# 段末（用户提问边界 / turn 收尾）：
|
||
store.commit("修复 stdout", hist_len=42)
|
||
# 浏览 / 恢复：
|
||
store.apply_code(node_id) # 把工作区还原到该节点
|
||
store.rewind_head(node_id) # 只移 HEAD（下次 commit 即从此分叉）
|
||
"""
|
||
|
||
def __init__(self, root: str, cwd: str) -> None:
|
||
self.root = os.path.abspath(root)
|
||
self.cwd = os.path.abspath(cwd)
|
||
self.objects_dir = os.path.join(self.root, "objects")
|
||
self.tree_path = os.path.join(self.root, "tree.json")
|
||
|
||
# id -> {parent, children[], title, created, kind, hist_len, files{rel: hash|None}}
|
||
self.nodes: dict[str, dict] = {}
|
||
self.head: str | None = None
|
||
self.root_id: str | None = None
|
||
self._seq = 0
|
||
# 全局：每个被追踪文件的「首次改前」状态（hash 或 None=当时不存在）。
|
||
self.baseline: dict[str, str | None] = {}
|
||
self.tracked: set[str] = set()
|
||
|
||
# 段内暂存：本段触碰过的相对路径（commit 时落定）。
|
||
self._touched: set[str] = set()
|
||
# 最近一次 restore 的 redo 点（内存级软保险②），结构同 node.files。
|
||
self._redo: dict | None = None
|
||
|
||
self.load()
|
||
if self.nodes:
|
||
self._ensure_origin() # 迁移:旧树(无 origin)挂一个「会话起点」根
|
||
|
||
# --------------------------------------------------------------- origin
|
||
def _ensure_origin(self) -> None:
|
||
"""保证存在一个虚拟根「会话起点」(空状态):所有真实 checkpoint 都是它的
|
||
后代。幂等。第一次 commit / 迁移旧树 / resume 时调用。
|
||
|
||
- 它代表「第一轮对话之前」,使「回退到第一个会话之前」= 普通地 rewind 到它;
|
||
- files={} + hist_len=0 → apply_code(origin) 经 baseline 回退 = 还原到空起点;
|
||
- 它是唯一真根 → 不会出现 forest / head=None。"""
|
||
if any(nd.get("kind") == "origin" for nd in self.nodes.values()):
|
||
if self.head is None:
|
||
self.head = self.root_id
|
||
return
|
||
oid = "origin"
|
||
while oid in self.nodes:
|
||
oid += "_"
|
||
tops = self.top_level_nodes()
|
||
self.nodes[oid] = {"parent": None, "children": list(tops),
|
||
"title": "会话起点", "created": time.time(),
|
||
"kind": "origin", "hist_len": 0, "files": {}}
|
||
for t in tops:
|
||
self.nodes[t]["parent"] = oid
|
||
self.root_id = oid
|
||
if self.head is None:
|
||
self.head = oid
|
||
|
||
# ------------------------------------------------------------------ paths
|
||
def key(self, abs_path: str) -> str:
|
||
"""绝对路径在 cwd 下 → 存相对（参考 CC maybeShortenFilePath）；否则原样。"""
|
||
ap = os.path.abspath(abs_path)
|
||
try:
|
||
if os.path.commonpath([ap, self.cwd]) == self.cwd:
|
||
return os.path.relpath(ap, self.cwd).replace(os.sep, "/")
|
||
except ValueError:
|
||
pass # 跨盘符（Windows）等：commonpath 抛错 → 用绝对路径
|
||
return ap.replace(os.sep, "/")
|
||
|
||
def _abs(self, rel: str) -> str:
|
||
if os.path.isabs(rel) or (len(rel) > 1 and rel[1] == ":"):
|
||
return os.path.normpath(rel)
|
||
return os.path.normpath(os.path.join(self.cwd, rel))
|
||
|
||
# ------------------------------------------------------------------ blobs
|
||
def _put_blob(self, data: bytes) -> str:
|
||
h = hashlib.sha256(data).hexdigest()
|
||
p = os.path.join(self.objects_dir, h[:2], h)
|
||
if not os.path.exists(p):
|
||
os.makedirs(os.path.dirname(p), exist_ok=True)
|
||
tmp = p + ".tmp"
|
||
with open(tmp, "wb") as f:
|
||
f.write(zlib.compress(data, 6))
|
||
os.replace(tmp, p)
|
||
return h
|
||
|
||
def _get_blob(self, h: str) -> bytes:
|
||
with open(os.path.join(self.objects_dir, h[:2], h), "rb") as f:
|
||
return zlib.decompress(f.read())
|
||
|
||
def _snapshot(self, abs_path: str) -> str | None:
|
||
"""当前内容 → blob hash；文件不存在 → None。"""
|
||
try:
|
||
with open(abs_path, "rb") as f:
|
||
data = f.read()
|
||
except (FileNotFoundError, NotADirectoryError, IsADirectoryError, PermissionError):
|
||
return _ABSENT
|
||
return self._put_blob(data)
|
||
|
||
# --------------------------------------------------------------- tracking
|
||
def track_pre_edit(self, abs_path: str) -> None:
|
||
"""改前埋点：首次见到该文件时记下它的「改前」状态作 baseline。
|
||
|
||
幂等——同一文件同段内多次调用只在首次抓 baseline（file_patch 连改不会
|
||
覆盖原始版本，对应 CC trackEdit 的 v1 保护）。"""
|
||
rel = self.key(abs_path)
|
||
self._touched.add(rel)
|
||
if rel not in self.baseline:
|
||
self.baseline[rel] = self._snapshot(abs_path) # 改前内容，可能是 None
|
||
self.tracked.add(rel)
|
||
|
||
def commit(self, title: str, hist_len: int | None = None, kind: str = "edit",
|
||
history: list | None = None) -> str:
|
||
"""段末落节点：继承 HEAD.files，把本段触碰文件的「改后」内容写进新节点。
|
||
|
||
`history` 给定时(对话树化)：切 `history[parent.hist_len:当前]` 为本轮对话增量，
|
||
存成 content-addressable conv blob，节点记其 hash + `hist_len=len(history)`。
|
||
此刻索引可靠 → 增量精确，恢复时沿路径拼回即可，不再靠会漂的绝对下标。
|
||
|
||
无触碰文件时也会落节点（纯对话推进也是 checkpoint）。返回新节点 id。"""
|
||
self._ensure_origin() # 第一次提交前先确保有「会话起点」根
|
||
parent = self.head # 至少是 origin,永不为 None
|
||
files = dict(self.nodes[parent]["files"]) if parent is not None else {}
|
||
for rel in self._touched:
|
||
files[rel] = self._snapshot(self._abs(rel)) # 改后内容，删除则 None
|
||
|
||
conv = None
|
||
if history is not None:
|
||
parent_len = self.nodes[parent].get("hist_len") if parent is not None else 0
|
||
parent_len = parent_len or 0 # None(旧节点) / origin → 0
|
||
delta = list(history[parent_len:])
|
||
hist_len = len(history)
|
||
conv = self._put_blob(
|
||
json.dumps(delta, ensure_ascii=False, default=str).encode("utf-8"))
|
||
|
||
nid = self._new_id()
|
||
self.nodes[nid] = {
|
||
"parent": parent,
|
||
"children": [],
|
||
"title": title,
|
||
"created": time.time(),
|
||
"kind": kind,
|
||
"hist_len": hist_len,
|
||
"files": files,
|
||
"conv": conv,
|
||
}
|
||
if parent is not None:
|
||
self.nodes[parent]["children"].append(nid)
|
||
elif self.root_id is None or self.root_id not in self.nodes:
|
||
self.root_id = nid # 只认第一个根;rewind-到-根之前再提交会产生第二个顶层节点(forest)
|
||
self.head = nid
|
||
self._touched.clear()
|
||
self.save()
|
||
return nid
|
||
|
||
# ------------------------------------------------------------- navigation
|
||
def rewind_head(self, node_id) -> None:
|
||
"""移动 HEAD(不动文件)。从非叶 HEAD 下次 commit 即 append 子节点 = fork。
|
||
node_id=None 表示「回到根之前」(空起点),下次 commit 产生新的顶层节点。"""
|
||
if node_id is not None and node_id not in self.nodes:
|
||
raise KeyError(node_id)
|
||
self.head = node_id
|
||
self.save()
|
||
|
||
def top_level_nodes(self) -> list[str]:
|
||
"""所有顶层节点(parent 为 None 或父已不存在)。通常一个 root,rewind-到-根
|
||
之前再提交会出现多个(forest)。"""
|
||
return [nid for nid, nd in self.nodes.items()
|
||
if nd.get("parent") is None or nd.get("parent") not in self.nodes]
|
||
|
||
def _target_state(self, node_id: str, rel: str):
|
||
"""文件 rel 在 node_id 时应有的状态：节点 files 里有就用它，否则回退到
|
||
baseline（CC 的 first-version）。返回 (known: bool, hash_or_None)。"""
|
||
files = self.nodes[node_id]["files"]
|
||
if rel in files:
|
||
return True, files[rel]
|
||
if rel in self.baseline:
|
||
return True, self.baseline[rel]
|
||
return False, None
|
||
|
||
def apply_code(self, node_id) -> list[tuple[str, str]]:
|
||
"""把工作区文件还原到 node_id 的状态。`node_id=None` → 还原到 baseline
|
||
(任何 checkpoint 之前的原始状态,用于「回退到根之前」)。
|
||
|
||
- **只动本 store 追踪过的路径**（self.tracked），绝不碰未记录文件（软保险①）。
|
||
- restore 前先把这些路径的当前内容存进 redo 点（软保险②）——误覆盖也能找回。
|
||
返回 [(rel, action)]，action ∈ {restored, deleted}。"""
|
||
if node_id is not None and node_id not in self.nodes:
|
||
raise KeyError(node_id)
|
||
|
||
# 软保险②：redo 点（记录当前 = 即将被覆盖的状态）。
|
||
redo_files: dict[str, str | None] = {}
|
||
for rel in self.tracked:
|
||
redo_files[rel] = self._snapshot(self._abs(rel))
|
||
self._redo = {"from": self.head, "files": redo_files}
|
||
|
||
changed: list[tuple[str, str]] = []
|
||
for rel in self.tracked:
|
||
if node_id is None:
|
||
known = rel in self.baseline
|
||
target = self.baseline.get(rel)
|
||
else:
|
||
known, target = self._target_state(node_id, rel)
|
||
if not known:
|
||
continue
|
||
ap = self._abs(rel)
|
||
cur = self._snapshot(ap)
|
||
if cur == target:
|
||
continue
|
||
if target is _ABSENT:
|
||
try:
|
||
os.remove(ap)
|
||
changed.append((rel, "deleted"))
|
||
except FileNotFoundError:
|
||
pass
|
||
else:
|
||
os.makedirs(os.path.dirname(ap) or ".", exist_ok=True)
|
||
with open(ap, "wb") as f:
|
||
f.write(self._get_blob(target))
|
||
changed.append((rel, "restored"))
|
||
return changed
|
||
|
||
def diff(self, node_id: str) -> list[dict]:
|
||
"""选中节点 vs 当前工作区，逐文件变更摘要（恢复到该节点会发生什么）。
|
||
|
||
返回 [{rel, action, insertions, deletions}]，仅列出会变的文件。
|
||
action: restore（覆盖）/ delete（删除）/ create（重建被删文件）。"""
|
||
if node_id not in self.nodes:
|
||
raise KeyError(node_id)
|
||
out: list[dict] = []
|
||
for rel in sorted(self.tracked):
|
||
known, target = self._target_state(node_id, rel)
|
||
if not known:
|
||
continue
|
||
cur = self._snapshot(self._abs(rel))
|
||
if cur == target:
|
||
continue
|
||
old = self._text(cur) # 当前工作区内容
|
||
new = self._text(target) # 该节点内容（恢复目标）
|
||
ins = dele = 0
|
||
for line in difflib.ndiff(old.splitlines(), new.splitlines()):
|
||
if line.startswith("+ "):
|
||
ins += 1
|
||
elif line.startswith("- "):
|
||
dele += 1
|
||
if target is _ABSENT:
|
||
action = "delete"
|
||
elif cur is _ABSENT:
|
||
action = "create"
|
||
else:
|
||
action = "restore"
|
||
out.append({"rel": rel, "action": action, "insertions": ins, "deletions": dele})
|
||
return out
|
||
|
||
def node_diff(self, node_id) -> list[dict]:
|
||
"""选中节点 vs **父节点**的逐文件内容 diff(= 这次 checkpoint 改了什么)。
|
||
返回 [{rel, old, new}](文本),仅含两者不同的文件;父缺失/文件不存在 → 空文本。
|
||
供 UI 渲染逐行 diff(区别于 `diff()` 的"节点 vs 当前工作区"计数摘要)。"""
|
||
if node_id not in self.nodes:
|
||
return []
|
||
nd = self.nodes[node_id]
|
||
par = nd.get("parent")
|
||
pf = self.nodes[par]["files"] if par in self.nodes else {}
|
||
nf = nd.get("files") or {}
|
||
out: list[dict] = []
|
||
for rel in sorted(set(nf) | set(pf)):
|
||
nh = nf.get(rel)
|
||
# 「上一个状态」:父节点记录过就用它;父节点没有(本节点首次触碰该文件)→
|
||
# 回退到 baseline(首次改前内容,CC first-version)。否则被 file_patch 修改的
|
||
# 已有文件会因父节点无记录而被误显示成整文件新增,而非真正的逐行 diff。
|
||
ph = pf[rel] if rel in pf else self.baseline.get(rel)
|
||
if nh == ph:
|
||
continue
|
||
out.append({"rel": rel, "old": self._text(ph), "new": self._text(nh)})
|
||
return out
|
||
|
||
def _text(self, h: str | None) -> str:
|
||
if h is _ABSENT:
|
||
return ""
|
||
try:
|
||
return self._get_blob(h).decode("utf-8", "replace")
|
||
except Exception:
|
||
return ""
|
||
|
||
# --------------------------------------------------------- delete + gc
|
||
def delete_subtree(self, node_id: str) -> list[str]:
|
||
"""删除节点及其整棵子树（不能删 HEAD 祖先路径上的节点 / root 由调用方把关）。
|
||
返回被删的节点 id 列表。删后做一次 blob GC。"""
|
||
if node_id not in self.nodes:
|
||
raise KeyError(node_id)
|
||
victims: list[str] = []
|
||
|
||
def collect(nid: str) -> None:
|
||
victims.append(nid)
|
||
for c in list(self.nodes[nid]["children"]):
|
||
collect(c)
|
||
|
||
collect(node_id)
|
||
parent = self.nodes[node_id]["parent"]
|
||
if parent is not None:
|
||
self.nodes[parent]["children"] = [
|
||
c for c in self.nodes[parent]["children"] if c != node_id
|
||
]
|
||
for nid in victims:
|
||
self.nodes.pop(nid, None)
|
||
if node_id == self.root_id:
|
||
self.root_id = None
|
||
if self.head in victims:
|
||
self.head = parent # HEAD 被删 → 退到父
|
||
self._gc()
|
||
self.save()
|
||
return victims
|
||
|
||
def _gc(self) -> int:
|
||
"""删除无人引用的 blob。返回删除数。"""
|
||
referenced: set[str] = set()
|
||
for node in self.nodes.values():
|
||
for h in node["files"].values():
|
||
if h is not _ABSENT:
|
||
referenced.add(h)
|
||
ch = node.get("conv") # 对话增量 blob 也是引用
|
||
if ch:
|
||
referenced.add(ch)
|
||
for h in self.baseline.values():
|
||
if h is not _ABSENT:
|
||
referenced.add(h)
|
||
if self._redo:
|
||
for h in self._redo["files"].values():
|
||
if h is not _ABSENT:
|
||
referenced.add(h)
|
||
removed = 0
|
||
if not os.path.isdir(self.objects_dir):
|
||
return 0
|
||
for sub in os.listdir(self.objects_dir):
|
||
d = os.path.join(self.objects_dir, sub)
|
||
if not os.path.isdir(d):
|
||
continue
|
||
for h in os.listdir(d):
|
||
if h not in referenced:
|
||
try:
|
||
os.remove(os.path.join(d, h))
|
||
removed += 1
|
||
except OSError:
|
||
pass
|
||
return removed
|
||
|
||
# ------------------------------------------------------------ persistence
|
||
def save(self) -> None:
|
||
os.makedirs(self.root, exist_ok=True)
|
||
payload = {
|
||
"nodes": self.nodes,
|
||
"head": self.head,
|
||
"root": self.root_id,
|
||
"seq": self._seq,
|
||
"baseline": self.baseline,
|
||
"tracked": sorted(self.tracked),
|
||
}
|
||
tmp = self.tree_path + ".tmp"
|
||
with open(tmp, "w", encoding="utf-8") as f:
|
||
json.dump(payload, f, ensure_ascii=False, default=str)
|
||
# 崩溃兜底：替换前把上一份完好的 tree.json 留作 .bak（load 损坏时回退它）。
|
||
if os.path.exists(self.tree_path):
|
||
try:
|
||
shutil.copyfile(self.tree_path, self.tree_path + ".bak")
|
||
except OSError:
|
||
pass
|
||
os.replace(tmp, self.tree_path)
|
||
|
||
def load(self) -> None:
|
||
"""加载树。主文件损坏/缺失 → 回退 .bak；都不行 → 留空 store（降级，不抛）。"""
|
||
for path in (self.tree_path, self.tree_path + ".bak"):
|
||
if not os.path.exists(path):
|
||
continue
|
||
try:
|
||
with open(path, encoding="utf-8") as f:
|
||
d = json.load(f)
|
||
except Exception:
|
||
continue # 损坏 → 尝试下一个候选
|
||
self.nodes = d.get("nodes", {})
|
||
self.head = d.get("head")
|
||
self.root_id = d.get("root")
|
||
self._seq = d.get("seq", 0)
|
||
self.baseline = d.get("baseline", {})
|
||
self.tracked = set(d.get("tracked", []))
|
||
return
|
||
|
||
def _new_id(self) -> str:
|
||
self._seq += 1
|
||
return f"v{self._seq}"
|
||
|
||
# ---------------------------------------------------------- session key
|
||
@staticmethod
|
||
def key_for_log(log_path: str) -> str:
|
||
"""会话 key = log_path 的 basename 去扩展名,如 `model_responses_123456`。
|
||
这是 GA 唯一稳定、/continue 据以识别会话的身份(task_dir 是 PID 级、重启即变,
|
||
不能用)。"""
|
||
return os.path.splitext(os.path.basename(log_path))[0]
|
||
|
||
@classmethod
|
||
def for_log(cls, temp_dir: str, log_path: str, cwd: str) -> "RewindStore":
|
||
"""按 GA 的 log_path 落到 `temp/.ga_rewind/<key>/`。"""
|
||
root = os.path.join(temp_dir, ".ga_rewind", cls.key_for_log(log_path))
|
||
return cls(root, cwd)
|
||
|
||
def resume_from(self, old_root: str) -> bool:
|
||
"""/continue 续接:把旧会话的 rewind 目录接到本 store(贴近 CC 的
|
||
copyFileHistoryForResume)。
|
||
|
||
GA 每个新 agent 生成新 log_path → 新 key,旧会话的树/blob 默认看不到;
|
||
这里在新 key 下重建:① 接管旧 tree.json(仅当本 store 尚为空);
|
||
② blob 内容寻址,直接 hardlink(同名即同内容,跳过;失败回退 copy)。
|
||
返回是否接管成功。"""
|
||
if not os.path.isdir(old_root) or os.path.abspath(old_root) == self.root:
|
||
return False
|
||
old_tree = os.path.join(old_root, "tree.json")
|
||
if not self.nodes and os.path.exists(old_tree):
|
||
try:
|
||
with open(old_tree, encoding="utf-8") as f:
|
||
d = json.load(f)
|
||
self.nodes = d.get("nodes", {})
|
||
self.head = d.get("head")
|
||
self.root_id = d.get("root")
|
||
self._seq = d.get("seq", 0)
|
||
self.baseline = d.get("baseline", {})
|
||
self.tracked = set(d.get("tracked", []))
|
||
except Exception:
|
||
return False
|
||
old_objects = os.path.join(old_root, "objects")
|
||
if os.path.isdir(old_objects):
|
||
for sub in os.listdir(old_objects):
|
||
sd = os.path.join(old_objects, sub)
|
||
if not os.path.isdir(sd):
|
||
continue
|
||
dd = os.path.join(self.objects_dir, sub)
|
||
os.makedirs(dd, exist_ok=True)
|
||
for h in os.listdir(sd):
|
||
dst = os.path.join(dd, h)
|
||
if os.path.exists(dst):
|
||
continue # 内容寻址:同名即同内容
|
||
src = os.path.join(sd, h)
|
||
try:
|
||
os.link(src, dst) # hardlink 不占额外空间
|
||
except OSError:
|
||
try:
|
||
shutil.copyfile(src, dst) # 跨盘/不支持 → 回退 copy
|
||
except OSError:
|
||
pass
|
||
if self.nodes:
|
||
self._ensure_origin() # 旧树若无 origin,迁移补上
|
||
self.save()
|
||
return True
|
||
|
||
# ------------------------------------------------------------------ views
|
||
def linear_path(self) -> list[str]:
|
||
"""root → HEAD 的线性路径（节点 id），供 /rewind 时间线展示。"""
|
||
return self.path_to(self.head)
|
||
|
||
def path_to(self, node_id) -> list[str]:
|
||
"""root → node_id 的节点 id 列表（含两端）。node_id 缺失 → []。"""
|
||
if node_id is None or node_id not in self.nodes:
|
||
return []
|
||
chain: list[str] = []
|
||
cur: str | None = node_id
|
||
while cur is not None and cur in self.nodes:
|
||
chain.append(cur)
|
||
cur = self.nodes[cur].get("parent")
|
||
chain.reverse()
|
||
return chain
|
||
|
||
# ----------------------------------------------------------- conversation
|
||
def _node_conv(self, node_id) -> list:
|
||
"""单节点的对话增量（list）。无 conv / 读失败 → []（容错降级）。"""
|
||
h = self.nodes.get(node_id, {}).get("conv")
|
||
if not h:
|
||
return []
|
||
try:
|
||
data = json.loads(self._get_blob(h).decode("utf-8"))
|
||
return data if isinstance(data, list) else []
|
||
except Exception:
|
||
return []
|
||
|
||
def rebuild_history(self, node_id) -> list:
|
||
"""沿 root→node_id 拼接各节点对话增量，整体重建 `backend.history`。
|
||
|
||
这是「树=真相源」的恢复入口：选任意节点（含旁支）都能精确还原其对话，
|
||
不依赖内存里那条线性 history。origin / 无 conv 的节点贡献空增量。"""
|
||
hist: list = []
|
||
for nid in self.path_to(node_id):
|
||
hist.extend(self._node_conv(nid))
|
||
return hist
|
||
|
||
# ----------------------------------------------------- 对账(防外部改写灾难)
|
||
@staticmethod
|
||
def _msg_user_text(msg) -> str:
|
||
"""真实用户提问文本;非 user / 纯 tool_result 回填 → ""(不算提问边界)。"""
|
||
if not isinstance(msg, dict) or msg.get("role") != "user":
|
||
return ""
|
||
c = msg.get("content")
|
||
if isinstance(c, str):
|
||
return c
|
||
if isinstance(c, list):
|
||
if any(isinstance(b, dict) and b.get("type") == "tool_result" for b in c):
|
||
return ""
|
||
return " ".join(b.get("text", "") for b in c
|
||
if isinstance(b, dict) and b.get("type") == "text")
|
||
return ""
|
||
|
||
@classmethod
|
||
def _turn_sig(cls, msgs) -> list:
|
||
"""一段消息的「对话身份」= 其中真实用户提问的文本序列(跳过 tool_result 轮,
|
||
故不含注入的易变内容)。用于 reconcile 判两段是否同一对话。"""
|
||
return [t for t in (cls._msg_user_text(m).strip() for m in msgs) if t]
|
||
|
||
def reconcile(self, history: list) -> int:
|
||
"""把 live history 里树尚未记录的尾部轮次吸收成 conv-only 节点,使树追平日志。
|
||
供「带 worldline 的 UI」在打开世界线 / 接管日志前调用 —— 别的(不更新树的)
|
||
UI 往同一日志追加后,若不对账,树会滞后;一旦在滞后状态 rewind,
|
||
`rewrite_projection` 会拿陈旧树回写、**抹掉那些外部轮次**(灾难性丢失)。
|
||
对账后树恒 ⊇ 日志,故 rewind 物理上不可能 clobber 未见过的轮次。
|
||
|
||
**安全闸**:仅当树 HEAD 路径与 history 前缀是「同一段对话」才在其上 append;
|
||
否则判为分歧(用户问了不同的问题),HEAD 退回 origin 另起一条顶层 trunk 容纳
|
||
live history,旧树保留为旁支、绝不拿它回写日志。空树 = n=0 特例(吸收全部,
|
||
即旧 _bootstrap_conv_tree)。幂等:无尾巴 → no-op,可任意重复调用。返回吸收条数。
|
||
|
||
判「同一段对话」用**用户提问序列**(`_turn_sig`)而非逐条全等:消息里会被注入
|
||
易变内容(如 `[WORKING MEMORY]`,只在 tool_result 轮),同一段对话在不同时刻
|
||
捕获并不逐字节相等;逐条全等会把它误判成分歧、复制出一整条重复 trunk。提问
|
||
文本在 tool_result 轮之外、稳定不变,是可靠的对话身份。"""
|
||
history = list(history or [])
|
||
self._ensure_origin()
|
||
head = self.head if self.head in self.nodes else self.root_id
|
||
known = self.rebuild_history(head)
|
||
n = len(known)
|
||
if self._turn_sig(history[:n]) != self._turn_sig(known):
|
||
# 分歧(提问序列不同):不信任陈旧树。退回空起点,后续 append 只反映 live history。
|
||
self.rewind_head(self.root_id)
|
||
n = len(self.rebuild_history(self.root_id)) # origin 无 conv → 0
|
||
if len(history) <= n:
|
||
return 0
|
||
# 按用户提问边界切尾巴 [n, len);逐段 commit(history=history[:段末]),
|
||
# commit 内部 delta = history[parent.hist_len:段末]。清 _touched 保证纯 conv-only。
|
||
starts = [i for i in range(n, len(history))
|
||
if self._msg_user_text(history[i]).strip()]
|
||
if not starts or starts[0] > n:
|
||
starts = [n] + starts
|
||
self._touched.clear()
|
||
absorbed = 0
|
||
for p, s in enumerate(starts):
|
||
end = starts[p + 1] if p + 1 < len(starts) else len(history)
|
||
title = next((self._msg_user_text(history[i]).strip()
|
||
for i in range(s, end)
|
||
if self._msg_user_text(history[i]).strip()), "")
|
||
title = (title or "（外部续接）").replace("\n", " ").strip()[:80]
|
||
self.commit(title or "（外部续接）", kind="edit", history=list(history[:end]))
|
||
absorbed += end - s
|
||
return absorbed
|
||
|
||
def first_user_message(self, node_id) -> dict | None:
|
||
"""node_id 对话增量里的首条 `role==user` 消息（原始 dict）；无则 None。
|
||
|
||
供 rewind 后精确 prefill（取「该提问本身」，不靠会漂的绝对下标）。文本提取
|
||
交调用方（消息内容格式是 GA backend 私有的）。"""
|
||
for msg in self._node_conv(node_id):
|
||
if isinstance(msg, dict) and msg.get("role") == "user":
|
||
return msg
|
||
return None
|
||
|
||
|
||
# ============================================================================
|
||
# 世界线视图模型 + 压缩 + 导航(从 rewind_tree_view.py 抽出的 UI 无关后端)
|
||
# ============================================================================
|
||
# 节点**语义**(字形 + 文字标签);**呈现色由前端决定**(颜色不属于后端)。
|
||
KIND_GLYPH = {"current": "◉", "rewind": "↩", "origin": "●"} # ◉=当前所在;●=普通/根
|
||
KIND_LABEL = {"current": "当前位置", "rewind": "回退来源", "origin": "会话起点"}
|
||
|
||
|
||
def kind_glyph(kind: str) -> str:
|
||
return KIND_GLYPH.get(kind, "●")
|
||
|
||
|
||
def kind_label(kind: str) -> str:
|
||
return KIND_LABEL.get(kind, "")
|
||
|
||
|
||
def rel_time(ago_s) -> str:
|
||
if ago_s is None:
|
||
return ""
|
||
if ago_s < 60:
|
||
return "刚刚"
|
||
m = ago_s // 60
|
||
if m < 60:
|
||
return f"{m} 分钟前"
|
||
h = m // 60
|
||
if h < 24:
|
||
return f"{h} 小时前"
|
||
return f"{h // 24} 天前"
|
||
|
||
|
||
def files_summary(files: List[str]) -> str:
|
||
if not files:
|
||
return "无变更"
|
||
if len(files) <= 3:
|
||
return "、".join(files)
|
||
return "、".join(files[:3]) + f" (+{len(files) - 3})"
|
||
|
||
|
||
_TITLE_MAX_CELLS = 40
|
||
|
||
|
||
def ellipsize(s: str, max_cells: int = _TITLE_MAX_CELLS) -> str:
|
||
"""按显示宽度(中文算 2 格)截断,超长加 `…`。用于左树/右上的长提问标题。"""
|
||
s = (s or "").replace("\n", " ").strip()
|
||
if cell_len(s) <= max_cells:
|
||
return s
|
||
return set_cell_size(s, max(1, max_cells - 1)).rstrip() + "…"
|
||
|
||
|
||
# --------------------------------------------------------------- 数据模型
|
||
@dataclass
|
||
class CheckpointNode:
|
||
id: str
|
||
title: str
|
||
parent_id: Optional[str] = None
|
||
children: List[str] = field(default_factory=list)
|
||
kind: str = "edit"
|
||
files: List[str] = field(default_factory=list)
|
||
ago: Optional[int] = 0
|
||
|
||
|
||
class CheckpointTree:
|
||
"""只读视图树(供压缩/渲染);从 RewindStore 构造。"""
|
||
|
||
def __init__(self) -> None:
|
||
self.nodes: Dict[str, CheckpointNode] = {}
|
||
self.root_id: Optional[str] = None
|
||
|
||
|
||
def _changed_files(store, node_id) -> List[str]:
|
||
import os
|
||
nd = store.nodes[node_id]
|
||
par = nd.get("parent")
|
||
pf = store.nodes[par]["files"] if par in store.nodes else {}
|
||
out: List[str] = []
|
||
for k, v in nd["files"].items():
|
||
if pf.get(k) != v:
|
||
b = os.path.basename(k)
|
||
if b not in out:
|
||
out.append(b)
|
||
return out
|
||
|
||
|
||
def tree_from_store(store, now: float) -> CheckpointTree:
|
||
"""把 RewindStore 的真实树投影成只读 CheckpointTree;HEAD 标 current。"""
|
||
t = CheckpointTree()
|
||
for nid, nd in store.nodes.items():
|
||
t.nodes[nid] = CheckpointNode(
|
||
id=nid,
|
||
title=nd.get("title") or "（空）",
|
||
parent_id=nd.get("parent"),
|
||
children=[c for c in nd.get("children", []) if c in store.nodes],
|
||
kind="current" if nid == store.head else nd.get("kind", "edit"),
|
||
files=_changed_files(store, nid),
|
||
ago=int(max(0, now - nd.get("created", now))),
|
||
)
|
||
t.root_id = store.root_id
|
||
return t
|
||
|
||
|
||
# --------------------------------------------- 压缩(强制折叠线性段,§4.2)
|
||
@dataclass
|
||
class DisplayNode:
|
||
key: int
|
||
kind: str # 'real' | 'fold'
|
||
depth: int = 0
|
||
parent_key: Optional[int] = None
|
||
children: List[int] = field(default_factory=list)
|
||
node_id: Optional[str] = None
|
||
seg: List[str] = field(default_factory=list)
|
||
end_id: Optional[str] = None
|
||
|
||
|
||
class CompressedTree:
|
||
def __init__(self, tree: CheckpointTree) -> None:
|
||
self.tree = tree
|
||
self.disp: Dict[int, DisplayNode] = {}
|
||
self.roots: List[int] = [] # 顶层显示节点 keys(通常一个 = origin)
|
||
self._k = 0
|
||
# 顶层节点(origin「会话起点」)常驻独立显示作锚点;它下面的真实对话链
|
||
# 照常折叠(§4.2)。origin 是「第一轮对话之前」的空状态,不是某轮对话,
|
||
# 所以独立显示语义清晰、不会像"第一轮对话被强行拎出"那样让人困惑。
|
||
for top in self._top_level():
|
||
self._build_root_node(top)
|
||
|
||
def _top_level(self) -> List[str]:
|
||
return [nid for nid, n in self.tree.nodes.items()
|
||
if n.parent_id is None or n.parent_id not in self.tree.nodes]
|
||
|
||
def _new_key(self) -> int:
|
||
k = self._k
|
||
self._k += 1
|
||
return k
|
||
|
||
def _build_root_node(self, root_id: str) -> int:
|
||
key = self._new_key()
|
||
self.disp[key] = DisplayNode(key=key, kind="real", node_id=root_id, depth=0)
|
||
self.roots.append(key)
|
||
for child_id in self.tree.nodes[root_id].children:
|
||
self._build_chain(child_id, 1, key)
|
||
return key
|
||
|
||
def _build_chain(self, start_id: str, depth: int, parent: Optional[int]) -> int:
|
||
seg = [start_id]
|
||
cur = start_id
|
||
while len(self.tree.nodes[cur].children) == 1:
|
||
cur = self.tree.nodes[cur].children[0]
|
||
seg.append(cur)
|
||
end = seg[-1]
|
||
key = self._new_key()
|
||
if len(seg) == 1:
|
||
self.disp[key] = DisplayNode(key=key, kind="real", node_id=end,
|
||
depth=depth, parent_key=parent)
|
||
else:
|
||
self.disp[key] = DisplayNode(key=key, kind="fold", depth=depth,
|
||
parent_key=parent, seg=list(seg), end_id=end)
|
||
if parent is None:
|
||
self.roots.append(key)
|
||
else:
|
||
self.disp[parent].children.append(key)
|
||
for child_id in self.tree.nodes[end].children:
|
||
self._build_chain(child_id, depth + 1, key)
|
||
return key
|
||
|
||
def flatten(self) -> List[int]:
|
||
order: List[int] = []
|
||
|
||
def dfs(k: int) -> None:
|
||
order.append(k)
|
||
for c in self.disp[k].children:
|
||
dfs(c)
|
||
|
||
for r in self.roots:
|
||
dfs(r)
|
||
return order
|
||
|
||
def end_node(self, key: int) -> CheckpointNode:
|
||
d = self.disp[key]
|
||
return self.tree.nodes[d.node_id if d.kind == "real" else d.end_id]
|
||
|
||
def label(self, key: int) -> str:
|
||
# 真实/折叠都只放末端标题(截断);折叠的 [N nodes] 标记放在 glyph 位(标题前)。
|
||
d = self.disp[key]
|
||
nid = d.node_id if d.kind == "real" else d.end_id
|
||
return ellipsize(self.tree.nodes[nid].title)
|
||
|
||
def glyph(self, key: int) -> str:
|
||
# 折叠节点开头用 [N nodes] 替代原来的 … 省略号(→ "[3 nodes] 标题");
|
||
# 真实节点仍用语义字形(● / ↩ / ◉)。
|
||
d = self.disp[key]
|
||
if d.kind == "fold":
|
||
return f"[{len(d.seg)} nodes]"
|
||
return kind_glyph(self.tree.nodes[d.node_id].kind)
|
||
|
||
|
||
# ----------------------------------------------------- 导航(纯计算)
|
||
def _order_depths(ct: CompressedTree) -> List[tuple[int, int]]:
|
||
return [(k, ct.disp[k].depth) for k in ct.flatten()]
|
||
|
||
|
||
def next_same_depth(order, sel, direction):
|
||
idx = next(i for i, (k, _) in enumerate(order) if k == sel)
|
||
depth = order[idx][1]
|
||
same = [i for i, (_, d) in enumerate(order) if d == depth]
|
||
pos = same.index(idx)
|
||
return order[same[(pos + direction) % len(same)]][0]
|
||
|
||
|
||
def nearest_depth_node(order, sel, delta):
|
||
idx = next(i for i, (k, _) in enumerate(order) if k == sel)
|
||
target = order[idx][1] + delta
|
||
if target < 0:
|
||
return sel
|
||
cand = [i for i, (_, d) in enumerate(order) if d == target]
|
||
if not cand:
|
||
return sel
|
||
if delta > 0:
|
||
best = min(cand, key=lambda i: (abs(i - idx), 0 if i >= idx else 1, i))
|
||
else:
|
||
best = min(cand, key=lambda i: (abs(i - idx), 0 if i <= idx else 1, -i))
|
||
return order[best][0]
|
||
|
||
|
||
def parent_sibling_first_child(ct, sel, direction):
|
||
d = ct.disp[sel]
|
||
if d.parent_key is None:
|
||
return sel
|
||
parent = ct.disp[d.parent_key]
|
||
if parent.parent_key is None:
|
||
return sel
|
||
siblings = ct.disp[parent.parent_key].children
|
||
if parent.key not in siblings or len(siblings) <= 1:
|
||
return sel
|
||
start = siblings.index(parent.key)
|
||
step = 1 if direction >= 0 else -1
|
||
for off in range(1, len(siblings)):
|
||
sib = ct.disp[siblings[(start + step * off) % len(siblings)]]
|
||
if sib.children:
|
||
return sib.children[0]
|
||
return sel
|
||
|
||
|
||
# ============================================================================
|
||
# 恢复编排(UI 无关):算出回退后的对话/文件/prefill 并落地,前端只刷新自己的显示
|
||
# ============================================================================
|
||
def user_msg_text(entry) -> str:
|
||
"""从一条 user 消息里取出用户可见文本(content 为 str 或 text block 列表)。"""
|
||
c = entry.get("content") if isinstance(entry, dict) else None
|
||
if isinstance(c, str):
|
||
return c
|
||
if isinstance(c, list):
|
||
parts = [b.get("text", "") for b in c
|
||
if isinstance(b, dict) and b.get("type") == "text"]
|
||
return "\n".join(p for p in parts if p)
|
||
return ""
|
||
|
||
|
||
def render_native_history(history, ts=None):
|
||
"""Inverse of `continue_cmd._parse_native_history`: render a `backend.history`
|
||
(alternating user/assistant) back into the native `=== Prompt/Response ===`
|
||
log text. Lives here (not continue_cmd) because the rewind tree is its only
|
||
caller — it rewrites `model_responses_<key>.txt` as a linear projection of
|
||
the current HEAD path so /continue + other UIs get a clean single-branch view
|
||
(the TUI itself reads only the tree). Round-trip holds:
|
||
`_parse_native_history(_pairs(render_native_history(h))) == h`.
|
||
|
||
Byte-for-byte matches llmcore's `_write_llm_log` for the tool backend:
|
||
Prompt = `json.dumps(merged, ensure_ascii=False, indent=2)`; Response =
|
||
`str(content_blocks)` which for a list equals `repr(...)` (survives
|
||
True/False/None that JSON would emit as true/false/null and break
|
||
`ast.literal_eval`). Only complete (user, assistant) pairs are emitted."""
|
||
ts = ts or datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||
out, i, n = [], 0, len(history)
|
||
while i < n:
|
||
u = history[i]
|
||
if isinstance(u, dict) and u.get('role') == 'user' and i + 1 < n:
|
||
a = history[i + 1]
|
||
if isinstance(a, dict) and a.get('role') == 'assistant':
|
||
out.append(f"=== Prompt === {ts}\n{json.dumps(u, ensure_ascii=False, indent=2)}\n\n")
|
||
out.append(f"=== Response === {ts}\n{a.get('content')!r}\n\n")
|
||
i += 2
|
||
continue
|
||
i += 1
|
||
return "".join(out)
|
||
|
||
|
||
def rewrite_projection(store, target, log_path: str) -> bool:
|
||
"""把 `model_responses_<key>.txt` 重写成 root→target 的线性投影(native
|
||
`=== Prompt/Response ===` 格式,`continue_cmd._parse_native_history` 可解析)。
|
||
|
||
给其他 UI + /continue 的兼容视角;消除 rewind/fork 后多分支混排。故障静默。"""
|
||
if store is None or not log_path:
|
||
return False
|
||
try:
|
||
text = render_native_history(store.rebuild_history(target))
|
||
# 回退到「会话起点」(origin)时 history 为空 → 日志被写空(0 字节)。这是【硬】回退:
|
||
# 日志恒等于树 HEAD,无错位,reconcile 不会复制分支。空日志会被普通 UI 的
|
||
# list_sessions(sz<32) 无害跳过;带 worldline 的 UI 靠 list_sessions(rewind_root=...)
|
||
# 的树感知发现 + continue_inplace(allow_empty=True) 仍能找回并恢复(空对话 + 重连树)。
|
||
tmp = log_path + ".tmp"
|
||
with open(tmp, "w", encoding="utf-8", errors="replace") as f:
|
||
f.write(text)
|
||
os.replace(tmp, log_path)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def restore_plan(store, node_id, mode: str = "both", to: str = "before",
|
||
log_path: str = "") -> Optional[dict]:
|
||
"""UI 无关的回退编排。`to`:
|
||
- **before**(选中段内某节点):回到该提问**之前**(连该节点一起清除),prefill 其提问;
|
||
target = parent(node) / origin。
|
||
- **at**(末尾占位项):在该节点处**继续**(HEAD→node,不清除、无 prefill)。
|
||
`mode`: both/conv/code。会就地:重建对话(返回 history,调用方自行赋给 backend)、
|
||
`apply_code`(还原文件,前自动留 redo 点)、移 HEAD、重写投影日志。
|
||
|
||
返回 None(无效) 或 dict:
|
||
{history: list|None(仅 conv 变更时), changed: [(rel,action)], prefill: str,
|
||
at_origin: bool, target: str, title: str, to: str}
|
||
—— 前端据此:赋 backend.history、重建界面消息、prefill 输入框、刷新。"""
|
||
if store is None or not node_id or node_id not in store.nodes:
|
||
return None
|
||
nd = store.nodes[node_id]
|
||
prefill = ""
|
||
if to == "at":
|
||
target = node_id
|
||
at_origin = nd.get("kind") == "origin"
|
||
else:
|
||
parent = nd.get("parent")
|
||
parent = parent if parent in store.nodes else None
|
||
target = parent if parent is not None else node_id
|
||
at_origin = parent is None
|
||
if not at_origin:
|
||
um = store.first_user_message(node_id)
|
||
prefill = (user_msg_text(um) if um else "") or nd.get("title", "")
|
||
|
||
history = None
|
||
if mode in ("both", "conv"):
|
||
history = store.rebuild_history(target)
|
||
changed = []
|
||
if mode in ("both", "code"):
|
||
try:
|
||
changed = store.apply_code(target)
|
||
except Exception:
|
||
changed = []
|
||
store.rewind_head(target)
|
||
if mode in ("both", "conv") and log_path:
|
||
rewrite_projection(store, target, log_path)
|
||
return {
|
||
"history": history,
|
||
"changed": changed,
|
||
"prefill": prefill,
|
||
"at_origin": at_origin,
|
||
"target": target,
|
||
"title": nd.get("title", ""),
|
||
"to": to,
|
||
}
|