WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/memory/recovery_fencing_test.go

@@ -0,0 +1,227 @@
+package memory
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/application/repository"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRecoveryMaintenanceRetainsPolicyAndLeaseChecks(t *testing.T) {
+	for _, maintenance := range []string{"topics", "consolidation"} {
+		for _, phase := range []string{"before", "after"} {
+			t.Run(maintenance+"/"+phase, func(t *testing.T) {
+				svc, db, ctx, messages, models, queue := recoveryHarness(t)
+				scope, err := ResolveScope(ctx)
+				require.NoError(t, err)
+				svc.ScheduleExtraction(ctx, "s", "m", "model")
+				if maintenance == "topics" {
+					_, err := svc.repo.BumpTopic(ctx, scope, "database tuning", types.NormalizeTopicKey("database tuning"), "database tuning")
+					require.NoError(t, err)
+					models.response = `{"memories":[],"topics":["deployment automation"]}`
+					models.responseFor = map[string]string{topicAdjudicationPrompt: `{"resolutions":[{"index":0,"same_as":0}]}`}
+				} else {
+					for i := range 6 {
+						require.NoError(t, svc.repo.CreateItem(ctx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID,
+							Kind: types.MemoryKindFact, Topic: "deployment automation", Content: fmt.Sprintf("deployment automation uses deterministic verification %d", i)}))
+					}
+					models.responseFor = map[string]string{consolidationSystemPrompt: `{"statement":"deployment automation uses deterministic verification"}`}
+				}
+				var beforeItems int64
+				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&beforeItems).Error)
+				before, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				mutate := func() { messages.owners = map[string]interfaces.MemoryScope{} }
+				models.onChat = func() {
+					if phase == "before" {
+						models.onGetChatModel = mutate
+					} else {
+						models.onChat = mutate
+					}
+				}
+				require.NoError(t, svc.Handle(context.Background(), queue.pop()))
+				calls := 1
+				if phase == "after" {
+					calls = 2
+				}
+				require.Equal(t, calls, models.callCount())
+				var afterItems int64
+				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&afterItems).Error)
+				require.Equal(t, beforeItems, afterItems)
+				after, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				require.Equal(t, before.SessionCursors, after.SessionCursors)
+				require.Equal(t, before.PendingSessions, after.PendingSessions)
+				require.Zero(t, after.ExtractFailures)
+				if maintenance == "topics" {
+					topics, err := svc.repo.TopTopics(ctx, scope, 10)
+					require.NoError(t, err)
+					require.Len(t, topics, 1)
+					require.Equal(t, 1, topics[0].Hits, "revoked maintenance cannot fall back to counting the rejected model output")
+				}
+			})
+		}
+	}
+}
+
+func TestRecoveryWithNoReadableTranscriptSkipsAutomaticMaintenance(t *testing.T) {
+	svc, _, ctx, messages, models, queue := recoveryHarness(t)
+	svc.ScheduleExtraction(ctx, "s", "m", "model")
+	scope, err := ResolveScope(ctx)
+	require.NoError(t, err)
+	for i := range 6 {
+		require.NoError(t, svc.repo.CreateItem(ctx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID,
+			Kind: types.MemoryKindFact, Content: fmt.Sprintf("deployment automation uses deterministic verification %d", i)}))
+	}
+	messages.owners = map[string]interfaces.MemoryScope{}
+	require.NoError(t, svc.Handle(context.Background(), queue.pop()))
+	require.Zero(t, models.callCount())
+	subject, err := svc.repo.GetSubject(ctx, scope)
+	require.NoError(t, err)
+	require.Nil(t, subject.ConsolidatedAt)
+	require.Nil(t, subject.ExtractCursor)
+}
+
+func TestRecoveryRechecksPolicyAndReservationAroundProvider(t *testing.T) {
+	for _, phase := range []string{"before", "after", "truncated"} {
+		for _, change := range []string{"deleted subject", "disabled subject", "disabled workspace", "revoked owner", "disabled Agent", "revoked context Agent", "replaced attempt", "expired lease"} {
+			t.Run(phase+"/"+change, func(t *testing.T) {
+				svc, db, ctx, messages, models, queue := recoveryHarness(t)
+				scope, err := ResolveScope(ctx)
+				require.NoError(t, err)
+				base := time.Now().Add(-time.Hour)
+				prior := userMessage("s", "earlier preference", base)
+				input := userMessage("s", "current preference", base.Add(time.Minute))
+				messages.set("s", []*types.Message{prior, input})
+				svc.ScheduleExtraction(ctx, "s", "m", "model")
+				cursor := base.Add(time.Second)
+				require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).
+					Updates(map[string]interface{}{"session_cursors": types.MemorySessionCursors{"s": cursor}, "extract_cursor": cursor}).Error)
+				before, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				models.response = `{"memories":[{"action":"add","kind":"preference","topic":"working style","content":"prefers clear examples"}],"topics":["software engineering"]}`
+				mutate := func() {
+					switch change {
+					case "deleted subject":
+						require.NoError(t, db.Where("tenant_id = ?", 1).Delete(&types.MemorySubject{}).Error)
+					case "disabled subject":
+						require.NoError(t, svc.repo.UpdateSubjectEnabled(ctx, scope, false))
+					case "disabled workspace":
+						svc.tenantRepo.(*stubTenantRepo).set(1, &types.MemoryConfig{Enabled: false})
+					case "revoked owner":
+						messages.owners = map[string]interfaces.MemoryScope{}
+					case "disabled Agent":
+						messages.deniedMessages = map[string]bool{input.ID: true}
+					case "revoked context Agent":
+						messages.deniedMessages = map[string]bool{prior.ID: true}
+					case "replaced attempt":
+						require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_attempt_id", "replacement").Error)
+					case "expired lease":
+						require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_lease_until", time.Now().Add(-time.Second)).Error)
+					}
+				}
+				if phase == "before" {
+					models.onGetChatModel = mutate
+				} else {
+					models.onChat = mutate
+					if phase == "truncated" {
+						models.truncateUntilCall = 1
+					}
+				}
+				err = svc.Handle(context.Background(), queue.pop())
+				stale := change == "replaced attempt" || change == "expired lease"
+				if stale {
+					require.ErrorIs(t, err, errExtractionAttemptLost)
+				} else {
+					require.NoError(t, err)
+				}
+				wantCalls := 1
+				if phase == "before" {
+					wantCalls = 0
+				}
+				require.Equal(t, wantCalls, models.callCount(), "revocation must also block the truncation retry and maintenance")
+				var count int64
+				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
+				require.Zero(t, count, "discard provider output after revocation")
+				after, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				if change == "deleted subject" {
+					require.Nil(t, after, "worker cannot recreate a deleted subject")
+					return
+				}
+				require.Equal(t, before.PendingSessions, after.PendingSessions)
+				require.Equal(t, before.SessionCursors, after.SessionCursors)
+				require.True(t, before.ExtractCursor.Equal(*after.ExtractCursor))
+				if !stale {
+					require.Zero(t, after.ExtractFailures, "policy changes pause without charging the provider retry budget")
+				}
+				if change == "replaced attempt" {
+					require.Equal(t, "replacement", after.ExtractAttemptID)
+				}
+			})
+		}
+	}
+}
+
+func TestRecoveryRechecksWebAccountLifecycle(t *testing.T) {
+	for _, phase := range []string{"queued", "before", "after"} {
+		for _, change := range []string{"UPDATE users SET is_active = false", "UPDATE users SET deleted_at = CURRENT_TIMESTAMP", "UPDATE tenant_members SET status = 'suspended'"} {
+			t.Run(phase+"/"+change, func(t *testing.T) {
+				svc, db, ctx, _, models, queue := recoveryHarness(t)
+				require.NoError(t, db.AutoMigrate(&types.Session{}, &types.Message{}))
+				for _, sql := range []string{
+					`CREATE TABLE users (id TEXT PRIMARY KEY, is_active BOOLEAN, is_system_admin BOOLEAN, deleted_at DATETIME)`,
+					`CREATE TABLE tenant_members (user_id TEXT, tenant_id INTEGER, status TEXT, deleted_at DATETIME)`,
+					`INSERT INTO users (id, is_active, is_system_admin) VALUES ('alice', true, false)`,
+					`INSERT INTO tenant_members (user_id, tenant_id, status) VALUES ('alice', 1, 'active')`,
+				} {
+					require.NoError(t, db.Exec(sql).Error)
+				}
+				base := time.Now().Add(-time.Hour).UTC().Truncate(time.Microsecond)
+				session := &types.Session{TenantID: 1, UserID: "alice"}
+				require.NoError(t, db.Create(session).Error)
+				input := &types.Message{SessionID: session.ID, RequestID: "r", Role: "user", Content: "prefers simple examples", CreatedAt: base}
+				answer := &types.Message{SessionID: session.ID, RequestID: "r", Role: "assistant", Content: "okay", IsCompleted: true, CreatedAt: base.Add(time.Second)}
+				require.NoError(t, db.Create(input).Error)
+				require.NoError(t, db.Create(answer).Error)
+				svc.messageRepo = repository.NewMessageRepository(db)
+				svc.ScheduleExtraction(ctx, session.ID, answer.ID, "model")
+				task := queue.pop()
+				require.NotNil(t, task)
+				scope, err := ResolveScope(ctx)
+				require.NoError(t, err)
+				before, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				mutate := func() { require.NoError(t, db.Exec(change).Error) }
+				models.response = `{"memories":[{"action":"add","content":"prefers simple examples"}],"topics":["examples"]}`
+				switch phase {
+				case "queued":
+					mutate()
+				case "before":
+					models.onGetChatModel = mutate
+				case "after":
+					models.onChat = mutate
+				}
+				require.NoError(t, svc.Handle(context.Background(), task))
+				wantCalls := 0
+				if phase == "after" {
+					wantCalls = 1
+				}
+				require.Equal(t, wantCalls, models.callCount())
+				after, err := svc.repo.GetSubject(ctx, scope)
+				require.NoError(t, err)
+				require.Equal(t, before.SessionCursors, after.SessionCursors, "revoked owners never replay or advance historical cursors")
+				require.Nil(t, after.ExtractCursor)
+				require.Zero(t, after.ExtractFailures)
+				var count int64
+				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
+				require.Zero(t, count)
+			})
+		}
+	}
+}