xinmi/internal/application/repository/wiki_recovery_budget_test.go

package repository

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os/exec"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func openRecoveryBudgetSQLite(t *testing.T, path string) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	return db
}

func seedRecoveryBudget(t *testing.T, db *gorm.DB, max int64) types.WikiRecoveryBudget {
	t.Helper()
	require.NoError(t, db.AutoMigrate(&types.WikiRecoveryBudget{}, &types.TaskPendingOp{}))
	for _, ddl := range []string{
		`CREATE TABLE IF NOT EXISTS knowledge_bases (id text PRIMARY KEY, tenant_id bigint, deleted_at timestamp)`,
		`CREATE TABLE IF NOT EXISTS knowledges (id text PRIMARY KEY, tenant_id bigint, knowledge_base_id text, deleted_at timestamp)`,
		`ALTER TABLE knowledge_bases ADD COLUMN summary_model_id text DEFAULT ''`,
		`ALTER TABLE knowledge_bases ADD COLUMN embedding_model_id text DEFAULT ''`,
		`ALTER TABLE knowledge_bases ADD COLUMN wiki_config text DEFAULT '{}'`,
		`ALTER TABLE knowledge_bases ADD COLUMN indexing_strategy text DEFAULT '{}'`,
		`ALTER TABLE knowledges ADD COLUMN title text DEFAULT ''`,
		`CREATE TABLE chunks (id text PRIMARY KEY, tenant_id bigint, knowledge_base_id text, knowledge_id text, content text, content_revision integer, chunk_index integer, start_at integer, end_at integer, chunk_type text, parent_chunk_id text, image_info text, is_enabled boolean, deleted_at timestamp)`,
		`CREATE TABLE models (id text PRIMARY KEY, tenant_id bigint, name text, type text, source text, status text, parameters text, deleted_at timestamp)`,
	} {
		require.NoError(t, db.Exec(ddl).Error)
	}
	require.NoError(t, db.Exec(`INSERT INTO knowledge_bases (id, tenant_id, summary_model_id, embedding_model_id, wiki_config, indexing_strategy) VALUES ('kb',7,'model','embed','{"synthesis_model_id":"model"}','{"wiki_enabled":true}')`).Error)
	require.NoError(t, db.Exec(`INSERT INTO knowledges (id,tenant_id,knowledge_base_id,title) VALUES ('doc',7,'kb','title')`).Error)
	require.NoError(t, db.Exec(`INSERT INTO chunks (id,tenant_id,knowledge_base_id,knowledge_id,content,content_revision,chunk_index,start_at,end_at,chunk_type,parent_chunk_id,image_info,is_enabled) VALUES ('chunk',7,'kb','doc','source',0,0,0,6,'text','','',?)`, true).Error)
	for _, id := range []string{"model", "embed"} {
		require.NoError(t, db.Exec(`INSERT INTO models (id,tenant_id,name,type,source,status,parameters) VALUES (?,7,?,'KnowledgeQA','remote','active','{}')`, id, id).Error)
	}
	if db.Dialector.Name() == "postgres" {
		for _, ddl := range []string{
			`ALTER TABLE knowledge_bases ALTER COLUMN wiki_config DROP DEFAULT`,
			`ALTER TABLE knowledge_bases ALTER COLUMN wiki_config TYPE jsonb USING wiki_config::jsonb`,
			`ALTER TABLE knowledge_bases ALTER COLUMN indexing_strategy DROP DEFAULT`,
			`ALTER TABLE knowledge_bases ALTER COLUMN indexing_strategy TYPE jsonb USING indexing_strategy::jsonb`,
			`ALTER TABLE models ALTER COLUMN parameters TYPE jsonb USING parameters::jsonb`,
		} {
			require.NoError(t, db.Exec(ddl).Error)
		}
	}
	row := types.WikiRecoveryBudget{TenantID: 7, DeadLetterID: 10, AttemptNumber: 5, PendingOpID: 42, KnowledgeBaseID: "kb", KnowledgeID: "doc", ModelID: "model", EmbeddingModelID: "embed", ReviewedRelease: "fixture", ReviewedAPIImage: "sha256:fixture", MaxProviderRequests: max, DeadlineAt: time.Now().UTC().Add(time.Hour).Truncate(time.Microsecond), Status: "active", PausedPayload: "[]", VerifiedProgress: "[]"}
	var err error
	row.SourceFingerprint, row.ModelFingerprint, err = WikiRecoveryFingerprints(db, row)
	require.NoError(t, err)
	require.NoError(t, db.Create(&row).Error)
	require.NoError(t, db.Exec(`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)`).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)
	return row
}

