xinmi/internal/application/service/wiki_recovery_budget_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"

	apprepo "github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/hibiken/asynq"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func seedServiceRecovery(t *testing.T) (*wikiIngestService, types.WikiRecoveryBudget, WikiIngestPayload) {
	t.Helper()
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "recovery.sqlite"))
	require.NoError(t, db.AutoMigrate(&types.WikiRecoveryBudget{}, &types.TaskPendingOp{}, &types.Chunk{}))
	for _, ddl := range []string{
		`CREATE TABLE knowledge_bases (id text, tenant_id bigint, summary_model_id text, embedding_model_id text, wiki_config text, indexing_strategy text, deleted_at datetime)`,
		`CREATE TABLE knowledges (id text, tenant_id bigint, knowledge_base_id text, title text, deleted_at datetime)`,
		`CREATE TABLE models (id text, tenant_id bigint, name text, type text, source text, status text, parameters text, deleted_at datetime)`,
		`CREATE TABLE wiki_dead_letter_recoveries (tenant_id bigint, dead_letter_id bigint, attempt_number integer, pending_op_id bigint, knowledge_base_id text, knowledge_id text, model_id text, deployed_fix_review text)`,
		`INSERT INTO knowledge_bases VALUES ('kb',7,'model','','{"synthesis_model_id":"model"}','{"wiki_enabled":true}',NULL)`,
		`INSERT INTO knowledges VALUES ('doc',7,'kb','title',NULL)`,
		`INSERT INTO models VALUES ('model',7,'model','KnowledgeQA','remote','active','{}',NULL)`,
	} {
		require.NoError(t, db.Exec(ddl).Error)
	}
	chunks := checkpointChunks(1)
	require.NoError(t, db.Create(chunks).Error)
	row := types.WikiRecoveryBudget{TenantID: 7, DeadLetterID: 10, AttemptNumber: 5, PendingOpID: 42, KnowledgeBaseID: "kb", KnowledgeID: "doc", ModelID: "model",
		ReviewedRelease: "fixture", ReviewedAPIImage: "sha256:fixture", MaxProviderRequests: 3, DeadlineAt: time.Now().UTC().Add(time.Hour).Truncate(time.Microsecond), Status: "active", PausedPayload: "[]", VerifiedProgress: "{}"}
	var err error
	row.SourceFingerprint, row.ModelFingerprint, err = apprepo.WikiRecoveryFingerprints(db, row)
	require.NoError(t, err)
	require.NoError(t, db.Create(&row).Error)
	review, err := json.Marshal(map[string]any{"tenant_id": row.TenantID, "dead_letter_id": row.DeadLetterID, "max_attempt": row.AttemptNumber, "budget_version": 1,
		"knowledge_base_id": row.KnowledgeBaseID, "knowledge_id": row.KnowledgeID, "model_id": row.ModelID, "embedding_model_id": row.EmbeddingModelID,
		"source_fingerprint": row.SourceFingerprint, "model_fingerprint": row.ModelFingerprint, "fix_release": row.ReviewedRelease, "api_image_id": row.ReviewedAPIImage,
		"max_provider_requests": row.MaxProviderRequests, "deadline_at": row.DeadlineAt})
	require.NoError(t, err)
	require.NoError(t, db.Exec(`INSERT INTO wiki_dead_letter_recoveries VALUES (?,?,?,?,?,?,?,?)`, row.TenantID, row.DeadLetterID, row.AttemptNumber, row.PendingOpID, row.KnowledgeBaseID, row.KnowledgeID, row.ModelID, string(review)).Error)
	body, err := json.Marshal(WikiPendingOp{WikiRecoveryIdentity: row.Identity(), KnowledgeID: "doc", DocTitle: "title", Op: WikiOpIngest, Language: "en-US"})
	require.NoError(t, err)
	require.NoError(t, db.Create(&types.TaskPendingOp{ID: row.PendingOpID, TenantID: row.TenantID, TaskType: wikiTaskType, Scope: wikiTaskScope, ScopeID: row.KnowledgeBaseID, DedupKey: row.KnowledgeID, Op: WikiOpIngest, Payload: body}).Error)
	t.Setenv("WEKNORA_REVIEWED_RELEASE", row.ReviewedRelease)
	t.Setenv("WEKNORA_REVIEWED_API_IMAGE", row.ReviewedAPIImage)
	svc := &wikiIngestService{checkpointDB: db, pendingRepo: apprepo.NewTaskPendingOpsRepository(db), task: &wikiGuardTaskQueue{}}
	return svc, row, WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
}

