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

package store

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"
)

// Secret 是一把带标签的服务密钥。列表接口不返回 Value。
type Secret struct {
	ID        string `json:"id"`
	Owner     string `json:"owner,omitempty"`
	Service   string `json:"service"`
	Label     string `json:"label"`
	Active    bool   `json:"active"`
	Mask      string `json:"mask"`
	CreatedAt string `json:"createdAt"`
	Value     string `json:"-"`
}

// ContextEntry 是一条条件注入。常驻或关键词命中后附在上下文末尾。
type ContextEntry struct {
	ID             string   `json:"id"`
	Owner          string   `json:"owner,omitempty"`
	ProjectID      string   `json:"projectId,omitempty"`
	ConversationID string   `json:"conversationId,omitempty"`
	Keys           []string `json:"keys"`
	Secondary      string   `json:"secondary,omitempty"`
	Logic          string   `json:"logic"`
	Content        string   `json:"content"`
	Constant       bool     `json:"constant"`
	ScanDepth      int      `json:"scanDepth"`
	SortOrder      int      `json:"sortOrder"`
	BudgetChars    int      `json:"budgetChars"`
	Enabled        bool     `json:"enabled"`
}

// UserCommand 是用户自定义的 / 命令，展开成提示词，不自动触发。
type UserCommand struct {
	ID        string `json:"id"`
	Owner     string `json:"owner,omitempty"`
	Name      string `json:"name"`
	Title     string `json:"title"`
	Template  string `json:"template"`
	CreatedAt string `json:"createdAt"`
}

func maskSecret(v string) string {
	v = strings.TrimSpace(v)
	if len(v) <= 4 {
		return "••••••"
	}
	return "••••" + v[len(v)-4:]
}

