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

internal/application/repository/memory_recovery.go

@@ -0,0 +1,186 @@
+package repository
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/google/uuid"
+	"gorm.io/gorm"
+)
+
+func currentExtractionSubject(subject *types.MemorySubject, attempt string) bool {
+	return subject != nil && subject.Enabled && subject.ExtractAttemptID == attempt &&
+		subject.ExtractStartedAt != nil && subject.ExtractLeaseUntil != nil && subject.ExtractLeaseUntil.After(time.Now())
+}
+
+// Lock the reservation with each model-output mutation. A recovered task or
+// subject deletion cannot interleave between this check and its item write.
+func (r *memoryRepository) withExtractionWrite(ctx context.Context, scope interfaces.MemoryScope, write func(*gorm.DB) error) error {
+	attempt := types.MemoryExtractionAttempt(ctx)
+	if attempt == "" {
+		return write(r.db.WithContext(ctx))
+	}
+	return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var subject types.MemorySubject
+		if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID).
+			Clauses(forUpdateClause()).First(&subject).Error; err != nil {
+			return err
+		}
+		if !currentExtractionSubject(&subject, attempt) {
+			return errMemoryAttemptLost
+		}
+		return write(tx)
+	})
+}
+
+var errMemoryAttemptLost = errors.New("memory extraction attempt is no longer current")
+
+func extractionRetryDelay(failures int) time.Duration {
+	// Persist the counter; a process restart or a new chat must not reset it.
+	if failures > 9 {
+		return 6 * time.Hour
+	}
+	if failures < 1 {
+		failures = 1
+	}
+	delay := time.Minute * time.Duration(1<<uint(failures-1))
+	if delay > 6*time.Hour {
+		return 6 * time.Hour
+	}
+	return delay
+}
+
+func failedExtractionUpdates(subject *types.MemorySubject, now time.Time) map[string]interface{} {
+	failures := subject.ExtractFailures + 1
+	if failures > types.MaxMemoryExtractionFailures {
+		failures = types.MaxMemoryExtractionFailures
+	}
+	var nextAt *time.Time
+	if failures < types.MaxMemoryExtractionFailures {
+		next := now.Add(extractionRetryDelay(failures))
+		nextAt = &next
+	}
+	return map[string]interface{}{
+		"extract_attempt_id": "", "extract_lease_until": nil, "extract_started_at": nil,
+		"extract_scheduled_at": nil, "extract_failures": failures,
+		"extract_next_at": nextAt, "updated_at": now,
+	}
+}
+
+// ReserveExtraction commits the task identity before sending its wake-up to
+// Redis or Lite. A lost wake-up remains discoverable after the lease expires.
+func (r *memoryRepository) ReserveExtraction(ctx context.Context, scope interfaces.MemoryScope, payload types.MemoryExtractPayload, delay, lease time.Duration) (*types.MemoryExtractPayload, error) {
+	var reserved *types.MemoryExtractPayload
+	err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var subject types.MemorySubject
+		query := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID)
+		if err := query.Clauses(forUpdateClause()).First(&subject).Error; err != nil {
+			return err
+		}
+		now := time.Now()
+		updates := map[string]interface{}{}
+		if payload.ChatModelID != "" {
+			subject.ExtractChatModelID = payload.ChatModelID
+			updates["extract_chat_model_id"] = payload.ChatModelID
+		}
+		if payload.Language != "" {
+			subject.ExtractLanguage = payload.Language
+			updates["extract_language"] = payload.Language
+		}
+		active := subject.ExtractLeaseUntil != nil && subject.ExtractLeaseUntil.After(now)
+		backingOff := subject.ExtractNextAt != nil && subject.ExtractNextAt.After(now)
+		if subject.Enabled && subject.ExtractFailures < types.MaxMemoryExtractionFailures && len(subject.PendingSessions) > 0 && !active && !backingOff {
+			if subject.ExtractAttemptID != "" {
+				for key, value := range failedExtractionUpdates(&subject, now) {
+					updates[key] = value
+				}
+			} else {
+				payload.AttemptID = uuid.NewString()
+				payload.TenantID, payload.SubjectID = scope.TenantID, scope.SubjectID
+				payload.SessionID = subject.PendingSessions[0]
+				payload.ChatModelID, payload.Language = subject.ExtractChatModelID, subject.ExtractLanguage
+				updates["extract_attempt_id"] = payload.AttemptID
+				updates["extract_scheduled_at"] = now
+				updates["extract_started_at"] = nil
+				updates["extract_lease_until"] = now.Add(delay + lease)
+				updates["extract_next_at"] = now.Add(delay)
+				reserved = &payload
+			}
+		}
+		if len(updates) == 0 {
+			return nil
+		}
+		return tx.Model(&types.MemorySubject{}).Where("id = ?", subject.ID).Updates(updates).Error
+	})
+	return reserved, err
+}
+
+func (r *memoryRepository) StartExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string, lease time.Duration) (bool, error) {
+	if attemptID == "" {
+		return false, nil
+	}
+	now := time.Now()
+	result := r.scoped(ctx, scope).Model(&types.MemorySubject{}).
+		Where("extract_failures < ?", types.MaxMemoryExtractionFailures).
+		Where("enabled = ? AND extract_attempt_id = ? AND extract_started_at IS NULL AND extract_lease_until > ?", true, attemptID, now).
+		Updates(map[string]interface{}{"extract_started_at": now, "extract_lease_until": now.Add(lease), "updated_at": now})
+	return result.RowsAffected == 1, result.Error
+}
+
+func (r *memoryRepository) FailExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string) error {
+	if attemptID == "" {
+		return nil
+	}
+	return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var subject types.MemorySubject
+		if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID).
+			Clauses(forUpdateClause()).First(&subject).Error; err != nil {
+			return err
+		}
+		if subject.ExtractAttemptID != attemptID {
+			return nil
+		}
+		return tx.Model(&types.MemorySubject{}).Where("id = ?", subject.ID).
+			Updates(failedExtractionUpdates(&subject, time.Now())).Error
+	})
+}
+
+func (r *memoryRepository) ListDueExtractions(ctx context.Context, limit int) ([]*types.MemorySubject, error) {
+	if limit < 1 || limit > 32 {
+		limit = 32
+	}
+	now := time.Now()
+	var subjects []*types.MemorySubject
+	err := r.db.WithContext(ctx).Where("enabled = ? AND pending_sessions IS NOT NULL AND pending_sessions <> ?", true, "[]").
+		Where("extract_failures < ?", types.MaxMemoryExtractionFailures).
+		Where("extract_next_at IS NULL OR extract_next_at <= ?", now).
+		Where("extract_lease_until IS NULL OR extract_lease_until <= ?", now).
+		// Give tasks queued by the previous binary time to finish during rollout.
+		Where("extract_attempt_id <> '' OR extract_scheduled_at IS NULL OR extract_scheduled_at <= ?", now.Add(-10*time.Minute)).
+		Order("extract_next_at ASC NULLS FIRST, updated_at ASC, id ASC").Limit(limit).Find(&subjects).Error
+	return subjects, err
+}
+
+func (r *memoryRepository) DeferExtraction(ctx context.Context, scope interfaces.MemoryScope, delay time.Duration) error {
+	return r.scoped(ctx, scope).Model(&types.MemorySubject{}).
+		Where("extract_attempt_id = ''").
+		Updates(map[string]interface{}{"extract_next_at": time.Now().Add(delay), "extract_scheduled_at": nil}).Error
+}
+
+func (r *memoryRepository) PauseExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string, delay time.Duration) error {
+	return r.scoped(ctx, scope).Model(&types.MemorySubject{}).Where("extract_attempt_id = ?", attemptID).
+		Updates(map[string]interface{}{
+			"extract_attempt_id": "", "extract_lease_until": nil, "extract_started_at": nil,
+			"extract_scheduled_at": nil, "extract_next_at": time.Now().Add(delay), "updated_at": time.Now(),
+		}).Error
+}
+
+func (r *memoryRepository) ResetExhaustedExtraction(ctx context.Context, scope interfaces.MemoryScope) (bool, error) {
+	result := r.scoped(ctx, scope).Model(&types.MemorySubject{}).
+		Where("enabled = ? AND extract_failures >= ? AND extract_attempt_id = ''", true, types.MaxMemoryExtractionFailures).
+		Updates(map[string]interface{}{"extract_failures": 0, "extract_next_at": nil, "extract_scheduled_at": nil, "updated_at": time.Now()})
+	return result.RowsAffected == 1, result.Error
+}