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

internal/router/sync_task.go

@@ -0,0 +1,232 @@
+package router
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/google/uuid"
+	"github.com/hibiken/asynq"
+	"go.uber.org/dig"
+)
+
+// SyncTaskExecutor executes tasks synchronously (in a goroutine) without Redis.
+// Used in Lite mode as a drop-in replacement for *asynq.Client.
+type SyncTaskExecutor struct {
+	mu              sync.RWMutex
+	handlers        map[string]func(context.Context, *asynq.Task) error
+	taskIDs         map[string]time.Time
+	activeKnowledge map[string]int
+}
+
+func NewSyncTaskExecutor() *SyncTaskExecutor {
+	return &SyncTaskExecutor{
+		handlers:        make(map[string]func(context.Context, *asynq.Task) error),
+		taskIDs:         make(map[string]time.Time),
+		activeKnowledge: make(map[string]int),
+	}
+}
+
+// RegisterHandler registers a handler for a given task type pattern.
+func (e *SyncTaskExecutor) RegisterHandler(pattern string, handler func(context.Context, *asynq.Task) error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.handlers[pattern] = handler
+}
+
+// HasActiveTaskID includes execution and retry backoff in Lite mode, where the
+// queue inspector has no external backend to inspect.
+func (e *SyncTaskExecutor) HasActiveTaskID(id string) bool {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	expiry, exists := e.taskIDs[id]
+	return exists && expiry.IsZero()
+}
+
+func (e *SyncTaskExecutor) HasActiveKnowledgeTask(id string) bool {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+	return e.activeKnowledge[id] > 0
+}
+
+// Enqueue satisfies interfaces.TaskEnqueuer.
+// Instead of queuing to Redis, it dispatches the task to a goroutine.
+// Supports ProcessIn (delay) and MaxRetry options for parity with asynq.
+func (e *SyncTaskExecutor) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
+	e.mu.RLock()
+	handler, ok := e.handlers[task.Type()]
+	e.mu.RUnlock()
+
+	if !ok {
+		return nil, fmt.Errorf("sync task executor: no handler registered for type %q", task.Type())
+	}
+
+	var delay time.Duration
+	var taskID string
+	var retention time.Duration
+	maxRetry := 25 // asynq default
+	maxRetrySet := false
+	for _, opt := range opts {
+		switch opt.Type() {
+		case asynq.TaskIDOpt:
+			taskID, _ = opt.Value().(string)
+		case asynq.RetentionOpt:
+			retention, _ = opt.Value().(time.Duration)
+		case asynq.ProcessInOpt:
+			if d, ok := opt.Value().(time.Duration); ok {
+				delay = d
+			}
+		case asynq.MaxRetryOpt:
+			if n, ok := opt.Value().(int); ok {
+				maxRetry = n
+				maxRetrySet = true
+			}
+		}
+	}
+	// Callers that explicitly pass MaxRetry(0) want no retries.
+	// Without the flag we can't distinguish "not set" from "set to 0".
+	if maxRetrySet && maxRetry < 0 {
+		maxRetry = 0
+	}
+
+	reservedID := taskID != ""
+	if reservedID {
+		e.mu.Lock()
+		now := time.Now()
+		for id, expiry := range e.taskIDs {
+			if !expiry.IsZero() && expiry.Before(now) {
+				delete(e.taskIDs, id)
+			}
+		}
+		if _, exists := e.taskIDs[taskID]; exists {
+			e.mu.Unlock()
+			return nil, asynq.ErrTaskIDConflict
+		}
+		e.taskIDs[taskID] = time.Time{}
+		e.mu.Unlock()
+	} else {
+		taskID = uuid.New().String()
+	}
+	info := &asynq.TaskInfo{
+		ID:    taskID,
+		Queue: "sync",
+		Type:  task.Type(),
+	}
+	var payload struct {
+		KnowledgeID string `json:"knowledge_id"`
+	}
+	_ = json.Unmarshal(task.Payload(), &payload)
+	if payload.KnowledgeID != "" {
+		e.mu.Lock()
+		e.activeKnowledge[payload.KnowledgeID]++
+		e.mu.Unlock()
+	}
+
+	go func() {
+		defer func() {
+			if payload.KnowledgeID != "" {
+				e.mu.Lock()
+				e.activeKnowledge[payload.KnowledgeID]--
+				if e.activeKnowledge[payload.KnowledgeID] == 0 {
+					delete(e.activeKnowledge, payload.KnowledgeID)
+				}
+				e.mu.Unlock()
+			}
+			if reservedID {
+				e.mu.Lock()
+				if retention > 0 {
+					e.taskIDs[taskID] = time.Now().Add(retention)
+				} else {
+					delete(e.taskIDs, taskID)
+				}
+				e.mu.Unlock()
+			}
+		}()
+		if delay > 0 {
+			time.Sleep(delay)
+		}
+
+		// Tag as a background worker execution so the per-model concurrency
+		// governor throttles Lite-mode ingestion/enrichment LLM calls, mirroring
+		// the asynq backgroundTaskMiddleware in the Redis path.
+		ctx := types.WithBackgroundTask(context.Background())
+		start := time.Now()
+		logger.Infof(ctx, "[SyncTask] Executing task type=%s id=%s", task.Type(), taskID)
+
+		var lastErr error
+		for attempt := 0; attempt <= maxRetry; attempt++ {
+			if attempt > 0 {
+				backoff := time.Duration(attempt) * 5 * time.Second
+				if backoff > 30*time.Second {
+					backoff = 30 * time.Second
+				}
+				logger.Infof(ctx, "[SyncTask] Retrying task type=%s id=%s attempt=%d/%d backoff=%s",
+					task.Type(), taskID, attempt, maxRetry, backoff)
+				time.Sleep(backoff)
+			}
+
+			attemptCtx := types.WithTaskRetryMetadata(ctx, attempt, maxRetry)
+			lastErr = handler(attemptCtx, task)
+			if lastErr == nil {
+				logger.Infof(ctx, "[SyncTask] Task completed type=%s id=%s elapsed=%v",
+					task.Type(), taskID, time.Since(start))
+				return
+			}
+		}
+
+		logger.Errorf(ctx, "[SyncTask] Task failed (exhausted retries) type=%s id=%s elapsed=%v err=%v",
+			task.Type(), taskID, time.Since(start), lastErr)
+	}()
+
+	return info, nil
+}
+
+type SyncTaskParams struct {
+	dig.In
+
+	Executor             *SyncTaskExecutor
+	KnowledgeService     interfaces.KnowledgeService
+	KnowledgeBaseService interfaces.KnowledgeBaseService
+	TagService           interfaces.KnowledgeTagService
+	DataSourceService    interfaces.DataSourceService
+	ChunkExtractor       interfaces.TaskHandler `name:"chunkExtractor"`
+	DataTableSummary     interfaces.TaskHandler `name:"dataTableSummary"`
+	ImageMultimodal      interfaces.TaskHandler `name:"imageMultimodal"`
+	KnowledgePostProcess interfaces.TaskHandler `name:"knowledgePostProcess"`
+	KnowledgeAutoTag     interfaces.TaskHandler `name:"knowledgeAutoTag"`
+	WikiIngest           interfaces.TaskHandler `name:"wikiIngest"`
+	TemporaryDocument    interfaces.TemporaryDocumentService
+	MemoryService        interfaces.MemoryService
+}
+
+// RegisterSyncHandlers registers all task handlers on the SyncTaskExecutor.
+// Used in Lite mode instead of RunAsynqServer.
+func RegisterSyncHandlers(params SyncTaskParams) {
+	params.Executor.RegisterHandler(types.TypeChunkExtract, params.ChunkExtractor.Handle)
+	params.Executor.RegisterHandler(types.TypeDataTableSummary, params.DataTableSummary.Handle)
+	params.Executor.RegisterHandler(types.TypeDocumentProcess, params.KnowledgeService.ProcessDocument)
+	params.Executor.RegisterHandler(types.TypeTemporaryDocumentProcess, params.TemporaryDocument.Process)
+	params.Executor.RegisterHandler(types.TypeManualProcess, params.KnowledgeService.ProcessManualUpdate)
+	params.Executor.RegisterHandler(types.TypeFAQImport, params.KnowledgeService.ProcessFAQImport)
+	params.Executor.RegisterHandler(types.TypeQuestionGeneration, params.KnowledgeService.ProcessQuestionGeneration)
+	params.Executor.RegisterHandler(types.TypeSummaryGeneration, params.KnowledgeService.ProcessSummaryGeneration)
+	params.Executor.RegisterHandler(types.TypeKBClone, params.KnowledgeService.ProcessKBClone)
+	params.Executor.RegisterHandler(types.TypeKnowledgeMove, params.KnowledgeService.ProcessKnowledgeMove)
+	params.Executor.RegisterHandler(types.TypeKnowledgeListDelete, params.KnowledgeService.ProcessKnowledgeListDelete)
+	params.Executor.RegisterHandler(types.TypeKnowledgeListReparse, params.KnowledgeService.ProcessKnowledgeListReparse)
+	params.Executor.RegisterHandler(types.TypeIndexDelete, params.TagService.ProcessIndexDelete)
+	params.Executor.RegisterHandler(types.TypeKBDelete, params.KnowledgeBaseService.ProcessKBDelete)
+	params.Executor.RegisterHandler(types.TypeImageMultimodal, params.ImageMultimodal.Handle)
+	params.Executor.RegisterHandler(types.TypeKnowledgePostProcess, params.KnowledgePostProcess.Handle)
+	params.Executor.RegisterHandler(types.TypeKnowledgeAutoTag, params.KnowledgeAutoTag.Handle)
+	params.Executor.RegisterHandler(types.TypeDataSourceSync, params.DataSourceService.ProcessSync)
+	params.Executor.RegisterHandler(types.TypeWikiIngest, params.WikiIngest.Handle)
+	params.Executor.RegisterHandler(types.TypeWikiFinalize, params.WikiIngest.Handle)
+	params.Executor.RegisterHandler(types.TypeMemoryExtract, params.MemoryService.Handle)
+	logger.Infof(context.Background(), "[SyncTask] All task handlers registered (Lite mode, no Redis)")
+}