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

package store

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

	core "agent-system-core"
)

// ---- 工件（元数据在库，内容在对象目录 data/artifacts/<id>）----
// 文件写走 临时文件 → rename → 落库：崩溃残留只可能产生无主孤立文件，
// 不产生库里有行、对象缺失的悬空引用。

func (s *Store) CreateArtifact(a *core.Artifact, content []byte) error {
	if a.ID == "" {
		a.ID = newID("art")
	}
	now := time.Now().UTC()
	a.CreatedAt, a.UpdatedAt = now, now
	a.Size = int64(len(content))
	ref, err := s.blobs.Put(a.ID, content)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(s.q(`INSERT INTO artifacts(id,project_id,owner,name,kind,mime,size,source_run,content_path,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)`),
		a.ID, a.ProjectID, a.Owner, a.Name, a.Kind, a.MIME, a.Size, a.SourceRun, ref, nowISO(), nowISO())
	if err != nil {
		_ = s.blobs.Delete(ref) // 落库失败：回收对象，不留悬空
		return err
	}
	return nil
}

func (s *Store) UpdateArtifactContent(owner, id string, content []byte) error {
	var path string
	err := s.db.QueryRow(s.q(`SELECT content_path FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(&path)
	if errors.Is(err, sql.ErrNoRows) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}
	// 同 ID 原子覆写（临时文件+rename）；行内引用不变，避免行/对象漂移。
	if _, err := s.blobs.Put(id, content); err != nil {
		return err
	}
	res, err := s.db.Exec(s.q(`UPDATE artifacts SET size=?, updated_at=? WHERE id=? AND owner=?`), len(content), nowISO(), id, owner)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

const artifactColumns = `id,project_id,owner,name,kind,mime,size,source_run,created_at,updated_at`

func scanArtifact(scan func(...any) error) (core.Artifact, error) {
	var a core.Artifact
	var owner, mime, sourceRun string
	var ca, ua string
	if err := scan(&a.ID, &a.ProjectID, &owner, &a.Name, &a.Kind, &mime, &a.Size, &sourceRun, &ca, &ua); err != nil {
		return a, err
	}
	a.Owner, a.MIME, a.SourceRun = owner, mime, sourceRun
	a.CreatedAt, _ = parseTime(ca)
	a.UpdatedAt, _ = parseTime(ua)
	return a, nil
}

func (s *Store) GetArtifact(owner, id string) (core.Artifact, error) {
	a, err := scanArtifact(func(dst ...any) error {
		return s.db.QueryRow(s.q(`SELECT `+artifactColumns+` FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(dst...)
	})
	if errors.Is(err, sql.ErrNoRows) {
		return a, ErrNotFound
	}
	return a, err
}

