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

package store

import (
	"database/sql"
	"fmt"
	"time"

	core "agent-system-core"
)

// ---- 插件授权与自有数据版本 ----

func (s *Store) SetPluginEnabled(owner, pluginID string, enabled bool) error {
	flag := 0
	if enabled {
		flag = 1
	}
	_, err := s.db.Exec(s.q(`INSERT INTO plugin_grants(owner,plugin_id,enabled) VALUES(?,?,?)
		ON CONFLICT(owner, plugin_id) DO UPDATE SET enabled=excluded.enabled`), owner, pluginID, flag)
	return err
}

func (s *Store) PluginEnabled(owner, pluginID string) (bool, error) {
	var flag int
	err := s.db.QueryRow(s.q(`SELECT enabled FROM plugin_grants WHERE owner=? AND plugin_id=?`), owner, pluginID).Scan(&flag)
	if err == sql.ErrNoRows {
		return true, nil
	}
	if err != nil {
		return false, err
	}
	return flag != 0, nil
}

func (s *Store) PluginSchemaVersion(pluginID string) (int, error) {
	var v int
	err := s.db.QueryRow(s.q(`SELECT version FROM plugin_schema WHERE plugin_id=?`), pluginID).Scan(&v)
	if err == sql.ErrNoRows {
		return 0, nil
	}
	return v, err
}

func (s *Store) SetPluginSchemaVersion(pluginID string, version int) error {
	now := time.Now().UTC().Format(time.RFC3339)
	_, err := s.db.Exec(s.q(`INSERT INTO plugin_schema(plugin_id,version,updated_at) VALUES(?,?,?)
		ON CONFLICT(plugin_id) DO UPDATE SET version=excluded.version, updated_at=excluded.updated_at`), pluginID, version, now)
	return err
}

func (s *Store) PutPluginKV(pluginID, owner, key, value string) error {
	_, err := s.db.Exec(s.q(`INSERT INTO plugin_kv(plugin_id,owner,k,v) VALUES(?,?,?,?)
		ON CONFLICT(plugin_id, owner, k) DO UPDATE SET v=excluded.v`), pluginID, owner, key, value)
	return err
}

func (s *Store) GetPluginKV(pluginID, owner, key string) (string, bool, error) {
	var v string
	err := s.db.QueryRow(s.q(`SELECT v FROM plugin_kv WHERE plugin_id=? AND owner=? AND k=?`), pluginID, owner, key).Scan(&v)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	return v, err == nil, err
}

// ---- 后台任务（与会话锁无关）----

type BackgroundTask struct {
	ID             string `json:"id"`
	Owner          string `json:"owner"`
	Kind           string `json:"kind"`
	Title          string `json:"title"`
	Status         string `json:"status"`
	Progress       int    `json:"progress"`
	Error          string `json:"error,omitempty"`
	ConversationID string `json:"conversationId,omitempty"`
	CreatedAt      string `json:"createdAt"`
	UpdatedAt      string `json:"updatedAt"`
}

func (s *Store) CreateBackgroundTask(owner, kind, title, conversationID string) (BackgroundTask, error) {
	now := time.Now().UTC().Format(time.RFC3339)
	t := BackgroundTask{
		ID: NewID("task"), Owner: owner, Kind: kind, Title: title, Status: "queued",
		ConversationID: conversationID, CreatedAt: now, UpdatedAt: now,
	}
	_, err := s.db.Exec(s.q(`INSERT INTO background_tasks(id,owner,kind,title,status,progress,error,conversation_id,created_at,updated_at)
		VALUES(?,?,?,?,?,?,?,?,?,?)`), t.ID, t.Owner, t.Kind, t.Title, t.Status, t.Progress, t.Error, t.ConversationID, t.CreatedAt, t.UpdatedAt)
	return t, err
}

