304 lines
12 KiB
Go
304 lines
12 KiB
Go
|
package memory
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type extractionFailureRepo struct {
|
||
|
|
interfaces.MemoryRepository
|
||
|
|
createFailures int
|
||
|
|
finishFailures int
|
||
|
|
beforeFinish func()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *extractionFailureRepo) CreateItem(ctx context.Context, item *types.MemoryItem) error {
|
||
|
|
if r.createFailures > 0 {
|
||
|
|
r.createFailures--
|
||
|
|
return errors.New("temporary memory write failure")
|
||
|
|
}
|
||
|
|
return r.MemoryRepository.CreateItem(ctx, item)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *extractionFailureRepo) FinishExtraction(ctx context.Context, scope interfaces.MemoryScope, cursors types.MemorySessionCursors, completed ...string) error {
|
||
|
|
if r.beforeFinish != nil {
|
||
|
|
callback := r.beforeFinish
|
||
|
|
r.beforeFinish = nil
|
||
|
|
callback()
|
||
|
|
}
|
||
|
|
if r.finishFailures > 0 {
|
||
|
|
r.finishFailures--
|
||
|
|
return errors.New("temporary checkpoint failure")
|
||
|
|
}
|
||
|
|
return r.MemoryRepository.FinishExtraction(ctx, scope, cursors, completed...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionRetryPreservesSiblingSessionsAndFailedWrites(t *testing.T) {
|
||
|
|
for _, failure := range []string{"model", "write", "checkpoint"} {
|
||
|
|
t.Run(failure, func(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
messages.set("a", []*types.Message{userMessage("a", "采用中文回答", base)})
|
||
|
|
messages.set("b", []*types.Message{userMessage("b", "采用简短回答", base.Add(time.Second))})
|
||
|
|
models.response = `{"memories":[{"action":"add","kind":"preference","topic":"语言","content":"采用中文回答"}]}`
|
||
|
|
repo := &extractionFailureRepo{MemoryRepository: svc.repo}
|
||
|
|
svc.repo = repo
|
||
|
|
switch failure {
|
||
|
|
case "model":
|
||
|
|
models.failNext = true
|
||
|
|
case "write":
|
||
|
|
repo.createFailures = 1
|
||
|
|
case "checkpoint":
|
||
|
|
repo.finishFailures = 1
|
||
|
|
}
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
svc.ScheduleExtraction(ctx, "b", "b1", "model")
|
||
|
|
task := queue.pop()
|
||
|
|
require.Error(t, svc.Handle(context.Background(), task))
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
subject, err := svc.repo.GetSubject(ctx, scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.ElementsMatch(t, []string{"a", "b"}, subject.PendingSessions)
|
||
|
|
require.Nil(t, subject.ExtractCursor)
|
||
|
|
require.NoError(t, svc.repo.DeferExtraction(ctx, scope, -time.Second))
|
||
|
|
require.NoError(t, svc.RecoverExtractions(ctx))
|
||
|
|
require.NoError(t, svc.Handle(context.Background(), queue.pop()))
|
||
|
|
require.Contains(t, models.seenTranscripts(), "采用简短回答")
|
||
|
|
items, count, err := svc.ListItems(ctx, types.MemoryStatusActive, 10, 0)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, int64(1), count)
|
||
|
|
require.Equal(t, "采用中文回答", items[0].Content)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionCommonCursorDoesNotSkipOverlappingSessions(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
for i := 0; i < 5; i++ {
|
||
|
|
id := fmt.Sprintf("session-%d", i)
|
||
|
|
rows := []*types.Message{userMessage(id, fmt.Sprintf("begin-%d", i), base.Add(time.Duration(i)*time.Second))}
|
||
|
|
if i == 0 {
|
||
|
|
rows = append(rows, userMessage(id, "late-first-session", base.Add(time.Minute)))
|
||
|
|
}
|
||
|
|
messages.set(id, rows)
|
||
|
|
svc.ScheduleExtraction(ctx, id, rows[0].ID, "model")
|
||
|
|
}
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
for i := 0; i < 5; i++ {
|
||
|
|
require.Contains(t, models.seenTranscripts(), fmt.Sprintf("begin-%d", i))
|
||
|
|
}
|
||
|
|
require.Contains(t, models.seenTranscripts(), "late-first-session")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionLateRegisteredSessionKeepsItsOwnCheckpoint(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
base := time.Now().Add(-time.Hour).Truncate(time.Microsecond)
|
||
|
|
messages.set("fast", []*types.Message{userMessage("fast", "fast-completed", base.Add(time.Minute))})
|
||
|
|
svc.ScheduleExtraction(ctx, "fast", "fast-answer", "model")
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
|
||
|
|
// This user turn predates the other session's completed extraction.
|
||
|
|
messages.set("slow", []*types.Message{
|
||
|
|
userMessage("slow", "old-history-not-replayed", base.Add(-time.Hour)),
|
||
|
|
userMessage("slow", "late-completed-turn", base),
|
||
|
|
})
|
||
|
|
messages.turnStarts = map[string]time.Time{"slow-answer": base}
|
||
|
|
svc.ScheduleExtraction(ctx, "slow", "slow-answer", "model")
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
var extracted string
|
||
|
|
for _, prompt := range models.prompts {
|
||
|
|
extracted += transcriptBlock(prompt)
|
||
|
|
}
|
||
|
|
require.Contains(t, extracted, "late-completed-turn")
|
||
|
|
require.NotContains(t, extracted, "old-history-not-replayed")
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
subject, err := svc.repo.GetSubject(ctx, scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, subject.SessionCursors["slow"].Equal(base))
|
||
|
|
require.True(t, subject.ExtractCursor.Equal(base.Add(time.Minute)))
|
||
|
|
|
||
|
|
// Further slow-session work must not inherit the fast session's watermark.
|
||
|
|
messages.set("slow", []*types.Message{userMessage("slow", "slow-next-turn", base.Add(time.Second))})
|
||
|
|
svc.ScheduleExtraction(ctx, "slow", "next-answer", "model")
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
require.Contains(t, models.seenTranscripts(), "slow-next-turn")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionRejectsForeignAndUnownedSessions(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
messages.owners = map[string]interfaces.MemoryScope{
|
||
|
|
"mine": scope,
|
||
|
|
"other-user": {TenantID: 1, SubjectID: "web_user:bob"},
|
||
|
|
"other-workspace": {TenantID: 2, SubjectID: scope.SubjectID},
|
||
|
|
}
|
||
|
|
for _, id := range []string{"mine", "other-user", "other-workspace", "unowned"} {
|
||
|
|
messages.set(id, []*types.Message{userMessage(id, id+"-private-text", time.Now().Add(-time.Minute))})
|
||
|
|
svc.ScheduleExtraction(ctx, id, "answer", "model")
|
||
|
|
}
|
||
|
|
// A stale or forged queued payload must be checked again by the worker.
|
||
|
|
_, _, err = svc.repo.EnqueuePendingSession(ctx, scope, "other-user", time.Minute)
|
||
|
|
require.NoError(t, err)
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
extracted := models.seenTranscripts()
|
||
|
|
require.Contains(t, extracted, "mine-private-text")
|
||
|
|
for _, id := range []string{"other-user", "other-workspace", "unowned"} {
|
||
|
|
require.NotContains(t, extracted, id+"-private-text")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionMessageCapRetainsSiblingBacklogs(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
for _, id := range []string{"a", "b"} {
|
||
|
|
var rows []*types.Message
|
||
|
|
for i := 0; i < 95; i++ {
|
||
|
|
rows = append(rows, userMessage(id, fmt.Sprintf("%s-message-%03d", id, i), base.Add(time.Duration(i)*time.Second)))
|
||
|
|
}
|
||
|
|
messages.set(id, rows)
|
||
|
|
svc.ScheduleExtraction(ctx, id, rows[0].ID, "model")
|
||
|
|
}
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
var extracted string
|
||
|
|
for _, prompt := range models.prompts {
|
||
|
|
extracted += transcriptBlock(prompt)
|
||
|
|
}
|
||
|
|
for _, id := range []string{"a", "b"} {
|
||
|
|
for i := 0; i < 95; i++ {
|
||
|
|
require.Equal(t, 1, strings.Count(extracted, fmt.Sprintf("%s-message-%03d", id, i)))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionTimestampBoundaryIncludesAllTiedMessages(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
var rows []*types.Message
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
for i := 0; i < 50; i++ {
|
||
|
|
rows = append(rows, userMessage("a", fmt.Sprintf("tied-%02d", i), base))
|
||
|
|
}
|
||
|
|
messages.set("a", rows)
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
for _, row := range rows {
|
||
|
|
require.Contains(t, models.seenTranscripts(), row.Content)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionCompletionPreservesNewArrivalInSameSession(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
rows := []*types.Message{userMessage("a", "before-run", base)}
|
||
|
|
messages.set("a", rows)
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
svc.repo = &extractionFailureRepo{MemoryRepository: svc.repo, beforeFinish: func() {
|
||
|
|
messages.set("a", append(rows, userMessage("a", "during-run", base.Add(time.Second))))
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a2", "model")
|
||
|
|
}}
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
require.Contains(t, models.seenTranscripts(), "during-run")
|
||
|
|
require.Equal(t, 2, models.calls)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionClaimIsDurableAndScoped(t *testing.T) {
|
||
|
|
svc, tenants, _, _, _ := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
for i := 0; i < 2; i++ {
|
||
|
|
pending, _, err := svc.repo.ClaimPendingSessions(ctx, scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, []string{"a"}, pending)
|
||
|
|
}
|
||
|
|
for _, other := range []interfaces.MemoryScope{{TenantID: 1, SubjectID: "web_user:bob"}, {TenantID: 2, SubjectID: scope.SubjectID}} {
|
||
|
|
_, err := svc.repo.EnsureSubject(ctx, other)
|
||
|
|
require.NoError(t, err)
|
||
|
|
pending, _, err := svc.repo.ClaimPendingSessions(ctx, other)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Empty(t, pending)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionInvalidOutputLeavesSiblingSessionsRetryable(t *testing.T) {
|
||
|
|
for _, output := range []string{"nil", "malformed", "missing"} {
|
||
|
|
t.Run(output, func(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
messages.set("a", []*types.Message{userMessage("a", "first", base)})
|
||
|
|
messages.set("b", []*types.Message{userMessage("b", "second", base.Add(time.Second))})
|
||
|
|
switch output {
|
||
|
|
case "nil":
|
||
|
|
models.nilResponses = 2
|
||
|
|
case "malformed":
|
||
|
|
models.response = "invalid JSON"
|
||
|
|
case "missing":
|
||
|
|
models.response = `{}`
|
||
|
|
}
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
svc.ScheduleExtraction(ctx, "b", "b1", "model")
|
||
|
|
task := queue.pop()
|
||
|
|
require.Error(t, svc.Handle(context.Background(), task))
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
subject, err := svc.repo.GetSubject(ctx, scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Nil(t, subject.ExtractCursor)
|
||
|
|
require.ElementsMatch(t, []string{"a", "b"}, subject.PendingSessions)
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
require.NoError(t, svc.repo.DeferExtraction(ctx, scope, -time.Second))
|
||
|
|
require.NoError(t, svc.RecoverExtractions(ctx))
|
||
|
|
require.NoError(t, svc.Handle(context.Background(), queue.pop()))
|
||
|
|
require.Contains(t, models.seenTranscripts(), "second")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestExtractionFollowUpEnqueueFailureReturnsErrorWithBacklog(t *testing.T) {
|
||
|
|
svc, tenants, messages, models, queue := newExtractionHarness(t)
|
||
|
|
ctx := enabledCtx(t, tenants, 1, "alice")
|
||
|
|
models.response = `{"memories":[]}`
|
||
|
|
base := time.Now().Add(-time.Hour)
|
||
|
|
var rows []*types.Message
|
||
|
|
for i := 0; i < 65; i++ {
|
||
|
|
rows = append(rows, userMessage("a", fmt.Sprintf("backlog-%02d", i), base.Add(time.Duration(i)*time.Second)))
|
||
|
|
}
|
||
|
|
messages.set("a", rows)
|
||
|
|
svc.ScheduleExtraction(ctx, "a", "a1", "model")
|
||
|
|
task := queue.pop()
|
||
|
|
queue.failNext = true
|
||
|
|
require.Error(t, svc.Handle(context.Background(), task))
|
||
|
|
scope, err := ResolveScope(ctx)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, svc.repo.DeferExtraction(ctx, scope, -time.Second))
|
||
|
|
require.NoError(t, svc.RecoverExtractions(ctx))
|
||
|
|
drainExtractions(t, svc, queue)
|
||
|
|
require.Contains(t, models.seenTranscripts(), "backlog-64")
|
||
|
|
}
|