228 lines
9.9 KiB
Go
228 lines
9.9 KiB
Go
|
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)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|