func (s *Store) UpdateBackgroundTask(owner, id, status string, progress int, errText string) error {
	now := time.Now().UTC().Format(time.RFC3339)
	res, err := s.db.Exec(s.q(`UPDATE background_tasks SET status=?, progress=?, error=?, updated_at=? WHERE id=? AND owner=?`),
		status, progress, errText, now, id, owner)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) GetBackgroundTask(owner, id string) (BackgroundTask, error) {
	var t BackgroundTask
	err := s.db.QueryRow(s.q(`SELECT id,owner,kind,title,status,progress,error,conversation_id,created_at,updated_at
		FROM background_tasks WHERE id=? AND owner=?`), id, owner).Scan(
		&t.ID, &t.Owner, &t.Kind, &t.Title, &t.Status, &t.Progress, &t.Error, &t.ConversationID, &t.CreatedAt, &t.UpdatedAt)
	if err == sql.ErrNoRows {
		return t, ErrNotFound
	}
	return t, err
}

func (s *Store) ListBackgroundTasks(owner string) ([]BackgroundTask, error) {
	rows, err := s.db.Query(s.q(`SELECT id,owner,kind,title,status,progress,error,conversation_id,created_at,updated_at
		FROM background_tasks WHERE owner=? ORDER BY created_at`), owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []BackgroundTask
	for rows.Next() {
		var t BackgroundTask
		if err := rows.Scan(&t.ID, &t.Owner, &t.Kind, &t.Title, &t.Status, &t.Progress, &t.Error, &t.ConversationID, &t.CreatedAt, &t.UpdatedAt); err != nil {
			return nil, err
		}
		out = append(out, t)
	}
	return out, rows.Err()
}

// ---- 记忆 ----

func (s *Store) PutMemory(owner string, e core.MemoryEntry) (core.MemoryEntry, error) {
	if e.ID == "" {
		e.ID = NewID("mem")
	}
	now := time.Now().UTC()
	if e.CreatedAt.IsZero() {
		e.CreatedAt = now
	}
	e.LastAccessed = now
	exp := ""
	if e.ExpiresAt != nil {
		exp = e.ExpiresAt.UTC().Format(time.RFC3339)
	}
	_, err := s.db.Exec(s.q(`INSERT INTO memory_entries(id,owner,scope,project_id,typ,content,source,confidence,created_at,last_accessed,expires_at)
		VALUES(?,?,?,?,?,?,?,?,?,?,?)`),
		e.ID, owner, e.Scope, string(e.ProjectID), e.Type, e.Content, e.Source, e.Confidence,
		e.CreatedAt.UTC().Format(time.RFC3339), e.LastAccessed.UTC().Format(time.RFC3339), exp)
	return e, err
}

func (s *Store) ListMemory(owner, scope, projectID string) ([]core.MemoryEntry, error) {
	rows, err := s.db.Query(s.q(`SELECT id,scope,project_id,typ,content,source,confidence,created_at,last_accessed,expires_at
		FROM memory_entries WHERE owner=? AND (scope=? OR ?='') AND (project_id=? OR ?='') ORDER BY created_at`),
		owner, scope, scope, projectID, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []core.MemoryEntry
	for rows.Next() {
		var e core.MemoryEntry
		var project, created, accessed, exp string
		if err := rows.Scan(&e.ID, &e.Scope, &project, &e.Type, &e.Content, &e.Source, &e.Confidence, &created, &accessed, &exp); err != nil {
			return nil, err
		}
		e.ProjectID = core.ID(project)
		e.CreatedAt, _ = time.Parse(time.RFC3339, created)
		e.LastAccessed, _ = time.Parse(time.RFC3339, accessed)
		if exp != "" {
			if t, err := time.Parse(time.RFC3339, exp); err == nil {
				e.ExpiresAt = &t
			}
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

// ---- 通用审计 ----

type AuditEvent struct {
	ID       int64  `json:"id"`
	Time     string `json:"time"`
	Owner    string `json:"owner"`
	Action   string `json:"action"`
	Actor    string `json:"actor,omitempty"`
	Target   string `json:"target,omitempty"`
	Evidence string `json:"evidence,omitempty"`
	TraceID  string `json:"traceId,omitempty"`
}

func (s *Store) AppendAuditEvent(ev AuditEvent) error {
	if ev.Time == "" {
		ev.Time = time.Now().UTC().Format(time.RFC3339)
	}
	_, err := s.db.Exec(s.q(`INSERT INTO audit_events(time,owner,action,actor,target,evidence,trace_id) VALUES(?,?,?,?,?,?,?)`),
		ev.Time, ev.Owner, ev.Action, ev.Actor, ev.Target, ev.Evidence, ev.TraceID)
	return err
}

func (s *Store) ListAuditEvents(owner string, limit int) ([]AuditEvent, error) {
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	rows, err := s.db.Query(s.q(`SELECT id,time,owner,action,actor,target,evidence,trace_id FROM audit_events
		WHERE owner=? OR ?='' ORDER BY id DESC LIMIT ?`), owner, owner, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []AuditEvent
	for rows.Next() {
		var ev AuditEvent
		if err := rows.Scan(&ev.ID, &ev.Time, &ev.Owner, &ev.Action, &ev.Actor, &ev.Target, &ev.Evidence, &ev.TraceID); err != nil {
			return nil, err
		}
		out = append(out, ev)
	}
	return out, rows.Err()
}

// ---- 危险命令确认挂起态 ----

func (s *Store) SavePendingConfirm(cardID, owner, payload string) error {
	now := time.Now().UTC().Format(time.RFC3339)
	_, err := s.db.Exec(s.q(`INSERT INTO pending_confirms(card_id,owner,payload,created_at) VALUES(?,?,?,?)
		ON CONFLICT(card_id) DO UPDATE SET owner=excluded.owner, payload=excluded.payload, created_at=excluded.created_at`),
		cardID, owner, payload, now)
	return err
}

func (s *Store) TakePendingConfirm(cardID string) (owner, payload string, ok bool, err error) {
	tx, err := s.db.Begin()
	if err != nil {
		return "", "", false, err
	}
	defer tx.Rollback()
	err = tx.QueryRow(s.q(`SELECT owner, payload FROM pending_confirms WHERE card_id=?`), cardID).Scan(&owner, &payload)
	if err == sql.ErrNoRows {
		return "", "", false, nil
	}
	if err != nil {
		return "", "", false, err
	}
	if _, err := tx.Exec(s.q(`DELETE FROM pending_confirms WHERE card_id=?`), cardID); err != nil {
		return "", "", false, err
	}
	if err := tx.Commit(); err != nil {
		return "", "", false, err
	}
	return owner, payload, true, nil
}

// ---- LLM 日预算 ----

func (s *Store) AddDailyTokens(owner, day string, tokens int) (int, error) {
	if day == "" {
		day = time.Now().UTC().Format("2006-01-02")
	}
	if _, err := s.db.Exec(s.q(`INSERT INTO llm_usage_daily(owner,day,tokens) VALUES(?,?,?)
		ON CONFLICT(owner, day) DO UPDATE SET tokens=llm_usage_daily.tokens + excluded.tokens`), owner, day, tokens); err != nil {
		return 0, err
	}
	var total int
	err := s.db.QueryRow(s.q(`SELECT tokens FROM llm_usage_daily WHERE owner=? AND day=?`), owner, day).Scan(&total)
	return total, err
}

// BudgetStatus 是一次记账的结果。
type BudgetStatus struct {
	Total int
	Over  bool
}

func (s *Store) AccountTokens(owner string, tokens, budget int) (BudgetStatus, error) {
	total, err := s.AddDailyTokens(owner, "", tokens)
	if err != nil {
		return BudgetStatus{}, fmt.Errorf("llm usage: %w", err)
	}
	return BudgetStatus{Total: total, Over: budget > 0 && total > budget}, nil
}