AutoZero/session/session.go

// Package session 实现 JSONL 会话持久化（借鉴 picoclaw 的 jsonl_backend）：
// 每条消息一行，追加写；resume 时整体重放。状态文件用临时文件 + rename 原子写。
package session

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// dirOverride 非空时 Dir() 返回它，供宿主自定义会话存储布局；
// 默认空，走 ~/.autogo/sessions（与来源仓库 AutoGo 保持镜像）。
var dirOverride string

// SetDir 覆盖会话存储目录。须在进程启动时调用一次；并发调用不设防。
func SetDir(d string) { dirOverride = d }

// Dir 返回会话存储目录：~/.autogo/sessions。
func Dir() string {
	if dirOverride != "" {
		return dirOverride
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return ".autogo-sessions"
	}
	return filepath.Join(home, ".autogo", "sessions")
}

// Session 是一个可持久化的会话。
type Session struct {
	Path     string
	Messages []llm.Message
	// Retired 是被游标退役、仍原样留在文件里的消息条数。它们不进上下文，但**没被删**
	// ——这正是回退敢做的原因：活跃视图可以缩，文件不缩。
	Retired int
	// retired 是那些消息本身。留着它，Rewrite 才能把历史原样写回去：
	// collectToolImages 那类"更新已有消息字段"的调用会整体重写文件，如果不保留
	// 退役内容，一次图片缓存就把刚做的回退历史抹掉了。
	retired []llm.Message
	// TurnEnds 是文件里的回合结束记录（按发生顺序）。回合"为什么停"必须能在会话
	// 文件本身看到：只导出 JSONL 取证时，光看最后一条消息分不清是模型说完、
	// 用户停止、报错，还是被步数上限切断。
	TurnEnds []TurnEnd
}

// turnEndKind 是回合结束记录行的判别值（与消息不撞：消息一定有 role）。
const turnEndKind = "autogo.turn_end"

// TurnEnd 是一次回合的退出原因。
type TurnEnd struct {
	Reason string `json:"reason"` // completed | stopped | failed | max_turns
	Error  string `json:"error,omitempty"`
	Step   int    `json:"step,omitempty"` // 本回合跑到第几步时结束
	Ms     int64  `json:"ms,omitempty"`
	At     int64  `json:"at,omitempty"` // unix ms
}

type turnEndRecord struct {
	Kind   string `json:"kind"`
	Reason string `json:"reason"`
	Error  string `json:"error,omitempty"`
	Step   int    `json:"step,omitempty"`
	Ms     int64  `json:"ms,omitempty"`
	At     int64  `json:"at,omitempty"`
}

// parseTurnEnd 判断这一行是不是一条回合结束记录。同样先用字节查找筛掉消息行。
func parseTurnEnd(line []byte) (TurnEnd, bool) {
	if !bytes.Contains(line, []byte(`"`+turnEndKind+`"`)) {
		return TurnEnd{}, false
	}
	var r turnEndRecord
	if json.Unmarshal(line, &r) != nil || r.Kind != turnEndKind {
		return TurnEnd{}, false
	}
	return TurnEnd{Reason: r.Reason, Error: r.Error, Step: r.Step, Ms: r.Ms, At: r.At}, true
}

// AppendTurnEnd 追加一条回合结束记录。best-effort：这是账，不该拦住回合收尾。
func (s *Session) AppendTurnEnd(t TurnEnd) {
	if s.Path == "" {
		return
	}
	t.At = time.Now().UnixMilli()
	line, err := json.Marshal(turnEndRecord{Kind: turnEndKind, Reason: t.Reason, Error: t.Error, Step: t.Step, Ms: t.Ms, At: t.At})
	if err != nil {
		return
	}
	if err := s.appendLine(append(line, '\n')); err != nil {
		return
	}
	s.TurnEnds = append(s.TurnEnds, t)
}

// cursorKind 是游标记录在 JSONL 里的判别值。它不可能与消息撞车：消息一定有 role，
// 游标一定没有。
const cursorKind = "autogo.cursor"

