WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/repository/wiki_recovery_budget.go

@@ -0,0 +1,376 @@
+package repository
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/binary"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+var ErrWikiRecoveryBudgetStopped = requestbudget.ErrStopped
+
+type WikiRecoveryBudgetRepository struct{ DB *gorm.DB }
+
+func recoveryScope(db *gorm.DB, expected types.WikiRecoveryBudget) *gorm.DB {
+	return db.Model(&types.WikiRecoveryBudget{}).Where(
+		"tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ? AND knowledge_base_id = ? AND knowledge_id = ? AND model_id = ?",
+		expected.TenantID, expected.DeadLetterID, expected.AttemptNumber, expected.PendingOpID,
+		expected.KnowledgeBaseID, expected.KnowledgeID, expected.ModelID)
+}
+
+func (r WikiRecoveryBudgetRepository) Load(ctx context.Context, tenant uint64, identity types.WikiRecoveryIdentity) (*types.WikiRecoveryBudget, error) {
+	if r.DB == nil || tenant == 0 || !identity.Valid() {
+		return nil, fmt.Errorf("%w: incomplete recovery identity", ErrWikiRecoveryBudgetStopped)
+	}
+	var row types.WikiRecoveryBudget
+	err := r.DB.WithContext(ctx).Where("tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ?",
+		tenant, identity.DeadLetterID, identity.Attempt, identity.PendingOpID).Take(&row).Error
+	if err != nil {
+		return nil, fmt.Errorf("%w: load budget: %v", ErrWikiRecoveryBudgetStopped, err)
+	}
+	return &row, nil
+}
+
+func (r WikiRecoveryBudgetRepository) Check(ctx context.Context, expected types.WikiRecoveryBudget) error {
+	return r.checkOrReserve(ctx, expected, "", "", false)
+}
+
+func (r WikiRecoveryBudgetRepository) Reserve(ctx context.Context, expected types.WikiRecoveryBudget, modelID, phase string) error {
+	return r.checkOrReserve(ctx, expected, modelID, phase, true)
+}
+
+// Row locking serializes the final slot across workers and processes. A failed
+// send never reaches this repository again to refund its reservation.
+func (r WikiRecoveryBudgetRepository) checkOrReserve(ctx context.Context, expected types.WikiRecoveryBudget, modelID, phase string, reserve bool) error {
+	if r.DB == nil {
+		return fmt.Errorf("%w: budget database unavailable", ErrWikiRecoveryBudgetStopped)
+	}
+	var stopped error
+	err := r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var row types.WikiRecoveryBudget
+		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&row).Error; err != nil {
+			return err
+		}
+		reason := ""
+		now := time.Now().UTC()
+		switch {
+		case row.Status != "active":
+			reason = row.StopReason
+			if reason == "" {
+				reason = row.Status
+			}
+		case row.MaxProviderRequests <= 0 || row.DeadlineAt.IsZero():
+			reason = "invalid_budget"
+		case !now.Before(row.DeadlineAt):
+			reason = "deadline"
+		case reserve && row.ReservedProviderRequests >= row.MaxProviderRequests:
+			reason = "request_limit"
+		case row.SourceFingerprint != expected.SourceFingerprint || row.ModelFingerprint != expected.ModelFingerprint ||
+			row.ReviewedRelease != expected.ReviewedRelease || row.ReviewedAPIImage != expected.ReviewedAPIImage ||
+			row.MaxProviderRequests != expected.MaxProviderRequests || !row.DeadlineAt.Equal(expected.DeadlineAt) ||
+			row.EmbeddingModelID != expected.EmbeddingModelID:
+			reason = "budget_identity_changed"
+		case reserve && modelID != row.ModelID && (row.EmbeddingModelID == "" || modelID != row.EmbeddingModelID):
+			reason = "model_mismatch"
+		}
+		if reason == "" {
+			if err := validateRecoveryReceipt(tx, row); err != nil {
+				reason = "review_receipt_mismatch"
+			}
+		}
+		if reason == "" {
+			source, model, err := WikiRecoveryFingerprints(tx, row)
+			if err != nil {
+				return err
+			}
+			if source != row.SourceFingerprint {
+				reason = "source_changed"
+			}
+			if model != row.ModelFingerprint {
+				reason = "model_configuration_changed"
+			}
+		}
+		if reason != "" {
+			stopped = fmt.Errorf("%w: %s", ErrWikiRecoveryBudgetStopped, reason)
+			if row.Status == "active" {
+				return recoveryScope(tx, expected).Updates(map[string]any{"status": "stopped_budget", "stop_reason": reason, "last_phase": phase, "updated_at": now}).Error
+			}
+			return nil
+		}
+		if !reserve {
+			return nil
+		}
+		result := recoveryScope(tx, expected).Where("status = ? AND reserved_provider_requests < max_provider_requests AND deadline_at > ?", "active", now).
+			Updates(map[string]any{"reserved_provider_requests": gorm.Expr("reserved_provider_requests + 1"), "last_phase": phase, "updated_at": now})
+		if result.Error != nil {
+			return result.Error
+		}
+		if result.RowsAffected != 1 {
+			return errors.New("reservation lost its active budget")
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("%w: budget transaction: %v", ErrWikiRecoveryBudgetStopped, err)
+	}
+	return stopped
+}
+
+// StopAndArchive removes only rows carrying this exact recovery identity. The
+// source documents, page drafts, history and validated checkpoints are retained.
+func (r WikiRecoveryBudgetRepository) StopAndArchive(ctx context.Context, expected types.WikiRecoveryBudget, reason, phase string) error {
+	return r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var budget types.WikiRecoveryBudget
+		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&budget).Error; err != nil {
+			return err
+		}
+		if budget.Status == "completed" {
+			return nil
+		}
+		var rows []types.TaskPendingOp
+		if err := tx.Where("tenant_id = ? AND scope_id = ? AND task_type IN ?", budget.TenantID, budget.KnowledgeBaseID,
+			[]string{types.TypeWikiIngest, types.TypeWikiFinalize}).Find(&rows).Error; err != nil {
+			return err
+		}
+		var retained []types.TaskPendingOp
+		var ids []int64
+		for _, row := range rows {
+			var identity types.WikiRecoveryIdentity
+			payloadMatches := json.Unmarshal(row.Payload, &identity) == nil && identity == budget.Identity()
+			keyMatches := row.TaskType == types.TypeWikiFinalize && strings.HasPrefix(row.DedupKey, "recovery:"+budget.Identity().Key()+":")
+			if (payloadMatches || keyMatches) && (row.TaskType != types.TypeWikiIngest || row.ID == budget.PendingOpID) {
+				retained = append(retained, row)
+				ids = append(ids, row.ID)
+			}
+		}
+		payload, err := json.Marshal(retained)
+		if err != nil {
+			return err
+		}
+		var progress []struct {
+			Fingerprint string
+			Stage       string
+			ItemKey     string
+			UpdatedAt   time.Time
+		}
+		if tx.Migrator().HasTable("wiki_ingest_checkpoints") {
+			if err := tx.Table("wiki_ingest_checkpoints").Select("fingerprint, stage, item_key, updated_at").
+				Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id = ?", budget.TenantID, budget.KnowledgeBaseID, budget.KnowledgeID).Find(&progress).Error; err != nil {
+				return err
+			}
+		}
+		progressJSON, err := json.Marshal(map[string]any{
+			"source_fingerprint":   budget.SourceFingerprint,
+			"validation_state":     "retained_metadata_requires_fingerprint_revalidation",
+			"retained_checkpoints": progress,
+		})
+		if err != nil {
+			return err
+		}
+		if budget.StopReason != "" {
+			reason = budget.StopReason
+		}
+		updates := map[string]any{"status": "stopped_budget", "stop_reason": reason, "updated_at": time.Now().UTC()}
+		if phase != "" && budget.LastPhase == "" {
+			updates["last_phase"] = phase
+		}
+		if len(retained) > 0 {
+			updates["paused_payload"] = string(payload)
+			updates["verified_progress"] = string(progressJSON)
+		}
+		if err := recoveryScope(tx, expected).Updates(updates).Error; err != nil {
+			return err
+		}
+		if len(ids) > 0 {
+			return tx.Where("tenant_id = ? AND id IN ?", budget.TenantID, ids).Delete(&types.TaskPendingOp{}).Error
+		}
+		return nil
+	})
+}
+
+func (r WikiRecoveryBudgetRepository) Complete(ctx context.Context, expected types.WikiRecoveryBudget, finalizeIDs []int64) (bool, error) {
+	completed := false
+	err := r.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var row types.WikiRecoveryBudget
+		if err := recoveryScope(tx, expected).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&row).Error; err != nil {
+			return err
+		}
+		if row.Status != "active" || !time.Now().Before(row.DeadlineAt) {
+			return fmt.Errorf("%w: finalization no longer active", ErrWikiRecoveryBudgetStopped)
+		}
+		if len(finalizeIDs) == 0 {
+			return errors.New("controlled recovery has no finalize receipt")
+		}
+		if err := validateRecoveryReceipt(tx, row); err != nil {
+			return fmt.Errorf("%w: finalization receipt mismatch", ErrWikiRecoveryBudgetStopped)
+		}
+		var selected []types.TaskPendingOp
+		if err := tx.Where("tenant_id = ? AND task_type = ? AND scope_id = ? AND id IN ?", row.TenantID, types.TypeWikiFinalize, row.KnowledgeBaseID, finalizeIDs).Find(&selected).Error; err != nil {
+			return err
+		}
+		if len(selected) != len(finalizeIDs) {
+			return fmt.Errorf("%w: finalization row scope mismatch", ErrWikiRecoveryBudgetStopped)
+		}
+		for _, op := range selected {
+			var identity types.WikiRecoveryIdentity
+			if json.Unmarshal(op.Payload, &identity) != nil || identity != row.Identity() {
+				return fmt.Errorf("%w: finalization row identity mismatch", ErrWikiRecoveryBudgetStopped)
+			}
+		}
+		if err := tx.Where("tenant_id = ? AND task_type = ? AND id IN ?", row.TenantID, types.TypeWikiFinalize, finalizeIDs).Delete(&types.TaskPendingOp{}).Error; err != nil {
+			return err
+		}
+		var remaining []types.TaskPendingOp
+		if err := tx.Where("tenant_id = ? AND scope_id = ? AND task_type IN ?", row.TenantID, row.KnowledgeBaseID, []string{types.TypeWikiIngest, types.TypeWikiFinalize}).Find(&remaining).Error; err != nil {
+			return err
+		}
+		for _, op := range remaining {
+			var identity types.WikiRecoveryIdentity
+			if (json.Unmarshal(op.Payload, &identity) == nil && identity == row.Identity()) ||
+				(op.TaskType == types.TypeWikiFinalize && strings.HasPrefix(op.DedupKey, "recovery:"+row.Identity().Key()+":")) {
+				return nil
+			}
+		}
+		if err := recoveryScope(tx, expected).Updates(map[string]any{"status": "completed", "last_phase": "finalize", "updated_at": time.Now().UTC()}).Error; err != nil {
+			return err
+		}
+		completed = true
+		return nil
+	})
+	return completed, err
+}
+
+// The receipt is the independent, persisted review. Reopening a process must
+// not accept an edited budget row as a fresh immutable snapshot.
+func validateRecoveryReceipt(db *gorm.DB, row types.WikiRecoveryBudget) error {
+	var receipt struct {
+		KnowledgeBaseID   string
+		KnowledgeID       string
+		ModelID           string
+		DeployedFixReview string
+	}
+	if err := db.Table("wiki_dead_letter_recoveries").Select("knowledge_base_id, knowledge_id, model_id, deployed_fix_review").
+		Where("tenant_id = ? AND dead_letter_id = ? AND attempt_number = ? AND pending_op_id = ?", row.TenantID, row.DeadLetterID, row.AttemptNumber, row.PendingOpID).Take(&receipt).Error; err != nil {
+		return err
+	}
+	var review struct {
+		TenantID            uint64    `json:"tenant_id"`
+		DeadLetterID        int64     `json:"dead_letter_id"`
+		AttemptNumber       int       `json:"max_attempt"`
+		BudgetVersion       int       `json:"budget_version"`
+		KnowledgeBaseID     string    `json:"knowledge_base_id"`
+		KnowledgeID         string    `json:"knowledge_id"`
+		ModelID             string    `json:"model_id"`
+		EmbeddingModelID    string    `json:"embedding_model_id"`
+		SourceFingerprint   string    `json:"source_fingerprint"`
+		ModelFingerprint    string    `json:"model_fingerprint"`
+		FixRelease          string    `json:"fix_release"`
+		APIImageID          string    `json:"api_image_id"`
+		MaxProviderRequests int64     `json:"max_provider_requests"`
+		DeadlineAt          time.Time `json:"deadline_at"`
+	}
+	if err := json.Unmarshal([]byte(receipt.DeployedFixReview), &review); err != nil {
+		return err
+	}
+	if receipt.KnowledgeBaseID != row.KnowledgeBaseID || receipt.KnowledgeID != row.KnowledgeID || receipt.ModelID != row.ModelID ||
+		review.TenantID != row.TenantID || review.DeadLetterID != row.DeadLetterID || review.AttemptNumber < row.AttemptNumber || review.BudgetVersion != 1 ||
+		review.KnowledgeBaseID != row.KnowledgeBaseID || review.KnowledgeID != row.KnowledgeID || review.ModelID != row.ModelID ||
+		review.EmbeddingModelID != row.EmbeddingModelID || review.SourceFingerprint != row.SourceFingerprint || review.ModelFingerprint != row.ModelFingerprint ||
+		review.FixRelease != row.ReviewedRelease || review.APIImageID != row.ReviewedAPIImage ||
+		review.MaxProviderRequests != row.MaxProviderRequests || !review.DeadlineAt.Equal(row.DeadlineAt) {
+		return errors.New("recovery differs from its review receipt")
+	}
+	return nil
+}
+
+// Fingerprints use unsigned 64-bit big-endian UTF-8 byte lengths followed by
+// each field's bytes. SQL-side callers can reproduce this without JSON ordering
+// or escaping ambiguities. Parameters remain encrypted and are only hashed.
+func recoveryFingerprint(fields []string) string {
+	h := sha256.New()
+	var size [8]byte
+	for _, field := range fields {
+		binary.BigEndian.PutUint64(size[:], uint64(len(field)))
+		h.Write(size[:])
+		h.Write([]byte(field))
+	}
+	return hex.EncodeToString(h.Sum(nil))
+}
+
+func WikiRecoveryFingerprints(db *gorm.DB, budget types.WikiRecoveryBudget) (string, string, error) {
+	var doc struct{ Title string }
+	if err := db.Table("knowledges").Select("title").Where("tenant_id = ? AND id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeID, budget.KnowledgeBaseID).Take(&doc).Error; err != nil {
+		return "", "", err
+	}
+	var kb struct {
+		SummaryModelID   string
+		EmbeddingModelID string
+		WikiConfig       string
+		IndexingStrategy string
+	}
+	if err := db.Table("knowledge_bases").Select("summary_model_id, embedding_model_id, wiki_config, indexing_strategy").Where("tenant_id = ? AND id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeBaseID).Take(&kb).Error; err != nil {
+		return "", "", err
+	}
+	var config struct {
+		SynthesisModelID string `json:"synthesis_model_id"`
+	}
+	if err := json.Unmarshal([]byte(kb.WikiConfig), &config); err != nil {
+		return "", "", err
+	}
+	if config.SynthesisModelID == "" {
+		config.SynthesisModelID = kb.SummaryModelID
+	}
+	var strategy struct {
+		WikiEnabled bool `json:"wiki_enabled"`
+	}
+	if err := json.Unmarshal([]byte(kb.IndexingStrategy), &strategy); err != nil {
+		return "", "", err
+	}
+	if !strategy.WikiEnabled || config.SynthesisModelID != budget.ModelID || kb.EmbeddingModelID != budget.EmbeddingModelID {
+		return "", "", errors.New("recovery knowledge base configuration changed")
+	}
+	var chunks []types.Chunk
+	if err := db.Where("tenant_id = ? AND knowledge_id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", budget.TenantID, budget.KnowledgeID, budget.KnowledgeBaseID).Order("id ASC").Find(&chunks).Error; err != nil {
+		return "", "", err
+	}
+	if len(chunks) == 0 {
+		return "", "", errors.New("recovery source has no chunks")
+	}
+	fields := []string{"wiki-recovery-source-v1", strconv.FormatUint(budget.TenantID, 10), budget.KnowledgeBaseID, budget.KnowledgeID, doc.Title}
+	for _, c := range chunks {
+		fields = append(fields, c.ID, c.Content, strconv.Itoa(c.ContentRevision), strconv.Itoa(c.ChunkIndex), strconv.Itoa(c.StartAt), strconv.Itoa(c.EndAt), c.ChunkType, c.ParentChunkID, c.ImageInfo, strconv.FormatBool(c.IsEnabled))
+	}
+	modelFields := []string{"wiki-recovery-model-v1", kb.WikiConfig, kb.IndexingStrategy}
+	for _, id := range []string{budget.ModelID, budget.EmbeddingModelID} {
+		if id == "" {
+			continue
+		}
+		var m struct {
+			ID         string
+			Name       string
+			Type       string
+			Source     string
+			Status     string
+			Parameters string
+		}
+		if err := db.Table("models").Select("id, name, type, source, status, parameters").Where("tenant_id = ? AND id = ? AND deleted_at IS NULL", budget.TenantID, id).Take(&m).Error; err != nil {
+			return "", "", err
+		}
+		if m.Status != "active" {
+			return "", "", errors.New("recovery model is not active")
+		}
+		modelFields = append(modelFields, m.ID, m.Name, m.Type, m.Source, m.Status, m.Parameters)
+	}
+	return recoveryFingerprint(fields), recoveryFingerprint(modelFields), nil
+}