func TestWikiRecoveryBudgetServiceHandoffResumesSameBudget(t *testing.T) {
	svc, row, payload := seedServiceRecovery(t)
	ctx, cancel, err := svc.bindWikiRecovery(context.Background(), payload, row.Identity(), row.KnowledgeID, row.ModelID)
	require.NoError(t, err)
	require.NoError(t, requestbudget.FromContext(ctx).Reserve(ctx, "model", "candidate"))
	require.NoError(t, svc.enqueueFinalize(ctx, payload, []string{"entity/a"}, map[string]string{"entity/a": "A"}, []wikiFinalizeChange{{Action: wikiFinalizeAdded, DocTitle: "title"}}, nil))
	cancel()
	var rows []*types.TaskPendingOp
	require.NoError(t, svc.checkpointDB.Order("id").Find(&rows).Error)
	require.Len(t, rows, 2)
	for _, op := range rows {
		require.Equal(t, wikiFinalizeTaskType, op.TaskType)
		var identity types.WikiRecoveryIdentity
		require.NoError(t, json.Unmarshal(op.Payload, &identity))
		require.Equal(t, row.Identity(), identity)
	}
	ctx, cancel, err = svc.bindWikiRecovery(context.Background(), payload, row.Identity(), "", row.ModelID)
	require.NoError(t, err)
	defer cancel()
	deadline, ok := ctx.Deadline()
	require.True(t, ok)
	require.Equal(t, row.DeadlineAt, deadline)
	guard := wikiRecoveryFromContext(ctx)
	require.Equal(t, int64(1), guard.row.ReservedProviderRequests)
	require.NoError(t, guard.Reserve(ctx, "model", "finalize"))
	completed, err := guard.repo.Complete(ctx, guard.row, []int64{rows[0].ID})
	require.NoError(t, err)
	require.False(t, completed)
	completed, err = guard.repo.Complete(ctx, guard.row, []int64{rows[1].ID})
	require.NoError(t, err)
	require.True(t, completed)
	saved, err := guard.repo.Load(ctx, 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, "completed", saved.Status)
	require.Equal(t, int64(2), saved.ReservedProviderRequests)
}

func TestWikiRecoveryBudgetHandoffFailureKeepsIngestAndRejectsOrdinaryOverlap(t *testing.T) {
	svc, row, payload := seedServiceRecovery(t)
	ctx, cancel, err := svc.bindWikiRecovery(context.Background(), payload, row.Identity(), row.KnowledgeID, row.ModelID)
	require.NoError(t, err)
	defer cancel()
	require.ErrorIs(t, svc.rejectOrdinaryRecoveryOverlap(ctx, payload, []WikiPendingOp{{KnowledgeID: "doc"}}), ErrWikiRecoveryBudgetStopped)
	require.NoError(t, svc.rejectOrdinaryRecoveryOverlap(ctx, payload, []WikiPendingOp{{KnowledgeID: "other"}}))
	require.NoError(t, svc.checkpointDB.Callback().Create().Before("gorm:create").Register("reject_finalize", func(tx *gorm.DB) {
		if tx.Statement.Table == "task_pending_ops" {
			tx.AddError(errors.New("fixture insert failure"))
		}
	}))
	require.Error(t, svc.enqueueFinalize(ctx, payload, []string{"entity/a"}, nil, []wikiFinalizeChange{{Action: wikiFinalizeAdded}}, nil))
	var remaining []types.TaskPendingOp
	require.NoError(t, svc.checkpointDB.Find(&remaining).Error)
	require.Len(t, remaining, 1)
	require.Equal(t, int64(42), remaining[0].ID)
}