func (s *Store) ArtifactContent(owner, id string) ([]byte, error) {
	var ref string
	err := s.db.QueryRow(s.q(`SELECT content_path FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(&ref)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, err
	}
	return s.blobs.Get(ref)
}

func (s *Store) ListArtifacts(owner, projectID string) ([]core.Artifact, error) {
	rows, err := s.db.Query(s.q(`SELECT `+artifactColumns+` FROM artifacts WHERE project_id=? AND owner=? ORDER BY updated_at DESC`), projectID, owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.Artifact{}
	for rows.Next() {
		a, err := scanArtifact(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

// DeleteArtifact 先删行后删对象：崩溃间隙只可能残留无主文件（可 GC），不悬空引用。
func (s *Store) DeleteArtifact(owner, id string) error {
	var path string
	err := s.db.QueryRow(s.q(`SELECT content_path FROM artifacts WHERE id=? AND owner=?`), id, owner).Scan(&path)
	if errors.Is(err, sql.ErrNoRows) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}
	if _, err := s.db.Exec(s.q(`DELETE FROM artifacts WHERE id=? AND owner=?`), id, owner); err != nil {
		return err
	}
	// 联动清理检索索引（知识行 + 向量；失败不阻断删除，残留可由重建索引收敛）。
	_ = s.DeleteKnowledgeByRef(owner, id)
	_ = s.DeleteKnowledgeVectors(context.Background(), owner, id)
	return s.blobs.Delete(path)
}

// ---- HITL 卡片（resolve 原子化：条件更新消灭先查后改窗口）----

func (s *Store) SaveHITLCard(owner string, c *core.HITLCard) error {
	_, err := s.db.Exec(s.q(`INSERT INTO hitl_cards(id,owner,run_id,conversation_id,payload,created_at,resolved,resolution) VALUES(?,?,?,?,?,?,?,?)
		ON CONFLICT(id) DO UPDATE SET payload=excluded.payload, resolved=excluded.resolved, resolution=excluded.resolution`),
		c.ID, owner, c.RunID, c.ConversationID, marshalJSON(c), nowISO(), boolInt(c.Resolved), "")
	return err
}

func (s *Store) GetHITLCard(owner, id string) (core.HITLCard, error) {
	var payload string
	err := s.db.QueryRow(s.q(`SELECT payload FROM hitl_cards WHERE id=? AND owner=?`), id, owner).Scan(&payload)
	if errors.Is(err, sql.ErrNoRows) {
		return core.HITLCard{}, ErrNotFound
	}
	if err != nil {
		return core.HITLCard{}, err
	}
	return unmarshalJSON[core.HITLCard](payload)
}

// ResolveHITLCard 原子置为已解决；已解决的重复请求返回明确错误（先查后改竞态修复）。
func (s *Store) ResolveHITLCard(owner, id string, resolution core.HITLResolvedPayload) error {
	c, err := s.GetHITLCard(owner, id)
	if err != nil {
		return err
	}
	if c.Resolved {
		return fmt.Errorf("store: card %s already resolved", id)
	}
	c.Resolved = true
	res, err := s.db.Exec(s.q(`UPDATE hitl_cards SET resolved=1, payload=?, resolution=? WHERE id=? AND owner=? AND resolved=0`),
		marshalJSON(c), marshalJSON(resolution), id, owner)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("store: card %s already resolved", id)
	}
	return nil
}

func (s *Store) ActiveHITLCards(owner, conversationID string) ([]core.HITLCard, error) {
	rows, err := s.db.Query(s.q(`SELECT payload FROM hitl_cards WHERE conversation_id=? AND owner=? AND resolved=0 ORDER BY created_at`), conversationID, owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.HITLCard{}
	for rows.Next() {
		var payload string
		if err := rows.Scan(&payload); err != nil {
			return nil, err
		}
		c, _ := unmarshalJSON[core.HITLCard](payload)
		out = append(out, c)
	}
	return out, rows.Err()
}

func boolInt(b bool) int {
	if b {
		return 1
	}
	return 0
}

// ---- 计划 / 待办（会话域；归属由上游 GetConversation 校验）----

func (s *Store) SavePlan(conversationID string, p core.PlanSnapshot) error {
	_, err := s.db.Exec(s.q(`INSERT INTO plans(conversation_id,payload) VALUES(?,?)
		ON CONFLICT(conversation_id) DO UPDATE SET payload=excluded.payload`), conversationID, marshalJSON(p))
	return err
}

func (s *Store) GetPlan(conversationID string) (core.PlanSnapshot, error) {
	var payload string
	err := s.db.QueryRow(s.q(`SELECT payload FROM plans WHERE conversation_id=?`), conversationID).Scan(&payload)
	if errors.Is(err, sql.ErrNoRows) {
		return core.PlanSnapshot{}, ErrNotFound
	}
	if err != nil {
		return core.PlanSnapshot{}, err
	}
	return unmarshalJSON[core.PlanSnapshot](payload)
}

func (s *Store) SaveTodos(conversationID string, l core.TodoList) error {
	_, err := s.db.Exec(s.q(`INSERT INTO todos(conversation_id,payload) VALUES(?,?)
		ON CONFLICT(conversation_id) DO UPDATE SET payload=excluded.payload`), conversationID, marshalJSON(l))
	return err
}

func (s *Store) GetTodos(conversationID string) (core.TodoList, error) {
	var payload string
	err := s.db.QueryRow(s.q(`SELECT payload FROM todos WHERE conversation_id=?`), conversationID).Scan(&payload)
	if errors.Is(err, sql.ErrNoRows) {
		return core.TodoList{}, ErrNotFound
	}
	if err != nil {
		return core.TodoList{}, err
	}
	return unmarshalJSON[core.TodoList](payload)
}

// ---- 知识（全文检索：SQLite FTS5 / PG tsvector+trigram；owner='' 为共享知识）----

// AddKnowledge 单事务批量写入：中途失败整体回滚，不再产生半批导入。
func (s *Store) AddKnowledge(docs ...core.KnowledgeDoc) error {
	if len(docs) == 0 {
		return nil
	}
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	ts := nowISO()
	// 按下标写回生成的 ID：调用方（rag 管线）用 ID 关联向量。
	for i := range docs {
		if docs[i].ID == "" {
			docs[i].ID = newID("knw")
		}
		if _, err := tx.Exec(s.q(`INSERT INTO knowledge(id,owner,project_id,source,title,ref,chunk,created_at) VALUES(?,?,?,?,?,?,?,?)`),
			docs[i].ID, docs[i].Owner, docs[i].ProjectID, docs[i].Source, docs[i].Title, docs[i].Ref, docs[i].Chunk, ts); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// DeleteKnowledgeByRef 删除引用对应的知识（工件删除时联动清索引）。
func (s *Store) DeleteKnowledgeByRef(owner, ref string) error {
	_, err := s.db.Exec(s.q(`DELETE FROM knowledge WHERE ref=? AND owner=?`), ref, owner)
	return err
}

func (s *Store) SearchKnowledge(owner string, q core.RetrievalQuery) ([]core.RetrievalHit, error) {
	if s.dialect.name() == "postgres" {
		return s.searchKnowledgePG(owner, q)
	}
	return s.searchKnowledgeSQLite(owner, q)
}

func (s *Store) searchKnowledgeSQLite(owner string, q core.RetrievalQuery) ([]core.RetrievalHit, error) {
	topK := q.TopK
	if topK <= 0 {
		topK = 5
	}
	rows, err := s.db.Query(s.q(`
		SELECT k.id,k.owner,k.project_id,k.source,k.title,k.ref,k.chunk,k.created_at, rank
		FROM knowledge_fts f JOIN knowledge k ON k.rowid=f.rowid
		WHERE knowledge_fts MATCH ? ORDER BY rank LIMIT ?`), ftsQuery(q.Text), topK*2)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.RetrievalHit{}
	for rows.Next() {
		h, err := scanKnowledgeHit(rows)
		if err != nil {
			return nil, err
		}
		if !knowledgeVisible(owner, q, h.Doc) {
			continue
		}
		out = append(out, h)
		if len(out) >= topK {
			break
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	// CJK 兜底：unicode61 分词把中文整句当单 token，FTS 前缀匹配失效时用 LIKE 子串。
	if len(out) == 0 && q.Text != "" {
		return s.searchKnowledgeLike(owner, q, topK)
	}
	return out, nil
}

func (s *Store) searchKnowledgeLike(owner string, q core.RetrievalQuery, topK int) ([]core.RetrievalHit, error) {
	rows, err := s.db.Query(s.q(`SELECT id,owner,project_id,source,title,ref,chunk,created_at FROM knowledge WHERE (owner=? OR owner='') AND (chunk LIKE ? OR title LIKE ?) LIMIT ?`),
		owner, "%"+q.Text+"%", "%"+q.Text+"%", topK)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []core.RetrievalHit{}
	for rows.Next() {
		var d core.KnowledgeDoc
		var own, ca string
		if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca); err != nil {
			return nil, err
		}
		d.Owner = own
		d.CreatedAt, _ = parseTime(ca)
		if !knowledgeVisible(owner, q, d) {
			continue
		}
		out = append(out, core.RetrievalHit{Doc: d, Score: 1.0, Kind: "fulltext"})
	}
	return out, rows.Err()
}

// knowledgeVisible 共同过滤：租户（本人 + 共享）、项目、来源。
func knowledgeVisible(owner string, q core.RetrievalQuery, d core.KnowledgeDoc) bool {
	if d.Owner != "" && d.Owner != owner {
		return false
	}
	if q.ProjectID != "" && d.ProjectID != "" && d.ProjectID != q.ProjectID {
		return false
	}
	if len(q.Sources) > 0 && !containsStr(q.Sources, d.Source) {
		return false
	}
	return true
}

type rowScanner interface{ Scan(dest ...any) error }

func scanKnowledgeHit(rows rowScanner) (core.RetrievalHit, error) {
	var d core.KnowledgeDoc
	var own, ca string
	var score float64
	if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca, &score); err != nil {
		return core.RetrievalHit{}, err
	}
	d.Owner = own
	d.CreatedAt, _ = parseTime(ca)
	return core.RetrievalHit{Doc: d, Score: -score, Kind: "fulltext"}, nil
}

// ftsQuery 做最小安全化：按空白切词加引号，FTS5 前缀匹配。
func ftsQuery(text string) string {
	words := splitWords(text)
	if len(words) == 0 {
		return `""`
	}
	out := ""
	for i, w := range words {
		if i > 0 {
			out += " "
		}
		out += `"` + w + `"*`
	}
	return out
}

func splitWords(s string) []string {
	out := []string{}
	cur := []rune{}
	for _, r := range s {
		if r == ' ' || r == '\t' || r == '\n' {
			if len(cur) > 0 {
				out = append(out, string(cur))
				cur = nil
			}
			continue
		}
		cur = append(cur, r)
	}
	if len(cur) > 0 {
		out = append(out, string(cur))
	}
	return out
}

func containsStr(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}