// cursorRecord 一行 {"kind":"autogo.cursor","keep":3}：把此刻已积累的活跃视图
// 截到前 keep 条，之后的读取继续往（缩短后的）视图上追加。
//
// 为什么不是直接改短文件（pigo/ZCode 的同一选择）：会话文件是唯一事实源，靠重写
// 维护历史就意味着「重写失败 = 历史没了」。回退这一步尤其危险——被截掉的那些消息
// 是用户唯一可能想撤销的内容。改成追加一条游标之后，回退在磁盘上只是一个
// 单行 append：写一半被杀，最坏结果是这行作废、回退没生效，而不是历史丢失。
type cursorRecord struct {
	Kind string `json:"kind"`
	Keep int    `json:"keep"`
}

// parseCursor 判断这一行是不是游标。先用一次便宜的字节查找筛掉绝大多数消息行
// （消息行可能几百 KB，逐行再解一遍 JSON 太浪费）。
func parseCursor(line []byte) (int, bool) {
	if !bytes.Contains(line, []byte(`"`+cursorKind+`"`)) {
		return 0, false
	}
	var c cursorRecord
	if json.Unmarshal(line, &c) != nil || c.Kind != cursorKind {
		return 0, false
	}
	if c.Keep < 0 {
		c.Keep = 0
	}
	return c.Keep, true
}

// New 创建带时间戳的新会话文件名。
func New() (*Session, error) {
	dir := Dir()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, err
	}
	return &Session{Path: newName(dir, time.Now().Format("20060102-150405"))}, nil
}

// newMu 保护 issued：本进程已发出、可能还没落盘的会话路径。
var (
	newMu  sync.Mutex
	issued = map[string]bool{}
)

// newName 在 dir 里挑一个没被占用的会话文件名。
//
// 名字只精确到秒，同一秒内连续新建会拿到同一个名字。会话文件是**第一条消息**才
// 写盘的，所以只查磁盘会漏掉"两份都还没落盘"那种撞名——那时光看 Stat 分不出来，
// 第二次"新建"经 LoadOrCreate 会把第一份（甚至某份已经写好消息的旧会话）当成
// 自己的。因此两道都查：盘上已有的文件、本进程已发出的名字，撞上就依次加 -2/-3 后缀。
func newName(dir, ts string) string {
	newMu.Lock()
	defer newMu.Unlock()
	path := filepath.Join(dir, ts+".jsonl")
	for i := 2; ; i++ {
		if _, err := os.Stat(path); os.IsNotExist(err) && !issued[path] {
			issued[path] = true
			return path
		}
		path = filepath.Join(dir, fmt.Sprintf("%s-%d.jsonl", ts, i))
	}
}

// sanitizeName 清理会话文件名（bot 渠道按 chat 命名，chat id 可能来自外部）。
func sanitizeName(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(name) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '-', r == '_':
			b.WriteRune(r)
		default:
			b.WriteRune('-')
		}
	}
	s := b.String()
	if len(s) > 80 {
		s = s[:80]
	}
	return s
}

// LoadOrCreate 打开具名会话；不存在则创建（bot 渠道按 chat 持久化记忆）。
func LoadOrCreate(name string) (*Session, error) {
	dir := Dir()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, err
	}
	path := filepath.Join(dir, sanitizeName(name)+".jsonl")
	if _, err := os.Stat(path); err == nil {
		return Load(path)
	}
	return &Session{Path: path}, nil
}

// Exists 报告具名会话是否已经落盘。后台通知等"会话之外的事"用它避免把
// 已删除的会话重新创建出来。
func Exists(name string) bool {
	if _, ok := SafeID(name); !ok {
		return false
	}
	_, err := os.Stat(filepath.Join(Dir(), name+".jsonl"))
	return err == nil
}

// Latest 返回最近修改的会话；目录为空时新建。
func Latest() (*Session, error) {
	dir := Dir()
	entries, err := os.ReadDir(dir)
	if err != nil || len(entries) == 0 {
		return New()
	}
	best := ""
	var bestMod time.Time
	for _, e := range entries {
		if e.IsDir() || filepath.Ext(e.Name()) != ".jsonl" {
			continue
		}
		info, err := e.Info()
		if err == nil && info.ModTime().After(bestMod) {
			bestMod, best = info.ModTime(), e.Name()
		}
	}
	if best == "" {
		return New()
	}
	return Load(filepath.Join(dir, best))
}

