xinmi/internal/application/service/memory/recovery_test.go

package memory

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func recoveryHarness(t *testing.T) (*Service, *gorm.DB, context.Context, *stubMessageRepo, *stubModelService, *stubEnqueuer) {
	t.Helper()
	svc, db, tenants := newMemoryHarness(t)
	messages, models, queue := &stubMessageRepo{}, &stubModelService{response: `{"memories":[]}`}, &stubEnqueuer{}
	svc.messageRepo, svc.modelService, svc.enqueuer = messages, models, queue
	ctx := enabledCtx(t, tenants, 1, "alice")
	messages.set("s", []*types.Message{userMessage("s", "recover this input", time.Now().Add(-time.Hour))})
	return svc, db, ctx, messages, models, queue
}

func retryDue(t *testing.T, db *gorm.DB) {
	t.Helper()
	require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).
		Update("extract_next_at", time.Now().Add(-time.Second)).Error)
}

func TestRecoveryAfterRestartNeedsNoNewChatAndRejectsDuplicateTasks(t *testing.T) {
	svc, db, ctx, _, models, queue := recoveryHarness(t)
	ctx = context.WithValue(ctx, types.LanguageContextKey, "Chinese")
	svc.ScheduleExtraction(ctx, "s", "m", "original-chat-model")
	oldTask := queue.pop()
	models.failNext = true
	require.Error(t, svc.Handle(context.Background(), oldTask))
	scope, err := ResolveScope(ctx)
	require.NoError(t, err)
	failed, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.Equal(t, 1, failed.ExtractFailures)
	require.True(t, failed.ExtractNextAt.After(time.Now()))
	require.Nil(t, failed.ExtractCursor)
	require.Equal(t, types.MemoryPendingSessions{"s"}, failed.PendingSessions)

	// A replacement process has no in-memory retries or queued triggers.
	restarted := &Service{repo: svc.repo, tenantRepo: svc.tenantRepo, messageRepo: svc.messageRepo, modelService: models, enqueuer: queue}
	require.NoError(t, restarted.RecoverExtractions(context.Background()))
	require.Nil(t, queue.pop(), "database deadline survives restart")
	require.NoError(t, restarted.Handle(context.Background(), oldTask))
	retryDue(t, db)
	require.NoError(t, restarted.RecoverExtractions(context.Background()))
	fresh := queue.pop()
	require.NotNil(t, fresh)
	var payload types.MemoryExtractPayload
	require.NoError(t, json.Unmarshal(fresh.Payload(), &payload))
	require.Equal(t, "original-chat-model", payload.ChatModelID)
	require.Equal(t, "Chinese", payload.Language)
	require.NoError(t, restarted.Handle(context.Background(), fresh))
	calls := models.calls
	require.NoError(t, restarted.Handle(context.Background(), fresh))
	require.NoError(t, restarted.Handle(context.Background(), oldTask))
	require.Equal(t, calls, models.calls, "duplicate/stale queue deliveries must not spend again")
	completed, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.Empty(t, completed.PendingSessions)
	require.NotNil(t, completed.ExtractCursor)
	require.Equal(t, 0, completed.ExtractFailures)
}

func TestNewChatDoesNotResetDurableBackoff(t *testing.T) {
	svc, _, ctx, _, models, queue := recoveryHarness(t)
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	models.failNext = true
	require.Error(t, svc.Handle(context.Background(), queue.pop()))
	scope, err := ResolveScope(ctx)
	require.NoError(t, err)
	before, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	svc.ScheduleExtraction(ctx, "s", "new-message", "model")
	after, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.True(t, before.ExtractNextAt.Equal(*after.ExtractNextAt))
	require.Equal(t, before.ExtractFailures, after.ExtractFailures)
	require.Nil(t, after.ExtractScheduledAt)
	require.Len(t, after.PendingSessions, 2)
	require.Nil(t, queue.pop())
}

func TestRecoveryHonorsActiveLeaseAndFencesExpiredWorker(t *testing.T) {
	svc, db, ctx, _, models, queue := recoveryHarness(t)
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	oldTask := queue.pop()
	var old types.MemoryExtractPayload
	require.NoError(t, json.Unmarshal(oldTask.Payload(), &old))
	scope, err := ResolveScope(ctx)
	require.NoError(t, err)
	started, err := svc.repo.StartExtraction(ctx, scope, old.AttemptID, extractInFlightGrace)
	require.NoError(t, err)
	require.True(t, started)
	retryDue(t, db)
	require.NoError(t, svc.RecoverExtractions(ctx))
	require.Nil(t, queue.pop())
	require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).
		Update("extract_lease_until", time.Now().Add(-time.Minute)).Error)
	require.NoError(t, svc.RecoverExtractions(ctx))
	require.Nil(t, queue.pop(), "crashed attempts receive durable backoff")
	retryDue(t, db)
	require.NoError(t, svc.RecoverExtractions(ctx))
	fresh := queue.pop()
	require.NotNil(t, fresh)
	require.Error(t, svc.repo.FinishExtraction(types.WithMemoryExtractionAttempt(ctx, old.AttemptID), scope, types.MemorySessionCursors{"s": time.Now()}, "s"))
	require.NoError(t, svc.repo.FailExtraction(ctx, scope, old.AttemptID))
	require.NoError(t, svc.Handle(context.Background(), oldTask))
	require.Zero(t, models.calls)
	require.NoError(t, svc.Handle(context.Background(), fresh))
	require.NotZero(t, models.calls)
}

