xinmi/internal/application/service/memory/recovery.go

package memory

import (
	"context"
	"time"

	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

// RecoverExtractions wakes only durable pending scopes. It never enumerates
// conversation history or changes a session checkpoint.
func (s *Service) RecoverExtractions(ctx context.Context) error {
	subjects, err := s.repo.ListDueExtractions(ctx, 32)
	if err != nil {
		return err
	}
	for _, subject := range subjects {
		if err := ctx.Err(); err != nil {
			return err
		}
		scope := interfaces.MemoryScope{TenantID: subject.TenantID, SubjectID: subject.SubjectID}
		cfg := s.workspaceConfig(ctx, scope.TenantID)
		if !cfg.AutoExtractEnabled() {
			// Release expired task identities before deferring disabled scopes,
			// so they do not monopolize the bounded recovery scan.
			if err := s.repo.PauseExtraction(ctx, scope, subject.ExtractAttemptID, 15*time.Minute); err != nil {
				return err
			}
			continue
		}
		delay := extractFollowUpDelay
		if subject.LastExtractedAt != nil {
			if remaining := cfg.ExtractMinInterval() - time.Since(*subject.LastExtractedAt); remaining > delay {
				delay = remaining
			}
		}
		recoveryCtx := context.WithValue(ctx, types.TenantIDContextKey, scope.TenantID)
		if subject.ExtractLanguage != "" {
			recoveryCtx = context.WithValue(recoveryCtx, types.LanguageContextKey, subject.ExtractLanguage)
		}
		if err := s.enqueueExtraction(recoveryCtx, scope, "", "", subject.ExtractChatModelID, delay); err != nil {
			logger.Warnf(ctx, "memory: recover pending extraction failed: %v", err)
		}
	}
	return nil
}

func (s *Service) RetryExtraction(ctx context.Context) error {
	scope, cfg, enabled := s.enabledScope(ctx)
	if !enabled || !cfg.AutoExtractEnabled() {
		return ErrMemoryDisabled
	}
	reset, err := s.repo.ResetExhaustedExtraction(ctx, scope)
	if err != nil || !reset {
		return err
	}
	return s.enqueueExtraction(ctx, scope, "", "", "", extractFollowUpDelay)
}

// StartExtractionRecovery is called after both Redis and Lite register the
// handler. The returned stop waits for the bounded sweep to exit.
func (s *Service) StartExtractionRecovery(parent context.Context) func() {
	ctx, cancel := context.WithCancel(parent)
	done := make(chan struct{})
	go func() {
		defer close(done)
		ticker := time.NewTicker(time.Minute)
		defer ticker.Stop()
		for {
			sweepCtx, sweepCancel := context.WithTimeout(ctx, 30*time.Second)
			if err := s.RecoverExtractions(sweepCtx); err != nil && ctx.Err() == nil {
				logger.Warnf(ctx, "memory: extraction recovery sweep failed: %v", err)
			}
			sweepCancel()
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
			}
		}
	}()
	return func() { cancel(); <-done }
}