// Load 重放 JSONL 文件恢复会话。
func Load(path string) (*Session, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	s := &Session{Path: path}
	// 用 jsonl.Each 而不是 bufio.Scanner：后者有 64KB 行长上限，超限会让**整个会话
	// 读不出来**（曾因一条 204KB 的 reasoning 消息导致会话在 WebUI 里消失）。
	err = jsonl.Each(f, func(line []byte) {
		if keep, ok := parseCursor(line); ok {
			// 游标只缩短活跃视图，被越过的消息收进 retired 原样留着（Rewrite 要写回去）。
			if keep < len(s.Messages) {
				s.Retired += len(s.Messages) - keep
				s.retired = append(s.retired, s.Messages[keep:]...)
				s.Messages = s.Messages[:keep]
			}
			return
		}
		if te, ok := parseTurnEnd(line); ok {
			s.TurnEnds = append(s.TurnEnds, te)
			return
		}
		var m llm.Message
		if json.Unmarshal(line, &m) != nil || m.Role == "" {
			// 坏行、或崩溃留下的半行：跳过。没有 role 的行不是一条消息——
			// 把它当空消息收进视图，会让下一次请求带着一条非法消息去撞 API。
			return
		}
		s.Messages = append(s.Messages, m)
	})
	return s, err
}

// Append 追加一条消息并落盘（追加写，兼容 tail）。
func (s *Session) Append(m llm.Message) error {
	s.Messages = append(s.Messages, m)
	b, err := json.Marshal(m)
	if err != nil {
		return err
	}
	// 换行是内容的一部分：少了它，多条消息会被写成**一行**，整份会话解析不出来
	// （曾经就是这么把自己改坏的）。
	return s.appendLine(append(b, '\n'))
}

// AppendAll 一次打开文件追加多条（撤销回退时接回一段历史用）。
//
// 只有 Append 可用时也能写，但撤销一次回退要往文件里塞回几十条消息，
// 逐条 open/close 既慢又给中途崩溃留更多口子。
func (s *Session) AppendAll(msgs []llm.Message) error {
	if len(msgs) == 0 {
		return nil
	}
	var b bytes.Buffer
	for _, m := range msgs {
		line, err := json.Marshal(m)
		if err != nil {
			return err
		}
		b.Write(line)
		b.WriteByte('\n')
	}
	if err := s.appendLine(b.Bytes()); err != nil {
		return err
	}
	s.Messages = append(s.Messages, msgs...)
	return nil
}