func (s *Store) ListSecrets(owner string) ([]Secret, error) {
	rows, err := s.db.Query(s.q(`SELECT id,service,label,value,active,created_at FROM secret_keys WHERE owner=? ORDER BY service, created_at`), owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []Secret{}
	for rows.Next() {
		var sec Secret
		var active int
		if err := rows.Scan(&sec.ID, &sec.Service, &sec.Label, &sec.Value, &active, &sec.CreatedAt); err != nil {
			return nil, err
		}
		sec.Active = active == 1
		sec.Mask = maskSecret(sec.Value)
		sec.Value = ""
		out = append(out, sec)
	}
	return out, rows.Err()
}

// ActiveSecret 返回该服务当前生效的明文。没有记录时 ok 为 false。
func (s *Store) ActiveSecret(owner, service string) (string, bool) {
	var v string
	err := s.db.QueryRow(s.q(`SELECT value FROM secret_keys WHERE owner=? AND service=? AND active=1 ORDER BY created_at DESC LIMIT 1`), owner, service).Scan(&v)
	if err != nil || strings.TrimSpace(v) == "" {
		return "", false
	}
	return v, true
}

// AddSecret 追加一把密钥并把它设为该服务唯一生效项。
func (s *Store) AddSecret(owner, service, label, value string) (Secret, error) {
	service = strings.TrimSpace(service)
	value = strings.TrimSpace(value)
	label = strings.TrimSpace(label)
	if service == "" || value == "" {
		return Secret{}, fmt.Errorf("store: secret service and value are required")
	}
	if label == "" {
		label = "未命名"
	}
	tx, err := s.db.Begin()
	if err != nil {
		return Secret{}, err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(s.q(`UPDATE secret_keys SET active=0 WHERE owner=? AND service=?`), owner, service); err != nil {
		return Secret{}, err
	}
	sec := Secret{ID: newID("key"), Service: service, Label: label, Active: true, CreatedAt: nowISO(), Mask: maskSecret(value)}
	if _, err := tx.Exec(s.q(`INSERT INTO secret_keys(id,owner,service,label,value,active,created_at) VALUES(?,?,?,?,?,?,?)`),
		sec.ID, owner, service, label, value, 1, sec.CreatedAt); err != nil {
		return Secret{}, err
	}
	return sec, tx.Commit()
}

func (s *Store) ActivateSecret(owner, id string) error {
	var service string
	err := s.db.QueryRow(s.q(`SELECT service FROM secret_keys WHERE id=? AND owner=?`), id, owner).Scan(&service)
	if err != nil {
		return ErrNotFound
	}
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(s.q(`UPDATE secret_keys SET active=0 WHERE owner=? AND service=?`), owner, service); err != nil {
		return err
	}
	if _, err := tx.Exec(s.q(`UPDATE secret_keys SET active=1 WHERE id=? AND owner=?`), id, owner); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *Store) RenameSecret(owner, id, label string) error {
	label = strings.TrimSpace(label)
	if label == "" {
		return fmt.Errorf("store: label required")
	}
	res, err := s.db.Exec(s.q(`UPDATE secret_keys SET label=? WHERE id=? AND owner=?`), label, id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) DeleteSecret(owner, id string) error {
	res, err := s.db.Exec(s.q(`DELETE FROM secret_keys WHERE id=? AND owner=?`), id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) HasSecret(owner, service string) (bool, error) {
	var n int
	err := s.db.QueryRow(s.q(`SELECT COUNT(1) FROM secret_keys WHERE owner=? AND service=?`), owner, service).Scan(&n)
	return n > 0, err
}

// SaveSnapshotIfChanged 在正文与最近一份不同时留档。相同则跳过。
func (s *Store) SaveSnapshotIfChanged(owner, scope, body, reason string) (bool, error) {
	var prev string
	err := s.db.QueryRow(s.q(`SELECT body FROM settings_snapshots WHERE owner=? AND scope=? ORDER BY created_at DESC LIMIT 1`), owner, scope).Scan(&prev)
	if err == nil && prev == body {
		return false, nil
	}
	if err != nil && err != sql.ErrNoRows {
		return false, err
	}
	_, err = s.db.Exec(s.q(`INSERT INTO settings_snapshots(id,owner,scope,body,reason,created_at) VALUES(?,?,?,?,?,?)`),
		newID("snap"), owner, scope, body, reason, nowISO())
	return err == nil, err
}

// LatestSnapshotAge 返回最近一份快照的时间。没有快照时 ok 为 false。
func (s *Store) LatestSnapshotAt(owner, scope string) (time.Time, bool) {
	var raw string
	err := s.db.QueryRow(s.q(`SELECT created_at FROM settings_snapshots WHERE owner=? AND scope=? ORDER BY created_at DESC LIMIT 1`), owner, scope).Scan(&raw)
	if err != nil {
		return time.Time{}, false
	}
	t, err := time.Parse(time.RFC3339Nano, raw)
	if err != nil {
		return time.Time{}, false
	}
	return t, true
}

func (s *Store) SetSessionNote(owner, id, note string) error {
	res, err := s.db.Exec(s.q(`UPDATE conversations SET session_note=?, updated_at=? WHERE id=? AND owner=?`), note, nowISO(), id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) SetMessageFlags(id string, excluded bool, group string, index int) error {
	ex := 0
	if excluded {
		ex = 1
	}
	_, err := s.db.Exec(s.q(`UPDATE messages SET excluded=?, variant_group=?, variant_index=? WHERE id=?`), ex, group, index, id)
	return err
}

// ExcludeRunGroup 把这条消息以及同一 run 的消息一起标成排除或恢复。
// 工具调用和结果落在同一条助手消息里，同一 run 的其它落库消息一并处理。
func (s *Store) ExcludeRunGroup(id string, excluded bool) error {
	var runID string
	err := s.db.QueryRow(s.q(`SELECT run_id FROM messages WHERE id=?`), id).Scan(&runID)
	if err != nil {
		return ErrNotFound
	}
	ex := 0
	if excluded {
		ex = 1
	}
	if _, err := s.db.Exec(s.q(`UPDATE messages SET excluded=? WHERE id=?`), ex, id); err != nil {
		return err
	}
	if runID == "" {
		return nil
	}
	_, err = s.db.Exec(s.q(`UPDATE messages SET excluded=? WHERE run_id=?`), ex, runID)
	return err
}

// SelectVariant 让同一组里只有这条候选进入上下文。
func (s *Store) SelectVariant(id string) error {
	var group, conv string
	err := s.db.QueryRow(s.q(`SELECT variant_group, conversation_id FROM messages WHERE id=?`), id).Scan(&group, &conv)
	if err != nil {
		return ErrNotFound
	}
	if group == "" {
		return s.ExcludeRunGroup(id, false)
	}
	if _, err := s.db.Exec(s.q(`UPDATE messages SET excluded=1 WHERE conversation_id=? AND variant_group=?`), conv, group); err != nil {
		return err
	}
	_, err = s.db.Exec(s.q(`UPDATE messages SET excluded=0 WHERE id=?`), id)
	return err
}

func (s *Store) SetRegenPending(conversationID, group string, index int) error {
	res, err := s.db.Exec(s.q(`UPDATE regen_pending SET variant_group=?, variant_index=? WHERE conversation_id=?`), group, index, conversationID)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n > 0 {
		return nil
	}
	_, err = s.db.Exec(s.q(`INSERT INTO regen_pending(conversation_id, variant_group, variant_index) VALUES(?,?,?)`), conversationID, group, index)
	return err
}

func (s *Store) TakeRegenPending(conversationID string) (group string, index int, ok bool) {
	err := s.db.QueryRow(s.q(`SELECT variant_group, variant_index FROM regen_pending WHERE conversation_id=?`), conversationID).Scan(&group, &index)
	if err != nil {
		return "", 0, false
	}
	_, _ = s.db.Exec(s.q(`DELETE FROM regen_pending WHERE conversation_id=?`), conversationID)
	return group, index, true
}

func (s *Store) ListContextEntries(owner, projectID, conversationID string) ([]ContextEntry, error) {
	rows, err := s.db.Query(s.q(`SELECT id,project_id,conversation_id,keys_json,secondary,logic,content,constant,scan_depth,sort_order,budget_chars,enabled
		FROM context_entries WHERE owner=? ORDER BY sort_order, created_at`), owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []ContextEntry{}
	for rows.Next() {
		e, err := scanEntry(rows)
		if err != nil {
			return nil, err
		}
		if e.ProjectID != "" && projectID != "" && e.ProjectID != projectID {
			continue
		}
		if e.ConversationID != "" && conversationID != "" && e.ConversationID != conversationID {
			continue
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

func scanEntry(rows interface{ Scan(...any) error }) (ContextEntry, error) {
	var e ContextEntry
	var keys, logic string
	var constant, enabled int
	err := rows.Scan(&e.ID, &e.ProjectID, &e.ConversationID, &keys, &e.Secondary, &logic, &e.Content, &constant, &e.ScanDepth, &e.SortOrder, &e.BudgetChars, &enabled)
	if err != nil {
		return e, err
	}
	_ = json.Unmarshal([]byte(keys), &e.Keys)
	if e.Keys == nil {
		e.Keys = []string{}
	}
	e.Logic = logic
	if e.Logic == "" {
		e.Logic = "and"
	}
	e.Constant = constant == 1
	e.Enabled = enabled == 1
	return e, nil
}

func (s *Store) SaveContextEntry(owner string, e ContextEntry) (ContextEntry, error) {
	if strings.TrimSpace(e.Content) == "" {
		return ContextEntry{}, fmt.Errorf("store: entry content required")
	}
	if e.Logic != "and" && e.Logic != "or" && e.Logic != "not" {
		e.Logic = "and"
	}
	if e.ScanDepth <= 0 {
		e.ScanDepth = 4
	}
	if e.BudgetChars <= 0 {
		e.BudgetChars = 2000
	}
	raw, _ := json.Marshal(e.Keys)
	constant, enabled := 0, 0
	if e.Constant {
		constant = 1
	}
	if e.Enabled || e.ID == "" {
		enabled = 1
		e.Enabled = true
	}
	if e.ID == "" {
		e.ID = newID("ctx")
		_, err := s.db.Exec(s.q(`INSERT INTO context_entries(id,owner,project_id,conversation_id,keys_json,secondary,logic,content,constant,scan_depth,sort_order,budget_chars,enabled,created_at)
			VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)`),
			e.ID, owner, e.ProjectID, e.ConversationID, string(raw), e.Secondary, e.Logic, e.Content, constant, e.ScanDepth, e.SortOrder, e.BudgetChars, enabled, nowISO())
		return e, err
	}
	res, err := s.db.Exec(s.q(`UPDATE context_entries SET project_id=?, conversation_id=?, keys_json=?, secondary=?, logic=?, content=?, constant=?, scan_depth=?, sort_order=?, budget_chars=?, enabled=? WHERE id=? AND owner=?`),
		e.ProjectID, e.ConversationID, string(raw), e.Secondary, e.Logic, e.Content, constant, e.ScanDepth, e.SortOrder, e.BudgetChars, enabled, e.ID, owner)
	if err != nil {
		return e, err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return e, ErrNotFound
	}
	return e, nil
}

func (s *Store) DeleteContextEntry(owner, id string) error {
	res, err := s.db.Exec(s.q(`DELETE FROM context_entries WHERE id=? AND owner=?`), id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) ListUserCommands(owner string) ([]UserCommand, error) {
	rows, err := s.db.Query(s.q(`SELECT id,name,title,template,created_at FROM user_commands WHERE owner=? ORDER BY name`), owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []UserCommand{}
	for rows.Next() {
		var c UserCommand
		if err := rows.Scan(&c.ID, &c.Name, &c.Title, &c.Template, &c.CreatedAt); err != nil {
			return nil, err
		}
		out = append(out, c)
	}
	return out, rows.Err()
}

func (s *Store) SaveUserCommand(owner string, c UserCommand) (UserCommand, error) {
	c.Name = strings.TrimSpace(c.Name)
	if !strings.HasPrefix(c.Name, "/") {
		c.Name = "/" + c.Name
	}
	if c.Title == "" {
		c.Title = c.Name
	}
	if strings.TrimSpace(c.Template) == "" {
		return UserCommand{}, fmt.Errorf("store: command template required")
	}
	if c.ID == "" {
		c.ID = newID("cmd")
		c.CreatedAt = nowISO()
		_, err := s.db.Exec(s.q(`INSERT INTO user_commands(id,owner,name,title,template,created_at) VALUES(?,?,?,?,?,?)`),
			c.ID, owner, c.Name, c.Title, c.Template, c.CreatedAt)
		return c, err
	}
	res, err := s.db.Exec(s.q(`UPDATE user_commands SET name=?, title=?, template=? WHERE id=? AND owner=?`), c.Name, c.Title, c.Template, c.ID, owner)
	if err != nil {
		return c, err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return c, ErrNotFound
	}
	return c, nil
}

func (s *Store) DeleteUserCommand(owner, id string) error {
	res, err := s.db.Exec(s.q(`DELETE FROM user_commands WHERE id=? AND owner=?`), id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

// MessageByID 读取一条消息。不存在时返回 ErrNotFound。
func (s *Store) MessageByID(id string) (core.Message, error) {
	var m core.Message
	var parts, ca string
	var excluded, index int
	err := s.db.QueryRow(s.q(`SELECT id,conversation_id,run_id,role,parts,created_at,excluded,variant_group,variant_index FROM messages WHERE id=?`), id).
		Scan(&m.ID, &m.ConversationID, &m.RunID, &m.Role, &parts, &ca, &excluded, &m.VariantGroup, &index)
	if err != nil {
		return m, ErrNotFound
	}
	m.Role = core.Role(m.Role)
	m.Parts, _ = unmarshalJSON[[]core.Part](parts)
	m.CreatedAt, _ = parseTime(ca)
	m.Excluded = excluded == 1
	m.VariantIndex = index
	return m, nil
}