func TestWikiRecoveryBudgetLastSlotPostgresIsAtomic(t *testing.T) {
	db := setupKBDeletePostgres(t)
	row := seedRecoveryBudget(t, db, 1)
	repo := WikiRecoveryBudgetRepository{DB: db}
	var wg sync.WaitGroup
	var accepted atomic.Int32
	for i := 0; i < 24; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if repo.Reserve(context.Background(), row, "model", "candidate") == nil {
				accepted.Add(1)
			}
		}()
	}
	wg.Wait()
	require.Equal(t, int32(1), accepted.Load())
	saved, err := repo.Load(context.Background(), 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, int64(1), saved.ReservedProviderRequests)
	require.Equal(t, "stopped_budget", saved.Status)
	require.Equal(t, "request_limit", saved.StopReason)
	require.Equal(t, row.DeadlineAt, saved.DeadlineAt.UTC())
}

func TestWikiRecoveryBudgetRestartSharesAllPhasesAndDeadline(t *testing.T) {
	path := filepath.Join(t.TempDir(), "budget.sqlite")
	db := openRecoveryBudgetSQLite(t, path)
	row := seedRecoveryBudget(t, db, 3)
	repo := WikiRecoveryBudgetRepository{DB: db}
	require.NoError(t, repo.Reserve(context.Background(), row, "model", "candidate"))
	require.NoError(t, repo.Reserve(context.Background(), row, "embed", "taxonomy"))
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())
	repo.DB = openRecoveryBudgetSQLite(t, path)
	saved, err := repo.Load(context.Background(), 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, row.DeadlineAt, saved.DeadlineAt)
	require.Equal(t, int64(2), saved.ReservedProviderRequests)
	require.NoError(t, repo.Reserve(context.Background(), *saved, "model", "finalize"))
	require.ErrorIs(t, repo.Reserve(context.Background(), *saved, "model", "retry"), ErrWikiRecoveryBudgetStopped)
	reloaded, err := repo.Load(context.Background(), 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, int64(3), reloaded.ReservedProviderRequests)
	require.Equal(t, row.DeadlineAt, reloaded.DeadlineAt)
	require.Equal(t, "request_limit", reloaded.StopReason)
}

