AutoZero/agent/checkpoint.go

package agent

import (
	"fmt"
	"strings"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
)

const (
	checkpointPrefix = "Here is a summary of the conversation so far:\n\n"
	checkpointSuffix = "\n\nContinue the task based on this summary."
)

// checkpointMessage 把交接文档装成工作集里的第一条 user 消息。
// 形状保持不变（测试与旧会话缓存都依赖这个前缀）。
func checkpointMessage(text string) llm.Message {
	return llm.Message{Role: llm.RoleUser, Content: checkpointPrefix + text + checkpointSuffix}
}

func isCheckpointMessage(m llm.Message) bool {
	return m.Role == llm.RoleUser && strings.HasPrefix(m.Content, checkpointPrefix)
}

// LoadCheckpoint 从 sidecar 恢复压缩交接文档，并把工作集重建成
// [checkpoint] + transcript[Covered:]。覆盖边界用消息指纹校验：历史被回退或
// 改写后指纹对不上，就当作没有 checkpoint——绝不能拿覆盖了别的消息的旧摘要继续跑。
func (a *Agent) LoadCheckpoint() bool {
	if a.Session == nil || a.Session.Path == "" {
		return false
	}
	cp, ok := session.LoadCheckpoint(a.sessionID())
	if !ok || cp.Covered <= 0 || cp.Covered > len(a.Session.Messages) {
		return false
	}
	if cp.TailHash != session.MessageFingerprint(a.Session.Messages[cp.Covered-1]) {
		return false
	}
	a.compactMu.Lock()
	defer a.compactMu.Unlock()
	a.checkpoint = cp.Text
	a.checkpointCovered = cp.Covered
	a.checkpointHash = cp.TailHash
	a.compacted = []llm.Message{checkpointMessage(cp.Text)}
	a.compactFrom = cp.Covered
	a.compactGaveUp = false
	a.resetInputUsage()
	return true
}

// validateCheckpointLocked 在 transcript 被回退/改写后作废旧 checkpoint。
// 覆盖边界的指纹对不上就整体退回完整历史，绝不拿一份覆盖了别的消息的摘要继续跑。
// 调用方持 compactMu。
func (a *Agent) validateCheckpointLocked() {
	if a.checkpoint == "" || a.Session == nil {
		return
	}
	valid := a.checkpointCovered > 0 && a.checkpointCovered <= len(a.Session.Messages) &&
		a.checkpointHash == session.MessageFingerprint(a.Session.Messages[a.checkpointCovered-1])
	if valid {
		return
	}
	a.checkpoint = ""
	a.checkpointCovered = 0
	a.checkpointHash = ""
	a.compacted = nil
	a.compactFrom = 0
	a.compactGaveUp = false
}

// saveCheckpointLocked 把当前 checkpoint 写进 sidecar。调用方持 compactMu。
func (a *Agent) saveCheckpointLocked() {
	id := a.sessionID()
	if id == "" || a.checkpoint == "" {
		return
	}
	err := session.SaveCheckpoint(id, session.Checkpoint{
		Text:     a.checkpoint,
		Covered:  a.checkpointCovered,
		TailHash: a.checkpointHash,
	})
	if err != nil {
		fmt.Fprintf(a.Opt.Stderr, "[context] save checkpoint failed: %v\n", err)
	}
}

// resetInputUsage 让下一次压缩判据退回估算：视图刚被替换，旧请求的真实 token
// 数已经不对应当前工作集了。
func (a *Agent) resetInputUsage() {
	a.usageMu.Lock()
	a.lastInputTokens = 0
	a.lastInputMsgs = 0
	a.lastSystemBytes = 0
	a.usageMu.Unlock()
}

// noteRequestUsage 记下一次成功请求的真实 prompt token（含缓存命中/写入）。
// 下一个 step 的压缩判据用它 + 新增消息估算，而不是纯 bytes/3。
func (a *Agent) noteRequestUsage(u *llm.Usage, view []llm.Message) {
	if u == nil || u.InputTokens <= 0 {
		return
	}
	sysLen := len(a.sysNow())
	a.usageMu.Lock()
	a.lastInputTokens = u.InputTokens
	a.lastInputMsgs = len(view)
	a.lastSystemBytes = sysLen
	a.usageMu.Unlock()
}

// observedTokens 返回当前工作集的 token 观测值：有真实请求数据时以它为基准，
// 加上那次请求之后新增消息与 system 动态段的估算；没有就退回 bytes/3。
func (a *Agent) observedTokens(view []llm.Message) int {
	a.usageMu.Lock()
	last, lastN, lastSys := a.lastInputTokens, a.lastInputMsgs, a.lastSystemBytes
	a.usageMu.Unlock()
	sys := a.sysNow()
	if last <= 0 || lastN < 0 || lastN > len(view) {
		return estimateTokens(view, sys)
	}
	added := 0
	if lastN < len(view) {
		added = estimateTokens(view[lastN:], "")
	}
	if delta := len(sys) - lastSys; delta > 0 {
		added += delta / 3
	}
	return last + added
}