GopherGraph/checkpoint.go

package GopherGraph

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// Checkpointer 定义了持久化状态的接口，S 是工作流的强类型状态。
type Checkpointer[S any] interface {
	// Save 保存当前线程快照。
	Save(ctx context.Context, threadID string, thread *Thread[S]) error
	// Load 读取之前保存的线程快照，用于恢复执行。
	Load(ctx context.Context, threadID string) (*Thread[S], error)
	// Delete 删除指定线程的快照文件，用于工作流完结后的清理。
	Delete(ctx context.Context, threadID string) error
}

// FileCheckpointer 实现了基于本地 JSON 文件的状态持久化。
type FileCheckpointer[S any] struct {
	dir    string // 存放状态文件的文件夹路径
	pretty bool   // 是否使用缩进格式（方便人工阅读，但体积更大）
}

// FileCheckpointerOption 是 FileCheckpointer 的可选配置项。
type FileCheckpointerOption func(*fileCheckpointerConfig)

type fileCheckpointerConfig struct {
	pretty bool
}

// WithPrettyJSON 开启 JSON 缩进格式（方便人工阅读，适合开发调试阶段）。
// 默认关闭，使用紧凑格式以减少磁盘 I/O 和文件体积。
func WithPrettyJSON() FileCheckpointerOption {
	return func(c *fileCheckpointerConfig) {
		c.pretty = true
	}
}

// NewFileCheckpointer 创建一个文件检查点管理器。
// 支持通过 opts 传入可选配置（如 WithPrettyJSON()）。
func NewFileCheckpointer[S any](dir string, opts ...FileCheckpointerOption) (*FileCheckpointer[S], error) {
	// 确保文件夹存在，不存在则创建
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create checkpoint dir: %w", err)
	}

	cfg := &fileCheckpointerConfig{}
	for _, opt := range opts {
		opt(cfg)
	}

	return &FileCheckpointer[S]{dir: dir, pretty: cfg.pretty}, nil
}

// Save 将 Thread 序列化为 JSON 并写入本地文件。
// 采用"先写临时文件再原子重命名"的方式，防止写入中途程序崩溃导致文件损坏。
// 若存在上次崩溃残留的临时文件，会在写入前自动清理。
func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error {
	if err := validateThreadID(threadID); err != nil {
		return err
	}

	path := filepath.Join(fc.dir, threadID+".json")
	tmpPath := path + ".tmp"

	// 序列化：生产环境默认使用紧凑格式以减少磁盘 I/O；开发阶段可用 WithPrettyJSON() 开启缩进
	var data []byte
	var err error
	if fc.pretty {
		data, err = json.MarshalIndent(thread, "", "  ")
	} else {
		data, err = json.Marshal(thread)
	}
	if err != nil {
		return fmt.Errorf("failed to marshal thread: %w", err)
	}

	// 清理上次崩溃可能残留的临时文件（忽略不存在的错误）
	_ = os.Remove(tmpPath)

	// 原子写入：先写临时文件，再 Rename 替换，确保断电/崩溃安全
	if err := os.WriteFile(tmpPath, data, 0644); err != nil {
		return fmt.Errorf("failed to write temp file: %w", err)
	}
	if err := os.Rename(tmpPath, path); err != nil {
		// Rename 失败时，尝试清理临时文件，避免残留
		_ = os.Remove(tmpPath)
		return fmt.Errorf("failed to rename temp file: %w", err)
	}
	return nil
}

// Load 从本地文件中读取 JSON 并反序列化为 Thread。
// 若加载到一个已完结（IsFinished == true）的线程，会包装 ErrAlreadyFinished 作为额外信息返回，
// 方便调用方做幂等处理（注意：此时 thread 依然有效，err 可通过 errors.Is 判断）。
func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) {
	if err := validateThreadID(threadID); err != nil {
		return nil, err
	}

	path := filepath.Join(fc.dir, threadID+".json")

	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("failed to read thread file: %w", err)
	}

	var thread Thread[S]
	if err := json.Unmarshal(data, &thread); err != nil {
		return nil, fmt.Errorf("failed to unmarshal thread: %w", err)
	}

	// 如果加载到已完结的线程，附带哨兵错误提示，方便调用方感知
	if thread.IsFinished {
		return &thread, fmt.Errorf("thread %q has already finished: %w", threadID, ErrAlreadyFinished)
	}

	return &thread, nil
}

// Delete 删除指定 threadID 对应的快照文件。
// 通常在工作流完结后调用，防止快照文件无限积累占用磁盘空间。
// 若文件不存在，静默返回 nil（幂等操作）。
func (fc *FileCheckpointer[S]) Delete(ctx context.Context, threadID string) error {
	if err := validateThreadID(threadID); err != nil {
		return err
	}

	path := filepath.Join(fc.dir, threadID+".json")
	if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to delete checkpoint file: %w", err)
	}
	return nil
}

// validateThreadID 校验 threadID 的合法性，防止路径穿越（path traversal）攻击。
// threadID 中不允许出现目录分隔符或 ".." 序列。
func validateThreadID(id string) error {
	if id == "" {
		return fmt.Errorf("threadID must not be empty")
	}
	if strings.ContainsAny(id, "/\\") || strings.Contains(id, "..") {
		return fmt.Errorf("threadID %q contains invalid characters", id)
	}
	return nil
}