func TestWikiRecoveryBudgetScopeAndLiveConfigurationFailClosed(t *testing.T) {
	for _, tc := range []string{"tenant", "pending_op", "model", "request_model", "source", "model_config", "expired", "extended_limit", "extended_deadline", "edited_persisted_limit", "edited_persisted_deadline", "missing_receipt", "receipt_scope"} {
		t.Run(tc, func(t *testing.T) {
			db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
			row := seedRecoveryBudget(t, db, 3)
			original := row
			model := "model"
			switch tc {
			case "tenant":
				row.TenantID++
			case "pending_op":
				row.PendingOpID++
			case "model":
				row.ModelID = "foreign"
			case "request_model":
				model = "foreign"
			case "source":
				require.NoError(t, db.Exec("UPDATE chunks SET content = 'changed'").Error)
			case "model_config":
				require.NoError(t, db.Exec("UPDATE models SET parameters = '{\"changed\":true}'").Error)
			case "expired":
				row.DeadlineAt = time.Now().Add(-time.Second)
				require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("deadline_at", row.DeadlineAt).Error)
			case "extended_limit":
				row.MaxProviderRequests++
			case "extended_deadline":
				row.DeadlineAt = row.DeadlineAt.Add(time.Hour)
			case "edited_persisted_limit":
				row.MaxProviderRequests++
				require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("max_provider_requests", row.MaxProviderRequests).Error)
			case "edited_persisted_deadline":
				row.DeadlineAt = row.DeadlineAt.Add(time.Hour)
				require.NoError(t, db.Model(&types.WikiRecoveryBudget{}).Where("tenant_id = 7").Update("deadline_at", row.DeadlineAt).Error)
			case "missing_receipt":
				require.NoError(t, db.Exec("DELETE FROM wiki_dead_letter_recoveries").Error)
			case "receipt_scope":
				require.NoError(t, db.Exec("UPDATE wiki_dead_letter_recoveries SET knowledge_id = 'foreign'").Error)
			}
			repo := WikiRecoveryBudgetRepository{DB: db}
			require.ErrorIs(t, repo.Reserve(context.Background(), row, model, "candidate"), ErrWikiRecoveryBudgetStopped)
			saved, err := repo.Load(context.Background(), 7, original.Identity())
			require.NoError(t, err)
			require.Zero(t, saved.ReservedProviderRequests)
			if tc == "expired" {
				require.Equal(t, "deadline", saved.StopReason)
			}
		})
	}
}

func TestWikiRecoveryBudgetArchiveExactRowsAndKeepEvidence(t *testing.T) {
	db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
	row := seedRecoveryBudget(t, db, 1)
	repo := WikiRecoveryBudgetRepository{DB: db}
	identity := row.Identity()
	for _, item := range []struct {
		id       int64
		tenant   uint64
		kind     string
		identity types.WikiRecoveryIdentity
	}{
		{42, 7, types.TypeWikiIngest, identity}, {43, 7, types.TypeWikiFinalize, identity},
		{44, 7, types.TypeWikiFinalize, types.WikiRecoveryIdentity{}}, {45, 8, types.TypeWikiFinalize, identity},
		{46, 7, types.TypeWikiIngest, identity},
	} {
		body, _ := json.Marshal(item.identity)
		require.NoError(t, db.Create(&types.TaskPendingOp{ID: item.id, TenantID: item.tenant, TaskType: item.kind, ScopeID: "kb", Payload: body}).Error)
	}
	require.NoError(t, repo.Reserve(context.Background(), row, "model", "candidate"))
	require.ErrorIs(t, repo.Reserve(context.Background(), row, "model", "summary"), ErrWikiRecoveryBudgetStopped)
	require.NoError(t, repo.StopAndArchive(context.Background(), row, "other", "stop"))
	require.NoError(t, repo.StopAndArchive(context.Background(), row, "other", "stop"))
	var ids []int64
	require.NoError(t, db.Model(&types.TaskPendingOp{}).Order("id").Pluck("id", &ids).Error)
	require.Equal(t, []int64{44, 45, 46}, ids)
	saved, err := repo.Load(context.Background(), 7, identity)
	require.NoError(t, err)
	var paused []types.TaskPendingOp
	require.NoError(t, json.Unmarshal([]byte(saved.PausedPayload), &paused))
	require.Len(t, paused, 2)
	require.Equal(t, "request_limit", saved.StopReason)
	for _, table := range []string{"knowledges", "chunks"} {
		var count int64
		require.NoError(t, db.Table(table).Count(&count).Error)
		require.Equal(t, int64(1), count)
	}
}

type repositoryHTTPGuard struct {
	repo WikiRecoveryBudgetRepository
	row  types.WikiRecoveryBudget
}

func (g repositoryHTTPGuard) Key() string                     { return g.row.Identity().Key() }
func (g repositoryHTTPGuard) Check(ctx context.Context) error { return g.repo.Check(ctx, g.row) }
func (g repositoryHTTPGuard) Reserve(ctx context.Context, model, phase string) error {
	return g.repo.Reserve(ctx, g.row, model, phase)
}

