Cyrene_v2/internal/runtime/loop/checkpoint.go

package loop

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
)

// checkpointStore 是循环持久化持久快照并恢复被中断
// run 所需的最小表面（design §1 durable execution）。
// 生产环境由 *persist.Store 满足；测试替换为假实现，以在
// 没有真实数据库的情况下断言恢复路径。
type checkpointStore interface {
	// Save 插入或更新一个 checkpoint。
	Save(ctx context.Context, cp *persist.Checkpoint) error
	// LoadLatest 返回 agentID 最近更新的 checkpoint；
	// 该 agent 没有时返回 ErrNotFound。
	LoadLatest(ctx context.Context, agentID string) (*persist.Checkpoint, error)
}

// checkpointState 是每步持久化的持久快照。携带消息
// 窗口与步索引，使长任务能从最新 checkpoint 恢复而不
// 丢失上下文（design §1 durable execution）。
type checkpointState struct {
	Step       int             `json:"step"`
	Messages   []model.Message `json:"messages"`
	StopReason string          `json:"stop_reason,omitempty"`
}

// checkpointAgentID 返回本回合 checkpoint 的归属键（design/15 §2.5.9）。
//
// 单路（ctx 上未挂并发路标识）时它逐字等于 sessionID —— 与加入并发路之前
// 完全一致，归属语义不变；真并行路（ctx 上挂了 legID）时升为
// `{sessionID}#{legID}`，使 LoadLatest 按路隔离：并行的两路各写各的窗口，
// 绝不会把对方的快照当成自己的续接（并发写同一个 key 会互相覆盖）。
//
// 归属键同时用作 checkpoint 行 id 的前缀，因此两路的行在 `checkpoints` 表里
// 天然不冲突。
func checkpointAgentID(ctx context.Context, sessionID string) string {
	if leg := ParallelLegFrom(ctx); leg != "" {
		return sessionID + "#" + leg
	}
	return sessionID
}

// saveCheckpoint 在 step 处持久化一个快照。它绝不让 run 失败：
// store 为 nil 时为空操作，marshal/save 出错则记录日志。
func (l *Loop) saveCheckpoint(ctx context.Context, sessionID string, step int, messages []model.Message, stopReason string) {
	if l.checkpoint == nil {
		return
	}
	owner := checkpointAgentID(ctx, sessionID)
	state := checkpointState{
		Step:       step,
		Messages:   messages,
		StopReason: stopReason,
	}
	buf, err := json.Marshal(state)
	if err != nil {
		l.log.Warn("checkpoint marshal", "step", step, "error", err)
		return
	}
	cp := &persist.Checkpoint{
		ID:      fmt.Sprintf("%s:%d", owner, step),
		AgentID: owner,
		Step:    step,
		State:   buf,
	}
	if err := l.checkpoint.Save(ctx, cp); err != nil {
		l.log.Warn("checkpoint save", "step", step, "error", err)
	}
}

// loadCheckpoint 从该会话最新的持久 checkpoint 恢复一个**崩溃遗留**的
// run（design/04 §1 durable execution）。它返回恢复后的消息
// 窗口、下一个要运行的循环步（checkpoint 步 + 1），以及
// 是否找到了可恢复的 checkpoint。不存在 checkpoint，或最新的一个
// 反映的是已完成/被阻塞的 run、或任何已落定的未完成退出（取消 / 报错 /
// 预算耗尽 / 护栏拦截）、或它的窗口已被调用方本次带来的消息取代
// （见 supersededCheckpointWindow）时，它返回调用方的消息并置步为 1，使循环
// 从新开始（主人 2026-10-04 裁决：只有崩溃才允许恢复；见 resumableStopReason）。
//
// 恢复是幂等的：checkpoint 精确保存了上一个已完成步结束时
// 的消息窗口（assistant 工具调用及其结果已融合在内），
// 因此从下一步恢复会由历史重建模型
// 上下文，绝不重跑已完成步的工具调用，也不
// 重新生成已完成的答案。加载是只读的 —— 恢复绝不删除
// checkpoint，因此每一步的快照始终可查询。
func (l *Loop) loadCheckpoint(ctx context.Context, sessionID string, messages []model.Message) ([]model.Message, int, bool) {
	if l.checkpoint == nil {
		return messages, 1, false
	}
	// 归属键按路取（design/15 §2.5.9）：单路 = sessionID（逐字不变），真并行路
	// = {sessionID}#{legID}，因此两路的崩溃快照互不可见，绝不会互相顶掉。
	owner := checkpointAgentID(ctx, sessionID)
	cp, err := l.checkpoint.LoadLatest(ctx, owner)
	if err != nil {
		if !errors.Is(err, persist.ErrNotFound) {
			l.log.Warn("checkpoint load", "session", sessionID, "agent", owner, "error", err)
		}
		return messages, 1, false
	}
	st := checkpointState{}
	if err := json.Unmarshal(cp.State, &st); err != nil {
		l.log.Warn("checkpoint unmarshal", "session", sessionID, "step", cp.Step, "error", err)
		return messages, 1, false
	}
	if !resumableStopReason(st.StopReason) {
		l.log.Debug("checkpoint not resumable", "session", sessionID, "step", cp.Step, "stop_reason", st.StopReason)
		return messages, 1, false
	}
	// 纵深防御第三层（前两层：driveTurn 的收尾 defer 把未落定退出落定成不可恢复；
	// 网关的 session_turn_fence.go 让新回合等上一回合真正返回）：快照的 stop reason
	// 说明它是崩溃遗留的在途 run，还要它的窗口与调用方本次带来的消息指向**同一条
	// 最新用户输入**才续接，否则当作全新回合。
	if supersededCheckpointWindow(st.Messages, messages) {
		l.log.Info("checkpoint window superseded", "session", sessionID, "step", cp.Step,
			"next_step", 1, "messages", len(messages))
		return messages, 1, false
	}
	l.log.Info("checkpoint resumed", "session", sessionID, "step", cp.Step, "next_step", cp.Step+1, "messages", len(st.Messages))
	return st.Messages, cp.Step + 1, true
}