func TestWikiRecoveryBudgetStopArchivesWithoutFailureCountOrCompletion(t *testing.T) {
	svc, row, payload := seedServiceRecovery(t)
	kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: "model", WikiConfig: &types.WikiConfig{}}
	kb.IndexingStrategy.WikiEnabled = true
	svc.kbService = &checkpointKBService{kb: kb}
	svc.modelService = &checkpointModelService{model: &types.Model{ID: "model", TenantID: 7}, chat: &wikiScriptedChat{id: "model"}}
	finalization := &checkpointFinalization{}
	svc.knowledgeRepo = finalization
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), &wikiCheckpointScope{7, "kb", "doc", "prior"}, "candidates", "complete", combinedExtraction{}))
	t.Setenv("WEKNORA_REVIEWED_RELEASE", "different-build")
	body, err := json.Marshal(payload)
	require.NoError(t, err)
	require.ErrorIs(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(wikiTaskType, body)), ErrWikiRecoveryBudgetStopped)
	repo := apprepo.WikiRecoveryBudgetRepository{DB: svc.checkpointDB}
	saved, err := repo.Load(context.Background(), 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, "stopped_budget", saved.Status)
	var paused []types.TaskPendingOp
	require.NoError(t, json.Unmarshal([]byte(saved.PausedPayload), &paused))
	require.Len(t, paused, 1)
	require.Zero(t, paused[0].FailCount)
	require.Zero(t, finalization.calls)
	var count int64
	require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Count(&count).Error)
	require.Equal(t, int64(1), count)
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(wikiTaskType, body)))
	require.Zero(t, finalization.calls)
}

func TestWikiRecoveryPhaseFailureKeepsCauseAndArchivesExactOperation(t *testing.T) {
	for _, test := range []struct {
		name string
		err  error
		want string
	}{
		{"incomplete", validateWikiResponse(&types.ChatResponse{Content: "partial", FinishReason: "length"}), "model_output_incomplete"},
		{"invalid", &wikiOutputError{reason: "invalid extraction JSON"}, "model_output_invalid"},
		{"timeout", context.DeadlineExceeded, "phase_timeout"},
		{"cancelled", context.Canceled, "phase_cancelled"},
		{"provider", errors.New("private provider response and credentials"), "phase_failed"},
		{"guard", ErrWikiRecoveryBudgetStopped, "guard_rejected"},
	} {
		t.Run(test.name, func(t *testing.T) {
			svc, row, payload := seedServiceRecovery(t)
			ctx, cancel, err := svc.bindWikiRecovery(context.Background(), payload, row.Identity(), row.KnowledgeID, row.ModelID)
			require.NoError(t, err)
			defer cancel()
			require.NoError(t, requestbudget.FromContext(ctx).Reserve(ctx, "model", "wiki_candidate_slug"))
			require.NoError(t, svc.checkpointDB.Create(&types.TaskPendingOp{ID: 43, TenantID: 7, TaskType: wikiTaskType, Scope: wikiTaskScope, ScopeID: "other", DedupKey: "other", Op: WikiOpIngest, Payload: []byte(`{"knowledge_id":"other"}`)}).Error)
			require.NoError(t, svc.saveWikiCheckpoint(ctx, &wikiCheckpointScope{7, "kb", "doc", "prior"}, "candidates", "complete", combinedExtraction{}))
			phaseErr := fmt.Errorf("candidate window 2/7: %w", test.err)
			result := svc.settleWikiRecoveryStop(ctx, row.Identity(), row.TenantID, phaseErr)
			require.ErrorIs(t, result, ErrWikiRecoveryBudgetStopped)
			require.ErrorIs(t, result, test.err)
			repo := apprepo.WikiRecoveryBudgetRepository{DB: svc.checkpointDB}
			saved, err := repo.Load(context.Background(), 7, row.Identity())
			require.NoError(t, err)
			require.Equal(t, "stopped_budget", saved.Status)
			require.Equal(t, test.want, saved.StopReason)
			require.Equal(t, "wiki_candidate_slug", saved.LastPhase)
			require.Equal(t, int64(1), saved.ReservedProviderRequests)
			require.Equal(t, row.MaxProviderRequests, saved.MaxProviderRequests)
			require.Equal(t, row.DeadlineAt, saved.DeadlineAt)
			var paused []types.TaskPendingOp
			require.NoError(t, json.Unmarshal([]byte(saved.PausedPayload), &paused))
			require.Len(t, paused, 1)
			require.Equal(t, row.PendingOpID, paused[0].ID)
			require.Zero(t, paused[0].FailCount)
			var remaining []types.TaskPendingOp
			require.NoError(t, svc.checkpointDB.Find(&remaining).Error)
			require.Len(t, remaining, 1)
			require.Equal(t, int64(43), remaining[0].ID)
			var count int64
			require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Count(&count).Error)
			require.Equal(t, int64(1), count)
			// A later sweep must preserve the original reason and archive.
			require.NoError(t, repo.StopAndArchive(context.Background(), *saved, "guard_rejected", "preflight"))
			after, err := repo.Load(context.Background(), 7, row.Identity())
			require.NoError(t, err)
			require.Equal(t, saved.StopReason, after.StopReason)
			require.Equal(t, saved.PausedPayload, after.PausedPayload)
			require.ErrorIs(t, repo.Reserve(context.Background(), *after, "model", "candidate"), ErrWikiRecoveryBudgetStopped)
		})
	}
}

