Agent-Native/runtime/checkpoint/checkpoint.go

// Package checkpoint 记录写文件之前的内容，并按用户消息把文件退回去。
// 只追踪显式交给它的写入。shell 与 python 的改动不经过这里。
package checkpoint

import (
	"bytes"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"sync"
)

// Record 是一次被追踪的写入。Before 是写入前的字节，After 是写入后的字节。
type Record struct {
	Seq     int    `json:"seq"`
	Anchor  string `json:"anchor,omitempty"`
	Root    string `json:"root"`
	Path    string `json:"path"`
	Existed bool   `json:"existed"`
	Before  []byte `json:"before,omitempty"`
	After   []byte `json:"after,omitempty"`
}

// Result 是一次回滚的分类。Skipped 是当前内容既不是写入后也不是写入前的文件
// （含目标是符号链接的情况——恢复会跟出根外，按外部改动处理）。
// Failed 是恢复动作本身失败的文件（Rewind 不因单文件失败中止，继续其余）。
type Result struct {
	Restored []string
	Skipped  []string
	Failed   []string
}

// Log 是追加写的检查点文件。
type Log struct {
	path    string
	mu      sync.Mutex
	lastSeq int // Open 时一次读出的最大 Seq；Record 递增（不再全读）
}

func Open(path string) (*Log, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, err
	}
	f, err := os.OpenFile(path, os.O_CREATE, 0o644)
	if err != nil {
		return nil, err
	}
	_ = f.Close()
	l := &Log{path: path}
	// 初始化序号（一次读，之后递增）。
	if all, err := l.read(); err == nil && len(all) > 0 {
		l.lastSeq = all[len(all)-1].Seq
	}
	return l, nil
}

func (l *Log) Path() string { return l.path }

