xinmi/internal/application/repository/memory_recovery.go

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
}