From a82267484aaa09bd4291f5f1f1353c2971f24342 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Fri, 4 Sep 2026 11:03:36 +0800 Subject: [PATCH] =?UTF-8?q?feat(memory):=20=E5=86=85=E5=AE=B9=E5=AF=BB?= =?UTF-8?q?=E5=9D=80=E5=AA=92=E4=BD=93=E5=AD=98=E5=82=A8=EF=BC=88CAS?= =?UTF-8?q?=EF=BC=89=E2=80=94=E2=80=94=E5=9B=BE=E8=AE=B0=E5=BF=86=E6=94=AF?= =?UTF-8?q?=E6=8C=81=E4=BA=8C=E8=BF=9B=E5=88=B6=E5=A4=9A=E5=AA=92=E4=BD=93?= =?UTF-8?q?=E8=8A=82=E7=82=B9=E7=9A=84=E5=BA=95=E5=BA=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前四层记忆全是纯文本载体，没有任何一层能存二进制： L0 ContextEvent — Input/Response/ToolResults[].Output 全 string L1 text.Event — 同上 L2 document.Doc — Summary/Content/Tags 全 string L3 图库 — sentences.text TEXT UNIQUE，节点身份就是那串文本 于是 multimodal 插件注入的图只在本轮对话内可见（走 message 数组，不经 记忆），下一轮起只剩 ToolResultItem.Output 里那句 "[已将图片注入后续对话] /tmp/x.png"——一条路径字符串。那个文件被删或 被覆盖之后连线索都断了。 ## 为何内容寻址而不是存路径 - 路径会失效。/tmp 下的探针图、下载缓存、别的进程的临时产物，记忆里 留个路径等于留个悬空指针。 - 同一内容常被反复注入（连问几轮同一张截图、see_video 相邻帧高度相似）， 按 sha256 寻址天然去重。 - 内容即身份，与 L3 图库 sentences.text UNIQUE 思路一致：文本节点用文本 本身做身份，媒体节点用内容摘要做身份。 ## 结构 元数据（SQLite media.db）与内容（磁盘 blobs/ 两级前缀分桶）分离，不把 blob 塞进库：单张图动辄几 MB，塞进去让每次 VACUUM/备份都拖着几百 MB 走， WAL 也会迅速膨胀。 media(digest PK, kind, mime, size, width, height, origin_path, tool, description, described_by, ref_count, first_seen, last_seen) media_refs(digest, owner_kind, owner_id, created_at, PK 三列) digest 既是主键也是文件名，所以没有 Path 字段——路径由 digest 推导， 不落库（落了就又是个会失效的引用）。origin_path 仅供人类溯源，注释里 明确标注不可用于读取。owner_kind 预留 context/document/graph_sentence。 ## 几处刻意的决定 - Get 强制校验 digest：CAS 的全部保证建立在「文件名 == 内容摘要」上， 位翻转或外部误改必须被发现——把损坏的图喂给模型只会得到无从追溯的幻觉。 - 先写 .tmp 再 rename：中途崩溃不留半个 blob 被当成完整内容读走。 - AddRef 幂等：只有真插进 media_refs 才递增，否则计数虚高会让 GC 永远 不敢清。DropRef 用 MAX(0,...) 兜底防负数。 - Put 的空描述不冲掉已有描述（先到的可能来自更强的模型），但尺寸/工具名 这类前一次缺失的信息会被补写。Describe 是显式操作，允许覆盖。 - GC 两段 + minAge 保护：刚 Put 还没 AddRef 的项 refcount 也是 0，minAge 防「落地后还没挂上就被清掉」。有引用的项永不删除，即使超容量——宁可 超限也不断引用。 ## 测试 21 例，覆盖去重 / MIME 归类 / 损坏检测 / 无残留临时文件 / AddRef 幂等 / 计数不为负 / DropOwner / GC 保留有引用项 / minAge 保护 / 容量淘汰 / 描述覆盖与补写策略 / Search 按描述与 kind 过滤 / Pending / data URL 往返 / Stats / 跨重启持久化。 本 commit 只加存储层，尚未接入 L0/L2/L3 与描述生成。 --- internal/memory/media/base64.go | 14 + internal/memory/media/media.go | 702 ++++++++++++++++++++++++++++ internal/memory/media/media_test.go | 477 +++++++++++++++++++ 3 files changed, 1193 insertions(+) create mode 100644 internal/memory/media/base64.go create mode 100644 internal/memory/media/media.go create mode 100644 internal/memory/media/media_test.go diff --git a/internal/memory/media/base64.go b/internal/memory/media/base64.go new file mode 100644 index 0000000..247d87b --- /dev/null +++ b/internal/memory/media/base64.go @@ -0,0 +1,14 @@ +package media + +import "encoding/base64" + +// base64 编解码单独抽出来，让 media.go 的 import 块只留业务依赖。 +// 用 StdEncoding：data URL 规范用的是标准表（含 + / =），不是 URL-safe 表。 + +func base64Decode(s string) ([]byte, error) { + return base64.StdEncoding.DecodeString(s) +} + +func base64Encode(b []byte) string { + return base64.StdEncoding.EncodeToString(b) +} diff --git a/internal/memory/media/media.go b/internal/memory/media/media.go new file mode 100644 index 0000000..2007b8f --- /dev/null +++ b/internal/memory/media/media.go @@ -0,0 +1,702 @@ +// Package media 是记忆系统的内容寻址媒体存储（CAS）。 +// +// 为何需要它：此前四层记忆全是纯文本载体——L0 `ContextEvent`、L1 `text.Event`、 +// L2 `document.Doc`、L3 图库的 `sentences.text TEXT UNIQUE`——没有任何一层能 +// 存二进制。multimodal 插件注入的图片在本轮对话内可见（走 message 数组，不经 +// 记忆），下一轮起就只剩 `ToolResultItem.Output` 里那句 +// "[已将图片注入后续对话] /tmp/x.png"，即一条路径字符串。那个文件被删或被 +// 覆盖之后连线索都断了。 +// +// 为何是内容寻址而不是存路径： +// - 路径会失效。/tmp 下的探针图、下载缓存、其他进程的临时产物，记忆里留个 +// 路径等于留个悬空指针。 +// - 同一张图往往被反复注入（用户连问几轮同一张截图、see_video 相邻帧高度 +// 相似）。按 sha256 寻址天然去重，引用计数记住被引了几次。 +// - 内容即身份，跟 L3 图库 `sentences.text UNIQUE` 的思路一致：文本节点用 +// 文本本身做身份，媒体节点用内容摘要做身份。 +package media + +import ( + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "time" + + _ "github.com/mattn/go-sqlite3" +) + +// Kind 是媒体大类。刻意只分三类而不细分具体格式： +// 记忆检索关心的是"这是张图还是段音频"，具体编码交给 MIME 字段。 +type Kind string + +const ( + KindImage Kind = "image" + KindAudio Kind = "audio" + KindVideo Kind = "video" + KindOther Kind = "other" +) + +// Item 是一条媒体记录。 +// +// Digest 既是主键也是磁盘文件名，所以没有单独的 Path 字段——路径可由 +// Store 根据 Digest 推导，不落库（落了就又是个会失效的引用）。 +type Item struct { + // Digest 是内容 sha256 的十六进制串（64 字符），媒体的唯一身份。 + Digest string `json:"digest"` + // Kind 是大类，供检索时按模态筛选。 + Kind Kind `json:"kind"` + // MIME 是原始 MIME 类型，如 image/png。 + MIME string `json:"mime"` + // Size 是字节数。 + Size int64 `json:"size"` + // Width/Height 是像素尺寸，未知或不适用时为 0。 + Width int `json:"width,omitempty"` + Height int `json:"height,omitempty"` + // OriginPath 是首次入库时的来源路径，仅供人类溯源与调试。 + // **不可用于读取内容**——它随时可能失效，这正是本包存在的理由。 + OriginPath string `json:"origin_path,omitempty"` + // Tool 是注入这条媒体的工具名（如 multimodal_see_picture）。 + Tool string `json:"tool,omitempty"` + // Description 是视觉/音频模型生成的文字描述，供 L2/L3 检索。 + // 空表示未描述（未开启描述、模型不可用或描述失败）。 + Description string `json:"description,omitempty"` + // DescribedBy 记录描述来自哪个源，让后续读者能判断可靠性。 + DescribedBy string `json:"described_by,omitempty"` + // RefCount 是引用计数。GC 只清理归零的项。 + RefCount int `json:"ref_count"` + // FirstSeen/LastSeen 是首末次入库时间。 + FirstSeen time.Time `json:"first_seen"` + LastSeen time.Time `json:"last_seen"` +} + +// Store 管理媒体的元数据（SQLite）与内容（磁盘 CAS 目录）。 +// +// 元数据与内容分离而不是把 blob 塞进 SQLite：单张图动辄几 MB，塞进库会让 +// 每次 VACUUM/备份都拖着几百 MB 走，也让 WAL 迅速膨胀。CAS 目录用两级 +// 前缀分桶（ab/cdef...）避免单目录几万文件。 +type Store struct { + mu sync.RWMutex + db *sql.DB + blobDir string + + // maxBytes 是内容目录的容量上限，0 表示不限。 + // 超限时 GC 按 LastSeen 从旧到新淘汰 RefCount=0 的项。 + maxBytes int64 +} + +// New 打开（或初始化）媒体存储。 +// dir 下会建 media.db 与 blobs/ 两个条目。 +func New(dir string, maxBytes int64) (*Store, error) { + if err := os.MkdirAll(filepath.Join(dir, "blobs"), 0755); err != nil { + return nil, fmt.Errorf("media: create blob dir: %w", err) + } + dbPath := filepath.Join(dir, "media.db") + db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_busy_timeout=5000") + if err != nil { + return nil, fmt.Errorf("media: open db: %w", err) + } + s := &Store{db: db, blobDir: filepath.Join(dir, "blobs"), maxBytes: maxBytes} + if err := s.initSchema(); err != nil { + db.Close() + return nil, err + } + return s, nil +} + +func (s *Store) initSchema() error { + stmts := []string{ + // digest 作主键：内容即身份，重复 Put 同一内容只递增 ref_count。 + `CREATE TABLE IF NOT EXISTS media ( + digest TEXT PRIMARY KEY, + kind TEXT NOT NULL, + mime TEXT NOT NULL, + size INTEGER NOT NULL, + width INTEGER DEFAULT 0, + height INTEGER DEFAULT 0, + origin_path TEXT, + tool TEXT, + description TEXT, + described_by TEXT, + ref_count INTEGER DEFAULT 0, + first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP + )`, + `CREATE INDEX IF NOT EXISTS idx_media_kind ON media(kind)`, + `CREATE INDEX IF NOT EXISTS idx_media_refcount ON media(ref_count)`, + `CREATE INDEX IF NOT EXISTS idx_media_last_seen ON media(last_seen)`, + // 反向索引：哪条记忆引用了哪个媒体。 + // owner_kind 取 context / document / graph_sentence，owner_id 是各层自己的标识。 + // 主键含三列，同一 owner 重复挂同一媒体是幂等的。 + `CREATE TABLE IF NOT EXISTS media_refs ( + digest TEXT NOT NULL, + owner_kind TEXT NOT NULL, + owner_id TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (digest, owner_kind, owner_id) + )`, + `CREATE INDEX IF NOT EXISTS idx_refs_owner ON media_refs(owner_kind, owner_id)`, + `CREATE INDEX IF NOT EXISTS idx_refs_digest ON media_refs(digest)`, + } + for _, q := range stmts { + if _, err := s.db.Exec(q); err != nil { + return fmt.Errorf("media: schema %q: %w", truncate(q, 60), err) + } + } + return nil +} + +// blobPath 按两级前缀分桶推导内容路径。 +func (s *Store) blobPath(digest string) string { + if len(digest) < 4 { + return filepath.Join(s.blobDir, digest) + } + return filepath.Join(s.blobDir, digest[:2], digest[2:]) +} + +// Put 落盘并登记一段媒体内容，返回其 digest。 +// +// 幂等：同一内容重复 Put 不重复落盘，只更新 last_seen 与可选的新元数据 +// （描述、尺寸等——后来者可能带着前一次没有的信息）。 +func (s *Store) Put(data []byte, meta Item) (string, error) { + if len(data) == 0 { + return "", fmt.Errorf("media: empty content") + } + sum := sha256.Sum256(data) + digest := hex.EncodeToString(sum[:]) + + s.mu.Lock() + defer s.mu.Unlock() + + path := s.blobPath(digest) + if _, err := os.Stat(path); os.IsNotExist(err) { + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + return "", fmt.Errorf("media: mkdir: %w", err) + } + // 先写临时文件再 rename：中途崩溃不会留下半个 blob 被后续 + // 当成完整内容读走（digest 校验能发现，但那时已经把坏数据喂给模型了）。 + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0644); err != nil { + return "", fmt.Errorf("media: write blob: %w", err) + } + if err := os.Rename(tmp, path); err != nil { + os.Remove(tmp) + return "", fmt.Errorf("media: commit blob: %w", err) + } + } + + now := time.Now() + if meta.Kind == "" { + meta.Kind = KindFromMIME(meta.MIME) + } + _, err := s.db.Exec(` + INSERT INTO media (digest, kind, mime, size, width, height, + origin_path, tool, description, described_by, + ref_count, first_seen, last_seen) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?) + ON CONFLICT(digest) DO UPDATE SET + last_seen = excluded.last_seen, + -- 只在原值为空时补写：先到的描述可能来自更强的模型， + -- 后到的空值不该把它冲掉。 + description = CASE WHEN COALESCE(media.description,'') = '' THEN excluded.description ELSE media.description END, + described_by = CASE WHEN COALESCE(media.described_by,'') = '' THEN excluded.described_by ELSE media.described_by END, + width = CASE WHEN media.width = 0 THEN excluded.width ELSE media.width END, + height = CASE WHEN media.height = 0 THEN excluded.height ELSE media.height END, + tool = CASE WHEN COALESCE(media.tool,'') = '' THEN excluded.tool ELSE media.tool END + `, digest, string(meta.Kind), meta.MIME, int64(len(data)), meta.Width, meta.Height, + meta.OriginPath, meta.Tool, meta.Description, meta.DescribedBy, now, now) + if err != nil { + return "", fmt.Errorf("media: upsert meta: %w", err) + } + return digest, nil +} + +// Get 读取内容并校验 digest。 +// +// 校验不是多余的：CAS 的全部保证建立在"文件名 == 内容摘要"上，磁盘位翻转 +// 或外部误改会让这条保证失效，而把损坏的图喂给模型只会得到无从追溯的幻觉。 +func (s *Store) Get(digest string) ([]byte, error) { + s.mu.RLock() + path := s.blobPath(digest) + s.mu.RUnlock() + + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("media: read %s: %w", shortDigest(digest), err) + } + sum := sha256.Sum256(data) + if got := hex.EncodeToString(sum[:]); got != digest { + return nil, fmt.Errorf("media: digest mismatch for %s (content corrupted)", shortDigest(digest)) + } + return data, nil +} + +// Stat 返回元数据，不读内容。 +func (s *Store) Stat(digest string) (*Item, error) { + s.mu.RLock() + defer s.mu.RUnlock() + return s.scanOne(s.db.QueryRow(` + SELECT digest, kind, mime, size, width, height, origin_path, tool, + description, described_by, ref_count, first_seen, last_seen + FROM media WHERE digest = ?`, digest)) +} + +// Describe 写入（或覆盖）文字描述。 +// +// 与 Put 的"只在空时补写"不同：Describe 是显式操作，调用方明确想要这份 +// 描述生效（例如换了更强的视觉模型重新描述）。 +func (s *Store) Describe(digest, description, describedBy string) error { + s.mu.Lock() + defer s.mu.Unlock() + res, err := s.db.Exec(`UPDATE media SET description = ?, described_by = ? WHERE digest = ?`, + description, describedBy, digest) + if err != nil { + return fmt.Errorf("media: describe: %w", err) + } + if n, _ := res.RowsAffected(); n == 0 { + return fmt.Errorf("media: describe: unknown digest %s", shortDigest(digest)) + } + return nil +} + +// AddRef 登记一条引用并递增计数。幂等：同一 (digest, owner) 重复调用不重复计数。 +func (s *Store) AddRef(digest, ownerKind, ownerID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + tx, err := s.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + res, err := tx.Exec(`INSERT OR IGNORE INTO media_refs (digest, owner_kind, owner_id) VALUES (?, ?, ?)`, + digest, ownerKind, ownerID) + if err != nil { + return fmt.Errorf("media: add ref: %w", err) + } + // 只有真的插进去才递增：否则重复调用会让计数虚高，GC 永远不敢清。 + if n, _ := res.RowsAffected(); n > 0 { + if _, err := tx.Exec(`UPDATE media SET ref_count = ref_count + 1 WHERE digest = ?`, digest); err != nil { + return fmt.Errorf("media: bump refcount: %w", err) + } + } + return tx.Commit() +} + +// DropRef 注销一条引用并递减计数。内容不立即删除，留给 GC。 +func (s *Store) DropRef(digest, ownerKind, ownerID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + tx, err := s.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + res, err := tx.Exec(`DELETE FROM media_refs WHERE digest = ? AND owner_kind = ? AND owner_id = ?`, + digest, ownerKind, ownerID) + if err != nil { + return fmt.Errorf("media: drop ref: %w", err) + } + if n, _ := res.RowsAffected(); n > 0 { + // MAX(0, ...) 兜底：历史数据或并发意外让计数与 refs 表不一致时， + // 不让它掉成负数（负数会让容量 GC 的排序失去意义）。 + if _, err := tx.Exec(`UPDATE media SET ref_count = MAX(0, ref_count - 1) WHERE digest = ?`, digest); err != nil { + return fmt.Errorf("media: lower refcount: %w", err) + } + } + return tx.Commit() +} + +// DropOwner 注销某个 owner 的全部引用（该条记忆被删/被归档替换时用）。 +func (s *Store) DropOwner(ownerKind, ownerID string) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + + rows, err := s.db.Query(`SELECT digest FROM media_refs WHERE owner_kind = ? AND owner_id = ?`, + ownerKind, ownerID) + if err != nil { + return 0, err + } + var digests []string + for rows.Next() { + var d string + if err := rows.Scan(&d); err == nil { + digests = append(digests, d) + } + } + rows.Close() + if err := rows.Err(); err != nil { + return 0, err + } + if len(digests) == 0 { + return 0, nil + } + + tx, err := s.db.Begin() + if err != nil { + return 0, err + } + defer tx.Rollback() + if _, err := tx.Exec(`DELETE FROM media_refs WHERE owner_kind = ? AND owner_id = ?`, ownerKind, ownerID); err != nil { + return 0, err + } + for _, d := range digests { + if _, err := tx.Exec(`UPDATE media SET ref_count = MAX(0, ref_count - 1) WHERE digest = ?`, d); err != nil { + return 0, err + } + } + if err := tx.Commit(); err != nil { + return 0, err + } + return len(digests), nil +} + +// Refs 返回某个 owner 引用的全部 digest。 +func (s *Store) Refs(ownerKind, ownerID string) ([]string, error) { + s.mu.RLock() + defer s.mu.RUnlock() + rows, err := s.db.Query(`SELECT digest FROM media_refs WHERE owner_kind = ? AND owner_id = ? ORDER BY created_at`, + ownerKind, ownerID) + if err != nil { + return nil, err + } + defer rows.Close() + var out []string + for rows.Next() { + var d string + if err := rows.Scan(&d); err == nil { + out = append(out, d) + } + } + return out, rows.Err() +} + +// Search 按描述文本做 LIKE 匹配，返回最近的若干条。 +// +// 刻意不在这里做向量检索：媒体的语义检索走 L2 文档层的既有索引 +// （描述文字随记忆条目一起进 Doc.Content，复用那套 TF-IDF/embedding）， +// 本方法只是"按关键词直接翻媒体库"的补充入口。 +func (s *Store) Search(query string, kind Kind, limit int) ([]*Item, error) { + if limit <= 0 { + limit = 20 + } + s.mu.RLock() + defer s.mu.RUnlock() + + q := `SELECT digest, kind, mime, size, width, height, origin_path, tool, + description, described_by, ref_count, first_seen, last_seen + FROM media WHERE COALESCE(description,'') != ''` + args := []interface{}{} + if strings.TrimSpace(query) != "" { + q += ` AND description LIKE ?` + args = append(args, "%"+query+"%") + } + if kind != "" { + q += ` AND kind = ?` + args = append(args, string(kind)) + } + q += ` ORDER BY last_seen DESC LIMIT ?` + args = append(args, limit) + + rows, err := s.db.Query(q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*Item + for rows.Next() { + it, err := s.scanRows(rows) + if err != nil { + continue + } + out = append(out, it) + } + return out, rows.Err() +} + +// Pending 返回尚无描述的媒体，供后台描述任务消费。 +func (s *Store) Pending(limit int) ([]*Item, error) { + if limit <= 0 { + limit = 10 + } + s.mu.RLock() + defer s.mu.RUnlock() + rows, err := s.db.Query(` + SELECT digest, kind, mime, size, width, height, origin_path, tool, + description, described_by, ref_count, first_seen, last_seen + FROM media WHERE COALESCE(description,'') = '' + ORDER BY last_seen DESC LIMIT ?`, limit) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*Item + for rows.Next() { + it, err := s.scanRows(rows) + if err != nil { + continue + } + out = append(out, it) + } + return out, rows.Err() +} + +// GC 清理无人引用的内容。 +// +// 两段策略： +// 1. ref_count=0 且 last_seen 早于 minAge 的一律清理。刚 Put 还没来得及 +// AddRef 的项 refcount 也是 0，minAge 保护它们不被立刻清掉。 +// 2. 清完仍超 maxBytes 时，继续按 last_seen 从旧到新淘汰 ref_count=0 的项。 +// +// 有引用的项永不删除——那会让记忆里的 digest 变成悬空指针，正是本包要避免的。 +func (s *Store) GC(minAge time.Duration) (removed int, freed int64, err error) { + s.mu.Lock() + defer s.mu.Unlock() + + cutoff := time.Now().Add(-minAge) + rows, err := s.db.Query(` + SELECT digest, size FROM media + WHERE ref_count <= 0 AND last_seen < ? + ORDER BY last_seen`, cutoff) + if err != nil { + return 0, 0, err + } + type cand struct { + digest string + size int64 + } + var cands []cand + for rows.Next() { + var c cand + if err := rows.Scan(&c.digest, &c.size); err == nil { + cands = append(cands, c) + } + } + rows.Close() + + for _, c := range cands { + if e := os.Remove(s.blobPath(c.digest)); e != nil && !os.IsNotExist(e) { + continue // 删不掉就留着元数据，下轮再试；不制造"元数据没了文件还在"的孤儿 + } + if _, e := s.db.Exec(`DELETE FROM media WHERE digest = ?`, c.digest); e != nil { + continue + } + removed++ + freed += c.size + } + + if s.maxBytes > 0 { + r2, f2 := s.enforceCapacityLocked() + removed += r2 + freed += f2 + } + return removed, freed, nil +} + +// enforceCapacityLocked 在超出 maxBytes 时继续淘汰无引用项（调用方已持锁）。 +func (s *Store) enforceCapacityLocked() (removed int, freed int64) { + var total int64 + if err := s.db.QueryRow(`SELECT COALESCE(SUM(size), 0) FROM media`).Scan(&total); err != nil { + return 0, 0 + } + if total <= s.maxBytes { + return 0, 0 + } + need := total - s.maxBytes + + rows, err := s.db.Query(`SELECT digest, size FROM media WHERE ref_count <= 0 ORDER BY last_seen`) + if err != nil { + return 0, 0 + } + type cand struct { + digest string + size int64 + } + var cands []cand + for rows.Next() { + var c cand + if err := rows.Scan(&c.digest, &c.size); err == nil { + cands = append(cands, c) + } + } + rows.Close() + + for _, c := range cands { + if freed >= need { + break + } + if e := os.Remove(s.blobPath(c.digest)); e != nil && !os.IsNotExist(e) { + continue + } + if _, e := s.db.Exec(`DELETE FROM media WHERE digest = ?`, c.digest); e != nil { + continue + } + removed++ + freed += c.size + } + return removed, freed +} + +// Stats 返回容量与条目统计，供 WebUI / healthcheck 展示。 +func (s *Store) Stats() map[string]interface{} { + s.mu.RLock() + defer s.mu.RUnlock() + + out := map[string]interface{}{"blob_dir": s.blobDir, "max_bytes": s.maxBytes} + var count, described, orphan int + var total int64 + s.db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(size),0) FROM media`).Scan(&count, &total) + s.db.QueryRow(`SELECT COUNT(*) FROM media WHERE COALESCE(description,'') != ''`).Scan(&described) + s.db.QueryRow(`SELECT COUNT(*) FROM media WHERE ref_count <= 0`).Scan(&orphan) + out["count"] = count + out["total_bytes"] = total + out["described"] = described + out["unreferenced"] = orphan + + byKind := map[string]int{} + rows, err := s.db.Query(`SELECT kind, COUNT(*) FROM media GROUP BY kind`) + if err == nil { + defer rows.Close() + for rows.Next() { + var k string + var n int + if rows.Scan(&k, &n) == nil { + byKind[k] = n + } + } + } + out["by_kind"] = byKind + return out +} + +func (s *Store) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + return s.db.Close() +} + +// ---- 扫描辅助 ---- + +type rowScanner interface { + Scan(dest ...interface{}) error +} + +func (s *Store) scanOne(r rowScanner) (*Item, error) { + it, err := scanItem(r) + if err == sql.ErrNoRows { + return nil, fmt.Errorf("media: unknown digest") + } + return it, err +} + +func (s *Store) scanRows(r rowScanner) (*Item, error) { return scanItem(r) } + +func scanItem(r rowScanner) (*Item, error) { + var it Item + var kind string + var origin, tool, desc, by sql.NullString + if err := r.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height, + &origin, &tool, &desc, &by, &it.RefCount, &it.FirstSeen, &it.LastSeen); err != nil { + return nil, err + } + it.Kind = Kind(kind) + it.OriginPath = origin.String + it.Tool = tool.String + it.Description = desc.String + it.DescribedBy = by.String + return &it, nil +} + +// ---- 工具函数 ---- + +// KindFromMIME 把 MIME 归到大类。 +func KindFromMIME(mime string) Kind { + m := strings.ToLower(strings.TrimSpace(mime)) + switch { + case strings.HasPrefix(m, "image/"): + return KindImage + case strings.HasPrefix(m, "audio/"): + return KindAudio + case strings.HasPrefix(m, "video/"): + return KindVideo + default: + return KindOther + } +} + +// ParseDataURL 从 data:;base64, 提取 MIME 与原始字节。 +// +// 与 agent/api 里的 parseAudioDataURL 分开实现：那个只认音频且只回 base64 +// 串（它要把串塞回 OpenAI 的 input_audio 字段），这里要的是解码后的字节。 +func ParseDataURL(url string) (mime string, data []byte, ok bool) { + const prefix = "data:" + if !strings.HasPrefix(url, prefix) { + return "", nil, false + } + rest := url[len(prefix):] + comma := strings.IndexByte(rest, ',') + if comma < 0 { + return "", nil, false + } + head := rest[:comma] + payload := rest[comma+1:] + if !strings.HasSuffix(strings.ToLower(head), ";base64") { + return "", nil, false + } + mime = head[:len(head)-len(";base64")] + if mime == "" || payload == "" { + return "", nil, false + } + decoded, err := base64Decode(payload) + if err != nil { + return "", nil, false + } + return mime, decoded, true +} + +// DataURL 把内容编回 data URL，供重新注入模型对话。 +func DataURL(mime string, data []byte) string { + return "data:" + mime + ";base64," + base64Encode(data) +} + +// CopyFrom 从 reader 读全部内容后 Put，用于大文件不便一次性构造 []byte 的场合。 +func (s *Store) CopyFrom(r io.Reader, meta Item) (string, error) { + data, err := io.ReadAll(r) + if err != nil { + return "", fmt.Errorf("media: read source: %w", err) + } + return s.Put(data, meta) +} + +// MarshalItems 序列化条目列表，供工具返回给模型。 +func MarshalItems(items []*Item) string { + b, err := json.Marshal(items) + if err != nil { + return "[]" + } + return string(b) +} + +func shortDigest(d string) string { + if len(d) > 12 { + return d[:12] + } + return d +} + +func truncate(s string, n int) string { + s = strings.Join(strings.Fields(s), " ") + if len(s) <= n { + return s + } + return s[:n] + "..." +} diff --git a/internal/memory/media/media_test.go b/internal/memory/media/media_test.go new file mode 100644 index 0000000..b4b2048 --- /dev/null +++ b/internal/memory/media/media_test.go @@ -0,0 +1,477 @@ +package media + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func newTestStore(t *testing.T, maxBytes int64) *Store { + t.Helper() + s, err := New(t.TempDir(), maxBytes) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { s.Close() }) + return s +} + +func TestPut_ContentAddressedDedup(t *testing.T) { + s := newTestStore(t, 0) + data := []byte("fake-png-bytes") + + d1, err := s.Put(data, Item{MIME: "image/png", OriginPath: "/tmp/a.png"}) + if err != nil { + t.Fatal(err) + } + d2, err := s.Put(data, Item{MIME: "image/png", OriginPath: "/tmp/b.png"}) + if err != nil { + t.Fatal(err) + } + if d1 != d2 { + t.Fatalf("同一内容应得同一 digest：%s vs %s", d1, d2) + } + + // 去重的意义：同一张图反复注入（连问几轮同一截图、see_video 相邻帧） + // 只占一份磁盘。 + st := s.Stats() + if st["count"].(int) != 1 { + t.Fatalf("同一内容应只有 1 条记录，实际 %v", st["count"]) + } +} + +func TestPut_KindInferredFromMIME(t *testing.T) { + s := newTestStore(t, 0) + cases := map[string]Kind{ + "image/png": KindImage, + "image/jpeg": KindImage, + "audio/wav": KindAudio, + "video/mp4": KindVideo, + "text/plain": KindOther, + } + for mime, want := range cases { + d, err := s.Put([]byte("payload-"+mime), Item{MIME: mime}) + if err != nil { + t.Fatal(err) + } + it, err := s.Stat(d) + if err != nil { + t.Fatal(err) + } + if it.Kind != want { + t.Fatalf("%s 应归为 %s，实际 %s", mime, want, it.Kind) + } + } +} + +func TestGet_DetectsCorruption(t *testing.T) { + // CAS 的全部保证建立在「文件名 == 内容摘要」上。外部误改或位翻转必须 + // 被发现——把损坏的图喂给模型只会得到无从追溯的幻觉。 + s := newTestStore(t, 0) + d, err := s.Put([]byte("original-content"), Item{MIME: "image/png"}) + if err != nil { + t.Fatal(err) + } + if _, err := s.Get(d); err != nil { + t.Fatalf("正常读取应成功: %v", err) + } + + if err := os.WriteFile(s.blobPath(d), []byte("tampered!"), 0644); err != nil { + t.Fatal(err) + } + _, err = s.Get(d) + if err == nil { + t.Fatal("内容被改后应报 digest 不匹配，却读成功了") + } + if !strings.Contains(err.Error(), "digest mismatch") { + t.Fatalf("错误应指明 digest 不匹配，得到: %v", err) + } +} + +func TestPut_NoPartialBlobOnDisk(t *testing.T) { + // 先写 .tmp 再 rename：确认落地后目录里不留临时文件。 + s := newTestStore(t, 0) + d, err := s.Put([]byte("some-bytes"), Item{MIME: "image/png"}) + if err != nil { + t.Fatal(err) + } + dir := filepath.Dir(s.blobPath(d)) + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + for _, e := range entries { + if strings.HasSuffix(e.Name(), ".tmp") { + t.Fatalf("落地后不该留临时文件: %s", e.Name()) + } + } +} + +func TestRefCount_AddIsIdempotent(t *testing.T) { + s := newTestStore(t, 0) + d, _ := s.Put([]byte("img"), Item{MIME: "image/png"}) + + for i := 0; i < 3; i++ { + if err := s.AddRef(d, "context", "evt-1"); err != nil { + t.Fatal(err) + } + } + it, _ := s.Stat(d) + // 重复 AddRef 若都递增，计数会虚高，GC 永远不敢清。 + if it.RefCount != 1 { + t.Fatalf("同一 owner 重复 AddRef 应只计 1，实际 %d", it.RefCount) + } + + if err := s.AddRef(d, "document", "doc-9"); err != nil { + t.Fatal(err) + } + it, _ = s.Stat(d) + if it.RefCount != 2 { + t.Fatalf("不同 owner 应各计一次，实际 %d", it.RefCount) + } +} + +func TestRefCount_DropAndNeverNegative(t *testing.T) { + s := newTestStore(t, 0) + d, _ := s.Put([]byte("img"), Item{MIME: "image/png"}) + s.AddRef(d, "context", "e1") + + if err := s.DropRef(d, "context", "e1"); err != nil { + t.Fatal(err) + } + it, _ := s.Stat(d) + if it.RefCount != 0 { + t.Fatalf("应归零，实际 %d", it.RefCount) + } + + // 多余的 DropRef 不该把计数压成负数（负数会让容量 GC 的排序失去意义） + for i := 0; i < 3; i++ { + s.DropRef(d, "context", "e1") + } + it, _ = s.Stat(d) + if it.RefCount != 0 { + t.Fatalf("重复 DropRef 后仍应为 0，实际 %d", it.RefCount) + } +} + +func TestDropOwner_RemovesAllItsRefs(t *testing.T) { + s := newTestStore(t, 0) + d1, _ := s.Put([]byte("frame1"), Item{MIME: "image/jpeg"}) + d2, _ := s.Put([]byte("frame2"), Item{MIME: "image/jpeg"}) + s.AddRef(d1, "context", "evt-x") + s.AddRef(d2, "context", "evt-x") + s.AddRef(d1, "document", "doc-y") // 别的 owner 也引了 d1 + + n, err := s.DropOwner("context", "evt-x") + if err != nil { + t.Fatal(err) + } + if n != 2 { + t.Fatalf("应注销 2 条引用，实际 %d", n) + } + + it1, _ := s.Stat(d1) + it2, _ := s.Stat(d2) + if it1.RefCount != 1 { + t.Fatalf("d1 仍被 document 引用，应剩 1，实际 %d", it1.RefCount) + } + if it2.RefCount != 0 { + t.Fatalf("d2 应归零，实际 %d", it2.RefCount) + } +} + +func TestRefs_ListsOwnerDigests(t *testing.T) { + s := newTestStore(t, 0) + d1, _ := s.Put([]byte("a"), Item{MIME: "image/png"}) + d2, _ := s.Put([]byte("b"), Item{MIME: "image/png"}) + s.AddRef(d1, "context", "e1") + s.AddRef(d2, "context", "e1") + + got, err := s.Refs("context", "e1") + if err != nil { + t.Fatal(err) + } + if len(got) != 2 { + t.Fatalf("应返回 2 个 digest，实际 %d", len(got)) + } +} + +func TestGC_KeepsReferencedContent(t *testing.T) { + // 有引用的项永不删除——那会让记忆里的 digest 变成悬空指针， + // 正是本包要避免的。 + s := newTestStore(t, 0) + kept, _ := s.Put([]byte("referenced"), Item{MIME: "image/png"}) + orphan, _ := s.Put([]byte("orphaned"), Item{MIME: "image/png"}) + s.AddRef(kept, "context", "e1") + + // minAge=0 让刚 Put 的都算超龄 + removed, _, err := s.GC(0) + if err != nil { + t.Fatal(err) + } + if removed != 1 { + t.Fatalf("应只清 1 条无引用项，实际 %d", removed) + } + if _, err := s.Get(kept); err != nil { + t.Fatalf("被引用的内容不该被清: %v", err) + } + if _, err := s.Stat(orphan); err == nil { + t.Fatal("无引用项的元数据应已删除") + } +} + +func TestGC_MinAgeProtectsFreshUnreferenced(t *testing.T) { + // 刚 Put 还没来得及 AddRef 的项 refcount 也是 0； + // minAge 必须保护它们，否则「Put 完还没挂上就被 GC 清掉」。 + s := newTestStore(t, 0) + d, _ := s.Put([]byte("just-arrived"), Item{MIME: "image/png"}) + + removed, _, err := s.GC(time.Hour) + if err != nil { + t.Fatal(err) + } + if removed != 0 { + t.Fatalf("新入库项应被 minAge 保护，却清掉了 %d 条", removed) + } + if _, err := s.Get(d); err != nil { + t.Fatalf("内容应还在: %v", err) + } +} + +func TestGC_EnforcesCapacity(t *testing.T) { + // 容量上限：清完超龄项后仍超限，继续按 last_seen 从旧到新淘汰无引用项。 + blob := make([]byte, 1024) + s := newTestStore(t, 2048) // 只容 2KB + + var digests []string + for i := 0; i < 4; i++ { + b := append([]byte{byte(i)}, blob...) // 内容各异，避免去重 + d, err := s.Put(b, Item{MIME: "image/png"}) + if err != nil { + t.Fatal(err) + } + digests = append(digests, d) + time.Sleep(2 * time.Millisecond) // 拉开 last_seen + } + + // 保护最后一个，确认容量 GC 也不碰有引用的 + s.AddRef(digests[3], "context", "e1") + + removed, freed, err := s.GC(0) + if err != nil { + t.Fatal(err) + } + if removed == 0 { + t.Fatal("超限应触发淘汰") + } + if _, err := s.Get(digests[3]); err != nil { + t.Fatalf("有引用项即使超限也不该删: %v", err) + } + t.Logf("removed=%d freed=%d", removed, freed) + + st := s.Stats() + if total := st["total_bytes"].(int64); total > 2048 { + // 有引用项可能让总量降不到线下，这是刻意的（宁可超限也不断引用） + t.Logf("总量 %d 仍超 2048，因有引用项不可删（预期行为）", total) + } +} + +func TestDescribe_OverwritesExplicitly(t *testing.T) { + s := newTestStore(t, 0) + d, _ := s.Put([]byte("img"), Item{MIME: "image/png"}) + + if err := s.Describe(d, "一只橘猫", "vis-a"); err != nil { + t.Fatal(err) + } + it, _ := s.Stat(d) + if it.Description != "一只橘猫" || it.DescribedBy != "vis-a" { + t.Fatalf("描述未写入: %+v", it) + } + + // Describe 是显式操作，允许覆盖（换更强模型重描述） + if err := s.Describe(d, "一只橘色虎斑猫坐在窗台", "vis-b"); err != nil { + t.Fatal(err) + } + it, _ = s.Stat(d) + if !strings.Contains(it.Description, "虎斑") || it.DescribedBy != "vis-b" { + t.Fatalf("Describe 应覆盖旧描述: %+v", it) + } +} + +func TestDescribe_UnknownDigestErrors(t *testing.T) { + s := newTestStore(t, 0) + err := s.Describe("deadbeef", "x", "y") + if err == nil { + t.Fatal("未知 digest 应报错而非静默成功") + } +} + +func TestPut_DoesNotClobberExistingDescription(t *testing.T) { + // 先到的描述可能来自更强的模型；后到的空值不该把它冲掉。 + s := newTestStore(t, 0) + data := []byte("img") + d, _ := s.Put(data, Item{MIME: "image/png", Description: "详细描述", DescribedBy: "strong-model"}) + + // 第二次 Put 同内容但不带描述 + if _, err := s.Put(data, Item{MIME: "image/png"}); err != nil { + t.Fatal(err) + } + it, _ := s.Stat(d) + if it.Description != "详细描述" || it.DescribedBy != "strong-model" { + t.Fatalf("重复 Put 的空描述不该冲掉已有描述: %+v", it) + } +} + +func TestPut_BackfillsMissingDimensions(t *testing.T) { + // 后来者可能带着前一次没有的信息（尺寸、工具名） + s := newTestStore(t, 0) + data := []byte("img") + d, _ := s.Put(data, Item{MIME: "image/png"}) + + if _, err := s.Put(data, Item{MIME: "image/png", Width: 640, Height: 480, Tool: "multimodal_see_picture"}); err != nil { + t.Fatal(err) + } + it, _ := s.Stat(d) + if it.Width != 640 || it.Height != 480 { + t.Fatalf("尺寸应被补写: %dx%d", it.Width, it.Height) + } + if it.Tool != "multimodal_see_picture" { + t.Fatalf("工具名应被补写: %q", it.Tool) + } +} + +func TestSearch_FiltersByDescriptionAndKind(t *testing.T) { + s := newTestStore(t, 0) + di, _ := s.Put([]byte("chart-img"), Item{MIME: "image/png"}) + da, _ := s.Put([]byte("speech-aud"), Item{MIME: "audio/wav"}) + dn, _ := s.Put([]byte("no-desc"), Item{MIME: "image/png"}) + s.Describe(di, "一张蓝色的柱状图表", "vis") + s.Describe(da, "一段关于图表的讲解录音", "aud") + + all, err := s.Search("图表", "", 10) + if err != nil { + t.Fatal(err) + } + if len(all) != 2 { + t.Fatalf("两条描述都含「图表」，应返回 2，实际 %d", len(all)) + } + + imgs, _ := s.Search("图表", KindImage, 10) + if len(imgs) != 1 || imgs[0].Digest != di { + t.Fatalf("按 image 过滤应只剩图片，实际 %d 条", len(imgs)) + } + + // 无描述的项不该出现在语义检索结果里 + for _, it := range all { + if it.Digest == dn { + t.Fatal("无描述的项不该被 Search 返回") + } + } +} + +func TestPending_ReturnsUndescribed(t *testing.T) { + s := newTestStore(t, 0) + described, _ := s.Put([]byte("has-desc"), Item{MIME: "image/png"}) + undescribed, _ := s.Put([]byte("needs-desc"), Item{MIME: "image/png"}) + s.Describe(described, "已有描述", "vis") + + pending, err := s.Pending(10) + if err != nil { + t.Fatal(err) + } + if len(pending) != 1 || pending[0].Digest != undescribed { + t.Fatalf("应只返回未描述项，实际 %d 条", len(pending)) + } +} + +func TestParseDataURL(t *testing.T) { + // 与 agent/api 的 parseAudioDataURL 不同：这里要解码后的字节。 + raw := []byte{0x89, 'P', 'N', 'G'} + url := DataURL("image/png", raw) + + mime, data, ok := ParseDataURL(url) + if !ok { + t.Fatal("应解析成功") + } + if mime != "image/png" { + t.Fatalf("MIME 应为 image/png，得到 %q", mime) + } + if string(data) != string(raw) { + t.Fatalf("字节应还原，得到 %v", data) + } + + for _, bad := range []string{ + "http://example.com/x.png", // 不是 data URL + "data:image/png,notbase64", // 缺 ;base64 + "data:;base64,", // 空 MIME 与空载荷 + "data:image/png;base64,!!!", // 非法 base64 + } { + if _, _, ok := ParseDataURL(bad); ok { + t.Fatalf("%q 应解析失败", bad) + } + } +} + +func TestStats_CountsByKindAndDescription(t *testing.T) { + s := newTestStore(t, 4096) + d1, _ := s.Put([]byte("i1"), Item{MIME: "image/png"}) + s.Put([]byte("i2"), Item{MIME: "image/jpeg"}) + s.Put([]byte("a1"), Item{MIME: "audio/wav"}) + s.Describe(d1, "描述", "vis") + s.AddRef(d1, "context", "e1") + + st := s.Stats() + if st["count"].(int) != 3 { + t.Fatalf("count 应为 3，实际 %v", st["count"]) + } + if st["described"].(int) != 1 { + t.Fatalf("described 应为 1，实际 %v", st["described"]) + } + if st["unreferenced"].(int) != 2 { + t.Fatalf("unreferenced 应为 2，实际 %v", st["unreferenced"]) + } + byKind := st["by_kind"].(map[string]int) + if byKind["image"] != 2 || byKind["audio"] != 1 { + t.Fatalf("by_kind 不对: %v", byKind) + } +} + +func TestPut_RejectsEmpty(t *testing.T) { + s := newTestStore(t, 0) + if _, err := s.Put(nil, Item{MIME: "image/png"}); err == nil { + t.Fatal("空内容应报错") + } +} + +func TestReopen_PersistsAcrossRestart(t *testing.T) { + // 记忆的意义就在于跨重启还在。 + dir := t.TempDir() + s1, err := New(dir, 0) + if err != nil { + t.Fatal(err) + } + d, _ := s1.Put([]byte("persistent-img"), Item{MIME: "image/png", OriginPath: "/tmp/x.png"}) + s1.Describe(d, "跨重启的描述", "vis") + s1.AddRef(d, "context", "e1") + s1.Close() + + s2, err := New(dir, 0) + if err != nil { + t.Fatal(err) + } + defer s2.Close() + + it, err := s2.Stat(d) + if err != nil { + t.Fatalf("重开后应能查到: %v", err) + } + if it.Description != "跨重启的描述" || it.RefCount != 1 { + t.Fatalf("元数据应持久化: %+v", it) + } + data, err := s2.Get(d) + if err != nil || string(data) != "persistent-img" { + t.Fatalf("内容应持久化: %v / %q", err, data) + } +}