Agent-System-Backend/internal/store/messages.go

package store

import (
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	core "agent-system-core"
)

// isUniqueViolation 判定唯一约束冲突（PG 23505 / SQLite UNIQUE）。
func isUniqueViolation(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "23505") ||
		strings.Contains(msg, "UNIQUE constraint") ||
		strings.Contains(msg, "duplicate key value")
}

// ---- 消息 ----

// AppendMessage 以单语句原子分配会话内序号（INSERT ... SELECT MAX+1）；
// 并发同会话追加由 uq_messages_conv_seq 唯一约束兜底，冲突有界重试。
func (s *Store) AppendMessage(m *core.Message) error {
	if m.ID == "" {
		m.ID = newID("msg")
	}
	m.CreatedAt = m.CreatedAt.UTC()
	insert := `INSERT INTO messages(id,conversation_id,run_id,role,parts,seq,created_at)
		SELECT ?,?,?,?,?,COALESCE(MAX(seq),0)+1,? FROM messages WHERE conversation_id=?`
	for attempt := 0; attempt < 8; attempt++ {
		_, err := s.db.Exec(s.q(insert), m.ID, m.ConversationID, m.RunID, string(m.Role), marshalJSON(m.Parts), nowISO(), m.ConversationID)
		if err == nil {
			s.TouchConversation(m.ConversationID)
			return nil
		}
		if !isUniqueViolation(err) {
			return err
		}
		// 并发同会话追加冲突（PG 多连接竞态）：短抖动退避后重算序号。
		time.Sleep(time.Duration(attempt) * time.Millisecond)
	}
	return fmt.Errorf("store: message seq conflict retries exhausted (conversation %s)", m.ConversationID)
}

