Agent-Native/runtime/checkpoint/checkpoint.go

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

import (
	"bytes"
	"encoding/json"
	"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 是当前内容既不是写入后也不是写入前的文件。
type Result struct {
	Restored []string
	Skipped  []string
}

// Log 是追加写的检查点文件。
type Log struct {
	path string
	mu   sync.Mutex
}

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()
	return &Log{path: path}, 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()
	all, err := l.read()
	if err != nil {
		return err
	}
	rec := Record{
		Seq: len(all) + 1, 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
}

// 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)
	for _, rel := range plan.Restored {
		group := groupFor(selected(all, anchor, ordered), rel)
		if err := restore(group); err != nil {
			return Result{}, err
		}
	}
	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
}

func restore(group []Record) error {
	if len(group) == 0 {
		return nil
	}
	first := group[0]
	abs := filepath.Join(first.Root, filepath.FromSlash(first.Path))
	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)
}