// appendLine 把一段已经带换行的 JSONL 追加到会话文件。Path 为空 = 临时会话，只改内存。
func (s *Session) appendLine(payload []byte) error {
	if s.Path == "" {
		return nil
	}
	f, err := os.OpenFile(s.Path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.Write(payload)
	return err
}

// RetireTo 把活跃视图截到前 keep 条，并在文件末尾追加一条游标记录。
//
// 与 Truncate 的分工是刻意的：RetireTo 是**回退**（历史留在盘上、可撤销、不进上下文），
// Truncate/Rewrite 是**清空**（/clear 那种"这些内容我不要了"的语义，真的把文件改短）。
// 写盘失败时内存一点没动——界面与文件仍然一致。
func (s *Session) RetireTo(keep int) error {
	if keep < 0 {
		keep = 0
	}
	if keep >= len(s.Messages) {
		return nil // 没东西可退役，也不留一条无意义的游标
	}
	line, err := json.Marshal(cursorRecord{Kind: cursorKind, Keep: keep})
	if err != nil {
		return err
	}
	if err := s.appendLine(append(line, '\n')); err != nil {
		return err
	}
	s.Retired += len(s.Messages) - keep
	s.retired = append(s.retired, s.Messages[keep:]...)
	s.Messages = s.Messages[:keep]
	return nil
}

// Rewrite 用当前内存中的消息重写 JSONL 文件。
//
// 这是**破坏性**的一条：写完之后文件里只有活跃视图，游标与已退役的历史一起没了。
// 所以回退走 RetireTo，这里只留给"我确实不要这些内容了"的 /clear，以及更新已有
// 消息字段的场合。
//
// 必须 tmp+rename 原子替换：原来直接 os.Create 目标文件再逐条写，进程在中间被杀
// （或磁盘写满）会把整条会话截成半截——而会话文件是唯一事实源。
func (s *Session) Rewrite() error {
	if s.Path == "" {
		return nil
	}
	var b bytes.Buffer
	// 先写退役的历史，再写一条 keep=0 的游标把它们从活跃视图里隔开，最后写当前视图。
	// 不这么做的话，任何一次 Rewrite（例如给图片回填缓存文件名）都会把回退留下的
	// 历史抹掉——文件是"更新一条消息"的载体，不该顺手删掉别的东西。
	emit := func(m llm.Message) error {
		line, err := json.Marshal(m)
		if err != nil {
			return err
		}
		b.Write(line)
		b.WriteByte('\n')
		return nil
	}
	for _, m := range s.retired {
		if err := emit(m); err != nil {
			return err
		}
	}
	if len(s.retired) > 0 {
		cur, err := json.Marshal(cursorRecord{Kind: cursorKind, Keep: 0})
		if err != nil {
			return err
		}
		b.Write(cur)
		b.WriteByte('\n')
	}
	for _, m := range s.Messages {
		if err := emit(m); err != nil {
			return err
		}
	}
	for _, te := range s.TurnEnds {
		line, err := json.Marshal(turnEndRecord{Kind: turnEndKind, Reason: te.Reason, Error: te.Error, Step: te.Step, Ms: te.Ms, At: te.At})
		if err != nil {
			return err
		}
		b.Write(line)
		b.WriteByte('\n')
	}
	tmp := s.Path + ".tmp"
	if err := os.WriteFile(tmp, b.Bytes(), 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, s.Path)
}

// Reset 清空一切：活跃视图、退役历史、回合账。
//
// /clear 要的是"这些内容我不要了"，所以退役的那部分也得跟着丢——否则 Rewrite
// 会把它们原样写回去，"清空"就只剩清空视图、文件却越清越肥。
func (s *Session) Reset() {
	s.Messages = nil
	s.retired = nil
	s.Retired = 0
	s.TurnEnds = nil
}

// SaveState 原子写一个状态文件（借鉴 picoclaw state.go 的 tmp+rename）。
func SaveState(path string, v any) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	b, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, b, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

// LoadState 读回 SaveState 写的状态文件。文件缺失或损坏返回错误，
// 由调用方决定是否降级（如项目计划缺失 = 空计划，不是致命）。
func LoadState(path string, v any) error {
	b, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, v)
}

// String 摘要信息。
func (s *Session) String() string {
	return fmt.Sprintf("session %s (%d messages)", filepath.Base(s.Path), len(s.Messages))
}

// CountPrompts 返回用户消息条数，即"第几个 prompt"意义上的回合数。
//
// 口径只有一处定义（连同 PromptBoundary），因为留档的回合号与回退的切点
// 必须是同一套数法：否则 `a2go rewind 3` 回到的不是第三次提问之前。
// 数的是 RoleUser 全部消息（含只带图片的），与写入时的计数规则一致。
func CountPrompts(msgs []llm.Message) int {
	n := 0
	for _, m := range msgs {
		if m.Role == llm.RoleUser && !m.Synthetic {
			n++
		}
	}
	return n
}

// PromptBoundary 返回第 n 个用户消息的下标，也就是"回到第 n 个 prompt 之前"
// 应当保留的长度；找不到（n 越界）返回 -1。
//
// 切点落在用户消息上是有意的：切在 assistant 的 tool_use 与它的 tool_result
// 之间会留下没有配对的 tool_use，下一次请求被 API 直接拒掉。
func PromptBoundary(msgs []llm.Message, n int) int {
	count := 0
	for i := range msgs {
		if msgs[i].Role != llm.RoleUser || msgs[i].Synthetic {
			continue
		}
		count++
		if count == n {
			return i
		}
	}
	return -1
}

// TailFromPrompt 返回第 n 个用户 prompt 起的尾部**副本**与切点下标（找不到返回 -1）。
//
// 它只读不动任何东西，这是刻意的顺序约束：调用方必须先用这份副本归档、成功后
// 再 RetireTo。反过来就会出现"消息已经从活跃视图摘掉、归档却失败"的窗口——
// 那正是这次改造要关掉的东西。
func (s *Session) TailFromPrompt(n int) (idx int, dropped []llm.Message) {
	idx = PromptBoundary(s.Messages, n)
	if idx < 0 {
		return -1, nil
	}
	return idx, append([]llm.Message(nil), s.Messages[idx:]...)
}