From 6afe361804573d9d68ac476a28b7d428d475ac04 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Tue, 15 Sep 2026 06:32:36 +0800 Subject: [PATCH] =?UTF-8?q?fix(memory):=20=E8=AE=B0=E5=BF=86=E5=B1=82?= =?UTF-8?q?=E5=90=AF=E5=8A=A8=E6=8E=A5=E7=BA=BF/=E5=B9=B6=E5=8F=91/?= =?UTF-8?q?=E8=90=BD=E7=9B=98=E4=B8=80=E8=87=B4=E6=80=A7=E6=95=B4=E5=A4=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 按设计方案整顿记忆系统，收敛一批"单测照不出、只在长跑生产里暴露"的缺陷： - 启动接线：initMemoryStack 残留 `defer distiller.Stop()`，规则蒸馏 10min 心跳启动即死。改为由调用点 cleanup 停机，并补 Stopped() 探针 + TestInitMemoryStackKeepsDistillerRunning / TestStartKeepsLoopRunningUntilStop。 - L0 相关性上下文：SetDenseSpace 注入稠密空间时回填已有事件的稠密向量， 否则旧事件走稀疏余弦、新事件走稠密余弦，同一次 Prune 里两种尺度混排。 - 文档检索：QueryScored 访问计数从读锁内写移出（-race 竞争），更新后置脏， 优雅关停可落盘、FindColdDocs 冷度判据跨重启不再失真。 - 蒸馏管线：只 flush 未落盘记录（persisted 标记）、蒸馏成功后从 raw 文件 删除对应行、原子写文件，修重启重复蒸馏导致 mention_count 膨胀。 - 图库：全量 Recall 加实体上限（内部整备路径，防大图整表进内存）； ClearSentenceID 补写锁；Commit/upsertEntity 计数语义注释澄清。 - 索引器：recalled 去重集加 FIFO 上限，防长跑进程自动注入越来越沉默。 - 文档/媒体注释修正；README 记忆层流程对齐跨模态召回。 验证：go build ./...、go vet、go test -race ./internal/memory/... ./internal/agent/core/... ./cmd/homed/... 全绿。 --- README.md | 2 +- cmd/homed/bootstrap.go | 8 +- cmd/homed/bootstrap_test.go | 42 ++++++ internal/agent/core/context.go | 29 ++++ internal/agent/core/context_test.go | 36 +++++ internal/memory/document/document.go | 69 ++++++---- internal/memory/document/document_test.go | 54 ++++++++ internal/memory/graph.go | 28 +++- internal/memory/indexer.go | 19 +++ internal/memory/media/media.go | 2 +- internal/memory/pipeline/pipeline.go | 157 ++++++++++++++++++---- internal/memory/pipeline/pipeline_test.go | 28 +++- 12 files changed, 413 insertions(+), 61 deletions(-) create mode 100644 cmd/homed/bootstrap_test.go diff --git a/README.md b/README.md index 7570ebd..8ed8944 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ flowchart TB end subgraph D[② Document 文件记忆] DS[DocStore JSON+TF-IDF] - Q1[Query 摘要自动注入] -->|【相关记忆文档】| SP + Q1[QueryScored+crossModalMarkdown] -->|【跨模态相关记忆】| SP Q2[doc_query LLM主动召回] -->|Consume+删除源| DS Q2 -->|原始时间戳写入上下文| RC CD[FindColdDocs 72h] -->|docToTriples| G diff --git a/cmd/homed/bootstrap.go b/cmd/homed/bootstrap.go index 7a3944f..6f65ef5 100644 --- a/cmd/homed/bootstrap.go +++ b/cmd/homed/bootstrap.go @@ -147,9 +147,6 @@ func initMemoryStack(dataDir string) (*memoryStack, func()) { } else { log.Printf("[homed] graph memory initialized") } - if memDB != nil { - } - memIdx := memory.NewIndexer(memDB) memIdx.Sync() // 启动时立即同步，避免前30分钟空窗 socialStore := social.New(memDB) @@ -160,8 +157,11 @@ func initMemoryStack(dataDir string) (*memoryStack, func()) { BatchSize: 50, }) if memDB != nil { + // 这里**故意不写 defer distiller.Stop()**：本函数在 return 时即触发 + // defer，而 Stop() → cancel() 会让刚启动的 distillLoop 立刻退出， + // 规则蒸馏管线启动即死、10min 心跳从不运行（旧 main() 拆分时的残留）。 + // 停机由调用点注册的 cleanup 负责（见下方返回值）。 distiller.Start() - defer distiller.Stop() } return &memoryStack{db: memDB, indexer: memIdx, social: socialStore, distiller: distiller}, diff --git a/cmd/homed/bootstrap_test.go b/cmd/homed/bootstrap_test.go new file mode 100644 index 0000000..eddc87d --- /dev/null +++ b/cmd/homed/bootstrap_test.go @@ -0,0 +1,42 @@ +package main + +import ( + "os" + "path/filepath" + "testing" +) + +// TestInitMemoryStackKeepsDistillerRunning 锁死启动接线回归： +// initMemoryStack 必须返回一个**仍在运行**的蒸馏器。 +// +// 历史 bug：main() 拆分时函数体内残留一句 `defer distiller.Stop()`， +// 函数一 return 就 cancel 掉刚启动的循环，规则蒸馏 10min 心跳从不运行。 +// 该缺陷不会让任何单测变红——pipeline 的 TestDistillOnce* 直接调 +// distillOnce，绕过了 Start/Stop 接线；只有在这里按「启动阶段函数」的 +// 真实调用方式断言，才照得出来。 +func TestInitMemoryStackKeepsDistillerRunning(t *testing.T) { + dir := t.TempDir() + // NewGraphDB 需要父目录已存在（生产由 dataDir 初始化保证）。 + if err := os.MkdirAll(filepath.Join(dir, "memory"), 0755); err != nil { + t.Fatal(err) + } + st, cleanup := initMemoryStack(dir) + if st == nil || st.distiller == nil { + cleanup() + t.Fatal("initMemoryStack 未返回蒸馏器") + } + if st.db == nil { + cleanup() + t.Skip("图库未初始化，无法验证蒸馏接线") + } + if st.distiller.Stopped() { + cleanup() + t.Fatal("initMemoryStack 返回后蒸馏循环已被停掉（defer Stop 残留？）") + } + + // cleanup 是唯一的停机点：先停蒸馏器、再关图库。 + cleanup() + if !st.distiller.Stopped() { + t.Fatal("cleanup 之后蒸馏器应已停止") + } +} diff --git a/internal/agent/core/context.go b/internal/agent/core/context.go index 8a38bb4..ec90426 100644 --- a/internal/agent/core/context.go +++ b/internal/agent/core/context.go @@ -3,6 +3,7 @@ package core import ( "encoding/json" "fmt" + "log" "os" "path/filepath" "sort" @@ -68,10 +69,38 @@ func NewRelevanceContext(savePath string, embedder *memory.StaticEmbedder) *Rele // SetDenseSpace 注入稠密多模态向量空间。配置后 L0 相关性裁剪可用稠密向量 // 余弦（与媒体检索、文档检索共享同一空间），未配置时退化到稀疏词向量。 +// +// 注入时**回填已有事件**的稠密向量。为什么必须回填：NewRelevanceContext 先 +// load()、再 SetDenseSpace，载入时 c.denseSpace 还是 nil，旧事件只算了稀疏 +// 向量；若这里只赋值不回填，Prune 里旧事件因 DenseFP 为空、长度不符而全部 +// 走稀疏余弦，新事件走稠密余弦 —— 同一次排序里两种尺度混排，谁留下谁归档 +// 取决于事件新旧而非相关性。对齐 DocStore.BuildDenseIndex 的做法。 +// +// 注意 DenseVec/DenseFP 刻意不持久化（json:"-"）：这是每次启动一次性重算的 +// 缓存，不落盘，因此这里也不需要 Save。 func (c *RelevanceContext) SetDenseSpace(ds vector.MultimodalEmbedder) { c.mu.Lock() defer c.mu.Unlock() c.denseSpace = ds + if ds == nil || !ds.Loaded() { + return + } + fp := ds.Fingerprint() + dim := ds.Dim() + filled := 0 + for _, evt := range c.events { + if evt == nil { + continue + } + if evt.DenseFP == fp && len(evt.DenseVec) == dim { + continue + } + c.computeVector(evt) + filled++ + } + if filled > 0 { + log.Printf("[agent] context dense backfill: %d events", filled) + } } func (c *RelevanceContext) SetToolDefLookup(fn func(name string) *sdk.ToolDef) { diff --git a/internal/agent/core/context_test.go b/internal/agent/core/context_test.go index bbb6f97..dbf996b 100644 --- a/internal/agent/core/context_test.go +++ b/internal/agent/core/context_test.go @@ -2,6 +2,7 @@ package core import ( "os" + "path/filepath" "testing" "time" @@ -237,6 +238,41 @@ func containsStr(s, substr string) bool { return false } +// TestSetDenseSpaceBackfillsExistingEvents 锁死 L0 稠密回填： +// +// NewRelevanceContext 先 load()（此时 denseSpace 仍为 nil，旧事件只算了稀疏 +// 向量），SetDenseSpace 才注入稠密空间。若不回填已有事件，它们的 DenseFP +// 为空、DenseVec 长度不符，Prune 里旧事件走稀疏余弦、新事件走稠密余弦—— +// 同一次排序里混排两种尺度，谁留下只取决于事件新旧。 +func TestSetDenseSpaceBackfillsExistingEvents(t *testing.T) { + path := filepath.Join(t.TempDir(), "context.json") + + // run1：写入事件并落盘（不注入稠密空间）。 + c1 := NewRelevanceContext(path, memory.NewStaticEmbedder("")) + c1.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "昨天的决定"}) + c1.Append(ContextEvent{Timestamp: time.Now(), Source: "agent", Response: "记为待办"}) + if err := c1.Save(); err != nil { + t.Fatal(err) + } + + // run2：模拟重启——load() 发生在 SetDenseSpace 之前。 + c2 := NewRelevanceContext(path, memory.NewStaticEmbedder("")) + if c2.Len() == 0 { + t.Fatal("重启后未读回任何事件") + } + c2.SetDenseSpace(fakeSpace{}) + + events := c2.Recent(c2.Len()) + if len(events) == 0 { + t.Fatal("no events") + } + for i, e := range events { + if e.DenseFP != "fake-space" || len(e.DenseVec) != 2 { + t.Errorf("event %d 未回填稠密向量: fp=%q len=%d", i, e.DenseFP, len(e.DenseVec)) + } + } +} + func splitLines(s string) []string { var lines []string start := 0 diff --git a/internal/memory/document/document.go b/internal/memory/document/document.go index ec0c6a4..b008b07 100644 --- a/internal/memory/document/document.go +++ b/internal/memory/document/document.go @@ -59,8 +59,6 @@ type Store struct { denseSpace vector.MultimodalEmbedder } -const maxSummaries = 10000 - // NewStore 创建文档存储。tokenizer 由外层注入（如 jieba），核心不直接依赖分词库。 func NewStore(dir string, tokenizer tfidf.Tokenizer) *Store { return &Store{ @@ -359,44 +357,63 @@ type DocHit struct { } func (s *Store) QueryScored(text string, topK int) []DocHit { - s.mu.RLock() - defer s.mu.RUnlock() if topK <= 0 { topK = 5 } - // 主路径 + var hits []DocHit + + // 主路径：稠密检索。只读，持读锁。 + s.mu.RLock() if s.denseSpace != nil && s.denseSpace.Loaded() { if qv, err := s.denseSpace.VectorizeDense(text); err == nil { - results := s.denseSearchScored(qv, topK) - for i := range results { - if d, ok := s.docs[results[i].Doc.ID]; ok { - d.AccessCount++ - d.LastAccess = time.Now() - results[i].Doc = d - } - } - return results + hits = s.denseSearchScored(qv, topK) } } + s.mu.RUnlock() // Fallback（需要写锁来 ensureTFIDF） - s.mu.RUnlock() - s.mu.Lock() - s.ensureTFIDF() - s.mu.Unlock() - s.mu.RLock() + if hits == nil { + s.mu.Lock() + s.ensureTFIDF() + results := s.tfidfIdx.Search(text, topK) + for _, r := range results { + if d, ok := s.docs[r.ID]; ok { + hits = append(hits, DocHit{Doc: d, Score: r.Score}) + } + } + s.mu.Unlock() + } - results := s.tfidfIdx.Search(text, topK) - var out []DocHit - for _, r := range results { - if d, ok := s.docs[r.ID]; ok { + // 访问计数是**写**：必须离开读锁后再取写锁更新。 + // + // 此前在 RLock 下直接 `d.AccessCount++` / `d.LastAccess = time.Now()`， + // 与 FindColdDocs / flush 的读并发构成数据竞争（-race 实测），且不置脏 —— + // 计数只在内存里涨，优雅关停也不落盘，FindColdDocs 的「≤2 次访问」 + // 冷度判据跨重启失真。 + s.markAccess(hits) + return hits +} + +// markAccess 记录一次检索命中：累加访问计数、刷新最后访问时间并置脏等待落盘。 +// 调用方不得持有 s.mu。 +func (s *Store) markAccess(hits []DocHit) { + if len(hits) == 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + for _, h := range hits { + if h.Doc == nil { + continue + } + if d, ok := s.docs[h.Doc.ID]; ok { d.AccessCount++ - d.LastAccess = time.Now() - out = append(out, DocHit{Doc: d, Score: r.Score}) + d.LastAccess = now } } - return out + s.dirty = true } func (s *Store) denseSearchScored(queryVec []float64, topK int) []DocHit { diff --git a/internal/memory/document/document_test.go b/internal/memory/document/document_test.go index 399905e..347e6fa 100644 --- a/internal/memory/document/document_test.go +++ b/internal/memory/document/document_test.go @@ -1,7 +1,9 @@ package document import ( + "fmt" "os" + "sync" "testing" "time" @@ -652,3 +654,55 @@ func TestStopFlushesDirtyDocs(t *testing.T) { t.Fatalf("Stop 未落盘: %+v", got) } } + +// TestQueryScoredConcurrentAndDirty 锁死两点： +// 1. QueryScored 更新访问计数必须离开读锁（旧实现在 RLock 下写，-race 会报）； +// 2. 更新后必须置脏，否则优雅关停也不落盘、FindColdDocs 的冷度判据跨重启失真。 +func TestQueryScoredConcurrentAndDirty(t *testing.T) { + dir := t.TempDir() + s := NewStore(dir, memory.TokenizeWords) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + for i := 0; i < 8; i++ { + // 文档必须互相区分：若 8 篇文本完全相同，查询词出现在每一篇里， + // IDF = log(N/df) = 0，词向量全零、检索恒为空，测试就成了假阴性。 + summary, content := "无关的天气与散步记录", "今天适合出门" + if i < 3 { + summary, content = "并发检索目标", "并发内容" + } + if err := s.Insert(&Doc{ID: fmt.Sprintf("d%d", i), Summary: summary, Content: content}); err != nil { + t.Fatal(err) + } + } + // 清掉 Insert 置的脏位，验证 QueryScored 自己会置脏。 + s.mu.Lock() + s.dirty = false + s.mu.Unlock() + + var wg sync.WaitGroup + for g := 0; g < 8; g++ { + wg.Add(1) + go func() { + defer wg.Done() + for k := 0; k < 30; k++ { + s.QueryScored("并发检索目标", 5) + s.FindColdDocs(72*time.Hour, 2) + } + }() + } + wg.Wait() + + s.mu.RLock() + dirty := s.dirty + access := s.docs["d0"].AccessCount + s.mu.RUnlock() + if access <= 1 { + t.Errorf("AccessCount 未随检索累加: %d", access) + } + if !dirty { + t.Error("QueryScored 更新访问计数后未置脏") + } +} diff --git a/internal/memory/graph.go b/internal/memory/graph.go index 0c6a984..40664fc 100644 --- a/internal/memory/graph.go +++ b/internal/memory/graph.go @@ -4,6 +4,7 @@ import ( "database/sql" "encoding/json" "fmt" + "log" "strings" "sync" "time" @@ -11,6 +12,11 @@ import ( _ "github.com/mattn/go-sqlite3" ) +// maxFullRecallEntities 是「无关键词全量读取」路径的实体上限。 +// 该路径只服务于内部整备（Indexer.Sync / 实体合并检测），并非用户检索； +// 无上限时一张大图会被整表 read 进内存。超限时 GraphDB.Recall 会记日志。 +const maxFullRecallEntities = 10000 + type Entity struct { ID int64 `json:"id"` Name string `json:"name"` @@ -255,10 +261,14 @@ func (g *GraphDB) migrateRelationUnique(tx *sql.Tx) error { return nil } -// Commit 把三元组写入图库，返回新建的实体数与关系数。 // Path 返回本库的存储路径（父 agent 用它为驻留子打开**受限句柄**）。 func (g *GraphDB) Path() string { return g.dbPath } +// Commit 把三元组写入图库，返回通过实体名校验并写入/刷新的实体数与**新建**的关系数。 +// +// 两个计数的语义刻意不同，因为上游只用它们判断「有没有东西写进去」： +// 实体计数含已存在实体的 mention_count 刷新（见 upsertEntity）， +// 关系计数只统计真正新建的关系（已存在则仅刷新 confidence）。 func (g *GraphDB) Commit(triples []Triple, sessionID string, turnID int) (int, int, error) { _, ec, rc, err := g.commit(triples, sessionID, turnID, false) return ec, rc, err @@ -409,6 +419,12 @@ func validEntityName(name string) bool { return hasLetter } +// upsertEntity 写入/刷新一个实体，返回 1 表示该实体**通过名校验并被写入或刷新**， +// 0 表示名校验未通过。 +// +// 注意返回值语义不是「新建数」：`ON CONFLICT DO UPDATE` 在更新时 +// RowsAffected 同样为 1，所以返回值等于「通过校验的 upsert 次数」。 +// 调用方（Commit）把它当「写入了几个实体」用，不是「新建了几个」。 func (g *GraphDB) upsertEntity(tx *sql.Tx, name string, entityType string) (int, error) { if !validEntityName(name) { return 0, nil @@ -442,9 +458,12 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se result := &RecallResult{} if len(keywords) == 0 && len(seedEntities) == 0 { + // 全量读取仅用于内部整备（Indexer.Sync / 实体合并检测）， + // 必须加限额：无 LIMIT 时大图会被整表读进内存。 rows, err := g.db.Query( `SELECT id, name, type, mention_count, created_at, updated_at - FROM entities ORDER BY mention_count DESC`, + FROM entities ORDER BY mention_count DESC LIMIT ?`, + maxFullRecallEntities, ) if err != nil { return nil, err @@ -458,6 +477,9 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se } result.Entities = append(result.Entities, e) } + if len(result.Entities) >= maxFullRecallEntities { + log.Printf("[graph] full recall 命中实体上限 %d，可能有实体未纳入", maxFullRecallEntities) + } relRows, err := g.db.Query( `SELECT r.id, r.source_id, r.target_id, e1.name, e2.name, @@ -1047,6 +1069,8 @@ func (g *GraphDB) Archive(days int) (int, error) { // ClearSentenceID 清除指定关系的 sentence_id（LLM复审后解除句子引用） func (g *GraphDB) ClearSentenceID(relationID int64) error { + g.mu.Lock() + defer g.mu.Unlock() _, err := g.db.Exec( `UPDATE relations SET sentence_id = 0, updated_at = CURRENT_TIMESTAMP WHERE id = ?`, relationID, diff --git a/internal/memory/indexer.go b/internal/memory/indexer.go index 732d79f..e865994 100644 --- a/internal/memory/indexer.go +++ b/internal/memory/indexer.go @@ -16,8 +16,18 @@ type Indexer struct { mu sync.RWMutex trained bool recalled map[string]bool // 已通过工具调用显式召回的实体名，自动注入时跳过 + + // recalledOrder 记录 recalled 的插入顺序，用于超限时按 FIFO 淘汰。 + recalledOrder []string } +// maxRecalledEntities 是「已召回实体」去重集的上限。 +// +// 无上限时它只增不减：进程活得越久，被永久跳过的实体越多，自动注入 +// 越来越「沉默」——一个只在长跑进程里才暴露的隐蔽退化。超限后淘汰最旧的 +// 名字（允许重新注入），而不是丢弃整个集合。 +const maxRecalledEntities = 1024 + func NewIndexer(db *GraphDB) *Indexer { return &Indexer{ db: db, @@ -32,7 +42,16 @@ func (idx *Indexer) MarkRecalled(names ...string) { idx.mu.Lock() defer idx.mu.Unlock() for _, name := range names { + if idx.recalled[name] { + continue + } idx.recalled[name] = true + idx.recalledOrder = append(idx.recalledOrder, name) + } + for len(idx.recalledOrder) > maxRecalledEntities { + oldest := idx.recalledOrder[0] + idx.recalledOrder = idx.recalledOrder[1:] + delete(idx.recalled, oldest) } } diff --git a/internal/memory/media/media.go b/internal/memory/media/media.go index de6c5af..e5a5f3e 100644 --- a/internal/memory/media/media.go +++ b/internal/memory/media/media.go @@ -235,7 +235,7 @@ func (s *Store) Stat(digest string) (*Item, error) { FROM media WHERE digest = ?`, digest)) } -// Stat 返回元数据，不读内容。 +// Delete 删除媒体内容与元数据。 // // 这不是 GC，也不看引用计数：调用方是记忆系统本身——当它把一个记忆块 // 永久地从三层记忆中删掉（而非在层间迁移）时，媒体作为块的内容一并删除。 diff --git a/internal/memory/pipeline/pipeline.go b/internal/memory/pipeline/pipeline.go index ac6a8f6..ca40193 100644 --- a/internal/memory/pipeline/pipeline.go +++ b/internal/memory/pipeline/pipeline.go @@ -24,6 +24,12 @@ type RawRecord struct { Content string `json:"content"` CreatedAt time.Time `json:"created_at"` Distilled bool `json:"distilled"` + + // persisted 表示该记录已经写在磁盘 raw 文件里。 + // 未导出：只影响本进程的落盘行为，不进 JSON。 + // 作用：flush 只写**尚未落盘**的记录，否则 loadExisting 载入的记录 + // 会被再写一份，重启后同一批记录从新旧两个文件各读回一次。 + persisted bool } type DistillerConfig struct { @@ -33,16 +39,16 @@ type DistillerConfig struct { } type Distiller struct { - mu sync.Mutex - db *memory.GraphDB - rawPath string - records []RawRecord - nextID int64 - cfg DistillerConfig - ctx context.Context - cancel context.CancelFunc - onMemory func(input, response string) - embedder nlp.Vectorizer + mu sync.Mutex + db *memory.GraphDB + rawPath string + records []RawRecord + nextID int64 + cfg DistillerConfig + ctx context.Context + cancel context.CancelFunc + onMemory func(input, response string) + embedder nlp.Vectorizer } func (d *Distiller) SetEmbedder(ev nlp.Vectorizer) { d.embedder = ev } @@ -76,6 +82,22 @@ func (d *Distiller) Stop() { d.flush() } +// Stopped 报告蒸馏循环是否已被 Stop() 取消。 +// +// Stop() 里的 cancel() 是同步生效的，所以本方法在 Stop() 返回后立即为 true， +// 不受循环 goroutine 何时退出的影响。启动自检、健康检查用它确认 +// 「Start 之后没有被立即 Stop 掉」——历史回归：main() 拆分时 +// initMemoryStack 里残留一句 defer distiller.Stop()，函数一返回就把刚起的 +// 循环杀了，10min 心跳从不运行。Start() 之前返回 false（尚未被停）。 +func (d *Distiller) Stopped() bool { + select { + case <-d.ctx.Done(): + return true + default: + return false + } +} + func (d *Distiller) Append(sessionID string, role string, content string) { d.mu.Lock() defer d.mu.Unlock() @@ -89,19 +111,35 @@ func (d *Distiller) Append(sessionID string, role string, content string) { func (d *Distiller) flush() { d.mu.Lock() defer d.mu.Unlock() - if len(d.records) == 0 { + d.flushLocked() +} + +// flushLocked 把**尚未落盘**的记录追加写入一个新的 raw 文件（原子写）。 +// +// 只写 !persisted 的记录：loadExisting 载入的记录已经在磁盘上，若 flush 再把 +// 它们整体重写一份，重启后同一批记录会同时从旧文件与新文件被读回， +// 实体 mention_count 与关系被重复蒸馏。 +func (d *Distiller) flushLocked() { + var pending []RawRecord + for _, r := range d.records { + if !r.persisted { + pending = append(pending, r) + } + } + if len(pending) == 0 { return } path := filepath.Join(d.rawPath, fmt.Sprintf("raw_%d.tsv", time.Now().UnixNano())) - f, err := os.Create(path) - if err != nil { + var sb strings.Builder + for _, r := range pending { + fmt.Fprintf(&sb, "%d\t%s\t%s\t%s\t%d\n", r.ID, r.SessionID, r.Role, r.Content, r.CreatedAt.Unix()) + } + if err := writeFileAtomic(path, []byte(sb.String())); err != nil { log.Printf("[memory] flush error: %v", err) return } - defer f.Close() - for _, r := range d.records { - line := fmt.Sprintf("%d\t%s\t%s\t%s\t%d\n", r.ID, r.SessionID, r.Role, r.Content, r.CreatedAt.Unix()) - f.WriteString(line) + for i := range d.records { + d.records[i].persisted = true } } @@ -167,6 +205,7 @@ func (d *Distiller) loadExisting() { } d.records = append(d.records, RawRecord{ ID: d.nextID, SessionID: parts[1], Role: parts[2], Content: parts[3], CreatedAt: createdAt, + persisted: true, }) d.nextID++ loaded++ @@ -222,6 +261,9 @@ func (d *Distiller) distillOnce() { } if d.distillBatch(toDistill[i:end]) { distilled += end - i + // 成功即从磁盘 raw 文件里删除这些行，否则重启后 loadExisting + // 会把它们当未蒸馏记录重新读回，每次启动重蒸同一批历史。 + d.removeRawRecords(toDistill[i:end]) } else { // 蒸馏失败：记录写回待处理队列，下次 tick 重试 d.mu.Lock() @@ -279,6 +321,80 @@ func (d *Distiller) cleanupRawFiles() { } } +// writeFileAtomic 写临时文件再 rename，避免进程在写一半时崩溃留下半个文件。 +func writeFileAtomic(path string, data []byte) error { + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0644); err != nil { + return err + } + return os.Rename(tmp, path) +} + +// rawKey 唯一标识一条原始记录，用于在 raw 文件里按内容定位并删除。 +// +// 为什么按内容而不是 ID：loadExisting 读回记录时会重新分配连续 ID， +// 文件里的 ID 与内存 ID 并不一一对应。 +func rawKey(session, role, content string, ts int64) string { + return session + "\x00" + role + "\x00" + content + "\x00" + strconv.FormatInt(ts, 10) +} + +// removeRawRecords 从磁盘 raw 文件中删除已成功蒸馏的记录。 +// +// 蒸馏成功后记录若只从内存移除、磁盘文件不动，下次启动 loadExisting 会把 +// 它们当未蒸馏记录重新读回，导致每次重启都重蒸同一批历史（实体 +// mention_count 膨胀，且 distillBatch 的 sessionID 取自 map 首个键， +// 不确定性会放大重复）。 +func (d *Distiller) removeRawRecords(batch []RawRecord) { + if len(batch) == 0 { + return + } + drop := make(map[string]bool, len(batch)) + for _, r := range batch { + drop[rawKey(r.SessionID, r.Role, r.Content, r.CreatedAt.Unix())] = true + } + entries, err := os.ReadDir(d.rawPath) + if err != nil { + return + } + for _, entry := range entries { + ext := filepath.Ext(entry.Name()) + if ext != ".tsv" && ext != ".jsonl" { + continue + } + path := filepath.Join(d.rawPath, entry.Name()) + data, err := os.ReadFile(path) + if err != nil { + continue + } + var kept []string + removed := false + for _, line := range strings.Split(strings.TrimRight(string(data), "\n"), "\n") { + if line == "" { + continue + } + parts := splitLine(line) + if len(parts) >= 5 { + if ts, err := strconv.ParseInt(parts[4], 10, 64); err == nil && + drop[rawKey(parts[1], parts[2], parts[3], ts)] { + removed = true + continue + } + } + kept = append(kept, line) + } + if !removed { + continue + } + if len(kept) == 0 { + os.Remove(path) + continue + } + if err := writeFileAtomic(path, []byte(strings.Join(kept, "\n")+"\n")); err != nil { + log.Printf("[memory] rewrite raw %s: %v", entry.Name(), err) + } + } +} + func extractKeyTriples(userContent, assistantContent string, embedder nlp.Vectorizer) []memory.Triple { var triples []memory.Triple @@ -310,13 +426,6 @@ func truncate(s string, max int) string { return s } -func parseLines(data string) []string { - if data == "" { - return nil - } - return strings.Split(strings.TrimRight(data, "\n"), "\n") -} - func splitLine(line string) []string { if line == "" { return nil diff --git a/internal/memory/pipeline/pipeline_test.go b/internal/memory/pipeline/pipeline_test.go index 62f68de..63d6b85 100644 --- a/internal/memory/pipeline/pipeline_test.go +++ b/internal/memory/pipeline/pipeline_test.go @@ -185,9 +185,9 @@ func TestDistillOnceBatchLimit(t *testing.T) { func TestExtractKeyTriples(t *testing.T) { tests := []struct { - user string + user string assistant string - check func([]memory.Triple) bool + check func([]memory.Triple) bool }{ { user: "我住在北京", @@ -201,7 +201,7 @@ func TestExtractKeyTriples(t *testing.T) { }, }, { - user: "我在杭州读书", + user: "我在杭州读书", assistant: "好的", check: func(triples []memory.Triple) bool { for _, tr := range triples { @@ -264,3 +264,25 @@ func TestTruncate(t *testing.T) { t.Errorf("expected 'hi', got %q", truncate("hi", 10)) } } + +// TestStartKeepsLoopRunningUntilStop 锁死 Start/Stop 的接线契约： +// Start 之后蒸馏循环必须处于运行态，只有显式 Stop 才退出。 +// +// 这条回归的直接来源：main() 拆分时 initMemoryStack 里残留一句 +// defer distiller.Stop()，函数返回即 cancel，循环启动即死——而 +// TestDistillOnce* 直接调 distillOnce，绕过了 Start/Stop，照不出这个洞。 +func TestStartKeepsLoopRunningUntilStop(t *testing.T) { + d := NewDistiller(nil, t.TempDir(), DistillerConfig{ + Interval: time.Hour, // 不依赖 tick，只验证循环存活 + RetentionDays: 7, + BatchSize: 50, + }) + d.Start() + if d.Stopped() { + t.Fatal("Start 之后蒸馏循环必须处于运行态（不可被 defer Stop 杀掉）") + } + d.Stop() + if !d.Stopped() { + t.Fatal("Stop 之后应处于已停止态") + } +}