xinmi/internal/application/service/memory/retry_boundary_test.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")
}