diff --git a/docs/api/memory.md b/docs/api/memory.md index fb7b4150d..33fc9792d 100644 --- a/docs/api/memory.md +++ b/docs/api/memory.md @@ -87,6 +87,14 @@ curl --location --request PUT 'http://localhost:8080/api/v1/memory/settings' \ 推断出的记忆（`status=pending`）确认后才注入提示词；否决会留下 tombstone，避免下一轮蒸馏把同一句话再写回来。 +如果推断用于修改已有记忆，确认前旧记忆继续生效；确认会在同一事务中激活推断并替换旧条目。已经失效、到期或基于已修改或已删除内容的推断不能再次确认，接口返回 HTTP 409，客户端应刷新列表。编辑正文会同时使旧向量失效，并尝试生成新向量；生成失败时保留词面检索，后续整理可补齐向量。 + +自动提取在独立的 `memory_extraction_sessions` 表按会话保存 `(created_at, message_id)` 游标、排队状态和重试计数，每次最多领取 32 个会话。用户记录只保存工作进程租约，读取记忆时不加载全部会话历史。已完成的游标仍保留，避免旧会话再次提取全部历史。 + +单条无效决策或已失效的更新目标不会阻塞同批其他结果。无法解析或持续截断的模型输出按消息段最多尝试三轮（每轮截断后可增加预算重试一次），随后记录失败范围并推进该段游标；同会话后续消息及其他会话继续处理。进度表保留该会话最近一次失败的起止游标、错误类别、次数和跳过时间，不保存原始模型输出；被跳过的范围不会自动补提取。网络、模型配置和数据库错误不适用跳过策略，仍保留进度供重试。批次截断、工作进程重启和提取期间的新消息都保留待处理状态。 + +升级时沿用旧的用户级游标作为各会话的初始边界，之后该边界不再推进；旧待处理列表在首次调度时导入独立表。这样避免升级触发全量历史重放，但不会自动恢复旧实现已跳过、且位于旧游标之前的消息。 + ## GET `/memory/export` 返回 `{success, total, truncated, data}`，并带 `Content-Disposition: attachment; filename="weknora-memories.json"`。`truncated` 仅在触达导出上限（2 万条）时为 true。 diff --git a/internal/application/repository/memory.go b/internal/application/repository/memory.go index 4b10bb80b..6ca5c329e 100644 --- a/internal/application/repository/memory.go +++ b/internal/application/repository/memory.go @@ -97,108 +97,6 @@ func (r *memoryRepository) UpdateSubjectBlock( }).Error } -// EnqueuePendingSession is the whole "never drop a turn" mechanism, so it runs -// inside a transaction: reading the subject, appending the session and claiming -// the in-flight slot must not interleave with a concurrent turn, or two turns -// could both decide nobody is scheduled (two tasks) or both decide someone is -// (a turn recorded against a run that already read the queue). -func (r *memoryRepository) EnqueuePendingSession( - ctx context.Context, scope interfaces.MemoryScope, sessionID string, inFlightTimeout time.Duration, -) (*types.MemorySubject, bool, error) { - var ( - snapshot types.MemorySubject - shouldSend bool - ) - err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var subject types.MemorySubject - if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID). - Clauses(forUpdateClause()). - First(&subject).Error; err != nil { - return err - } - snapshot = subject - - now := time.Now() - updates := map[string]interface{}{"updated_at": now} - if pending := subject.PendingSessions.Append(sessionID); len(pending) != len(subject.PendingSessions) { - updates["pending_sessions"] = pending - } - // A stale marker (worker crashed, task lost) must not wedge the subject - // forever, so the claim expires. - inFlight := subject.ExtractScheduledAt != nil && - now.Sub(*subject.ExtractScheduledAt) < inFlightTimeout - if !inFlight { - updates["extract_scheduled_at"] = now - shouldSend = true - } - return tx.Model(&types.MemorySubject{}). - Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID). - Updates(updates).Error - }) - if err != nil { - return nil, false, err - } - return &snapshot, shouldSend, nil -} - -// ClaimPendingSessions empties the queue and returns it together with the -// watermark to walk forward from. Emptying it here (rather than after the run) -// is deliberate: turns arriving during the run land in a fresh queue and -// trigger a follow-up, instead of being erased by the run that never saw them. -func (r *memoryRepository) ClaimPendingSessions( - ctx context.Context, scope interfaces.MemoryScope, -) ([]string, time.Time, error) { - var ( - pending []string - cursor time.Time - ) - err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var subject types.MemorySubject - if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID). - Clauses(forUpdateClause()). - First(&subject).Error; err != nil { - return err - } - pending = append(pending, subject.PendingSessions...) - if subject.ExtractCursor != nil { - cursor = *subject.ExtractCursor - } - return tx.Model(&types.MemorySubject{}). - Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID). - Updates(map[string]interface{}{ - "pending_sessions": types.MemoryPendingSessions{}, - "updated_at": time.Now(), - }).Error - }) - if err != nil { - return nil, time.Time{}, err - } - return pending, cursor, nil -} - -func (r *memoryRepository) FinishExtraction( - ctx context.Context, scope interfaces.MemoryScope, cursor time.Time, -) error { - now := time.Now() - updates := map[string]interface{}{ - "last_extracted_at": now, - "extract_scheduled_at": nil, - "updated_at": now, - } - if !cursor.IsZero() { - updates["extract_cursor"] = cursor - } - return r.scoped(ctx, scope).Model(&types.MemorySubject{}).Updates(updates).Error -} - -func (r *memoryRepository) ReleaseExtractionSlot( - ctx context.Context, scope interfaces.MemoryScope, -) error { - return r.scoped(ctx, scope). - Model(&types.MemorySubject{}). - Updates(map[string]interface{}{"extract_scheduled_at": nil, "updated_at": time.Now()}).Error -} - func (r *memoryRepository) CreateItem(ctx context.Context, item *types.MemoryItem) error { if item.ID == "" { item.ID = uuid.New().String() @@ -351,37 +249,63 @@ func (r *memoryRepository) FindActiveByKey( func (r *memoryRepository) UpdateItemContent( ctx context.Context, scope interfaces.MemoryScope, id, content, normalizedKey string, importance int, ) error { - return r.scoped(ctx, scope). - Model(&types.MemoryItem{}). - Where("id = ?", id). - Updates(map[string]interface{}{ - "content": content, - "normalized_key": normalizedKey, - "importance": importance, - "origin": types.MemoryOriginManual, - "updated_at": time.Now(), + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + var current types.MemoryItem + if err := tx.Where("tenant_id = ? AND subject_id = ? AND id = ?", scope.TenantID, scope.SubjectID, id).First(¤t).Error; err != nil { + return err + } + if current.Content != content { + if err := tx.Where("tenant_id = ? AND subject_id = ? AND item_id = ?", scope.TenantID, scope.SubjectID, id).Delete(&types.MemoryItemEmbedding{}).Error; err != nil { + return err + } + // Editing a confirmed fact invalidates proposals based on its old wording. + if err := tx.Model(&types.MemoryItem{}). + Where("tenant_id = ? AND subject_id = ? AND replaces_id = ? AND status = ?", + scope.TenantID, scope.SubjectID, id, types.MemoryStatusPending). + Updates(map[string]interface{}{"status": types.MemoryStatusSuperseded, "invalid_at": time.Now(), "superseded_by": id}).Error; err != nil { + return err + } + } + return tx.Model(¤t).Updates(map[string]interface{}{ + "content": content, "normalized_key": normalizedKey, "importance": importance, + "origin": types.MemoryOriginManual, "updated_at": time.Now(), }).Error + }) } func (r *memoryRepository) SupersedeItem( ctx context.Context, scope interfaces.MemoryScope, id, supersededBy string, ) error { - now := time.Now() - return r.scoped(ctx, scope). - Model(&types.MemoryItem{}). - Where("id = ? AND status = ?", id, types.MemoryStatusActive). - Updates(map[string]interface{}{ - "status": types.MemoryStatusSuperseded, - "invalid_at": now, - "superseded_by": supersededBy, - "updated_at": now, - }).Error + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + return tx.Model(&types.MemoryItem{}). + Where("tenant_id = ? AND subject_id = ? AND ((id = ? AND status = ?) OR (replaces_id = ? AND status = ?))", + scope.TenantID, scope.SubjectID, id, types.MemoryStatusActive, id, types.MemoryStatusPending). + Updates(map[string]interface{}{ + "status": types.MemoryStatusSuperseded, "invalid_at": time.Now(), + "superseded_by": supersededBy, "updated_at": time.Now(), + }).Error + }) } func (r *memoryRepository) DeleteItem( ctx context.Context, scope interfaces.MemoryScope, id string, ) error { - return r.scoped(ctx, scope).Where("id = ?", id).Delete(&types.MemoryItem{}).Error + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + if err := tx.Model(&types.MemoryItem{}). + Where("tenant_id = ? AND subject_id = ? AND replaces_id = ? AND status = ?", + scope.TenantID, scope.SubjectID, id, types.MemoryStatusPending). + Updates(map[string]interface{}{ + "status": types.MemoryStatusSuperseded, "invalid_at": time.Now(), + }).Error; err != nil { + return err + } + if err := tx.Where("tenant_id = ? AND subject_id = ? AND item_id = ?", scope.TenantID, scope.SubjectID, id). + Delete(&types.MemoryItemEmbedding{}).Error; err != nil { + return err + } + return tx.Where("tenant_id = ? AND subject_id = ? AND id = ?", scope.TenantID, scope.SubjectID, id). + Delete(&types.MemoryItem{}).Error + }) } func (r *memoryRepository) DeleteAll( @@ -561,12 +485,27 @@ func (r *memoryRepository) UpsertItemEmbedding( if embedding.CreatedAt.IsZero() { embedding.CreatedAt = now } - return r.db.WithContext(ctx). - Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "item_id"}}, - DoUpdates: clause.AssignmentColumns([]string{"model_id", "dims", "vector", "updated_at"}), - }). - Create(embedding).Error + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + if embedding.SourceContent != "" { + var current types.MemoryItem + err := tx.Where("tenant_id = ? AND subject_id = ? AND id = ?", scope.TenantID, scope.SubjectID, embedding.ItemID).First(¤t).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil + } + if err != nil { + return err + } + if current.Content != embedding.SourceContent || current.Topic != embedding.SourceTopic { + return nil + } + } + return tx. + Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "item_id"}}, + DoUpdates: clause.AssignmentColumns([]string{"model_id", "dims", "vector", "updated_at"}), + }). + Create(embedding).Error + }) } // DeleteItemEmbedding drops one memory's vector so the backfill rebuilds it. @@ -651,15 +590,6 @@ func (r *memoryRepository) MarkForcedConsolidated( Updates(map[string]interface{}{"forced_consolidated_at": now, "updated_at": now}).Error } -func (r *memoryRepository) SetItemStatus( - ctx context.Context, scope interfaces.MemoryScope, id, status string, -) error { - return r.scoped(ctx, scope). - Model(&types.MemoryItem{}). - Where("id = ?", id). - Updates(map[string]interface{}{"status": status, "updated_at": time.Now()}).Error -} - // BumpTopic counts one more sighting. The insert-then-increment shape keeps two // concurrent turns from both deciding the topic is new. func (r *memoryRepository) BumpTopic( diff --git a/internal/application/repository/memory_extraction.go b/internal/application/repository/memory_extraction.go new file mode 100644 index 000000000..cd94ea7e7 --- /dev/null +++ b/internal/application/repository/memory_extraction.go @@ -0,0 +1,230 @@ +package repository + +import ( + "context" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +func (r *memoryRepository) withSubject(ctx context.Context, scope interfaces.MemoryScope, fn func(*gorm.DB, *types.MemorySubject) error) error { + return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var subject types.MemorySubject + if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID). + Clauses(forUpdateClause()).First(&subject).Error; err != nil { + return err + } + return fn(tx, &subject) + }) +} + +func saveExtractionState(tx *gorm.DB, subject *types.MemorySubject) error { + return tx.Model(subject).Updates(map[string]interface{}{ + "extraction_state": subject.ExtractionState, "pending_sessions": subject.PendingSessions, + "extract_scheduled_at": subject.ExtractScheduledAt, "updated_at": time.Now(), + }).Error +} + +func extractionRows(tx *gorm.DB, scope interfaces.MemoryScope) *gorm.DB { + return tx.Model(&types.MemoryExtractionSession{}). + Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID) +} + +func enqueueExtractionSession(tx *gorm.DB, subject *types.MemorySubject, id string, bump bool) error { + if id == "" { + return nil + } + row := types.MemoryExtractionSession{ + TenantID: subject.TenantID, SubjectID: subject.SubjectID, SessionID: id, Revision: 1, Pending: true, + } + // Preserve the pre-upgrade boundary instead of replaying every historical + // conversation. This value is frozen; only per-session cursors advance now. + if subject.ExtractCursor != nil { + row.Cursor.At = *subject.ExtractCursor + } + conflict := clause.OnConflict{ + Columns: []clause.Column{{Name: "tenant_id"}, {Name: "subject_id"}, {Name: "session_id"}}, + DoNothing: true, + } + if bump { + conflict.DoNothing = false + conflict.DoUpdates = clause.Assignments(map[string]interface{}{ + "revision": gorm.Expr("memory_extraction_sessions.revision + 1"), "pending": true, "updated_at": time.Now(), + }) + } + return tx.Clauses(conflict).Create(&row).Error +} + +func importLegacySessions(tx *gorm.DB, subject *types.MemorySubject) error { + for _, id := range subject.PendingSessions { + if err := enqueueExtractionSession(tx, subject, id, false); err != nil { + return err + } + } + subject.PendingSessions = types.MemoryPendingSessions{} + return nil +} + +func hasPendingExtraction(tx *gorm.DB, scope interfaces.MemoryScope) (bool, error) { + var rows []types.MemoryExtractionSession + result := extractionRows(tx, scope).Select("session_id").Where("pending = ?", true).Limit(1).Find(&rows) + return len(rows) > 0, result.Error +} + +func (r *memoryRepository) HasPendingExtraction(ctx context.Context, scope interfaces.MemoryScope) (bool, error) { + return hasPendingExtraction(r.db.WithContext(ctx), scope) +} + +func (r *memoryRepository) EnqueuePendingSession(ctx context.Context, scope interfaces.MemoryScope, sessionID string, timeout time.Duration) (*types.MemorySubject, bool, error) { + var snapshot types.MemorySubject + shouldSend := false + err := r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + snapshot = *subject + if err := importLegacySessions(tx, subject); err != nil { + return err + } + if err := enqueueExtractionSession(tx, subject, sessionID, true); err != nil { + return err + } + now := time.Now() + running := subject.ExtractionState.LeaseID != "" && subject.ExtractionState.LeaseUntil.After(now) + queued := subject.ExtractScheduledAt != nil && now.Sub(*subject.ExtractScheduledAt) < timeout + if !running && !queued { + pending, err := hasPendingExtraction(tx, scope) + if err != nil { + return err + } + if pending { + subject.ExtractScheduledAt = &now + shouldSend = true + } + } + return saveExtractionState(tx, subject) + }) + return &snapshot, shouldSend, err +} + +func (r *memoryRepository) ClaimPendingSessions(ctx context.Context, scope interfaces.MemoryScope, fallbackSession, leaseID string, ttl time.Duration) (*types.MemoryExtractionBatch, error) { + var batch *types.MemoryExtractionBatch + err := r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + now := time.Now() + if subject.ExtractionState.LeaseID != "" && subject.ExtractionState.LeaseUntil.After(now) { + batch = &types.MemoryExtractionBatch{RetryAt: subject.ExtractionState.LeaseUntil} + return nil + } + if err := importLegacySessions(tx, subject); err != nil { + return err + } + // Bootstrap legacy payloads once; a duplicate must not reactivate a drained row. + if err := enqueueExtractionSession(tx, subject, fallbackSession, false); err != nil { + return err + } + var sessions []types.MemoryExtractionSession + if err := extractionRows(tx, scope).Where("pending = ?", true). + Order("updated_at ASC, session_id ASC"). + Limit(types.MaxMemoryPendingSessions).Find(&sessions).Error; err != nil { + return err + } + if len(sessions) == 0 { + return saveExtractionState(tx, subject) + } + batch = &types.MemoryExtractionBatch{Sessions: sessions} + subject.ExtractionState.LeaseID = leaseID + subject.ExtractionState.LeaseUntil = now.Add(ttl) + return saveExtractionState(tx, subject) + }) + return batch, err +} + +func validExtractionLease(subject *types.MemorySubject, leaseID string) bool { + return subject.ExtractionState.LeaseID == leaseID && subject.ExtractionState.LeaseUntil.After(time.Now()) +} + +func (r *memoryRepository) CheckpointExtraction(ctx context.Context, scope interfaces.MemoryScope, leaseID string, session types.MemoryExtractionSession, cursor types.MemoryMessageCursor, drained bool) error { + return r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + if !validExtractionLease(subject, leaseID) { + return types.ErrMemoryExtractionLeaseLost + } + var progress types.MemoryExtractionSession + if err := extractionRows(tx, scope).Where("session_id = ?", session.SessionID). + First(&progress).Error; err != nil { + return err + } + if cursor.After(progress.Cursor) { + progress.Cursor = cursor + } + // Updating just this row rotates unfinished work without rewriting the + // subject's entire history. Completed cursors remain small indexed records. + updates := map[string]interface{}{ + "cursor_at": progress.Cursor.At, "cursor_id": progress.Cursor.ID, + "pending": !drained || progress.Revision != session.Revision, "updated_at": time.Now(), + } + if progress.FailedAt == nil { + updates["failure_count"] = 0 + updates["failure_code"] = "" + } + return extractionRows(tx, scope).Where("session_id = ?", session.SessionID).Updates(updates).Error + }) +} + +func (r *memoryRepository) RecordExtractionFailure( + ctx context.Context, scope interfaces.MemoryScope, leaseID string, failure interfaces.MemoryExtractionFailure, +) (bool, error) { + skip := false + err := r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + if !validExtractionLease(subject, leaseID) { + return types.ErrMemoryExtractionLeaseLost + } + var progress types.MemoryExtractionSession + if err := extractionRows(tx, scope).Where("session_id = ?", failure.Session.SessionID). + First(&progress).Error; err != nil { + return err + } + attempts := 1 + if progress.FailedFrom.At.Equal(progress.Cursor.At) && progress.FailedFrom.ID == progress.Cursor.ID { + attempts += progress.FailureCount + } + skip = attempts >= 3 + updates := map[string]interface{}{ + "failure_count": attempts, "failure_code": failure.Code, "updated_at": time.Now(), + "failed_at": nil, + "failed_from_at": progress.Cursor.At, "failed_from_id": progress.Cursor.ID, + "failed_to_at": failure.End.At, "failed_to_id": failure.End.ID, + } + if skip { + updates["failed_at"] = time.Now() + } + return extractionRows(tx, scope).Where("session_id = ?", failure.Session.SessionID).Updates(updates).Error + }) + return skip, err +} + +func (r *memoryRepository) FinishExtraction(ctx context.Context, scope interfaces.MemoryScope, leaseID string) error { + return r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + if subject.ExtractionState.LeaseID != leaseID { + return types.ErrMemoryExtractionLeaseLost + } + subject.ExtractionState.LeaseID = "" + subject.ExtractionState.LeaseUntil = time.Time{} + subject.ExtractScheduledAt = nil + if err := saveExtractionState(tx, subject); err != nil { + return err + } + return tx.Model(subject).Update("last_extracted_at", time.Now()).Error + }) +} + +func (r *memoryRepository) ReleaseExtractionSlot(ctx context.Context, scope interfaces.MemoryScope, leaseID string) error { + return r.withSubject(ctx, scope, func(tx *gorm.DB, subject *types.MemorySubject) error { + if subject.ExtractionState.LeaseID != leaseID { + return nil + } + subject.ExtractionState.LeaseID = "" + subject.ExtractionState.LeaseUntil = time.Time{} + subject.ExtractScheduledAt = nil + return saveExtractionState(tx, subject) + }) +} diff --git a/internal/application/repository/memory_lifecycle.go b/internal/application/repository/memory_lifecycle.go new file mode 100644 index 000000000..f5647e6d6 --- /dev/null +++ b/internal/application/repository/memory_lifecycle.go @@ -0,0 +1,125 @@ +package repository + +import ( + "context" + "errors" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "gorm.io/gorm" +) + +// SaveItem serializes replacements with confirmation and manual edits. A +// proposal can replace another proposal, but cannot retire an active fact. +func (r *memoryRepository) SaveItem(ctx context.Context, scope interfaces.MemoryScope, item *types.MemoryItem, replacesID string) error { + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + query := func() *gorm.DB { return tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID) } + var target types.MemoryItem + if replacesID != "" { + if err := query().First(&target, "id = ?", replacesID).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return types.ErrMemoryConflict + } + return err + } + if target.Status != types.MemoryStatusActive && target.Status != types.MemoryStatusPending { + // A successfully applied decision may be replayed after a failed + // checkpoint. Return its replacement instead of writing it twice. + var replacement types.MemoryItem + if target.SupersededBy != "" && query().First(&replacement, "id = ?", target.SupersededBy).Error == nil && + replacement.Status == item.Status && replacement.Content == item.Content { + *item = replacement + return nil + } + return types.ErrMemoryConflict + } + } + var live []*types.MemoryItem + if err := query().Where("normalized_key = ? AND status IN ?", item.NormalizedKey, + []string{types.MemoryStatusActive, types.MemoryStatusPending}).Find(&live).Error; err != nil { + return err + } + for _, old := range live { + if old.Content == item.Content && old.Status == item.Status { + *item = *old + return nil + } + } + if item.Status == types.MemoryStatusPending { + if target.Status == types.MemoryStatusActive { + item.ReplacesID = target.ID + } + if target.Status == types.MemoryStatusPending { + item.ReplacesID = target.ReplacesID + } + if item.ReplacesID == "" { + for _, old := range live { + if old.Status == types.MemoryStatusActive { + item.ReplacesID = old.ID + break + } + } + } + } + item.TenantID, item.SubjectID = scope.TenantID, scope.SubjectID + if err := tx.Create(item).Error; err != nil { + return err + } + ids := make([]string, 0, len(live)+2) + if target.ID != "" { + live = append(live, &target) + } + if target.ReplacesID != "" && item.Status == types.MemoryStatusActive { + ids = append(ids, target.ReplacesID) + } + for _, old := range live { + if item.Status == types.MemoryStatusPending && old.Status == types.MemoryStatusActive { + continue + } + ids = append(ids, old.ID) + } + if len(ids) == 0 { + return nil + } + return query().Model(&types.MemoryItem{}). + Where("id <> ? AND (id IN ? OR replaces_id IN ?) AND status IN ?", item.ID, ids, ids, + []string{types.MemoryStatusActive, types.MemoryStatusPending}). + Updates(map[string]interface{}{"status": types.MemoryStatusSuperseded, "invalid_at": time.Now(), "superseded_by": item.ID, "updated_at": time.Now()}).Error + }) +} + +func (r *memoryRepository) ConfirmPendingItem(ctx context.Context, scope interfaces.MemoryScope, id string) error { + return r.withSubject(ctx, scope, func(tx *gorm.DB, _ *types.MemorySubject) error { + query := func() *gorm.DB { return tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID) } + var item types.MemoryItem + if err := query().First(&item, "id = ?", id).Error; err != nil { + return err + } + if item.Status == types.MemoryStatusActive { + return nil + } + if item.Status != types.MemoryStatusPending || (item.ExpiresAt != nil && !item.ExpiresAt.After(time.Now())) { + return types.ErrMemoryConflict + } + if item.ReplacesID != "" { + var target types.MemoryItem + err := query().First(&target, "id = ?", item.ReplacesID).Error + if errors.Is(err, gorm.ErrRecordNotFound) || (err == nil && target.Status != types.MemoryStatusActive) { + return types.ErrMemoryConflict + } + if err != nil { + return err + } + } + now := time.Now() + if err := query().Model(&types.MemoryItem{}). + Where("id <> ? AND (normalized_key = ? OR id = ? OR (replaces_id <> '' AND replaces_id = ?)) AND status IN ?", + id, item.NormalizedKey, item.ReplacesID, item.ReplacesID, []string{types.MemoryStatusActive, types.MemoryStatusPending}). + Updates(map[string]interface{}{"status": types.MemoryStatusSuperseded, "invalid_at": now, "superseded_by": id, "updated_at": now}).Error; err != nil { + return err + } + return query().Model(&types.MemoryItem{}).Where("id = ?", id). + Updates(map[string]interface{}{"status": types.MemoryStatusActive, "updated_at": now}).Error + }) +} diff --git a/internal/application/repository/message.go b/internal/application/repository/message.go index 455b82fad..2f4ef09c5 100644 --- a/internal/application/repository/message.go +++ b/internal/application/repository/message.go @@ -124,7 +124,18 @@ func (r *messageRepository) ListMessagesBySessionAfterTime( return messages, nil } -// UpdateMessage updates an existing message +// ListMessagesBySessionAfterCursor uses a stable tie-breaker for memory paging. +func (r *messageRepository) ListMessagesBySessionAfterCursor(ctx context.Context, sessionID string, cursor types.MemoryMessageCursor, limit int) ([]*types.Message, error) { + var messages []*types.Message + query := r.db.WithContext(ctx).Where("session_id = ?", sessionID) + if !cursor.At.IsZero() || cursor.ID != "" { + query = query.Where("created_at > ? OR (created_at = ? AND id > ?)", cursor.At, cursor.At, cursor.ID) + } + err := query.Order("created_at ASC, id ASC").Limit(limit).Find(&messages).Error + return messages, err +} + +// UpdateMessage updates an existing message. func (r *messageRepository) UpdateMessage(ctx context.Context, message *types.Message) error { return r.db.WithContext(ctx).Model(&types.Message{}).Where( "id = ? AND session_id = ?", message.ID, message.SessionID, diff --git a/internal/application/service/memory/consistency_test.go b/internal/application/service/memory/consistency_test.go new file mode 100644 index 000000000..af4793040 --- /dev/null +++ b/internal/application/service/memory/consistency_test.go @@ -0,0 +1,265 @@ +package memory + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/require" +) + +func TestMemoryConsistencyPendingMustNotReplaceActive(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "我是后端工程师", Origin: types.MemoryOriginManual}) + require.NoError(t, err) + proposal, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "可能是产品经理", Origin: types.MemoryOriginExtracted, Inferred: true}) + require.NoError(t, err) + require.Equal(t, types.MemoryStatusPending, proposal.Status) + current, err := s.repo.GetItem(ctx, scope, old.ID) + require.NoError(t, err) + require.Equal(t, types.MemoryStatusActive, current.Status, "unconfirmed inference must not retire a confirmed fact") +} +func TestMemoryConsistencyUpdateTargetMustWinOverChangedTopic(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Topic: "生产数据库", Content: "线上使用 MySQL"}) + require.NoError(t, err) + zero := 0 + require.NoError(t, s.applyDecisions(ctx, scope, s.workspaceConfig(ctx, 1), transcriptSegment{lines: []transcriptLine{{sessionID: "s", messageID: "m", content: "迁移到 PostgreSQL"}}}, []*types.MemoryItem{old}, []extractionDecision{{Action: "update", Target: &zero, Kind: types.MemoryKindFact, Topic: "数据库选型", Content: "线上已迁移到 PostgreSQL"}})) + _, total, err := s.ListItems(ctx, types.MemoryStatusActive, 20, 0) + require.NoError(t, err) + require.Equal(t, int64(1), total, "indexed update must not leave contradictory old memory active") +} +func TestMemoryConsistencyCappedSessionsMustBeFollowedUp(t *testing.T) { + s, tr, msg, model, q := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + model.response = `{"memories":[]}` + base := time.Now().Add(-time.Hour) + for i := 0; i < 4; i++ { + id := fmt.Sprintf("session-%d", i) + msg.set(id, []*types.Message{userMessage(id, fmt.Sprintf("unique-session-marker-%d", i), base.Add(time.Duration(i)*time.Minute))}) + s.ScheduleExtraction(ctx, id, "m", "model") + } + drainExtractions(t, s, q) + require.True(t, strings.Contains(model.seenTranscripts(), "unique-session-marker-3"), "fourth claimed session must survive the three-segment cap") +} +func TestMemoryConsistencyRetryMustKeepAllClaimedSessions(t *testing.T) { + s, tr, msg, model, q := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + model.response = `{"memories":[]}` + base := time.Now().Add(-time.Hour) + for i := 0; i < 2; i++ { + id := fmt.Sprintf("session-%d", i) + msg.set(id, []*types.Message{userMessage(id, fmt.Sprintf("retry-session-marker-%d", i), base.Add(time.Duration(i)*time.Minute))}) + s.ScheduleExtraction(ctx, id, "m", "model") + } + task := q.pop() + require.NotNil(t, task) + model.failNext = true + require.Error(t, s.Handle(context.Background(), task)) + require.NoError(t, s.Handle(context.Background(), task)) + drainExtractions(t, s, q) + require.True(t, strings.Contains(model.seenTranscripts(), "retry-session-marker-1"), "retry must retain the second drained session") +} + +func TestMemoryConsistencyManualEditMustRefreshEmbedding(t *testing.T) { + s, tr, _ := newVectorHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Content: "回答直接给结论"}) + require.NoError(t, err) + _, err = s.UpdateItem(ctx, old.ID, "生产环境需要调大连接池", 3) + require.NoError(t, err) + s.backfillEmbeddings(ctx, scope, s.workspaceConfig(ctx, 1)) + vectors, err := s.repo.ItemEmbeddings(ctx, scope, []string{old.ID}, "embed-1") + require.NoError(t, err) + require.Equal(t, []float32{0, 1, 0}, vectors[old.ID], "edit and maintenance must not leave the old semantic vector") +} + +func TestMemoryConsistencyProposalConfirmationAndRejection(t *testing.T) { + for _, confirm := range []bool{false, true} { + t.Run(fmt.Sprint(confirm), func(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "我是工程师", Origin: types.MemoryOriginManual}) + require.NoError(t, err) + proposal, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "可能是经理", Inferred: true}) + require.NoError(t, err) + require.Equal(t, old.ID, proposal.ReplacesID) + if confirm { + _, err = s.ConfirmItem(ctx, proposal.ID) + require.NoError(t, err) + _, err = s.ConfirmItem(ctx, proposal.ID) + require.NoError(t, err, "confirmation is idempotent") + } else { + require.NoError(t, s.RejectItem(ctx, proposal.ID)) + } + current, err := s.repo.GetItem(ctx, scope, old.ID) + require.NoError(t, err) + if confirm { + require.Equal(t, types.MemoryStatusSuperseded, current.Status) + require.Equal(t, proposal.ID, current.SupersededBy) + } else { + require.Equal(t, types.MemoryStatusActive, current.Status) + } + _, total, err := s.ListItems(ctx, types.MemoryStatusActive, 20, 0) + require.NoError(t, err) + require.Equal(t, int64(1), total) + }) + } +} + +func TestMemoryConsistencyEditInvalidatesOldProposals(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "我是工程师"}) + require.NoError(t, err) + proposal, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "可能是经理", Inferred: true}) + require.NoError(t, err) + _, err = s.UpdateItem(ctx, old.ID, "我是设计师", 4) + require.NoError(t, err) + _, err = s.ConfirmItem(ctx, proposal.ID) + require.ErrorIs(t, err, types.ErrMemoryConflict) + items, _, err := s.ListItems(ctx, types.MemoryStatusActive, 20, 0) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, "我是设计师", items[0].Content) +} + +func TestMemoryConsistencyReplacingProposalKeepsOriginalTarget(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "我是工程师"}) + require.NoError(t, err) + first, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "可能是经理", Inferred: true}) + require.NoError(t, err) + second, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindProfile, Topic: "职业", Content: "可能是设计师", Inferred: true}) + require.NoError(t, err) + require.Equal(t, old.ID, second.ReplacesID) + _, err = s.ConfirmItem(ctx, first.ID) + require.ErrorIs(t, err, types.ErrMemoryConflict) + _, err = s.ConfirmItem(ctx, second.ID) + require.NoError(t, err) + _, total, err := s.ListItems(ctx, types.MemoryStatusActive, 20, 0) + require.NoError(t, err) + require.Equal(t, int64(1), total) +} + +func TestMemoryConsistencyTimestampTiesAndLargePendingQueue(t *testing.T) { + for _, manySessions := range []bool{false, true} { + t.Run(fmt.Sprint(manySessions), func(t *testing.T) { + s, tr, msg, model, q := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + model.response = `{"memories":[]}` + at := time.Now().Add(-time.Hour) + var rows []*types.Message + for i := 0; i < 85; i++ { + id := fmt.Sprintf("marker-%03d", i) + session := "single" + if manySessions { + session = id + } + row := userMessage(session, id, at) + if manySessions { + msg.set(session, []*types.Message{row}) + } else { + rows = append(rows, row) + msg.set(session, rows) + } + s.ScheduleExtraction(ctx, session, id, "model") + } + drainExtractions(t, s, q) + var seen string + for _, prompt := range model.prompts { + seen += transcriptBlock(prompt) + } + for i := 0; i < 85; i++ { + require.Equal(t, 1, strings.Count(seen, fmt.Sprintf("marker-%03d", i)), "each message must appear exactly once as extractable input") + } + }) + } +} + +func TestMemoryConsistencyGapDoesNotSkipNextSegment(t *testing.T) { + s, tr, msg, model, q := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + model.response = `{"memories":[]}` + at := time.Now().Add(-24 * time.Hour) + var rows []*types.Message + for i := 0; i < 7; i++ { + rows = append(rows, userMessage("s", fmt.Sprintf("gap-marker-%d", i), at.Add(time.Duration(i)*2*time.Hour))) + } + msg.set("s", rows) + s.ScheduleExtraction(ctx, "s", "m", "model") + drainExtractions(t, s, q) + var seen string + for _, prompt := range model.prompts { + seen += transcriptBlock(prompt) + } + for i := 0; i < 7; i++ { + require.Equal(t, 1, strings.Count(seen, fmt.Sprintf("gap-marker-%d", i))) + } +} + +func TestMemoryConsistencyLeaseAndConcurrentEnqueue(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + _, err := s.repo.EnsureSubject(ctx, scope) + require.NoError(t, err) + _, _, err = s.repo.EnqueuePendingSession(ctx, scope, "s", time.Minute) + require.NoError(t, err) + batch, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "first", time.Minute) + require.NoError(t, err) + require.Len(t, batch.Sessions, 1) + duplicate, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "second", time.Minute) + require.NoError(t, err) + require.NotNil(t, duplicate) + require.False(t, duplicate.RetryAt.IsZero()) + require.Empty(t, duplicate.Sessions) + _, queued, err := s.repo.EnqueuePendingSession(ctx, scope, "s", time.Minute) + require.NoError(t, err) + require.False(t, queued) + require.NoError(t, s.repo.CheckpointExtraction(ctx, scope, "first", batch.Sessions[0], types.MemoryMessageCursor{At: time.Now(), ID: "a"}, true)) + require.NoError(t, s.repo.FinishExtraction(ctx, scope, "first")) + next, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "next", time.Minute) + require.NoError(t, err) + require.Len(t, next.Sessions, 1, "a turn arriving while the batch runs must remain pending") + require.NoError(t, s.repo.ReleaseExtractionSlot(ctx, scope, "first")) + require.ErrorIs(t, s.repo.CheckpointExtraction(ctx, scope, "first", batch.Sessions[0], types.MemoryMessageCursor{}, true), types.ErrMemoryExtractionLeaseLost) + require.NoError(t, s.repo.CheckpointExtraction(ctx, scope, "next", next.Sessions[0], next.Sessions[0].Cursor, true)) + require.NoError(t, s.repo.FinishExtraction(ctx, scope, "next")) + empty, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "duplicate", time.Minute) + require.NoError(t, err) + require.Nil(t, empty, "redelivering a finished task must be a no-op") +} + +func TestMemoryConsistencyEmbeddingFailureAndStaleCompletion(t *testing.T) { + s, tr, model := newVectorHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Content: "回答直接给结论"}) + require.NoError(t, err) + model.embedder.fail = true + _, err = s.UpdateItem(ctx, old.ID, "生产环境需要调大连接池", 3) + require.NoError(t, err) + vectors, err := s.repo.ItemEmbeddings(ctx, scope, []string{old.ID}, "embed-1") + require.NoError(t, err) + require.Empty(t, vectors, "failed refresh must not leave a stale vector searchable") + model.embedder.fail = false + s.storeItemEmbedding(ctx, scope, s.workspaceConfig(ctx, 1), old) + vectors, err = s.repo.ItemEmbeddings(ctx, scope, []string{old.ID}, "embed-1") + require.NoError(t, err) + require.Empty(t, vectors, "a slow embedding of old content must be discarded") + require.Equal(t, 1, s.backfillEmbeddings(ctx, scope, s.workspaceConfig(ctx, 1))) + vectors, err = s.repo.ItemEmbeddings(ctx, scope, []string{old.ID}, "embed-1") + require.NoError(t, err) + require.Equal(t, []float32{0, 1, 0}, vectors[old.ID]) +} diff --git a/internal/application/service/memory/coverage_test.go b/internal/application/service/memory/coverage_test.go index 6d12d6031..8fc7627f5 100644 --- a/internal/application/service/memory/coverage_test.go +++ b/internal/application/service/memory/coverage_test.go @@ -356,8 +356,8 @@ func TestATruncatedRunDoesNotSwallowTheMessages(t *testing.T) { task := enqueuer.pop() require.NotNil(t, task) - require.Error(t, svc.Handle(context.Background(), task), - "a run that read nothing has to fail, or the messages are consumed for good") + require.NoError(t, svc.Handle(context.Background(), task)) + require.NotEmpty(t, enqueuer.tasks, "invalid output must schedule another attempt") scope, err := ResolveScope(ctx) require.NoError(t, err) diff --git a/internal/application/service/memory/extract.go b/internal/application/service/memory/extract.go index da9b2685b..97d14ec81 100644 --- a/internal/application/service/memory/extract.go +++ b/internal/application/service/memory/extract.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "sort" "strings" "time" @@ -14,6 +13,7 @@ import ( "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/google/uuid" "github.com/hibiken/asynq" ) @@ -101,8 +101,8 @@ func (s *Service) ScheduleExtraction(ctx context.Context, sessionID, messageID, return } - // The subject row carries the queue and the watermark, so it has to exist - // before the first turn is recorded. + // The subject row serializes queue mutations, so it must exist before + // the first per-session progress row is recorded. subject, err := s.repo.EnsureSubject(ctx, scope) if err != nil { logger.Warnf(ctx, "memory: ensure subject for extraction failed: %v", err) @@ -114,7 +114,7 @@ func (s *Service) ScheduleExtraction(ctx context.Context, sessionID, messageID, delay := cfg.ExtractDelay() previous, shouldEnqueue, err := s.repo.EnqueuePendingSession( - ctx, scope, sessionID, delay+extractInFlightGrace, + ctx, scope, sessionID, cfg.ExtractMinInterval()+delay+extractInFlightGrace, ) if err != nil { logger.Warnf(ctx, "memory: record pending session failed: %v", err) @@ -129,7 +129,7 @@ func (s *Service) ScheduleExtraction(ctx context.Context, sessionID, messageID, // The minimum interval only defers: if the previous run was recent, the // task is queued further out rather than the turn being discarded. if previous != nil && previous.LastExtractedAt != nil { - if remaining := cfg.ExtractMinInterval() - time.Since(*previous.LastExtractedAt); remaining > delay { + if remaining := cfg.ExtractMinInterval() - time.Now().Sub(*previous.LastExtractedAt); remaining > delay { delay = remaining } } @@ -145,7 +145,7 @@ func (s *Service) enqueueExtraction( scope interfaces.MemoryScope, sessionID, messageID, chatModelID string, delay time.Duration, -) { +) error { payload := types.MemoryExtractPayload{ TenantID: scope.TenantID, SubjectID: scope.SubjectID, @@ -159,7 +159,7 @@ func (s *Service) enqueueExtraction( if err != nil { logger.Warnf(ctx, "memory: marshal extraction payload failed: %v", err) s.releaseSlot(ctx, scope) - return + return err } task := asynq.NewTask(types.TypeMemoryExtract, body) @@ -170,11 +170,13 @@ func (s *Service) enqueueExtraction( ); err != nil { logger.Warnf(ctx, "memory: enqueue extraction failed: %v", err) s.releaseSlot(ctx, scope) + return err } + return nil } func (s *Service) releaseSlot(ctx context.Context, scope interfaces.MemoryScope) { - if err := s.repo.ReleaseExtractionSlot(ctx, scope); err != nil { + if err := s.repo.ReleaseExtractionSlot(ctx, scope, ""); err != nil { logger.Warnf(ctx, "memory: release extraction slot failed: %v", err) } } @@ -244,6 +246,8 @@ type extractionResponse struct { Topics []string `json:"topics"` } +var errInvalidExtractionOutput = errors.New("invalid memory extraction output") + // Handle runs one distillation pass. func (s *Service) Handle(ctx context.Context, task *asynq.Task) error { var payload types.MemoryExtractPayload @@ -283,102 +287,123 @@ func (s *Service) Handle(ctx context.Context, task *asynq.Task) error { return nil } - // Take the queue before reading anything. Turns arriving from here on land - // in a fresh queue and get their own follow-up run rather than being - // erased by this one. - pending, cursor, err := s.repo.ClaimPendingSessions(ctx, scope) + // A lease serializes duplicate tasks, while the durable session queue + // stays intact until each segment has actually been applied. + leaseID := uuid.NewString() + batch, err := s.repo.ClaimPendingSessions(ctx, scope, payload.SessionID, leaseID, extractInFlightGrace) if err != nil { return fmt.Errorf("claim pending sessions: %w", err) } - if payload.SessionID != "" && !containsString(pending, payload.SessionID) { - pending = append(pending, payload.SessionID) - } - - // Expired items are archived before the existing-notes list is built, so a - // finished task is not offered to the model as something still true. - if archived, err := s.repo.ExpireOverdue(ctx, scope); err != nil { - logger.Warnf(ctx, "memory: expire overdue items failed: %v", err) - } else if archived > 0 { - logger.Infof(ctx, "memory: archived %d expired items", archived) - } - - segments, truncated, err := s.collectSegments(ctx, pending, cursor) - if err != nil { - return err - } - if len(segments) == 0 { - // Nothing new to read. Advance nothing, but release the slot so the - // next turn can schedule a run immediately. - s.releaseSlot(ctx, scope) + if batch == nil { return nil } - - // One call per topic segment. Handing the model a flat pile of messages - // spanning two conversations and an hour-long gap is exactly where - // extraction quality falls apart, and it also destroys attribution. - var newCursor time.Time - for _, segment := range segments { - existing, err := s.repo.ListActiveByKinds(ctx, scope, types.MemoryKinds, extractCandidatePool) - if err != nil { - return fmt.Errorf("load existing memories: %w", err) + if !batch.RetryAt.IsZero() { + // A redelivery can arrive before a dead worker's lease expires. Simply + // acknowledging it here would strand the durable queue forever. + if s.enqueuer == nil { + return fmt.Errorf("memory extraction is leased until %s", batch.RetryAt) } - existing = s.narrowToRelevant(ctx, scope, cfg, segment, existing) - forgotten, err := s.repo.ListTombstones(ctx, scope, 30) - if err != nil { - logger.Warnf(ctx, "memory: load tombstones failed: %v", err) + return s.enqueueExtraction(ctx, scope, payload.SessionID, payload.MessageID, payload.ChatModelID, time.Until(batch.RetryAt)+time.Second) + } + defer func() { + cleanup, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := s.repo.ReleaseExtractionSlot(cleanup, scope, leaseID); err != nil { + logger.Warnf(cleanup, "memory: release worker lease failed: %v", err) } + }() + // Stop model calls before the lease can expire and another worker starts. + ctx, cancel := context.WithTimeout(ctx, extractInFlightGrace-time.Minute) + defer cancel() + if _, err := s.repo.ExpireOverdue(ctx, scope); err != nil { + return err + } - knownTopics, err := s.repo.TopTopics(ctx, scope, topicCandidateLimit) + processed := 0 + var retryErr error + for _, session := range batch.Sessions { + segments, more, err := s.collectSessionSegments(ctx, session) if err != nil { - logger.Warnf(ctx, "memory: load known topics failed: %v", err) - knownTopics = nil - } - - parsed, err := s.callExtractionModel(ctx, cfg, payload, segment, existing, forgotten, knownTopics) - if err != nil { - // Leave the watermark where it is: the messages this run failed on - // must be read again, not skipped. Advancing over the segments that - // did succeed keeps the failure from replaying them. - if !newCursor.IsZero() { - if err := s.repo.FinishExtraction(ctx, scope, newCursor); err != nil { - logger.Warnf(ctx, "memory: advance extraction cursor failed: %v", err) - } - } - s.releaseSlot(ctx, scope) return err } - s.applyDecisions(ctx, scope, cfg, segment, existing, parsed.Memories) - // Subjects are counted separately from memories: one question is noise, - // the same subject across conversations is an interest. - s.observeTopics(ctx, scope, cfg, s.extractionModelID(ctx, cfg, payload), parsed.Topics) - if segment.end.After(newCursor) { - newCursor = segment.end + if len(segments) == 0 { + if err := s.repo.CheckpointExtraction(ctx, scope, leaseID, session, session.Cursor, true); err != nil { + return err + } + continue + } + for i, segment := range segments { + cursor := types.MemoryMessageCursor{At: segment.end, ID: segment.endID} + if len(segment.lines) > 0 { + if err := s.extractSegment(ctx, scope, cfg, payload, segment); err != nil { + if !errors.Is(err, errInvalidExtractionOutput) { + return err + } + failure := interfaces.MemoryExtractionFailure{ + Session: session, End: cursor, Code: "invalid_model_output", + } + skip, recordErr := s.repo.RecordExtractionFailure(ctx, scope, leaseID, failure) + if recordErr != nil { + return recordErr + } + if !skip { + retryErr = err + processed++ + break // Retry this range later, while other sessions can progress. + } + logger.Warnf(ctx, "memory: skipping invalid segment; failure recorded for session %s", + session.SessionID) + } + } + if err := s.repo.CheckpointExtraction(ctx, scope, leaseID, session, cursor, !more && i == len(segments)-1); err != nil { + return err + } + session.Cursor = cursor + processed++ + if processed >= extractMaxSegmentsPerRun { + break + } + } + if processed >= extractMaxSegmentsPerRun { + break } } - - if err := s.repo.FinishExtraction(ctx, scope, newCursor); err != nil { - logger.Warnf(ctx, "memory: advance extraction cursor failed: %v", err) + if err := s.repo.FinishExtraction(ctx, scope, leaseID); err != nil { + return err + } + if err := s.scheduleFollowUpIfNeeded(ctx, scope, cfg, payload); err != nil { + return err + } + if retryErr != nil && s.enqueuer == nil { + return retryErr } - - // Either this run hit its message cap, or new turns arrived while it was - // working. Both mean there is more to read, and both are how the "every - // message is eventually considered" guarantee survives a busy user. - s.scheduleFollowUpIfNeeded(ctx, scope, cfg, payload, truncated) - - // Maintenance rides along on a background run that has already happened - // rather than needing its own scheduler, which keeps it working identically - // in Lite mode where there is no asynq worker to hold a periodic job. s.consolidateIfDue(ctx, scope, cfg, s.extractionModelID(ctx, cfg, payload)) return nil } -func containsString(values []string, target string) bool { - for _, value := range values { - if value == target { - return true - } +func (s *Service) extractSegment(ctx context.Context, scope interfaces.MemoryScope, cfg *types.MemoryConfig, payload types.MemoryExtractPayload, segment transcriptSegment) error { + existing, err := s.repo.ListActiveByKinds(ctx, scope, types.MemoryKinds, extractCandidatePool) + if err != nil { + return fmt.Errorf("load existing memories: %w", err) } - return false + existing = s.narrowToRelevant(ctx, scope, cfg, segment, existing) + forgotten, err := s.repo.ListTombstones(ctx, scope, 30) + if err != nil { + return fmt.Errorf("load memory tombstones: %w", err) + } + knownTopics, err := s.repo.TopTopics(ctx, scope, topicCandidateLimit) + if err != nil { + return fmt.Errorf("load known topics: %w", err) + } + parsed, err := s.callExtractionModel(ctx, cfg, payload, segment, existing, forgotten, knownTopics) + if err != nil { + return err + } + if err := s.applyDecisions(ctx, scope, cfg, segment, existing, parsed.Memories); err != nil { + return err + } + s.observeTopics(ctx, scope, cfg, s.extractionModelID(ctx, cfg, payload), parsed.Topics) + return nil } // scheduleFollowUpIfNeeded queues the next run when work remains. @@ -387,37 +412,26 @@ func (s *Service) scheduleFollowUpIfNeeded( scope interfaces.MemoryScope, cfg *types.MemoryConfig, payload types.MemoryExtractPayload, - truncated bool, -) { +) error { if s.enqueuer == nil { - return + return nil } - subject, err := s.repo.GetSubject(ctx, scope) + pending, err := s.repo.HasPendingExtraction(ctx, scope) if err != nil { - logger.Warnf(ctx, "memory: reload subject for follow-up failed: %v", err) - return + return fmt.Errorf("load pending memory extraction: %w", err) } - if subject == nil { - return - } - if !truncated && len(subject.PendingSessions) == 0 { - return + if !pending { + return nil } sessionID := payload.SessionID - if len(subject.PendingSessions) > 0 { - sessionID = subject.PendingSessions[0] - } // Claim the slot again for the successor; FinishExtraction just cleared it. if _, shouldEnqueue, err := s.repo.EnqueuePendingSession( - ctx, scope, sessionID, cfg.ExtractDelay()+extractInFlightGrace, + ctx, scope, "", cfg.ExtractMinInterval()+cfg.ExtractDelay()+extractInFlightGrace, ); err != nil || !shouldEnqueue { - if err != nil { - logger.Warnf(ctx, "memory: claim follow-up slot failed: %v", err) - } - return + return err } logger.Infof(ctx, "memory: queueing follow-up distillation for subject %s", scope.SubjectID) - s.enqueueExtraction(ctx, scope, sessionID, payload.MessageID, payload.ChatModelID, extractFollowUpDelay) + return s.enqueueExtraction(ctx, scope, sessionID, payload.MessageID, payload.ChatModelID, extractFollowUpDelay) } // transcriptLine is one thing the user said, kept with the identity of the @@ -439,112 +453,66 @@ type transcriptSegment struct { context []string // end is the newest message timestamp the segment covers, including the // assistant rows in between, and is what the watermark advances to. - end time.Time + end time.Time + endID string } -// collectSegments reads everything past the watermark and cuts it into -// segments. -// -// Walking forward from a watermark rather than reading "the most recent N -// messages" is what makes coverage a property of the data: a burst of turns, a -// second concurrent session, or a slow worker can delay a message but cannot -// make it invisible. Segmenting on top of that is what keeps quality: a run -// that spans two conversations and an hour-long gap is one where the model has -// to guess which statement belongs to which situation. -// -// Only role=user messages are extracted from. The assistant's own words are the -// model talking to itself, and feeding them back is how a prompt injection in a -// retrieved document ends up stored as a fact about the user. -func (s *Service) collectSegments( - ctx context.Context, sessions []string, cursor time.Time, -) (segments []transcriptSegment, truncated bool, err error) { - for _, sessionID := range sessions { - if strings.TrimSpace(sessionID) == "" { +// collectSessionSegments reads a bounded page using a session-local cursor. +// Every row, including assistant-only pages, belongs to a checkpoint. A gap +// is flushed before admitting the next message so it cannot advance over it. +func (s *Service) collectSessionSegments(ctx context.Context, session types.MemoryExtractionSession) ([]transcriptSegment, bool, error) { + messages, err := s.messageRepo.ListMessagesBySessionAfterCursor(ctx, session.SessionID, session.Cursor, extractMaxMessagesPerRun+1) + if err != nil { + return nil, false, fmt.Errorf("load session messages: %w", err) + } + more := len(messages) > extractMaxMessagesPerRun + if more { + messages = messages[:extractMaxMessagesPerRun] + } + var segments []transcriptSegment + current := transcriptSegment{sessionID: session.SessionID} + var lastAt time.Time + hasRows := false + flush := func() { + if !hasRows { + return + } + segments = append(segments, current) + current = transcriptSegment{sessionID: session.SessionID} + hasRows = false + } + for _, message := range messages { + if message == nil { continue } - messages, err := s.messageRepo.ListMessagesBySessionAfterTime( - ctx, sessionID, cursor, extractMaxMessagesPerRun+1, - ) - if err != nil { - return nil, false, fmt.Errorf("load session messages: %w", err) + content := strings.TrimSpace(message.Content) + isUser := message.Role == "user" && content != "" + if isUser && !lastAt.IsZero() && message.CreatedAt.Sub(lastAt) > extractSegmentGap { + flush() } - if len(messages) > extractMaxMessagesPerRun { - truncated = true - messages = messages[:extractMaxMessagesPerRun] + current.end, current.endID = message.CreatedAt, message.ID + hasRows = true + if !isUser { + continue } - - var ( - lines []transcriptLine - end time.Time - lastAt time.Time - flushed []transcriptSegment - ) - flush := func() { - if len(lines) == 0 { - return - } - flushed = append(flushed, transcriptSegment{ - sessionID: sessionID, - lines: lines, - end: end, - }) - lines = nil + lastAt = message.CreatedAt + if runes := []rune(content); len(runes) > extractMaxLineRunes { + content = string(runes[:extractMaxLineRunes]) } - for _, message := range messages { - if message == nil { - continue - } - // The watermark covers assistant rows too: they are not read, but - // leaving them behind it would make the cursor move back and forth - // around every turn. - if message.CreatedAt.After(end) { - end = message.CreatedAt - } - if message.Role != "user" { - continue - } - content := strings.TrimSpace(message.Content) - if content == "" { - continue - } - if !lastAt.IsZero() && message.CreatedAt.Sub(lastAt) > extractSegmentGap { - flush() - } - lastAt = message.CreatedAt - if runes := []rune(content); len(runes) > extractMaxLineRunes { - content = string(runes[:extractMaxLineRunes]) - } - lines = append(lines, transcriptLine{ - sessionID: sessionID, - messageID: message.ID, - at: message.CreatedAt, - content: content, - }) + current.lines = append(current.lines, transcriptLine{sessionID: session.SessionID, messageID: message.ID, at: message.CreatedAt, content: content}) + } + flush() + for i := range segments { + if len(segments[i].lines) == 0 { + continue } - flush() - - for i := range flushed { - // Every segment but the first in a session already has its lead-in - // inside this run; the first one needs it fetched. - if i == 0 { - flushed[i].context = s.priorContext(ctx, sessionID, flushed[i].lines) - } else { - flushed[i].context = tailContents(flushed[i-1].lines, extractContextLines) - } - segments = append(segments, flushed[i]) + if i == 0 { + segments[i].context = s.priorContext(ctx, session.SessionID, segments[i].lines) + } else { + segments[i].context = tailContents(segments[i-1].lines, extractContextLines) } } - - sort.SliceStable(segments, func(i, j int) bool { - return segments[i].lines[0].at.Before(segments[j].lines[0].at) - }) - if len(segments) > extractMaxSegmentsPerRun { - // The rest stay ahead of the watermark and are picked up by the - // follow-up run, so capping the work of one run never loses a message. - segments = segments[:extractMaxSegmentsPerRun] - truncated = true - } - return segments, truncated, nil + return segments, more, nil } // priorContext fetches the few user messages just before a segment. @@ -697,9 +665,9 @@ Lines: Existing notes: (none) {"memories":[ {"action":"add","target":null,"kind":"profile","topic":"可能的身份", - "content":"可能在负责连锁门店的排班","importance":2,"source":1, + "content":"可能在负责仓库单据管理","importance":2,"source":1, "expires_at":null,"inferred":true}], -"topics":["门店排班管理"]} +"topics":["仓库单据保留规则"]} The identity is a guess, so it is marked inferred and waits for confirmation. The subject is counted either way. @@ -961,7 +929,7 @@ func (s *Service) callExtractionModel( return extractionResponse{}, err } if response == nil { - return extractionResponse{}, nil + return extractionResponse{}, fmt.Errorf("%w: no response", errInvalidExtractionOutput) } // A truncated call is retried once with room to spare. Reasoning models @@ -977,24 +945,16 @@ func (s *Service) callExtractionModel( return extractionResponse{}, err } if response == nil || isTruncated(response) { - // Returning an error is what keeps the watermark where it is, so - // these messages are read again rather than silently consumed by a - // run that learned nothing. return extractionResponse{}, fmt.Errorf( - "extraction model returned no usable output within %d tokens; "+ + "%w: no usable output within %d tokens; "+ "if this is a reasoning model, its thinking is consuming the budget", - extractBudgetRetryTokens) + errInvalidExtractionOutput, extractBudgetRetryTokens) } } parsed, err := parseExtractionResponse(response.Content) if err != nil { - // A malformed but complete response is the model's fault, not a - // transient failure: the same prompt at temperature zero produces the - // same garbage, so retrying only burns the budget. Truncation is - // handled above precisely because it is *not* this case. - logger.Warnf(ctx, "memory: unparsable extraction response: %v", err) - return extractionResponse{}, nil + return extractionResponse{}, fmt.Errorf("%w: %v", errInvalidExtractionOutput, err) } return parsed, nil } @@ -1075,7 +1035,7 @@ func (s *Service) applyDecisions( segment transcriptSegment, existing []*types.MemoryItem, decisions []extractionDecision, -) { +) error { applied := 0 // Two decisions about the same topic inside one response would otherwise // supersede each other, leaving a superseded row from a single run. @@ -1092,12 +1052,24 @@ func (s *Service) applyDecisions( topic := types.SanitizeMemoryTopic(decision.Topic) // An update or delete says which note it means by index; fall back to - // the topic only when the index is absent or out of range. + // the topic only when the index is absent. var target *types.MemoryItem - if decision.Target != nil && *decision.Target >= 0 && *decision.Target < len(existing) { - target = existing[*decision.Target] - } - if target != nil && topic == "" { + if action == "update" || action == "delete" { + if decision.Target != nil { + if *decision.Target < 0 || *decision.Target >= len(existing) || existing[*decision.Target] == nil { + continue + } + target = existing[*decision.Target] + } else { + var err error + target, err = s.repo.FindActiveByKey(ctx, scope, types.MemoryItemKey(topic, decision.Content)) + if err != nil { + return fmt.Errorf("find memory decision target: %w", err) + } + if target == nil { + continue + } + } topic = target.Topic } @@ -1108,24 +1080,19 @@ func (s *Service) applyDecisions( switch action { case "delete": - if target == nil { - found, err := s.repo.FindActiveByKey(ctx, scope, key) - if err != nil || found == nil { - continue - } - target = found - } // Superseding with no replacement keeps the note visible in the // memory manager as something that stopped being true, which is // more useful than it disappearing without explanation. if err := s.repo.SupersedeItem(ctx, scope, target.ID, ""); err != nil { - logger.Warnf(ctx, "memory: delete decision failed: %v", err) - continue + return fmt.Errorf("delete memory decision: %w", err) } seenTopics[key] = struct{}{} applied++ s.rebuildBlock(ctx, scope) case "add", "update": + if types.SanitizeMemoryContent(decision.Content) == "" { + continue + } source := decision.resolveSource(segment) item := types.MemoryItem{ Kind: decision.Kind, @@ -1138,11 +1105,16 @@ func (s *Service) applyDecisions( ExpiresAt: parseExpiry(decision.ExpiresAt), Inferred: decision.Inferred, } - if _, err := s.write(ctx, scope, cfg, item); err != nil { - if !errors.Is(err, ErrPreviouslyForgotten) && !errors.Is(err, ErrSensitiveContent) { - logger.Warnf(ctx, "memory: apply extraction decision failed: %v", err) + var targetID string + if action == "update" && target != nil { + targetID = target.ID + } + if _, err := s.writeReplacing(ctx, scope, cfg, item, targetID); err != nil { + if errors.Is(err, ErrPreviouslyForgotten) || errors.Is(err, ErrSensitiveContent) || + errors.Is(err, types.ErrMemoryConflict) { + continue } - continue + return fmt.Errorf("apply memory decision: %w", err) } seenTopics[key] = struct{}{} applied++ @@ -1151,4 +1123,5 @@ func (s *Service) applyDecisions( if applied > 0 { logger.Infof(ctx, "memory: stored %d memories for subject %s", applied, scope.SubjectID) } + return nil } diff --git a/internal/application/service/memory/extract_test.go b/internal/application/service/memory/extract_test.go index 111da5b40..293230166 100644 --- a/internal/application/service/memory/extract_test.go +++ b/internal/application/service/memory/extract_test.go @@ -171,17 +171,17 @@ func TestExtractionAppliesUpdateAndDeleteDecisions(t *testing.T) { require.Equal(t, int64(2), superseded) } -func TestExtractionToleratesUnparsableModelOutput(t *testing.T) { - svc, tenantRepo, messages, models, _ := newExtractionHarness(t) +func TestExtractionPreservesWorkOnUnparsableModelOutput(t *testing.T) { + svc, tenantRepo, messages, models, queue := newExtractionHarness(t) tenantRepo.set(7, &types.MemoryConfig{Enabled: true, WriteMode: types.MemoryWriteAuto}) messages.messages = []*types.Message{{Role: "user", Content: "随便说点什么"}} models.response = "抱歉，我不太明白你的意思。" - // Garbage is the model's fault, not a transient failure, so returning an - // error would just re-run the same prompt until the retry budget is gone. + // The first malformed output retains the range and queues a retry. require.NoError(t, svc.Handle(context.Background(), extractTask(t, types.MemoryExtractPayload{ TenantID: 7, SubjectID: "web_user:alice", SessionID: "s", MessageID: "m", ChatModelID: "m1", }))) + require.NotNil(t, queue.pop()) _, total, err := svc.ListItems(enabledCtx(t, tenantRepo, 7, "alice"), "", 10, 0) require.NoError(t, err) require.Zero(t, total) diff --git a/internal/application/service/memory/persistence_consistency_test.go b/internal/application/service/memory/persistence_consistency_test.go new file mode 100644 index 000000000..121cbdf10 --- /dev/null +++ b/internal/application/service/memory/persistence_consistency_test.go @@ -0,0 +1,179 @@ +package memory + +import ( + "fmt" + "os" + "strings" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestMemoryConsistencyRealMessagePaging(t *testing.T) { + s, db, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + require.NoError(t, db.AutoMigrate(&types.Message{})) + at := time.Now().UTC().Truncate(time.Second) + for i := 0; i < 85; i++ { + require.NoError(t, db.Exec("INSERT INTO messages (id, session_id, role, content, created_at) VALUES (?, ?, ?, ?, ?)", fmt.Sprintf("m%03d", i), "s", "user", "hello", at).Error) + } + require.NoError(t, db.Exec("INSERT INTO messages (id, session_id, role, content, created_at) VALUES (?, ?, ?, ?, ?)", "other", "unrelated", "user", "private", at).Error) + require.NoError(t, db.Create(&types.Message{ + ID: "deleted", SessionID: "s", Role: "user", CreatedAt: at, DeletedAt: gorm.DeletedAt{Time: at, Valid: true}, + }).Error) + s.messageRepo = repository.NewMessageRepository(db) + var cursor types.MemoryMessageCursor + seen := map[string]bool{} + for { + rows, err := s.messageRepo.ListMessagesBySessionAfterCursor(ctx, "s", cursor, 40) + require.NoError(t, err) + if len(rows) == 0 { + break + } + for _, row := range rows { + require.False(t, seen[row.ID]) + require.Equal(t, "s", row.SessionID) + seen[row.ID] = true + cursor = types.MemoryMessageCursor{At: row.CreatedAt, ID: row.ID} + } + } + require.Len(t, seen, 85) +} + +func TestMemoryConsistencyReplacementRollsBackAsOneOperation(t *testing.T) { + s, db, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Topic: "数据库", Content: "使用 MySQL"}) + require.NoError(t, err) + require.NoError(t, db.Exec(`CREATE TRIGGER fail_supersede BEFORE UPDATE OF status ON memory_items + WHEN NEW.status = 'superseded' BEGIN SELECT RAISE(ABORT, 'injected failure'); END`).Error) + _, err = s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Topic: "数据库", Content: "已迁移到 PostgreSQL"}) + require.Error(t, err) + items, total, err := s.ListItems(ctx, "", 20, 0) + require.NoError(t, err) + require.Equal(t, int64(1), total, "failed replacement must roll back the inserted row") + require.Equal(t, old.ID, items[0].ID) + require.Equal(t, types.MemoryStatusActive, items[0].Status) +} + +func TestMemoryConsistencyLeaseRecoveryRetainsProgress(t *testing.T) { + s, db, tr := newMemoryHarness(t) + at := time.Now() + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + _, err := s.repo.EnsureSubject(ctx, scope) + require.NoError(t, err) + _, _, err = s.repo.EnqueuePendingSession(ctx, scope, "s", time.Minute) + require.NoError(t, err) + first, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "dead-worker", time.Minute) + require.NoError(t, err) + cursor := types.MemoryMessageCursor{At: at.Add(-time.Hour), ID: "completed"} + require.NoError(t, s.repo.CheckpointExtraction(ctx, scope, "dead-worker", first.Sessions[0], cursor, false)) + subject, err := s.repo.GetSubject(ctx, scope) + require.NoError(t, err) + subject.ExtractionState.LeaseUntil = time.Now().Add(-time.Minute) + require.NoError(t, db.Model(subject).Update("extraction_state", subject.ExtractionState).Error) + recovered, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "recovery", time.Minute) + require.NoError(t, err) + require.Len(t, recovered.Sessions, 1) + require.True(t, recovered.Sessions[0].Cursor.At.Equal(cursor.At)) + require.Equal(t, cursor.ID, recovered.Sessions[0].Cursor.ID) + require.ErrorIs(t, s.repo.FinishExtraction(ctx, scope, "dead-worker"), types.ErrMemoryExtractionLeaseLost) + require.NoError(t, s.repo.CheckpointExtraction(ctx, scope, "recovery", recovered.Sessions[0], cursor, true)) +} + +func TestMemoryConsistencyRedeliveryWaitsForCrashedWorkerLease(t *testing.T) { + s, tr, messages, models, queue := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + models.response = `{"memories":[]}` + messages.set("s", []*types.Message{userMessage("s", "must-survive-restart", time.Now().Add(-time.Hour))}) + s.ScheduleExtraction(ctx, "s", "m", "model") + original := queue.pop() + _, err := s.repo.ClaimPendingSessions(ctx, scope, "s", "crashed", time.Minute) + require.NoError(t, err) + require.NoError(t, s.Handle(ctx, original)) + require.Zero(t, models.callCount()) + retry := queue.pop() + require.NotNil(t, retry, "a busy lease must defer redelivery rather than consume the only surviving task") + require.Greater(t, queue.options[len(queue.options)-1].processIn, 50*time.Second) + // Simulate the crashed worker's lease being released by recovery. + require.NoError(t, s.repo.ReleaseExtractionSlot(ctx, scope, "crashed")) + require.NoError(t, s.Handle(ctx, retry)) + require.Equal(t, 1, models.callCount()) +} + +func TestMemoryConsistencyMigrationRestoresLegacyPendingTarget(t *testing.T) { + _, db, _ := newMemoryHarness(t) + testMemoryConsistencyMigration(t, db, "sqlite") +} + +func execMemoryMigration(t *testing.T, db *gorm.DB, path string) { + t.Helper() + contents, err := os.ReadFile(path) + require.NoError(t, err) + var sql strings.Builder + for _, line := range strings.Split(string(contents), "\n") { + if !strings.HasPrefix(strings.TrimSpace(line), "--") { + sql.WriteString(line) + sql.WriteByte('\n') + } + } + for _, stmt := range strings.Split(sql.String(), ";") { + if strings.TrimSpace(stmt) != "" { + require.NoError(t, db.Exec(stmt).Error, stmt) + } + } +} + +func testMemoryConsistencyMigration(t *testing.T, db *gorm.DB, dialect string) { + t.Helper() + for _, table := range []string{ + "memory_extraction_sessions", "memory_item_embeddings", "memory_doc_affinity", + "memory_topic_stats", "memory_tombstones", "memory_items", "memory_subjects", + } { + require.NoError(t, db.Exec("DROP TABLE IF EXISTS "+table).Error) + } + require.NoError(t, db.Exec("CREATE TABLE IF NOT EXISTS tenants (id BIGINT PRIMARY KEY)").Error) + require.NoError(t, db.Exec("CREATE TABLE IF NOT EXISTS messages (id VARCHAR(36) PRIMARY KEY)").Error) + baseline := "sqlite/000004_memory" + migration := "sqlite/000015_memory_consistency" + if dialect == "postgres" { + baseline = "versioned/000084_memory" + migration = "versioned/000094_memory_consistency" + } + execMemoryMigration(t, db, "../../../../migrations/"+baseline+".up.sql") + for _, row := range []struct{ id, key, status, by string }{ + {"old", "job", "superseded", "proposal"}, + {"proposal", "job", "pending", ""}, + {"old2", "database", "superseded", "proposal2"}, + {"proposal2", "database", "pending", ""}, + {"newer", "database", "active", ""}, + } { + require.NoError(t, db.Exec("INSERT INTO memory_items "+ + "(id, tenant_id, subject_id, normalized_key, status, superseded_by, valid_from, kind, content) "+ + "VALUES (?, 1, 'alice', ?, ?, ?, ?, 'fact', 'legacy fact')", + row.id, row.key, row.status, row.by, time.Now()).Error) + } + execMemoryMigration(t, db, "../../../../migrations/"+migration+".up.sql") + require.True(t, db.Migrator().HasTable(&types.MemoryExtractionSession{})) + require.True(t, db.Migrator().HasIndex(&types.MemoryItem{}, "idx_memory_replaces")) + var old, proposal, untouched types.MemoryItem + require.NoError(t, db.First(&old, "id = ?", "old").Error) + require.NoError(t, db.First(&proposal, "id = ?", "proposal").Error) + require.NoError(t, db.First(&untouched, "id = ?", "old2").Error) + require.Equal(t, types.MemoryStatusActive, old.Status) + require.Empty(t, old.SupersededBy) + require.Equal(t, "old", proposal.ReplacesID) + require.Equal(t, types.MemoryStatusSuperseded, untouched.Status, "do not override a newer active fact") + var active int64 + require.NoError(t, db.Model(&types.MemoryItem{}).Where("normalized_key = 'database' AND status = 'active'").Count(&active).Error) + require.Equal(t, int64(1), active) + execMemoryMigration(t, db, "../../../../migrations/"+migration+".down.sql") + require.False(t, db.Migrator().HasTable(&types.MemoryExtractionSession{})) + require.False(t, db.Migrator().HasColumn(&types.MemoryItem{}, "replaces_id")) +} diff --git a/internal/application/service/memory/postgres_consistency_test.go b/internal/application/service/memory/postgres_consistency_test.go new file mode 100644 index 000000000..44b6d1988 --- /dev/null +++ b/internal/application/service/memory/postgres_consistency_test.go @@ -0,0 +1,124 @@ +package memory + +import ( + "context" + "os" + "sync" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Opt in with a disposable test database. Each run uses an isolated schema. +func TestMemoryConsistencyPostgres(t *testing.T) { + dsn := os.Getenv("WEKNORA_MEMORY_TEST_POSTGRES_DSN") + if dsn == "" { + t.Skip("set WEKNORA_MEMORY_TEST_POSTGRES_DSN to run PostgreSQL integration tests") + } + admin, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + require.NoError(t, err) + sqlAdmin, err := admin.DB() + require.NoError(t, err) + defer func() { require.NoError(t, sqlAdmin.Close()) }() + require.NoError(t, admin.Exec(`CREATE EXTENSION IF NOT EXISTS "uuid-ossp"`).Error) + schema := "memory_test_" + uuid.NewString()[:8] + require.NoError(t, admin.Exec("CREATE SCHEMA "+schema).Error) + defer admin.Exec("DROP SCHEMA " + schema + " CASCADE") + db, err := gorm.Open(postgres.Open(dsn+" search_path="+schema+",public"), &gorm.Config{Logger: logger.Discard}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + defer func() { require.NoError(t, sqlDB.Close()) }() + + testMemoryConsistencyMigration(t, db, "postgres") + execMemoryMigration(t, db, "../../../../migrations/versioned/000094_memory_consistency.up.sql") + repo := repository.NewMemoryRepository(db) + ctx := context.Background() + scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "alice"} + _, err = repo.EnsureSubject(ctx, scope) + require.NoError(t, err) + var wg sync.WaitGroup + errs := make(chan error, 20) + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, _, err := repo.EnqueuePendingSession(ctx, scope, "session", time.Minute) + errs <- err + }() + } + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + batch, err := repo.ClaimPendingSessions(ctx, scope, "", "worker", time.Minute) + require.NoError(t, err) + require.Len(t, batch.Sessions, 1) + require.EqualValues(t, 20, batch.Sessions[0].Revision) + duplicate, err := repo.ClaimPendingSessions(ctx, scope, "", "duplicate", time.Minute) + require.NoError(t, err) + require.False(t, duplicate.RetryAt.IsZero()) + end := types.MemoryMessageCursor{At: time.Now().UTC().Truncate(time.Microsecond), ID: "last"} + for i := 1; i <= 3; i++ { + skip, err := repo.RecordExtractionFailure(ctx, scope, "worker", interfaces.MemoryExtractionFailure{ + Session: batch.Sessions[0], End: end, Code: "invalid_model_output", + }) + require.NoError(t, err) + require.Equal(t, i == 3, skip) + } + require.NoError(t, repo.CheckpointExtraction(ctx, scope, "worker", batch.Sessions[0], end, true)) + require.NoError(t, repo.FinishExtraction(ctx, scope, "worker")) + pending, err := repo.HasPendingExtraction(ctx, scope) + require.NoError(t, err) + require.False(t, pending) + + old := &types.MemoryItem{ + ID: uuid.NewString(), TenantID: scope.TenantID, SubjectID: scope.SubjectID, + Kind: types.MemoryKindFact, Content: "old", Topic: "job", NormalizedKey: "job", + Status: types.MemoryStatusActive, ValidFrom: time.Now(), + } + require.NoError(t, repo.SaveItem(ctx, scope, old, "")) + // Seed two competing proposals to exercise confirmation itself, including legacy duplicates. + ids := []string{uuid.NewString(), uuid.NewString()} + for _, id := range ids { + require.NoError(t, db.Create(&types.MemoryItem{ + ID: id, TenantID: scope.TenantID, SubjectID: scope.SubjectID, + Kind: types.MemoryKindFact, Content: id, Topic: "job", NormalizedKey: "job", + Status: types.MemoryStatusPending, ReplacesID: old.ID, ValidFrom: time.Now(), + }).Error) + } + results := make(chan error, 2) + for _, id := range ids { + wg.Add(1) + go func(id string) { + defer wg.Done() + results <- repo.ConfirmPendingItem(ctx, scope, id) + }(id) + } + wg.Wait() + close(results) + successes := 0 + for err := range results { + if err == nil { + successes++ + } else { + require.ErrorIs(t, err, types.ErrMemoryConflict) + } + } + require.Equal(t, 1, successes) + var active int64 + require.NoError(t, db.Model(&types.MemoryItem{}). + Where("tenant_id = ? AND subject_id = ? AND status = ?", + scope.TenantID, scope.SubjectID, types.MemoryStatusActive). + Count(&active).Error) + require.EqualValues(t, 1, active) +} diff --git a/internal/application/service/memory/review_consistency_test.go b/internal/application/service/memory/review_consistency_test.go new file mode 100644 index 000000000..785c771de --- /dev/null +++ b/internal/application/service/memory/review_consistency_test.go @@ -0,0 +1,174 @@ +package memory + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/stretchr/testify/require" +) + +func TestMemoryConsistencyInvalidDecisionsDoNotBlockValidOnes(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + old, err := s.Remember(ctx, types.MemoryItem{Kind: types.MemoryKindFact, Topic: "职业", Content: "我是工程师"}) + require.NoError(t, err) + require.NoError(t, s.repo.DeleteItem(ctx, scope, old.ID)) // Snapshot became stale during the model call. + zero, invalid := 0, 99 + require.NoError(t, s.applyDecisions(ctx, scope, s.workspaceConfig(ctx, 1), + transcriptSegment{}, []*types.MemoryItem{old}, []extractionDecision{ + {Action: "update", Target: &zero, Content: "我是经理"}, + {Action: "add", Content: " "}, + {Action: "update", Target: &invalid, Content: "错误索引"}, + {Action: "update", Topic: "不存在", Content: "不应该变成新增"}, + {Action: "add", Kind: types.MemoryKindFact, Topic: "编辑器", Content: "使用 Neovim"}, + })) + items, total, err := s.ListItems(ctx, types.MemoryStatusActive, 10, 0) + require.NoError(t, err) + require.EqualValues(t, 1, total) + require.Equal(t, "使用 Neovim", items[0].Content) +} + +func TestMemoryConsistencyDeleteInvalidatesPendingTarget(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + old, err := s.Remember(ctx, types.MemoryItem{Topic: "职业", Content: "我是工程师"}) + require.NoError(t, err) + proposal, err := s.Remember(ctx, types.MemoryItem{Topic: "职业", Content: "可能是经理", Inferred: true}) + require.NoError(t, err) + require.NoError(t, s.DeleteItem(ctx, old.ID)) + _, total, err := s.ListItems(ctx, types.MemoryStatusPending, 10, 0) + require.NoError(t, err) + require.Zero(t, total) + _, err = s.ConfirmItem(ctx, proposal.ID) + require.ErrorIs(t, err, types.ErrMemoryConflict) +} + +func TestMemoryConsistencyPoisonSegmentHasBoundedRetries(t *testing.T) { + s, tr, messages, models, queue := newExtractionHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + at := time.Now().Add(-24 * time.Hour) + // Separate sessions also prove a poison input cannot pin the other queue entries. + badMessage := userMessage("bad", "poison-marker", at) + messages.set("bad", []*types.Message{badMessage}) + messages.set("good", []*types.Message{userMessage("good", "valid-marker", at.Add(time.Hour))}) + models.response = `{"memories":[]}` + models.responseFor = map[string]string{"poison-marker": "not json"} + s.ScheduleExtraction(ctx, "bad", "m", "model") + s.ScheduleExtraction(ctx, "good", "m", "model") + drainExtractions(t, s, queue) + poisonCalls := 0 + for _, prompt := range models.prompts { + if containsTranscript(prompt, "poison-marker") { + poisonCalls++ + } + } + require.Equal(t, 3, poisonCalls) + require.Contains(t, models.seenTranscripts(), "valid-marker") + pending, err := s.repo.HasPendingExtraction(ctx, scopeFor(t, ctx)) + require.NoError(t, err) + require.False(t, pending) + + // Failure metadata is checked through a repository claim after a fresh turn. + _, _, err = s.repo.EnqueuePendingSession(ctx, scopeFor(t, ctx), "bad", time.Minute) + require.NoError(t, err) + batch, err := s.repo.ClaimPendingSessions(ctx, scopeFor(t, ctx), "", "inspect", time.Minute) + require.NoError(t, err) + require.Len(t, batch.Sessions, 1) + progress := batch.Sessions[0] + require.Equal(t, 3, progress.FailureCount) + require.NotNil(t, progress.FailedAt) + require.Equal(t, "invalid_model_output", progress.FailureCode) + require.True(t, progress.Cursor.At.Equal(progress.FailedTo.At)) + require.Equal(t, progress.Cursor.ID, progress.FailedTo.ID) + require.NoError(t, s.repo.FinishExtraction(ctx, scopeFor(t, ctx), "inspect")) + // Later messages in the same conversation still run, without re-reading the poison input. + messages.set("bad", []*types.Message{ + badMessage, + userMessage("bad", "later-valid-marker", at.Add(2*time.Hour)), + }) + models.responseFor = nil + s.ScheduleExtraction(ctx, "bad", "later", "model") + before := len(models.prompts) + drainExtractions(t, s, queue) + require.Greater(t, len(models.prompts), before) + require.Contains(t, transcriptBlock(models.prompts[before]), "later-valid-marker") + require.NotContains(t, transcriptBlock(models.prompts[before]), "poison-marker") +} + +func containsTranscript(prompt, marker string) bool { + // Inspect only extractable input, excluding context and prompt examples. + return strings.Contains(transcriptBlock(prompt), marker) +} + +func TestMemoryConsistencyProgressDoesNotGrowSubjectJSON(t *testing.T) { + s, db, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + scope := scopeFor(t, ctx) + subject, err := s.repo.EnsureSubject(ctx, scope) + require.NoError(t, err) + legacyBoundary := time.Now().Add(-time.Hour).UTC().Truncate(time.Second) + require.NoError(t, db.Model(subject).Updates(map[string]interface{}{ + "extract_cursor": legacyBoundary, + "pending_sessions": types.MemoryPendingSessions{"legacy"}, + }).Error) + for i := 0; i < 100; i++ { + _, _, err := s.repo.EnqueuePendingSession(ctx, scope, fmt.Sprintf("session-%03d", i), time.Minute) + require.NoError(t, err) + } + for i := 0; ; i++ { + lease := fmt.Sprintf("lease-%d", i) + batch, err := s.repo.ClaimPendingSessions(ctx, scope, "", lease, time.Minute) + require.NoError(t, err) + if batch == nil { + break + } + require.LessOrEqual(t, len(batch.Sessions), types.MaxMemoryPendingSessions) + for _, session := range batch.Sessions { + require.True(t, session.Cursor.At.Equal(legacyBoundary), "upgrade must not replay pre-cursor history") + require.NoError(t, s.repo.CheckpointExtraction(ctx, scope, lease, session, + types.MemoryMessageCursor{At: legacyBoundary.Add(time.Minute), ID: "done"}, true)) + } + require.NoError(t, s.repo.FinishExtraction(ctx, scope, lease)) + } + subject, err = s.repo.GetSubject(ctx, scope) + require.NoError(t, err) + raw, err := json.Marshal(subject.ExtractionState) + require.NoError(t, err) + require.Less(t, len(raw), 100) + require.Empty(t, subject.PendingSessions) + require.True(t, subject.ExtractCursor.Equal(legacyBoundary), "legacy baseline must never advance") + var count int64 + require.NoError(t, db.Model(&types.MemoryExtractionSession{}).Count(&count).Error) + require.EqualValues(t, 101, count) + + // A completed row remembers its cursor; a genuinely new row uses the frozen baseline. + _, _, err = s.repo.EnqueuePendingSession(ctx, scope, "legacy", time.Minute) + require.NoError(t, err) + batch, err := s.repo.ClaimPendingSessions(ctx, scope, "", "again", time.Minute) + require.NoError(t, err) + require.Equal(t, "done", batch.Sessions[0].Cursor.ID) +} + +type brokenDecisionLookup struct{ interfaces.MemoryRepository } + +func (brokenDecisionLookup) FindActiveByKey( + context.Context, interfaces.MemoryScope, string, +) (*types.MemoryItem, error) { + return nil, fmt.Errorf("database unavailable") +} + +func TestMemoryConsistencyDecisionDatabaseErrorsRemainRetryable(t *testing.T) { + s, _, tr := newMemoryHarness(t) + ctx := enabledCtx(t, tr, 1, "alice") + s.repo = brokenDecisionLookup{s.repo} + err := s.applyDecisions(ctx, scopeFor(t, ctx), s.workspaceConfig(ctx, 1), + transcriptSegment{}, nil, []extractionDecision{{Action: "update", Topic: "职业", Content: "经理"}}) + require.ErrorContains(t, err, "database unavailable") +} diff --git a/internal/application/service/memory/service.go b/internal/application/service/memory/service.go index d8b89e92c..ffe41853b 100644 --- a/internal/application/service/memory/service.go +++ b/internal/application/service/memory/service.go @@ -299,6 +299,10 @@ func (s *Service) write( cfg *types.MemoryConfig, item types.MemoryItem, ) (*types.MemoryItem, error) { + return s.writeReplacing(ctx, scope, cfg, item, "") +} + +func (s *Service) writeReplacing(ctx context.Context, scope interfaces.MemoryScope, cfg *types.MemoryConfig, item types.MemoryItem, targetID string) (*types.MemoryItem, error) { content := types.SanitizeMemoryContent(item.Content) if content == "" { return nil, errors.New("memory: empty content") @@ -347,11 +351,19 @@ func (s *Service) write( topic := types.SanitizeMemoryTopic(item.Topic) normalizedKey := types.MemoryItemKey(topic, content) - existing, err := s.repo.FindActiveByKey(ctx, scope, normalizedKey) + var existing *types.MemoryItem + if targetID != "" { + existing, err = s.repo.GetItem(ctx, scope, targetID) + if err == nil && existing == nil { + return nil, types.ErrMemoryConflict + } + } else { + existing, err = s.repo.FindActiveByKey(ctx, scope, normalizedKey) + } if err != nil { return nil, fmt.Errorf("find conflicting memory: %w", err) } - if existing != nil && types.SanitizeMemoryContent(existing.Content) == content { + if existing != nil && existing.Status == types.MemoryStatusActive && types.SanitizeMemoryContent(existing.Content) == content { // Same statement about the same topic: nothing changed, so keep the // original timestamps instead of churning the row on every turn. return existing, nil @@ -366,7 +378,7 @@ func (s *Service) write( if err != nil { return nil, err } - if duplicate != nil && !longer { + if duplicate != nil && !longer && (duplicate.Status == types.MemoryStatusActive || statusForWrite(item) == types.MemoryStatusPending) { return duplicate, nil } // The new statement subsumes the old one, so let it supersede. @@ -392,15 +404,11 @@ func (s *Service) write( if stored.Origin == "" { stored.Origin = types.MemoryOriginExtracted } - if err := s.repo.CreateItem(ctx, stored); err != nil { - return nil, fmt.Errorf("create memory item: %w", err) - } if existing != nil { - // Supersede rather than delete: the old statement keeps its content - // and gains invalid_at, so the memory manager can show what changed. - if err := s.repo.SupersedeItem(ctx, scope, existing.ID, stored.ID); err != nil { - logger.Warnf(ctx, "memory: supersede %s failed: %v", existing.ID, err) - } + targetID = existing.ID + } + if err := s.repo.SaveItem(ctx, scope, stored, targetID); err != nil { + return nil, fmt.Errorf("save memory item: %w", err) } s.enforceCapacity(ctx, scope, cfg) @@ -747,6 +755,12 @@ func (s *Service) UpdateItem( if sanitized == "" { return nil, errors.New("memory: empty content") } + if redacted, changed := types.RedactSensitive(sanitized); changed { + if types.IsMostlyRedacted(redacted) { + return nil, ErrSensitiveContent + } + sanitized = types.SanitizeMemoryContent(redacted) + } // Keep the original topic: the user is correcting the statement, not // re-filing it under a different subject, and reusing the topic is what // keeps the correction able to supersede a future extraction. @@ -756,7 +770,12 @@ func (s *Service) UpdateItem( return nil, err } s.rebuildBlock(ctx, scope) - return s.repo.GetItem(ctx, scope, id) + updated, err := s.repo.GetItem(ctx, scope, id) + if err != nil { + return nil, err + } + s.storeItemEmbedding(ctx, scope, s.workspaceConfig(ctx, scope.TenantID), updated) + return updated, nil } // DeleteItem forgets one memory permanently. @@ -1288,9 +1307,10 @@ func (s *Service) ConfirmItem(ctx context.Context, id string) (*types.MemoryItem if existing == nil { return nil, ErrItemNotFound } - if err := s.repo.SetItemStatus(ctx, scope, id, types.MemoryStatusActive); err != nil { + if err := s.repo.ConfirmPendingItem(ctx, scope, id); err != nil { return nil, err } + s.enforceCapacity(ctx, scope, s.workspaceConfig(ctx, scope.TenantID)) s.rebuildBlock(ctx, scope) return s.repo.GetItem(ctx, scope, id) } diff --git a/internal/application/service/memory/service_test.go b/internal/application/service/memory/service_test.go index d6981385b..d9484b7da 100644 --- a/internal/application/service/memory/service_test.go +++ b/internal/application/service/memory/service_test.go @@ -33,7 +33,7 @@ func newMemoryHarness(t *testing.T) (*Service, *gorm.DB, *stubTenantRepo) { sqlDB.SetMaxOpenConns(1) require.NoError(t, db.AutoMigrate(&types.MemorySubject{}, &types.MemoryItem{}, &types.MemoryTombstone{}, &types.MemoryTopicStat{}, &types.MemoryDocAffinity{}, - &types.MemoryItemEmbedding{})) + &types.MemoryItemEmbedding{}, &types.MemoryExtractionSession{})) tenantRepo := &stubTenantRepo{ configs: map[uint64]*types.MemoryConfig{}, diff --git a/internal/application/service/memory/stubs_test.go b/internal/application/service/memory/stubs_test.go index 1ed00f1c8..e3e4a0a2c 100644 --- a/internal/application/service/memory/stubs_test.go +++ b/internal/application/service/memory/stubs_test.go @@ -373,3 +373,24 @@ func (e *stubEmbedder) BatchEmbedWithPool( func (e *stubEmbedder) GetModelName() string { return "stub-embedder" } func (e *stubEmbedder) GetDimensions() int { return 3 } func (e *stubEmbedder) GetModelID() string { return "embed-1" } + +func (s *stubMessageRepo) ListMessagesBySessionAfterCursor(ctx context.Context, sessionID string, cursor types.MemoryMessageCursor, limit int) ([]*types.Message, error) { + messages, err := s.ListMessagesBySessionAfterTime(ctx, sessionID, time.Time{}, 0) + if err != nil { + return nil, err + } + sort.SliceStable(messages, func(i, j int) bool { + return (types.MemoryMessageCursor{At: messages[j].CreatedAt, ID: messages[j].ID}).After(types.MemoryMessageCursor{At: messages[i].CreatedAt, ID: messages[i].ID}) + }) + var out []*types.Message + for _, message := range messages { + next := types.MemoryMessageCursor{At: message.CreatedAt, ID: message.ID} + if (cursor.At.IsZero() && cursor.ID == "") || next.After(cursor) { + out = append(out, message) + } + if limit > 0 && len(out) >= limit { + break + } + } + return out, nil +} diff --git a/internal/application/service/memory/vector.go b/internal/application/service/memory/vector.go index 105dacc90..0cab4176c 100644 --- a/internal/application/service/memory/vector.go +++ b/internal/application/service/memory/vector.go @@ -100,10 +100,12 @@ func (s *Service) storeItemEmbedding( return } err := s.repo.UpsertItemEmbedding(ctx, scope, &types.MemoryItemEmbedding{ - ItemID: item.ID, - ModelID: modelID, - Dims: len(vector), - Vector: types.EncodeEmbedding(vector), + ItemID: item.ID, + SourceContent: item.Content, + SourceTopic: item.Topic, + ModelID: modelID, + Dims: len(vector), + Vector: types.EncodeEmbedding(vector), }) if err != nil { logger.Warnf(ctx, "memory: store embedding failed: %v", err) @@ -298,10 +300,12 @@ func (s *Service) backfillEmbeddings( break } err := s.repo.UpsertItemEmbedding(ctx, scope, &types.MemoryItemEmbedding{ - ItemID: item.ID, - ModelID: modelID, - Dims: len(vector), - Vector: types.EncodeEmbedding(vector), + ItemID: item.ID, + SourceContent: item.Content, + SourceTopic: item.Topic, + ModelID: modelID, + Dims: len(vector), + Vector: types.EncodeEmbedding(vector), }) if err != nil { logger.Warnf(ctx, "memory: backfill embedding failed: %v", err) diff --git a/internal/database/migration_sqlite_versioned_schema_test.go b/internal/database/migration_sqlite_versioned_schema_test.go index ed9434e15..5f858ce40 100644 --- a/internal/database/migration_sqlite_versioned_schema_test.go +++ b/internal/database/migration_sqlite_versioned_schema_test.go @@ -14,6 +14,7 @@ import ( // 000041 task queue, 000053 system settings, 000055 processing spans, // 000063 knowledge multi-tags, 000093 browser authorization. var versionedSQLiteTables = []string{ + "memory_extraction_sessions", "task_pending_ops", "task_dead_letters", "system_settings", @@ -27,6 +28,8 @@ var versionedSQLiteTables = []string{ // versionedSQLiteColumns maps each existing table to the columns that the // versioned migrations add and the SQLite baseline was missing. var versionedSQLiteColumns = map[string][]string{ + "memory_subjects": {"extraction_state"}, // 000094 + "memory_items": {"replaces_id"}, // 000094 "tenants": {"api_principal_config"}, // 000064 "users": {"is_system_admin"}, // 000053 "knowledges": {"pending_subtasks_count"}, // 000056 @@ -37,7 +40,7 @@ var versionedSQLiteColumns = map[string][]string{ "mcp_tool_approvals": {"enabled"}, // 000091 } -const expectedSQLiteMigrationVersion = 14 +const expectedSQLiteMigrationVersion = 15 func TestSQLiteMigrationsCreateVersionedSchema(t *testing.T) { repoRoot := sqliteRepoRoot(t) diff --git a/internal/handler/memory.go b/internal/handler/memory.go index 47b32eb7b..678cad00b 100644 --- a/internal/handler/memory.go +++ b/internal/handler/memory.go @@ -452,6 +452,10 @@ func (h *MemoryHandler) fail(c *gin.Context, err error, message string) { c.Error(apperrors.NewUnauthorizedError("no principal in request")) case errors.Is(err, memory.ErrItemNotFound): c.Error(apperrors.NewNotFoundError("memory not found")) + case errors.Is(err, types.ErrMemoryConflict): + c.Error(apperrors.NewConflictError(err.Error())) + case errors.Is(err, memory.ErrSensitiveContent): + c.Error(apperrors.NewBadRequestError(err.Error())) case errors.Is(err, memory.ErrMemoryDisabled): c.Error(apperrors.NewBadRequestError("memory is disabled")) default: diff --git a/internal/handler/memory_consistency_test.go b/internal/handler/memory_consistency_test.go new file mode 100644 index 000000000..c3928ca4a --- /dev/null +++ b/internal/handler/memory_consistency_test.go @@ -0,0 +1,51 @@ +package handler + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Tencent/WeKnora/internal/application/service/memory" + "github.com/Tencent/WeKnora/internal/middleware" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type memoryFailureService struct{ interfaces.MemoryService } + +func (memoryFailureService) ConfirmItem(context.Context, string) (*types.MemoryItem, error) { + return nil, fmt.Errorf("confirm: %w", types.ErrMemoryConflict) +} + +func (memoryFailureService) UpdateItem(context.Context, string, string, int) (*types.MemoryItem, error) { + return nil, fmt.Errorf("edit: %w", memory.ErrSensitiveContent) +} + +func TestMemoryConsistencyHTTPFailures(t *testing.T) { + router := gin.New() + router.Use(middleware.ErrorHandler()) + handler := NewMemoryHandler(memoryFailureService{}) + router.POST("/memory/items/:id/confirm", handler.ConfirmItem) + router.PUT("/memory/items/:id", handler.UpdateItem) + for _, tc := range []struct { + method, path, body string + status int + }{ + {http.MethodPost, "/memory/items/stale/confirm", "", http.StatusConflict}, + {http.MethodPut, "/memory/items/item", `{"content":"sensitive"}`, http.StatusBadRequest}, + } { + t.Run(tc.method, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + require.Equal(t, tc.status, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), `"success":false`) + }) + } +} diff --git a/internal/types/interfaces/memory.go b/internal/types/interfaces/memory.go index a5647574b..eec1083f7 100644 --- a/internal/types/interfaces/memory.go +++ b/internal/types/interfaces/memory.go @@ -41,15 +41,25 @@ type MemoryRepository interface { EnqueuePendingSession( ctx context.Context, scope MemoryScope, sessionID string, inFlightTimeout time.Duration, ) (*types.MemorySubject, bool, error) - // ClaimPendingSessions takes the pending queue for processing, returning - // the sessions to drain and the cursor to walk forward from. - ClaimPendingSessions(ctx context.Context, scope MemoryScope) ([]string, time.Time, error) - // FinishExtraction advances the watermark and releases the in-flight slot. - // A zero cursor leaves the watermark untouched. - FinishExtraction(ctx context.Context, scope MemoryScope, cursor time.Time) error - // ReleaseExtractionSlot clears the in-flight marker without advancing the - // watermark, used when a run ends early. - ReleaseExtractionSlot(ctx context.Context, scope MemoryScope) error + // ClaimPendingSessions leases a snapshot without removing durable work. + // A nil batch means no work remains; RetryAt defers a busy lease. + ClaimPendingSessions(ctx context.Context, scope MemoryScope, fallbackSession, leaseID string, ttl time.Duration) (*types.MemoryExtractionBatch, error) + // CheckpointExtraction acknowledges a processed segment (or a recorded + // skip). A concurrent enqueue changes Revision and keeps the session pending. + CheckpointExtraction(ctx context.Context, scope MemoryScope, leaseID string, session types.MemoryExtractionSession, cursor types.MemoryMessageCursor, drained bool) error + HasPendingExtraction(ctx context.Context, scope MemoryScope) (bool, error) + // RecordExtractionFailure returns true after the bounded invalid-output + // retry budget. It preserves a failure range without storing transcript text. + RecordExtractionFailure( + ctx context.Context, scope MemoryScope, leaseID string, session MemoryExtractionFailure, + ) (bool, error) + FinishExtraction(ctx context.Context, scope MemoryScope, leaseID string) error + // Empty leaseID only releases a queued task, never a running worker. + ReleaseExtractionSlot(ctx context.Context, scope MemoryScope, leaseID string) error + // SaveItem atomically inserts and resolves active/pending replacements. + SaveItem(ctx context.Context, scope MemoryScope, item *types.MemoryItem, replacesID string) error + // ConfirmPendingItem atomically activates a proposal and retires its target. + ConfirmPendingItem(ctx context.Context, scope MemoryScope, id string) error // CreateItem inserts one memory item. CreateItem(ctx context.Context, item *types.MemoryItem) error @@ -116,9 +126,6 @@ type MemoryRepository interface { // ListLive returns items of one kind that the user can see: in use plus // proposed and awaiting a decision. ListLive(ctx context.Context, scope MemoryScope, kind string, limit int) ([]*types.MemoryItem, error) - // SetItemStatus moves an item between statuses, used to confirm or reject - // something the system inferred. - SetItemStatus(ctx context.Context, scope MemoryScope, id, status string) error // BumpTopic records one more sighting of a topic and returns the running // total, so a caller can decide whether it has recurred enough to promote. @@ -174,6 +181,13 @@ type MemoryRepository interface { CountActive(ctx context.Context, scope MemoryScope) (int64, error) } +// MemoryExtractionFailure identifies an invalid-output range without retaining message text. +type MemoryExtractionFailure struct { + Session types.MemoryExtractionSession + End types.MemoryMessageCursor + Code string +} + // MemoryRecall is what one turn pulls in: the resident block plus any // situational items matched against the query. type MemoryRecall struct { diff --git a/internal/types/interfaces/message.go b/internal/types/interfaces/message.go index b6ad52487..f919c656a 100644 --- a/internal/types/interfaces/message.go +++ b/internal/types/interfaces/message.go @@ -87,6 +87,8 @@ type MessageRepository interface { ListMessagesBySessionAfterTime( ctx context.Context, sessionID string, afterTime time.Time, limit int, ) ([]*types.Message, error) + // ListMessagesBySessionAfterCursor uses (created_at, id) for lossless paging. + ListMessagesBySessionAfterCursor(ctx context.Context, sessionID string, cursor types.MemoryMessageCursor, limit int) ([]*types.Message, error) // UpdateMessage updates a message UpdateMessage(ctx context.Context, message *types.Message) error // UpdateMessageImages updates only the images JSONB column for a message diff --git a/internal/types/memory.go b/internal/types/memory.go index da2548f6b..c8f39f039 100644 --- a/internal/types/memory.go +++ b/internal/types/memory.go @@ -205,14 +205,13 @@ type MemorySubject struct { BlockUpdatedAt *time.Time `json:"block_updated_at" gorm:"column:block_updated_at"` ItemCount int `json:"item_count" gorm:"column:item_count;not null;default:0"` LastExtractedAt *time.Time `json:"last_extracted_at" gorm:"column:last_extracted_at"` - // ExtractCursor is the watermark: everything this subject said up to and - // including this instant has already been considered for distillation. - // Distillation walks forward from here, which is what makes "no message is - // skipped" a property of the data rather than of timing. - ExtractCursor *time.Time `json:"extract_cursor" gorm:"column:extract_cursor"` - // PendingSessions are the sessions with turns past the cursor. A turn that - // arrives while a task is already in flight is recorded here instead of - // being dropped, so it is picked up by the run that is already coming. + // ExtractCursor is the legacy subject-wide watermark, retained for + // the upgrade boundary for newly initialized session cursors. It is never + // advanced by new workers; each session has its own progress row. + ExtractCursor *time.Time `json:"extract_cursor" gorm:"column:extract_cursor"` + ExtractionState MemoryExtractionState `json:"-" gorm:"column:extraction_state;type:jsonb"` + // PendingSessions is the legacy queue, imported into indexed progress rows + // once on the next enqueue or claim. New workers never grow this array. PendingSessions MemoryPendingSessions `json:"pending_sessions" gorm:"column:pending_sessions;type:jsonb"` // ExtractScheduledAt marks a distillation task as in flight, so concurrent // turns enqueue one task rather than one per turn. @@ -235,10 +234,7 @@ type MemorySubject struct { // distillation for one subject. type MemoryPendingSessions []string -// MaxMemoryPendingSessions bounds the queue. A subject chatting in more -// sessions than this between two runs loses the oldest entry rather than -// growing the row without limit; the cursor still covers those sessions the -// next time they produce a turn. +// MaxMemoryPendingSessions bounds one processing batch, never the durable queue. const MaxMemoryPendingSessions = 32 func (p MemoryPendingSessions) Value() (driver.Value, error) { @@ -270,7 +266,7 @@ func (p *MemoryPendingSessions) Scan(value interface{}) error { return json.Unmarshal(b, p) } -// Append adds a session id, keeping the queue de-duplicated and bounded. +// Append adds a session id without dropping existing work. func (p MemoryPendingSessions) Append(sessionID string) MemoryPendingSessions { sessionID = strings.TrimSpace(sessionID) if sessionID == "" { @@ -282,9 +278,6 @@ func (p MemoryPendingSessions) Append(sessionID string) MemoryPendingSessions { } } updated := append(p, sessionID) - if len(updated) > MaxMemoryPendingSessions { - updated = updated[len(updated)-MaxMemoryPendingSessions:] - } return updated } @@ -320,6 +313,7 @@ type MemoryItem struct { // Without it an in-flight task stays in context forever and slowly turns // the memory into a list of things the user finished months ago. ExpiresAt *time.Time `json:"expires_at" gorm:"column:expires_at"` + ReplacesID string `json:"replaces_id,omitempty" gorm:"column:replaces_id;type:varchar(36);not null;default:''"` SupersededBy string `json:"superseded_by" gorm:"column:superseded_by;type:varchar(36)"` LastUsedAt *time.Time `json:"last_used_at" gorm:"column:last_used_at"` UseCount int `json:"use_count" gorm:"column:use_count;not null;default:0"` @@ -1371,9 +1365,12 @@ func TopicLooksLikeOneQuestion(topic string) bool { // a few kilobytes of float per row along for the ride. Only the code that // actually scores similarity loads these. type MemoryItemEmbedding struct { - ItemID string `json:"item_id" gorm:"primaryKey;type:varchar(36)"` - TenantID uint64 `json:"tenant_id" gorm:"not null;index:idx_mem_emb_scope,priority:1"` - SubjectID string `json:"subject_id" gorm:"type:varchar(512);not null;index:idx_mem_emb_scope,priority:2"` + // Input snapshot prevents a slow embedding call from overwriting a newer edit. + SourceContent string `json:"-" gorm:"-"` + SourceTopic string `json:"-" gorm:"-"` + ItemID string `json:"item_id" gorm:"primaryKey;type:varchar(36)"` + TenantID uint64 `json:"tenant_id" gorm:"not null;index:idx_mem_emb_scope,priority:1"` + SubjectID string `json:"subject_id" gorm:"type:varchar(512);not null;index:idx_mem_emb_scope,priority:2"` // ModelID records which model produced this vector. Vectors from different // models are not comparable, so a model change has to invalidate them // rather than silently score nonsense. diff --git a/internal/types/memory_extraction.go b/internal/types/memory_extraction.go new file mode 100644 index 000000000..649240400 --- /dev/null +++ b/internal/types/memory_extraction.go @@ -0,0 +1,76 @@ +package types + +import ( + "database/sql/driver" + "encoding/json" + "errors" + "fmt" + "time" +) + +var ErrMemoryConflict = errors.New("memory changed; reload before applying this proposal") +var ErrMemoryExtractionLeaseLost = errors.New("memory extraction lease lost") + +// MemoryMessageCursor breaks timestamp ties using the message primary key. +type MemoryMessageCursor struct { + At time.Time `json:"at"` + ID string `json:"id"` +} + +func (c MemoryMessageCursor) After(other MemoryMessageCursor) bool { + return c.At.After(other.At) || (c.At.Equal(other.At) && c.ID > other.ID) +} + +// MemoryExtractionSession is a small indexed row, not an entry in a growing +// subject JSON document. Completed cursors prevent historical re-extraction. +type MemoryExtractionSession struct { + TenantID uint64 `json:"-" gorm:"primaryKey;autoIncrement:false"` + SubjectID string `json:"-" gorm:"primaryKey;type:varchar(512)"` + SessionID string `json:"-" gorm:"primaryKey;type:varchar(36)"` + Revision uint64 `json:"revision" gorm:"not null;default:0"` + Cursor MemoryMessageCursor `json:"cursor" gorm:"embedded;embeddedPrefix:cursor_"` + Pending bool `json:"-" gorm:"not null;default:false"` + FailureCount int `json:"-" gorm:"not null;default:0"` + FailureCode string `json:"-" gorm:"type:varchar(64);not null;default:''"` + FailedFrom MemoryMessageCursor `json:"-" gorm:"embedded;embeddedPrefix:failed_from_"` + FailedTo MemoryMessageCursor `json:"-" gorm:"embedded;embeddedPrefix:failed_to_"` + FailedAt *time.Time `json:"-"` + UpdatedAt time.Time `json:"-"` +} + +// TableName selects the indexed session progress table. +func (MemoryExtractionSession) TableName() string { return "memory_extraction_sessions" } + +// MemoryExtractionState contains only the subject-wide worker lease. +// Cursors and queued work live in MemoryExtractionSession. +type MemoryExtractionState struct { + LeaseID string `json:"lease_id,omitempty"` + LeaseUntil time.Time `json:"lease_until,omitempty"` +} + +func (s MemoryExtractionState) Value() (driver.Value, error) { return json.Marshal(s) } + +func (s *MemoryExtractionState) Scan(value interface{}) error { + *s = MemoryExtractionState{} + var raw []byte + switch v := value.(type) { + case nil: + return nil + case []byte: + raw = v + case string: + raw = []byte(v) + default: + return fmt.Errorf("memory extraction state: unsupported value %T", value) + } + if len(raw) == 0 { + return nil + } + return json.Unmarshal(raw, s) +} + +type MemoryExtractionBatch struct { + // RetryAt keeps a redelivered task alive while a crashed worker's lease expires. + RetryAt time.Time + Sessions []MemoryExtractionSession +} diff --git a/migrations/sqlite/000015_memory_consistency.down.sql b/migrations/sqlite/000015_memory_consistency.down.sql new file mode 100644 index 000000000..a83f1f63a --- /dev/null +++ b/migrations/sqlite/000015_memory_consistency.down.sql @@ -0,0 +1,4 @@ +DROP INDEX IF EXISTS idx_memory_replaces; +DROP TABLE IF EXISTS memory_extraction_sessions; +ALTER TABLE memory_items DROP COLUMN replaces_id; +ALTER TABLE memory_subjects DROP COLUMN extraction_state; diff --git a/migrations/sqlite/000015_memory_consistency.up.sql b/migrations/sqlite/000015_memory_consistency.up.sql new file mode 100644 index 000000000..e57002dcd --- /dev/null +++ b/migrations/sqlite/000015_memory_consistency.up.sql @@ -0,0 +1,61 @@ +-- Per-session extraction progress and deferred proposal replacement. +ALTER TABLE memory_subjects ADD COLUMN extraction_state TEXT; +ALTER TABLE memory_items ADD COLUMN replaces_id VARCHAR(36) NOT NULL DEFAULT ''; + +-- Progress is an indexed row per conversation; subject rows only hold a lease. +CREATE TABLE memory_extraction_sessions ( + tenant_id BIGINT NOT NULL, + subject_id VARCHAR(512) NOT NULL, + session_id VARCHAR(36) NOT NULL, + revision BIGINT NOT NULL DEFAULT 0, + cursor_at DATETIME, cursor_id VARCHAR(36) NOT NULL DEFAULT '', + pending BOOLEAN NOT NULL DEFAULT false, + failure_count INTEGER NOT NULL DEFAULT 0, + failure_code VARCHAR(64) NOT NULL DEFAULT '', + failed_from_at DATETIME, failed_from_id VARCHAR(36) NOT NULL DEFAULT '', + failed_to_at DATETIME, failed_to_id VARCHAR(36) NOT NULL DEFAULT '', + failed_at DATETIME, updated_at DATETIME NOT NULL, + PRIMARY KEY (tenant_id, subject_id, session_id) +); +CREATE INDEX idx_memory_extraction_pending ON memory_extraction_sessions (tenant_id, subject_id, pending, updated_at, session_id); +CREATE INDEX idx_memory_replaces ON memory_items (tenant_id, subject_id, replaces_id, status); + +-- Recover the old active record when a legacy pending inference prematurely +-- superseded it. Keep an explicit target so confirmation can still replace it. +UPDATE memory_items +SET replaces_id = ( + SELECT old.id FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.superseded_by = memory_items.id AND old.status = 'superseded' + ORDER BY old.valid_from DESC, old.id DESC LIMIT 1 +) +WHERE status = 'pending' AND EXISTS ( + SELECT 1 FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.superseded_by = memory_items.id AND old.status = 'superseded' +); + +-- Restore at most one record per key, and never displace a newer active fact. +UPDATE memory_items +SET status = 'active', invalid_at = NULL, superseded_by = '' +WHERE status = 'superseded' + AND NOT EXISTS ( + SELECT 1 FROM memory_items AS active + WHERE active.tenant_id = memory_items.tenant_id + AND active.subject_id = memory_items.subject_id + AND active.normalized_key = memory_items.normalized_key AND active.status = 'active' + ) + AND id = ( + SELECT old.id FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.normalized_key = memory_items.normalized_key AND old.status = 'superseded' + AND EXISTS ( + SELECT 1 FROM memory_items AS pending + WHERE pending.tenant_id = old.tenant_id AND pending.subject_id = old.subject_id + AND pending.status = 'pending' AND pending.replaces_id = old.id + ) + ORDER BY old.valid_from DESC, old.id DESC LIMIT 1 + ); diff --git a/migrations/versioned/000094_memory_consistency.down.sql b/migrations/versioned/000094_memory_consistency.down.sql new file mode 100644 index 000000000..a83f1f63a --- /dev/null +++ b/migrations/versioned/000094_memory_consistency.down.sql @@ -0,0 +1,4 @@ +DROP INDEX IF EXISTS idx_memory_replaces; +DROP TABLE IF EXISTS memory_extraction_sessions; +ALTER TABLE memory_items DROP COLUMN replaces_id; +ALTER TABLE memory_subjects DROP COLUMN extraction_state; diff --git a/migrations/versioned/000094_memory_consistency.up.sql b/migrations/versioned/000094_memory_consistency.up.sql new file mode 100644 index 000000000..78f71100a --- /dev/null +++ b/migrations/versioned/000094_memory_consistency.up.sql @@ -0,0 +1,61 @@ +-- Per-session extraction progress and deferred proposal replacement. +ALTER TABLE memory_subjects ADD COLUMN extraction_state JSONB; +ALTER TABLE memory_items ADD COLUMN replaces_id VARCHAR(36) NOT NULL DEFAULT ''; + +-- Progress is an indexed row per conversation; subject rows only hold a lease. +CREATE TABLE memory_extraction_sessions ( + tenant_id BIGINT NOT NULL, + subject_id VARCHAR(512) NOT NULL, + session_id VARCHAR(36) NOT NULL, + revision BIGINT NOT NULL DEFAULT 0, + cursor_at TIMESTAMP WITH TIME ZONE, cursor_id VARCHAR(36) NOT NULL DEFAULT '', + pending BOOLEAN NOT NULL DEFAULT false, + failure_count INTEGER NOT NULL DEFAULT 0, + failure_code VARCHAR(64) NOT NULL DEFAULT '', + failed_from_at TIMESTAMP WITH TIME ZONE, failed_from_id VARCHAR(36) NOT NULL DEFAULT '', + failed_to_at TIMESTAMP WITH TIME ZONE, failed_to_id VARCHAR(36) NOT NULL DEFAULT '', + failed_at TIMESTAMP WITH TIME ZONE, updated_at TIMESTAMP WITH TIME ZONE NOT NULL, + PRIMARY KEY (tenant_id, subject_id, session_id) +); +CREATE INDEX idx_memory_extraction_pending ON memory_extraction_sessions (tenant_id, subject_id, pending, updated_at, session_id); +CREATE INDEX idx_memory_replaces ON memory_items (tenant_id, subject_id, replaces_id, status); + +-- Recover the old active record when a legacy pending inference prematurely +-- superseded it. Keep an explicit target so confirmation can still replace it. +UPDATE memory_items +SET replaces_id = ( + SELECT old.id FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.superseded_by = memory_items.id AND old.status = 'superseded' + ORDER BY old.valid_from DESC, old.id DESC LIMIT 1 +) +WHERE status = 'pending' AND EXISTS ( + SELECT 1 FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.superseded_by = memory_items.id AND old.status = 'superseded' +); + +-- Restore at most one record per key, and never displace a newer active fact. +UPDATE memory_items +SET status = 'active', invalid_at = NULL, superseded_by = '' +WHERE status = 'superseded' + AND NOT EXISTS ( + SELECT 1 FROM memory_items AS active + WHERE active.tenant_id = memory_items.tenant_id + AND active.subject_id = memory_items.subject_id + AND active.normalized_key = memory_items.normalized_key AND active.status = 'active' + ) + AND id = ( + SELECT old.id FROM memory_items AS old + WHERE old.tenant_id = memory_items.tenant_id + AND old.subject_id = memory_items.subject_id + AND old.normalized_key = memory_items.normalized_key AND old.status = 'superseded' + AND EXISTS ( + SELECT 1 FROM memory_items AS pending + WHERE pending.tenant_id = old.tenant_id AND pending.subject_id = old.subject_id + AND pending.status = 'pending' AND pending.replaces_id = old.id + ) + ORDER BY old.valid_from DESC, old.id DESC LIMIT 1 + );