// resumableStopReason 报告最后一个步以上述 stop reason 结束的
// checkpoint 是否代表一个**崩溃遗留**的在途 run。只有落在工具循环中途的步
// 算「崩溃遗留」：循环已把工具结果并入窗口并即将继续，进程却猝死了——
// 此时收尾 defer 根本没执行，快照是进程最后留下的证据，从下一步接上会由
// 历史重建上下文而不重复工作（design/04 §1 durable execution）。
//
// 其余一律落到全新开始（主人 2026-10-04 裁决：只有崩溃才允许恢复）：
//   - 携带最终答案、护栏拒绝、预算耗尽或安全网终止的 checkpoint 是
//     已完成/被阻塞的 run——恢复它会重新生成已产出的答案（重复副作用）；
//   - canceled / aborted 是**已落定**的未完成退出，下一回合必须自己组
//     上下文（裁决 13.4.5），不得被旧窗口顶掉；
//   - 空的 stop reason 无法证明是崩溃遗留（崩溃留下的快照必带
//     tool_calls），保守起见同样不恢复。
//
// 打开真正新会话的调用方使用不同的 sessionID，本就错过 LoadLatest。
func resumableStopReason(stopReason string) bool {
	return stopReason == agent.StopReasonToolCalls
}

// supersededCheckpointWindow 报告快照窗口的最后一条用户消息是否已被调用方
// 本次带来的**更新**用户消息取代——取代了就不恢复（纵深防御第三层）。
//
// 它挡的是「上一回合收尾落定之前就发起下一回合」这一类路径的最后一道口子。
// 网关播种的新回合窗口是「该会话最近的持久化消息 + 刚发出的那条」
// （gateway.startTurnMessages），被放弃/被中断回合的最后一条用户消息因此
// **出现在**新窗口里，只是不再是最新的一条；此时续接旧快照就等于把用户刚发的
// 消息整段吞掉（裁决 13.4.5 的最新性 / 完整性）。
//
// 判据刻意保守：只有**正面证据**——快照窗口的最后一条用户消息确实出现在本回合
// 的消息里、而本回合又带来了更新的用户消息——才判为「已被取代」。判不出来的一律
// 维持既有崩溃语义：快照窗口与调用方消息完全无交集（进程重启后用户直接发了条
// 全新消息、历史不可得）是主人 2026-10-04 裁决明确保留的崩溃恢复代价
// （resume_test.go 的 TestRun_Resume_FromCrashSnapshot_LoadLatestCalled），
// 本判据不据它改判。
//
// 「同一条消息被重复播种」不会误判：判据比对的是两侧**最后一条**用户消息，
// 重复的同一条消息让两侧的键相等，判为同源，照常恢复。
func supersededCheckpointWindow(window, turn []model.Message) bool {
	windowLast := lastUserMessageKey(window)
	if windowLast == "" {
		return false
	}
	turnLast := lastUserMessageKey(turn)
	if turnLast == "" || turnLast == windowLast {
		return false
	}
	for _, m := range turn {
		if m.Role == model.RoleUser && userMessageKey(m) == windowLast {
			return true
		}
	}
	return false
}

// lastUserMessageKey 返回 messages 中最后一条用户消息的比对键；
// 没有用户消息（或该条没有任何可比内容）时返回 ""。
func lastUserMessageKey(messages []model.Message) string {
	for i := len(messages) - 1; i >= 0; i-- {
		if messages[i].Role == model.RoleUser {
			return userMessageKey(messages[i])
		}
	}
	return ""
}

// userMessageKey 返回一条用户消息的比对键：文本、图片引用与文档引用按出现顺序
// 拼接。纯图片消息（无文本）因此也能区分——只比文本会把「用户刚发的一张图」判成
// 空键而漏判；纯文档消息同理：客户端「超限走文件」发的就是一条只有附件的消息，
// 漏掉它会让续接判定在附件回合上失准（design/24 §9.5 移交的静默吞掉）。角色与
// 说话人不进键：这里只问「是不是同一条用户输入」。
func userMessageKey(m model.Message) string {
	var b strings.Builder
	for _, part := range m.Content {
		switch part.Type {
		case model.ContentTypeImage:
			b.WriteString("\x00img:")
			b.WriteString(part.ImageURL)
		case model.ContentTypeDocument:
			b.WriteString("\x00doc:")
			b.WriteString(part.DocumentName)
		default:
			b.WriteString(part.Text)
		}
	}
	return b.String()
}