func TestWikiRecoveryFailureReasonPrefersGuardCancellation(t *testing.T) {
	ctx, cancel := context.WithCancelCause(context.Background())
	cancel(fmt.Errorf("%w: deadline", ErrWikiRecoveryBudgetStopped))
	require.Equal(t, "guard_rejected", wikiRecoveryFailureReason(ctx, context.Canceled))
}

func TestWikiRecoveryIngestPreservesIncompleteMapAndReduceErrors(t *testing.T) {
	for _, stage := range []string{"map", "reduce"} {
		t.Run(stage, func(t *testing.T) {
			svc, row, payload := seedServiceRecovery(t)
			chunks := checkpointChunks(1)
			kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: "model", WikiConfig: &types.WikiConfig{}}
			kb.IndexingStrategy.WikiEnabled = true
			var calls atomic.Int32
			model := &wikiScriptedChat{id: "model", call: func(ctx context.Context, _ []chat.Message) (*types.ChatResponse, error) {
				if err := requestbudget.FromContext(ctx).Reserve(ctx, "model", stage); err != nil {
					return nil, err
				}
				if calls.Add(1) == 1 && stage == "reduce" {
					return &types.ChatResponse{Content: `{}`, FinishReason: "stop"}, nil
				}
				return &types.ChatResponse{FinishReason: "length"}, nil
			}}
			svc.kbService = &checkpointKBService{kb: kb}
			svc.modelService = &checkpointModelService{model: &types.Model{ID: "model", TenantID: 7}, chat: model}
			svc.knowledgeSvc = &checkpointKnowledge{}
			svc.chunkRepo = &checkpointChunksRepo{chunks: chunks}
			svc.wikiService = &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
			finalization := &checkpointFinalization{}
			svc.knowledgeRepo = finalization
			if stage == "reduce" {
				op := WikiPendingOp{WikiRecoveryIdentity: row.Identity(), Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: row.PendingOpID}
				scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
				require.NoError(t, err)
				require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", checkpointMap(2)))
			}
			body, err := json.Marshal(payload)
			require.NoError(t, err)
			err = svc.ProcessWikiIngest(context.Background(), asynq.NewTask(wikiTaskType, body))
			require.ErrorIs(t, err, ErrWikiRecoveryBudgetStopped)
			var outputErr *wikiOutputError
			require.ErrorAs(t, err, &outputErr)
			require.True(t, outputErr.incomplete)
			repo := apprepo.WikiRecoveryBudgetRepository{DB: svc.checkpointDB}
			saved, err := repo.Load(context.Background(), 7, row.Identity())
			require.NoError(t, err)
			require.Equal(t, "model_output_incomplete", saved.StopReason)
			require.Equal(t, stage, saved.LastPhase)
			require.Equal(t, int64(calls.Load()), saved.ReservedProviderRequests)
			require.Zero(t, finalization.calls)
			var paused []types.TaskPendingOp
			require.NoError(t, json.Unmarshal([]byte(saved.PausedPayload), &paused))
			require.Len(t, paused, 1)
			require.Equal(t, row.PendingOpID, paused[0].ID)
			require.Zero(t, paused[0].FailCount)
			priorCalls := calls.Load()
			require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(wikiTaskType, body)))
			require.Equal(t, priorCalls, calls.Load())
		})
	}
}