func (s *Store) ListMessages(conversationID string) ([]core.Message, error) {
	rows, err := s.db.Query(s.q(`SELECT id,conversation_id,run_id,role,parts,created_at FROM messages WHERE conversation_id=? ORDER BY seq`), conversationID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.Message{}
	for rows.Next() {
		var m core.Message
		var parts, ca string
		if err := rows.Scan(&m.ID, &m.ConversationID, &m.RunID, &m.Role, &parts, &ca); err != nil {
			return nil, err
		}
		m.Role = core.Role(m.Role)
		m.Parts, _ = unmarshalJSON[[]core.Part](parts)
		m.CreatedAt, _ = parseTime(ca)
		out = append(out, m)
	}
	return out, rows.Err()
}

func (s *Store) LastMessages(conversationID string, n int) ([]core.Message, error) {
	all, err := s.ListMessages(conversationID)
	if err != nil {
		return nil, err
	}
	if len(all) > n {
		all = all[len(all)-n:]
	}
	return all, nil
}

// ---- Runs（owner 冗余自会话；GetRun 面向 API 带租户过滤，内部路径用 getRunByID）----

func (s *Store) CreateRun(r *core.Run) error {
	if r.ID == "" {
		r.ID = newID("run")
	}
	r.CreatedAt = r.CreatedAt.UTC()
	_, err := s.db.Exec(s.q(`INSERT INTO runs(id,parent_run_id,conversation_id,project_id,owner,kind,profile_id,status,phase,agent_name,created_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)`),
		r.ID, r.ParentRunID, r.ConversationID, r.ProjectID, r.Owner, string(r.Kind), r.ProfileID, string(r.Status), string(r.Phase), r.AgentName, nowISO())
	return err
}

func (s *Store) UpdateRun(r *core.Run) error {
	var ended any
	if r.EndedAt != nil {
		ended = nowISO()
	}
	_, err := s.db.Exec(s.q(`UPDATE runs SET status=?, phase=?, usage=?, error=?, ended_at=? WHERE id=?`),
		string(r.Status), string(r.Phase), marshalJSON(r.Usage), marshalJSON(r.Error), ended, r.ID)
	return err
}

const runColumns = `id,parent_run_id,conversation_id,project_id,owner,kind,profile_id,status,phase,agent_name,created_at,ended_at,usage,error`

func scanRun(scan func(...any) error) (core.Run, error) {
	var r core.Run
	var parent, owner, profile, phase, agent, ended, usage, errJSON *string
	var kind, status string
	var ca string
	if err := scan(&r.ID, &parent, &r.ConversationID, &r.ProjectID, &owner, &kind, &profile, &status, &phase, &agent, &ca, &ended, &usage, &errJSON); err != nil {
		return r, err
	}
	if parent != nil {
		r.ParentRunID = *parent
	}
	if owner != nil && *owner != "" {
		r.Owner = *owner
	}
	if profile != nil {
		r.ProfileID = *profile
	}
	r.Kind = core.RunKind(kind)
	r.Status = core.RunStatus(status)
	if phase != nil {
		r.Phase = core.RunPhase(*phase)
	}
	if agent != nil {
		r.AgentName = *agent
	}
	r.CreatedAt, _ = parseTime(ca)
	if ended != nil {
		t, _ := parseTime(*ended)
		r.EndedAt = &t
	}
	if usage != nil && *usage != "" && *usage != "null" {
		r.Usage, _ = unmarshalJSON[*core.Usage](*usage)
	}
	if errJSON != nil && *errJSON != "" && *errJSON != "null" {
		r.Error, _ = unmarshalJSON[*core.ErrorPart](*errJSON)
	}
	return r, nil
}

// GetRun 面向 API 的租户过滤读取（跨租户 = NotFound）。
func (s *Store) GetRun(owner, id string) (core.Run, error) {
	r, err := scanRun(func(dst ...any) error {
		return s.db.QueryRow(s.q(`SELECT `+runColumns+` FROM runs WHERE id=? AND owner=?`), id, owner).Scan(dst...)
	})
	if errors.Is(err, sql.ErrNoRows) {
		return r, ErrNotFound
	}
	return r, err
}

// GetRunByID 内部读取，无租户过滤（runmgr 生命周期管理；API 层必须用 GetRun）。
func (s *Store) GetRunByID(id string) (core.Run, error) {
	r, err := scanRun(func(dst ...any) error {
		return s.db.QueryRow(s.q(`SELECT `+runColumns+` FROM runs WHERE id=?`), id).Scan(dst...)
	})
	if errors.Is(err, sql.ErrNoRows) {
		return r, ErrNotFound
	}
	return r, err
}

func (s *Store) ListRunsByConversation(conversationID string) ([]core.Run, error) {
	rows, err := s.db.Query(s.q(`SELECT `+runColumns+` FROM runs WHERE conversation_id=? ORDER BY created_at`), conversationID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.Run{}
	for rows.Next() {
		r, err := scanRun(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// ChildRuns 返回直接子 run（任务树）。
func (s *Store) ChildRuns(runID string) ([]core.Run, error) {
	rows, err := s.db.Query(s.q(`SELECT `+runColumns+` FROM runs WHERE parent_run_id=? ORDER BY created_at`), runID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.Run{}
	for rows.Next() {
		r, err := scanRun(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// ---- 事件（会话内连续 seq 游标权威；批量落盘降低写放大）----

// AppendEvents 在单事务内追加一批事件并分配连续 seq（base+1..base+n），
// 成功后回填 evs[i].Seq。事务内锁定会话序列：PG 用事务级咨询锁（跨实例安全），
// SQLite 单写连接天然串行。空批次为空操作。
func (s *Store) AppendEvents(conversationID string, evs []core.StreamEvent) (int64, error) {
	if len(evs) == 0 {
		return 0, nil
	}
	tx, err := s.db.Begin()
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()
	if err := s.dialect.lockConversation(tx, conversationID); err != nil {
		return 0, err
	}
	var base sql.NullInt64
	if err := tx.QueryRow(s.q(`SELECT COALESCE(MAX(seq),0) FROM events WHERE conversation_id=?`), conversationID).Scan(&base); err != nil {
		return 0, err
	}
	var sb strings.Builder
	sb.WriteString(`INSERT INTO events(conversation_id,seq,run_id,ts,v,type,payload) VALUES `)
	args := make([]any, 0, len(evs)*7)
	for i := range evs {
		if i > 0 {
			sb.WriteString(",")
		}
		sb.WriteString("(?,?,?,?,?,?,?)")
		var payload any
		if len(evs[i].Payload) > 0 {
			payload = []byte(evs[i].Payload)
		}
		args = append(args, conversationID, base.Int64+int64(i)+1, evs[i].RunID, evs[i].TS, evs[i].V, string(evs[i].Type), payload)
	}
	if _, err := tx.Exec(s.q(sb.String()), args...); err != nil {
		return 0, err
	}
	if err := tx.Commit(); err != nil {
		return 0, err
	}
	for i := range evs {
		evs[i].Seq = base.Int64 + int64(i) + 1
	}
	return base.Int64, nil
}

func (s *Store) EventsAfter(conversationID string, since int64, limit int) ([]core.StreamEvent, error) {
	rows, err := s.db.Query(s.q(`SELECT seq,run_id,ts,v,type,payload FROM events WHERE conversation_id=? AND seq>? ORDER BY seq LIMIT ?`),
		conversationID, since, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.StreamEvent{}
	for rows.Next() {
		var e core.StreamEvent
		var typ string
		var payload []byte
		if err := rows.Scan(&e.Seq, &e.RunID, &e.TS, &e.V, &typ, &payload); err != nil {
			return nil, err
		}
		e.ConversationID = conversationID
		e.Type = core.StreamEventType(typ)
		if len(payload) > 0 {
			e.Payload = payload
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

// EventAt 读取单个事件（PG 档 NOTIFY 扇出的按序号取回路径）。
func (s *Store) EventAt(conversationID string, seq int64) (core.StreamEvent, error) {
	var e core.StreamEvent
	var typ string
	var payload []byte
	err := s.db.QueryRow(s.q(`SELECT seq,run_id,ts,v,type,payload FROM events WHERE conversation_id=? AND seq=?`), conversationID, seq).
		Scan(&e.Seq, &e.RunID, &e.TS, &e.V, &typ, &payload)
	if errors.Is(err, sql.ErrNoRows) {
		return e, ErrNotFound
	}
	if err != nil {
		return e, err
	}
	e.ConversationID = conversationID
	e.Type = core.StreamEventType(typ)
	if len(payload) > 0 {
		e.Payload = payload
	}
	return e, nil
}

func (s *Store) LastEventSeq(conversationID string) (int64, error) {
	var seq sql.NullInt64
	err := s.db.QueryRow(s.q(`SELECT MAX(seq) FROM events WHERE conversation_id=?`), conversationID).Scan(&seq)
	if err != nil {
		return 0, err
	}
	return seq.Int64, nil
}

func newID(kind string) string {
	b := make([]byte, 6)
	rand.Read(b)
	return kind + "_" + hex.EncodeToString(b)
}

func parseTime(s string) (time.Time, error) {
	return time.Parse(time.RFC3339Nano, s)
}

// NewID 生成带前缀的随机 ID（导出供 runmgr/server 使用）。
func NewID(kind string) string { return newID(kind) }