func TestWikiRecoveryBudgetDatabaseFailureSendsNothingAndFailedSendConsumes(t *testing.T) {
	db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
	row := seedRecoveryBudget(t, db, 2)
	guard := repositoryHTTPGuard{WikiRecoveryBudgetRepository{DB: db}, row}
	var sends atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sends.Add(1)
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer server.Close()
	client := &http.Client{Transport: requestbudget.Transport{Base: http.DefaultTransport}}
	send := func() error {
		ctx := requestbudget.WithModel(requestbudget.WithGuard(context.Background(), guard), "model")
		req, _ := http.NewRequestWithContext(ctx, "POST", server.URL, nil)
		resp, err := client.Do(req)
		if resp != nil {
			resp.Body.Close()
		}
		return err
	}
	require.NoError(t, send())
	require.NoError(t, send())
	require.Error(t, send())
	require.Equal(t, int32(2), sends.Load())
	saved, err := guard.repo.Load(context.Background(), 7, row.Identity())
	require.NoError(t, err)
	require.Equal(t, int64(2), saved.ReservedProviderRequests)
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())
	require.Error(t, send())
	require.Equal(t, int32(2), sends.Load())
}

func TestWikiRecoveryBudgetCompleteRejectsForeignFinalizeAndWaitsForOwnRows(t *testing.T) {
	db := openRecoveryBudgetSQLite(t, filepath.Join(t.TempDir(), "budget.sqlite"))
	row := seedRecoveryBudget(t, db, 3)
	repo := WikiRecoveryBudgetRepository{DB: db}
	controlled, _ := json.Marshal(row.Identity())
	for _, op := range []types.TaskPendingOp{
		{ID: 50, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: controlled},
		{ID: 51, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: []byte(`{}`)},
		{ID: 52, TenantID: 7, TaskType: types.TypeWikiFinalize, ScopeID: "kb", Payload: controlled},
	} {
		require.NoError(t, db.Create(&op).Error)
	}
	completed, err := repo.Complete(context.Background(), row, []int64{51})
	require.ErrorIs(t, err, ErrWikiRecoveryBudgetStopped)
	require.False(t, completed)
	var count int64
	require.NoError(t, db.Model(&types.TaskPendingOp{}).Count(&count).Error)
	require.Equal(t, int64(3), count)
	completed, err = repo.Complete(context.Background(), row, []int64{50})
	require.NoError(t, err)
	require.False(t, completed)
	completed, err = repo.Complete(context.Background(), row, []int64{52})
	require.NoError(t, err)
	require.True(t, completed)
	require.NoError(t, db.Model(&types.TaskPendingOp{}).Count(&count).Error)
	require.Equal(t, int64(1), count)
}

func TestWikiRecoveryBudgetFingerprintMatchesScriptPostgres(t *testing.T) {
	node, err := exec.LookPath("node")
	if err != nil {
		t.Skip("Node is required to verify the recovery script fingerprint contract")
	}
	db := setupKBDeletePostgres(t)
	row := seedRecoveryBudget(t, db, 3)
	goSource, goModel, err := WikiRecoveryFingerprints(db, row)
	require.NoError(t, err)
	output, err := exec.Command(node, "-e", `process.stdout.write(JSON.stringify(require('./../../../scripts/wiki_recovery_budget_sql.cjs').recoveryFingerprintsSQL()))`).Output()
	require.NoError(t, err)
	var sql struct{ Source, Model string }
	require.NoError(t, json.Unmarshal(output, &sql))
	var got struct{ Source, Model string }
	require.NoError(t, db.Raw("SELECT "+sql.Source+" AS source, "+sql.Model+" AS model FROM (SELECT 7::bigint AS tenant_id, 'kb' AS scope_id, 'doc' AS related_id) d JOIN knowledge_bases k ON k.id=d.scope_id JOIN knowledges n ON n.id=d.related_id").Scan(&got).Error)
	require.Equal(t, goSource, got.Source)
	require.Equal(t, goModel, got.Model)
}