func TestWikiRecoveryBudgetDedupAndFinalizeNeverMixIdentity(t *testing.T) {
	identity := types.WikiRecoveryIdentity{DeadLetterID: 10, Attempt: 5, PendingOpID: 42, Version: 1}
	controlled, _ := json.Marshal(WikiPendingOp{WikiRecoveryIdentity: identity, KnowledgeID: "doc"})
	ordinary, _ := json.Marshal(WikiPendingOp{KnowledgeID: "doc"})
	rows := []*types.TaskPendingOp{{ID: 42, Payload: controlled}, {ID: 43, Payload: ordinary}}
	svc := &wikiIngestService{}
	ops, ids := svc.decodePendingRows(context.Background(), rows)
	require.Len(t, ops, 2)
	require.Len(t, ids, 2)
	require.True(t, ops[0].recoveryConflict)
	finalize, _ := json.Marshal(wikiFinalizeRow{WikiRecoveryIdentity: identity, Slug: "entity/a"})
	rows[0] = &types.TaskPendingOp{ID: 42, Op: wikiFinalizeOpSlug, DedupKey: "recovery:" + identity.Key() + ":entity/a", Payload: finalize}
	selected, got, err := selectWikiFinalizeGroup(rows)
	require.NoError(t, err)
	require.Equal(t, identity, got)
	require.Equal(t, rows[:1], selected)
	selected, got, err = selectWikiFinalizeGroup([]*types.TaskPendingOp{rows[1], rows[0]})
	require.NoError(t, err)
	require.Equal(t, identity, got)
	require.Equal(t, rows[:1], selected)
}

func TestWikiRecoveryBudgetFinalizeRejectsLostOrMalformedPayloadIdentity(t *testing.T) {
	identity := types.WikiRecoveryIdentity{DeadLetterID: 10, Attempt: 5, PendingOpID: 42, Version: 1}
	for _, body := range []string{`{"change":{"action":"added"}}`, `{"recovery_dead_letter_id":10,"recovery_attempt":5,"recovery_pending_op_id":42,"recovery_budget_version":1,"change":"bad"}`} {
		row := &types.TaskPendingOp{ID: 50, Op: wikiFinalizeOpChange, DedupKey: "recovery:" + identity.Key() + ":change:0", Payload: []byte(body)}
		selected, got, err := selectWikiFinalizeGroup([]*types.TaskPendingOp{row})
		require.Nil(t, selected)
		require.Equal(t, identity, got)
		require.ErrorIs(t, err, ErrWikiRecoveryBudgetStopped)
	}
}

func TestWikiRecoveryBudgetRejectsMismatchedTaskAndReviewAfterReopen(t *testing.T) {
	for _, change := range []string{"tenant", "kb", "doc", "model", "receipt", "limit"} {
		t.Run(change, func(t *testing.T) {
			svc, row, payload := seedServiceRecovery(t)
			knowledgeID, modelID := row.KnowledgeID, row.ModelID
			switch change {
			case "tenant":
				payload.TenantID++
			case "kb":
				payload.KnowledgeBaseID = "other"
			case "doc":
				knowledgeID = "other"
			case "model":
				modelID = "other"
			case "receipt":
				require.NoError(t, svc.checkpointDB.Exec("DELETE FROM wiki_dead_letter_recoveries").Error)
			case "limit":
				require.NoError(t, svc.checkpointDB.Exec("UPDATE wiki_recovery_budgets SET max_provider_requests = 999").Error)
			}
			_, cancel, err := svc.bindWikiRecovery(context.Background(), payload, row.Identity(), knowledgeID, modelID)
			cancel()
			require.ErrorIs(t, err, ErrWikiRecoveryBudgetStopped)
		})
	}
}

type recoveryIndexPages struct {
	interfaces.WikiPageService
	page   *types.WikiPage
	err    error
	writes int
}

func (p *recoveryIndexPages) GetIndex(context.Context, string) (*types.WikiPage, error) {
	return p.page, p.err
}
func (p *recoveryIndexPages) UpdatePage(_ context.Context, page *types.WikiPage) (*types.WikiPage, error) {
	p.writes++
	return page, nil
}

func TestWikiRecoveryBudgetFinalizeDoesNotSwallowProviderStop(t *testing.T) {
	svc, row, payload := seedServiceRecovery(t)
	ctx, cancel, err := svc.bindWikiRecovery(context.Background(), payload, row.Identity(), row.KnowledgeID, row.ModelID)
	require.NoError(t, err)
	defer cancel()
	pages := &recoveryIndexPages{page: &types.WikiPage{Content: "Previous grounded index"}}
	svc.wikiService = pages
	model := &wikiScriptedChat{id: "model", call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		return nil, requestbudget.ErrStopped
	}}
	require.ErrorIs(t, svc.rebuildIndexPage(ctx, model, payload, "new document", "English", ""), ErrWikiRecoveryBudgetStopped)
	require.Zero(t, pages.writes)
	require.ErrorIs(t, requestbudget.Check(ctx), ErrWikiRecoveryBudgetStopped)
	require.ErrorIs(t, context.Cause(ctx), ErrWikiRecoveryBudgetStopped)
}