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

internal/application/service/wiki_recovery_budget.go

@@ -0,0 +1,331 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	apprepo "github.com/Tencent/WeKnora/internal/application/repository"
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+var ErrWikiRecoveryBudgetStopped = requestbudget.ErrStopped
+
+// Check before configuration short-circuits so a disabled/deleted source or
+// expired deadline still settles durably without waiting for a model call.
+func (s *wikiIngestService) checkWikiRecoveryBudgets(ctx context.Context, payload WikiIngestPayload) error {
+	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
+		return nil
+	}
+	var rows []types.WikiRecoveryBudget
+	if err := s.checkpointDB.WithContext(ctx).Where("tenant_id = ? AND knowledge_base_id = ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, []string{"active", "stopped_budget"}).Find(&rows).Error; err != nil {
+		return err
+	}
+	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
+	for _, row := range rows {
+		if err := repo.Check(ctx, row); err != nil {
+			if archiveErr := repo.StopAndArchive(ctx, row, "guard_rejected", "preflight"); archiveErr != nil {
+				return errors.Join(err, archiveErr)
+			}
+		}
+	}
+	return nil
+}
+
+func (s *wikiIngestService) rejectOrdinaryRecoveryOverlap(ctx context.Context, payload WikiIngestPayload, ops []WikiPendingOp) error {
+	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
+		return nil
+	}
+	ids := make([]string, 0, len(ops))
+	for _, op := range ops {
+		ids = append(ids, op.KnowledgeID)
+	}
+	var count int64
+	if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
+		Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id IN ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, ids, []string{"active", "stopped_budget"}).Count(&count).Error; err != nil {
+		return err
+	}
+	if count != 0 {
+		return fmt.Errorf("%w: ordinary operation overlaps controlled recovery", ErrWikiRecoveryBudgetStopped)
+	}
+	return nil
+}
+
+type wikiRecoveryGuard struct {
+	repo    apprepo.WikiRecoveryBudgetRepository
+	row     types.WikiRecoveryBudget
+	cancel  context.CancelCauseFunc
+	mu      sync.Mutex
+	stopped error
+}
+
+func (g *wikiRecoveryGuard) Key() string {
+	return fmt.Sprintf("%d:%s", g.row.TenantID, g.row.Identity().Key())
+}
+
+func (g *wikiRecoveryGuard) fail(err error) error {
+	if err != nil {
+		g.mu.Lock()
+		if g.stopped == nil {
+			g.stopped = err
+		}
+		err = g.stopped
+		g.mu.Unlock()
+		g.cancel(err)
+	}
+	return err
+}
+
+func (g *wikiRecoveryGuard) Check(ctx context.Context) error {
+	g.mu.Lock()
+	stopped := g.stopped
+	g.mu.Unlock()
+	if stopped != nil {
+		return stopped
+	}
+	if !time.Now().Before(g.row.DeadlineAt) {
+		cleanup, cancel := wikiIngestCleanupContext(ctx)
+		defer cancel()
+		return g.fail(g.repo.Check(cleanup, g.row))
+	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	return g.fail(g.repo.Check(ctx, g.row))
+}
+
+func (g *wikiRecoveryGuard) Reserve(ctx context.Context, model, phase string) error {
+	if err := g.Check(ctx); err != nil {
+		return err
+	}
+	return g.fail(g.repo.Reserve(ctx, g.row, model, phase))
+}
+
+func wikiRecoveryFromContext(ctx context.Context) *wikiRecoveryGuard {
+	guard, _ := requestbudget.FromContext(ctx).(*wikiRecoveryGuard)
+	return guard
+}
+
+func wikiRecoveryIdentityFromContext(ctx context.Context) types.WikiRecoveryIdentity {
+	if guard := wikiRecoveryFromContext(ctx); guard != nil {
+		return guard.row.Identity()
+	}
+	return types.WikiRecoveryIdentity{}
+}
+
+func (s *wikiIngestService) bindWikiRecovery(ctx context.Context, payload WikiIngestPayload, identity types.WikiRecoveryIdentity, knowledgeID, modelID string) (context.Context, context.CancelFunc, error) {
+	if !identity.Present() {
+		return ctx, func() {}, nil
+	}
+	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
+	row, err := repo.Load(ctx, payload.TenantID, identity)
+	if err != nil {
+		return ctx, func() {}, err
+	}
+	if row.KnowledgeBaseID != payload.KnowledgeBaseID || (knowledgeID != "" && row.KnowledgeID != knowledgeID) || row.ModelID != modelID {
+		return ctx, func() {}, fmt.Errorf("%w: task scope mismatch", ErrWikiRecoveryBudgetStopped)
+	}
+	// Deployment must explicitly attest its immutable release and image. A
+	// controlled op cannot execute in an unreviewed local or stale API build.
+	if row.ReviewedRelease == "" || row.ReviewedAPIImage == "" ||
+		os.Getenv("WEKNORA_REVIEWED_RELEASE") != row.ReviewedRelease || os.Getenv("WEKNORA_REVIEWED_API_IMAGE") != row.ReviewedAPIImage {
+		return ctx, func() {}, fmt.Errorf("%w: reviewed runtime mismatch", ErrWikiRecoveryBudgetStopped)
+	}
+	ctx, deadlineCancel := context.WithDeadlineCause(ctx, row.DeadlineAt, fmt.Errorf("%w: deadline", ErrWikiRecoveryBudgetStopped))
+	ctx, cancel := context.WithCancelCause(ctx)
+	guard := &wikiRecoveryGuard{repo: repo, row: *row, cancel: cancel}
+	ctx = requestbudget.WithGuard(ctx, guard)
+	if err := guard.Check(ctx); err != nil {
+		deadlineCancel()
+		cancel(err)
+		return ctx, func() {}, err
+	}
+	return ctx, func() { deadlineCancel(); cancel(nil) }, nil
+}
+
+func (s *wikiIngestService) settleWikiRecoveryStop(ctx context.Context, identity types.WikiRecoveryIdentity, tenant uint64, err error) error {
+	if !identity.Present() || err == nil || errors.Is(err, errWikiCheckpointYield) {
+		return err
+	}
+	reason := wikiRecoveryFailureReason(ctx, err)
+	if !errors.Is(err, ErrWikiRecoveryBudgetStopped) {
+		err = fmt.Errorf("%w: recovery phase failed: %w", ErrWikiRecoveryBudgetStopped, err)
+	}
+	cleanup, cancel := wikiIngestCleanupContext(ctx)
+	defer cancel()
+	repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
+	row, loadErr := repo.Load(cleanup, tenant, identity)
+	if loadErr != nil {
+		return errors.Join(err, loadErr)
+	}
+	phase := row.LastPhase
+	if phase == "" {
+		phase = requestbudget.Phase(ctx)
+	}
+	if phase == "" {
+		phase = "recovery"
+	}
+	return errors.Join(err, repo.StopAndArchive(cleanup, *row, reason, phase))
+}
+
+// Persist only a bounded category, never provider response bodies or secrets.
+// Phase failures still stop the same recovery; they are not budget exhaustion.
+func wikiRecoveryFailureReason(ctx context.Context, err error) string {
+	if errors.Is(err, ErrWikiRecoveryBudgetStopped) || errors.Is(context.Cause(ctx), ErrWikiRecoveryBudgetStopped) {
+		return "guard_rejected"
+	}
+	var outputErr *wikiOutputError
+	if errors.As(err, &outputErr) {
+		if outputErr.incomplete {
+			return "model_output_incomplete"
+		}
+		return "model_output_invalid"
+	}
+	if errors.Is(err, context.DeadlineExceeded) {
+		return "phase_timeout"
+	}
+	if errors.Is(err, context.Canceled) {
+		return "phase_cancelled"
+	}
+	return "phase_failed"
+}
+
+func selectWikiFinalizeGroup(rows []*types.TaskPendingOp) ([]*types.TaskPendingOp, types.WikiRecoveryIdentity, error) {
+	var first types.WikiRecoveryIdentity
+	var selected []*types.TaskPendingOp
+	identities := make([]types.WikiRecoveryIdentity, len(rows))
+	for i, row := range rows {
+		var identity types.WikiRecoveryIdentity
+		var keyIdentity types.WikiRecoveryIdentity
+		if strings.HasPrefix(row.DedupKey, "recovery:") {
+			parts := strings.SplitN(row.DedupKey, ":", 6)
+			if len(parts) != 6 {
+				return nil, first, fmt.Errorf("%w: invalid finalize dedup identity", ErrWikiRecoveryBudgetStopped)
+			}
+			keyIdentity.DeadLetterID, _ = strconv.ParseInt(parts[1], 10, 64)
+			keyIdentity.Attempt, _ = strconv.Atoi(parts[2])
+			keyIdentity.PendingOpID, _ = strconv.ParseInt(parts[3], 10, 64)
+			keyIdentity.Version, _ = strconv.Atoi(parts[4])
+		}
+		if len(row.Payload) > 0 {
+			if err := json.Unmarshal(row.Payload, &identity); err != nil {
+				return nil, keyIdentity, fmt.Errorf("%w: invalid finalize payload: %v", ErrWikiRecoveryBudgetStopped, err)
+			}
+		}
+		if keyIdentity.Present() || identity.Present() {
+			if !keyIdentity.Valid() || keyIdentity != identity {
+				if keyIdentity.Valid() {
+					identity = keyIdentity
+				}
+				return nil, identity, fmt.Errorf("%w: finalize payload and dedup identity mismatch", ErrWikiRecoveryBudgetStopped)
+			}
+			var payload wikiFinalizeRow
+			if err := json.Unmarshal(row.Payload, &payload); err != nil {
+				return nil, identity, fmt.Errorf("%w: invalid finalize content: %v", ErrWikiRecoveryBudgetStopped, err)
+			}
+			if (row.Op != wikiFinalizeOpSlug && row.Op != wikiFinalizeOpChange) ||
+				(row.Op == wikiFinalizeOpSlug && payload.Slug == "") || (row.Op == wikiFinalizeOpChange && payload.Change == nil) {
+				return nil, identity, fmt.Errorf("%w: incomplete finalize content", ErrWikiRecoveryBudgetStopped)
+			}
+		}
+		identities[i] = identity
+		// Controlled rows go first so an ordinary index update cannot absorb
+		// this recovery's work and leave its own finalize lane unbudgeted.
+		if !first.Present() && identity.Present() {
+			first = identity
+		}
+	}
+	for i, row := range rows {
+		if identities[i] == first {
+			selected = append(selected, row)
+		}
+	}
+	if first.Present() && !first.Valid() {
+		return nil, first, fmt.Errorf("%w: invalid finalize identity", ErrWikiRecoveryBudgetStopped)
+	}
+	return selected, first, nil
+}
+
+func (s *wikiIngestService) rejectOrdinaryFinalizeRecoveryOverlap(ctx context.Context, payload WikiIngestPayload) error {
+	if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
+		return nil
+	}
+	var count int64
+	if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
+		Where("tenant_id = ? AND knowledge_base_id = ? AND status = ?", payload.TenantID, payload.KnowledgeBaseID, "active").Count(&count).Error; err != nil {
+		return err
+	}
+	if count > 0 {
+		return fmt.Errorf("%w: controlled recovery still owns index finalization", ErrWikiIngestConcurrent)
+	}
+	return nil
+}
+
+func (s *wikiIngestService) enqueueControlledWikiFinalize(ctx context.Context, payload WikiIngestPayload, slugs []string, titles map[string]string, changes []wikiFinalizeChange) error {
+	guard := wikiRecoveryFromContext(ctx)
+	if guard == nil {
+		return fmt.Errorf("%w: missing recovery guard", ErrWikiRecoveryBudgetStopped)
+	}
+	if err := guard.Check(ctx); err != nil {
+		return err
+	}
+	identity := guard.row.Identity()
+	var rows []types.TaskPendingOp
+	appendRow := func(op, key string, value wikiFinalizeRow) error {
+		value.WikiRecoveryIdentity = identity
+		body, err := json.Marshal(value)
+		if err != nil {
+			return err
+		}
+		rows = append(rows, types.TaskPendingOp{TenantID: payload.TenantID, TaskType: wikiFinalizeTaskType, Scope: wikiTaskScope, ScopeID: payload.KnowledgeBaseID, Op: op, DedupKey: "recovery:" + identity.Key() + ":" + key, Payload: body, EnqueuedAt: time.Now().UTC()})
+		return nil
+	}
+	for _, slug := range slugs {
+		if err := appendRow(wikiFinalizeOpSlug, slug, wikiFinalizeRow{Slug: slug, Title: titles[slug]}); err != nil {
+			return err
+		}
+	}
+	for i := range changes {
+		if err := appendRow(wikiFinalizeOpChange, fmt.Sprintf("change:%d", i), wikiFinalizeRow{Change: &changes[i]}); err != nil {
+			return err
+		}
+	}
+	if len(changes) == 0 {
+		return fmt.Errorf("%w: recovery has no completed document to finalize", ErrWikiRecoveryBudgetStopped)
+	}
+	err := s.checkpointDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var budget types.WikiRecoveryBudget
+		if err := tx.Where("tenant_id = ? AND pending_op_id = ?", payload.TenantID, identity.PendingOpID).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&budget).Error; err != nil {
+			return err
+		}
+		if budget.Identity() != identity || budget.Status != "active" || !time.Now().Before(budget.DeadlineAt) {
+			return fmt.Errorf("%w: finalize handoff no longer active", ErrWikiRecoveryBudgetStopped)
+		}
+		var op types.TaskPendingOp
+		if err := tx.Where("tenant_id = ? AND id = ? AND task_type = ? AND scope_id = ?", payload.TenantID, identity.PendingOpID, wikiTaskType, payload.KnowledgeBaseID).Take(&op).Error; err != nil {
+			return err
+		}
+		var saved types.WikiRecoveryIdentity
+		if json.Unmarshal(op.Payload, &saved) != nil || saved != identity {
+			return fmt.Errorf("%w: handoff op mismatch", ErrWikiRecoveryBudgetStopped)
+		}
+		if err := tx.Create(&rows).Error; err != nil {
+			return err
+		}
+		return tx.Where("tenant_id = ? AND id = ?", payload.TenantID, identity.PendingOpID).Delete(&types.TaskPendingOp{}).Error
+	})
+	if err != nil {
+		return err
+	}
+	s.scheduleFinalize(ctx, payload)
+	return nil
+}