func TestRecoveryRecordsQueueOutageAndRespectsWorkspaceAndSubjectSwitches(t *testing.T) {
	svc, db, ctx, _, models, queue := recoveryHarness(t)
	svc.enqueuer = nil
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	scope, err := ResolveScope(ctx)
	require.NoError(t, err)
	subject, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.Len(t, subject.PendingSessions, 1)
	require.Equal(t, 1, subject.ExtractFailures)
	svc.enqueuer = queue
	tenant := svc.tenantRepo.(*stubTenantRepo)
	tenant.set(1, &types.MemoryConfig{Enabled: false})
	retryDue(t, db)
	require.NoError(t, svc.RecoverExtractions(ctx))
	require.Nil(t, queue.pop())
	tenant.set(1, &types.MemoryConfig{Enabled: true, WriteMode: types.MemoryWriteAuto})
	require.NoError(t, svc.repo.UpdateSubjectEnabled(ctx, scope, false))
	retryDue(t, db)
	require.NoError(t, svc.RecoverExtractions(ctx))
	require.Nil(t, queue.pop())
	require.NoError(t, svc.repo.UpdateSubjectEnabled(ctx, scope, true))
	require.NoError(t, svc.RecoverExtractions(ctx))
	fresh := queue.pop()
	require.NotNil(t, fresh)
	tenant.set(1, &types.MemoryConfig{Enabled: false})
	require.NoError(t, svc.Handle(context.Background(), fresh))
	require.Zero(t, models.calls, "a queued task rechecks current workspace configuration")
	paused, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.Equal(t, subject.ExtractFailures, paused.ExtractFailures, "ordinary disabled deferrals do not consume attempts")
}

func TestExhaustedRecoveryBudgetSurvivesRestartNewChatAndToggle(t *testing.T) {
	svc, db, ctx, _, models, queue := recoveryHarness(t)
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	for attempt := 1; attempt <= types.MaxMemoryExtractionFailures; attempt++ {
		task := queue.pop()
		require.NotNil(t, task)
		models.failNext = true
		require.Error(t, svc.Handle(context.Background(), task))
		retryDue(t, db)
		require.NoError(t, svc.RecoverExtractions(ctx))
	}
	require.Nil(t, queue.pop())
	settings, err := svc.GetSettings(ctx)
	require.NoError(t, err)
	require.True(t, settings.ExtractionExhausted)
	require.True(t, settings.ExtractionPending)
	require.Equal(t, types.MaxMemoryExtractionFailures, settings.ExtractionFailures)
	svc.ScheduleExtraction(ctx, "s", "new-turn", "model")
	require.NoError(t, svc.SetEnabled(ctx, false))
	require.NoError(t, svc.SetEnabled(ctx, true))
	restarted := &Service{repo: svc.repo, tenantRepo: svc.tenantRepo, messageRepo: svc.messageRepo, modelService: models, enqueuer: queue}
	require.NoError(t, restarted.RecoverExtractions(context.Background()))
	require.Nil(t, queue.pop())
	require.Equal(t, types.MaxMemoryExtractionFailures, models.callCount())
	otherCtx := enabledCtx(t, svc.tenantRepo.(*stubTenantRepo), 1, "bob")
	require.NoError(t, svc.RetryExtraction(otherCtx))
	require.Nil(t, queue.pop(), "one subject cannot renew another subject's budget")
	require.NoError(t, svc.RetryExtraction(ctx))
	fresh := queue.pop()
	require.NotNil(t, fresh)
	require.NoError(t, svc.RetryExtraction(ctx))
	require.Nil(t, queue.pop(), "manual retries cannot duplicate an active reservation")
	require.NoError(t, svc.Handle(context.Background(), fresh))
	settings, err = svc.GetSettings(ctx)
	require.NoError(t, err)
	require.False(t, settings.ExtractionExhausted)
	require.False(t, settings.ExtractionPending)
	require.Zero(t, settings.ExtractionFailures)
}

func TestRecoveryExcludesDisallowedAgentTurnsAndPriorContext(t *testing.T) {
	svc, db, ctx, messages, models, queue := recoveryHarness(t)
	base := time.Now().Add(-time.Hour)
	secret := userMessage("s", "disabled Agent private preference", base)
	visible := userMessage("s", "allowed Agent preference", base.Add(time.Second))
	messages.set("s", []*types.Message{secret, visible})
	messages.deniedMessages = map[string]bool{secret.ID: true}
	svc.enqueuer = nil
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	svc.enqueuer = queue
	retryDue(t, db)
	require.NoError(t, svc.RecoverExtractions(ctx))
	require.NoError(t, svc.Handle(context.Background(), queue.pop()))
	require.NotContains(t, models.lastPrompt, secret.Content)
	require.Contains(t, models.lastPrompt, visible.Content)
}

func TestRecoveryDoesNotRecreateDeletedSubjectOrTrustLegacyPayload(t *testing.T) {
	svc, db, ctx, _, models, queue := recoveryHarness(t)
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	task := queue.pop()
	require.NoError(t, db.Where("tenant_id = ?", 1).Delete(&types.MemorySubject{}).Error)
	require.NoError(t, svc.Handle(context.Background(), task))
	var count int64
	require.NoError(t, db.Model(&types.MemorySubject{}).Count(&count).Error)
	require.Zero(t, count)
	require.Zero(t, models.calls)
}