// Record 追加一条检查点。rel 是相对 root 的路径。
func (l *Log) Record(anchor, root, rel string, before []byte, existed bool, after []byte) error {
	if l == nil {
		return nil
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	l.lastSeq++
	rec := Record{
		Seq: l.lastSeq, Anchor: anchor, Root: root, Path: filepath.ToSlash(rel),
		Existed: existed, Before: append([]byte(nil), before...), After: append([]byte(nil), after...),
	}
	raw, err := json.Marshal(rec)
	if err != nil {
		return err
	}
	f, err := os.OpenFile(l.path, os.O_APPEND|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.Write(append(raw, '\n'))
	return err
}

// Compact 只保留最近 keepAnchors 个 anchor 的记录（按 Seq 序判定新旧，
// 保留原 Seq 不重排，序号单调性不破）。keepAnchors ≤ 0 不动。
// 重写是原子的（临时文件 + rename）。
func (l *Log) Compact(keepAnchors int) error {
	if l == nil || keepAnchors <= 0 {
		return nil
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	all, err := l.read()
	if err != nil {
		return err
	}
	var anchors []string
	seen := map[string]bool{}
	for _, rec := range all {
		if !seen[rec.Anchor] {
			seen[rec.Anchor] = true
			anchors = append(anchors, rec.Anchor)
		}
	}
	if len(anchors) <= keepAnchors {
		return nil
	}
	keep := map[string]bool{}
	for _, a := range anchors[len(anchors)-keepAnchors:] {
		keep[a] = true
	}
	tmp := l.path + ".compact.tmp"
	f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return err
	}
	enc := json.NewEncoder(f)
	for _, rec := range all {
		if keep[rec.Anchor] {
			if err := enc.Encode(rec); err != nil {
				_ = f.Close()
				_ = os.Remove(tmp)
				return err
			}
		}
	}
	if err := f.Close(); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, l.path)
}

// Plan 列出从 anchor 起（含）会被回滚的文件，不改磁盘。
// ordered 是用户消息 id 的时间顺序。
func (l *Log) Plan(anchor string, ordered []string) (Result, error) {
	if l == nil {
		return Result{}, nil
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	all, err := l.read()
	if err != nil {
		return Result{}, err
	}
	return classify(all, anchor, ordered), nil
}

// Rewind 把 anchor 及其后的被追踪写入退回写入前。
// 当前内容与最后一次记录的 After 不一致时跳过并列入 Skipped。
func (l *Log) Rewind(anchor string, ordered []string) (Result, error) {
	if l == nil {
		return Result{}, nil
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	all, err := l.read()
	if err != nil {
		return Result{}, err
	}
	plan := classify(all, anchor, ordered)
	// 单文件失败不中止：符号链接目标归入 Skipped（外部改动语义），
	// 其余恢复错误记入 Failed 并继续（部分恢复对调用方可读）。
	restored := plan.Restored[:0]
	for _, rel := range plan.Restored {
		group := groupFor(selected(all, anchor, ordered), rel)
		if err := restore(group); err != nil {
			if errors.Is(err, errSymlinkTarget) {
				plan.Skipped = append(plan.Skipped, rel)
				continue
			}
			plan.Failed = append(plan.Failed, rel)
			continue
		}
		restored = append(restored, rel)
	}
	plan.Restored = restored
	return plan, nil
}

func (l *Log) read() ([]Record, error) {
	b, err := os.ReadFile(l.path)
	if err != nil {
		return nil, err
	}
	var out []Record
	for _, line := range bytes.Split(b, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		var rec Record
		if err := json.Unmarshal(line, &rec); err != nil {
			return nil, err
		}
		out = append(out, rec)
	}
	return out, nil
}

func selected(all []Record, anchor string, ordered []string) []Record {
	keep := map[string]bool{}
	on := false
	for _, id := range ordered {
		if id == anchor {
			on = true
		}
		if on {
			keep[id] = true
		}
	}
	if anchor != "" && len(ordered) == 0 {
		keep[anchor] = true
	}
	var out []Record
	for _, rec := range all {
		if keep[rec.Anchor] {
			out = append(out, rec)
		}
	}
	return out
}

func classify(all []Record, anchor string, ordered []string) Result {
	sel := selected(all, anchor, ordered)
	seen := map[string]bool{}
	var restored, skipped []string
	for _, rec := range sel {
		if seen[rec.Path] {
			continue
		}
		seen[rec.Path] = true
		group := groupFor(sel, rec.Path)
		switch fileState(group) {
		case stateRestore:
			restored = append(restored, rec.Path)
		case stateSkip:
			skipped = append(skipped, rec.Path)
		}
	}
	return Result{Restored: restored, Skipped: skipped}
}

func groupFor(sel []Record, rel string) []Record {
	var group []Record
	for _, rec := range sel {
		if rec.Path == rel {
			group = append(group, rec)
		}
	}
	return group
}

const (
	stateAlready = iota
	stateRestore
	stateSkip
)

func fileState(group []Record) int {
	if len(group) == 0 {
		return stateAlready
	}
	first, last := group[0], group[len(group)-1]
	abs := filepath.Join(last.Root, filepath.FromSlash(last.Path))
	current, err := os.ReadFile(abs)
	missing := os.IsNotExist(err)
	if err != nil && !missing {
		return stateSkip
	}
	if missing {
		if !last.Existed && len(last.After) == 0 {
			return stateAlready
		}
		if !first.Existed && bytes.Equal(nil, first.Before) && len(last.After) > 0 {
			return stateSkip
		}
		return stateSkip
	}
	if bytes.Equal(current, last.After) {
		if !first.Existed {
			return stateRestore
		}
		if bytes.Equal(current, first.Before) {
			return stateAlready
		}
		return stateRestore
	}
	if first.Existed && bytes.Equal(current, first.Before) {
		return stateAlready
	}
	return stateSkip
}

var errSymlinkTarget = errors.New("checkpoint: restore target is a symlink")

// restoreForTesting 是测试缝：非空时在每次恢复前调用（返回错误即恢复失败）。
// 生产代码永不应设置它；恢复期的故障注入无法绕开分类层的预检
// （权限类失败对 root 不可重现，故用显式缝）。
var restoreForTesting func(abs string) error

func restore(group []Record) error {
	if len(group) == 0 {
		return nil
	}
	first := group[0]
	abs := filepath.Join(first.Root, filepath.FromSlash(first.Path))
	if restoreForTesting != nil {
		if err := restoreForTesting(abs); err != nil {
			return err
		}
	}
	// 符号链接防护：恢复会跟出根外（与 fs_write_file 的防护同语义）。
	// 按外部改动处理——由调用方归入 Skipped，不写不删。
	if info, err := os.Lstat(abs); err == nil && info.Mode()&os.ModeSymlink != 0 {
		return errSymlinkTarget
	}
	if !first.Existed {
		return os.Remove(abs)
	}
	if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil {
		return err
	}
	return os.WriteFile(abs, first.Before, 0o644)
}