feat(tests): enhance image multimodal service tests and runtime task handling

Added new tests for the ImageMultimodalService to verify the handling of orphan tasks and the finalization of pending counters. Improved runtime task inspection with better error handling and task cancellation logic, ensuring robust management of orphan tasks in the system.

internal/application/service/image_multimodal.go

@@ -151,6 +151,9 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e
 		logger.Infof(ctx,
 			"[ImageMultimodal] Dropping task chunk=%s knowledge=%s kb=%s image=%s",
 			payload.ChunkID, payload.KnowledgeID, payload.KnowledgeBaseID, payload.ImageURL)
+		// Still count this image toward the parent finalize gate so a batch
+		// of dropped orphans cannot strand multimodal:pending forever.
+		s.checkAndFinalizeAllImages(ctx, payload)
 		return nil
 	}
 

internal/application/service/image_multimodal_orphan_test.go

@@ -9,7 +9,9 @@ import (
 	"github.com/Tencent/WeKnora/internal/application/repository"
 	"github.com/Tencent/WeKnora/internal/types"
 	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/alicebob/miniredis/v2"
 	"github.com/hibiken/asynq"
+	"github.com/redis/go-redis/v9"
 )
 
 type orphanKnowledgeRepo struct {
@@ -96,6 +98,58 @@ func TestImageMultimodalHandleDropsMissingKnowledge(t *testing.T) {
 	}
 }
 
+type orphanTaskEnqueuer struct {
+	interfaces.TaskEnqueuer
+	enqueued []*asynq.Task
+}
+
+func (e *orphanTaskEnqueuer) Enqueue(task *asynq.Task, _ ...asynq.Option) (*asynq.TaskInfo, error) {
+	e.enqueued = append(e.enqueued, task)
+	return &asynq.TaskInfo{ID: "post-process"}, nil
+}
+
+func TestImageMultimodalHandleDropFinalizesPendingCounter(t *testing.T) {
+	t.Parallel()
+	mr, err := miniredis.Run()
+	if err != nil {
+		t.Fatalf("start miniredis: %v", err)
+	}
+	t.Cleanup(mr.Close)
+	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	t.Cleanup(func() { _ = rdb.Close() })
+
+	enqueuer := &orphanTaskEnqueuer{}
+	svc := &ImageMultimodalService{
+		knowledgeRepo: &orphanKnowledgeRepo{err: repository.ErrKnowledgeNotFound},
+		kbService:     &orphanKBService{kb: &types.KnowledgeBase{ID: "kb-1"}},
+		redisClient:   rdb,
+		taskEnqueuer:  enqueuer,
+	}
+	const knowledgeID = "missing"
+	redisKey := "multimodal:pending:" + knowledgeID
+	if err := rdb.Set(context.Background(), redisKey, 1, 0).Err(); err != nil {
+		t.Fatal(err)
+	}
+	payload, err := json.Marshal(types.ImageMultimodalPayload{
+		TenantID:        1,
+		KnowledgeID:     knowledgeID,
+		KnowledgeBaseID: "kb-1",
+		ImageURL:        "minio://bucket/img.png",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := svc.Handle(context.Background(), asynq.NewTask(types.TypeImageMultimodal, payload)); err != nil {
+		t.Fatalf("orphan drop should succeed: %v", err)
+	}
+	if mr.Exists(redisKey) {
+		t.Fatal("pending counter should be cleared after drop finalize")
+	}
+	if len(enqueuer.enqueued) != 1 || enqueuer.enqueued[0].Type() != types.TypeKnowledgePostProcess {
+		t.Fatalf("expected post-process enqueue, got %d tasks", len(enqueuer.enqueued))
+	}
+}
+
 func TestImageMultimodalHandlePropagatesTransientKnowledgeError(t *testing.T) {
 	t.Parallel()
 	dbErr := errors.New("db unavailable")

internal/handler/system.go

@@ -1827,29 +1827,31 @@ func isKnowledgeGone(err error) bool {
 }
 
 // purgeOrphanRuntimeTask removes queue records for a task whose knowledge row
-// is already gone. Best-effort: CancelTasksForKnowledge sweeps every queue
-// state; ForceDeleteRuntimeTask is the fallback for the specific task ID.
+// is already gone. CancelTasksForKnowledge sweeps sibling tasks best-effort;
+// success is reported only once the requested task ID is gone.
 func (h *SystemHandler) purgeOrphanRuntimeTask(
 	ctx context.Context,
 	inspector interfaces.RuntimeTaskInspector,
 	queue, taskID, knowledgeID string,
 ) bool {
 	if h.taskInspector != nil && knowledgeID != "" {
-		deleted, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID)
-		if err != nil {
+		if _, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID); err != nil {
 			logger.Warnf(ctx, "purge orphan tasks for knowledge %s: %v", knowledgeID, err)
-		} else if deleted > 0 {
-			return true
 		}
 	}
 	supported, err := inspector.ForceDeleteRuntimeTask(ctx, queue, taskID)
-	if !supported || err != nil {
-		if err != nil {
-			logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err)
-		}
+	if !supported {
 		return false
 	}
-	return true
+	if err == nil {
+		return true
+	}
+	// The knowledge-wide sweep may have already removed this task ID.
+	if _, available, getErr := inspector.GetRuntimeTask(ctx, queue, taskID); available && getErr != nil {
+		return true
+	}
+	logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err)
+	return false
 }
 
 // MutateRuntimeTask executes a backend-advertised action after re-reading the

internal/handler/system_runtime_test.go

@@ -3,6 +3,7 @@ package handler
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -85,6 +86,10 @@ type runtimeTaskTestInspector struct {
 	cancelKnowledge string
 	cancelDeleted   int
 	mutatedQueue    string
+	forceDeleteErr      error
+	getRuntimeErr       error
+	getRuntimeErrFrom   int
+	getRuntimeTaskCalls int
 }
 
 func (r *runtimeTaskTestInspector) CancelTasksForKnowledge(
@@ -134,6 +139,10 @@ func (r *runtimeTaskTestInspector) ListRuntimeTasks(
 func (r *runtimeTaskTestInspector) GetRuntimeTask(
 	_ context.Context, queue, taskID string,
 ) (*types.RuntimeTaskInfo, bool, error) {
+	r.getRuntimeTaskCalls++
+	if r.getRuntimeErr != nil && r.getRuntimeTaskCalls >= r.getRuntimeErrFrom {
+		return nil, true, r.getRuntimeErr
+	}
 	for i := range r.tasks {
 		if r.tasks[i].ID == taskID {
 			return &r.tasks[i], true, nil
@@ -169,7 +178,7 @@ func (r *runtimeTaskTestInspector) ForceDeleteRuntimeTask(
 ) (bool, error) {
 	r.mutatedQueue = queue
 	r.forceDeleted = taskID
-	return true, nil
+	return true, r.forceDeleteErr
 }
 
 func TestGetRuntimeQueuesReportsIsolatedPoolCapacity(t *testing.T) {
@@ -401,3 +410,67 @@ func TestRuntimeTaskCancelPurgesOrphanWhenKnowledgeGone(t *testing.T) {
 			inspector.deletedTask, inspector.forceDeleted, inspector.cancelKnowledge)
 	}
 }
+
+func TestRuntimeTaskCancelPurgesOrphanAfterSiblingSweep(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	inspector := &runtimeTaskTestInspector{
+		cancelDeleted: 2,
+		tasks: []types.RuntimeTaskInfo{{
+			ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal,
+			State: types.RuntimeTaskArchived, TenantID: 42, KnowledgeID: "knowledge-gone",
+			AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel},
+		}},
+	}
+	canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound}
+	handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller}
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Params = gin.Params{
+		{Key: "queue", Value: types.QueueMultimodal},
+		{Key: "task_id", Value: "task-orphan"},
+		{Key: "action", Value: string(types.RuntimeTaskActionCancel)},
+	}
+	ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil)
+
+	handler.MutateRuntimeTask(ctx)
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String())
+	}
+	if inspector.cancelKnowledge != "knowledge-gone" || inspector.forceDeleted != "task-orphan" {
+		t.Fatalf("expected sibling sweep then force delete, got cancel=%q force=%q",
+			inspector.cancelKnowledge, inspector.forceDeleted)
+	}
+}
+
+func TestRuntimeTaskCancelPurgesOrphanWhenSweepAlreadyRemovedTask(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	inspector := &runtimeTaskTestInspector{
+		cancelDeleted:     1,
+		forceDeleteErr:    errors.New("task not found"),
+		getRuntimeErr:     errors.New("task not found"),
+		getRuntimeErrFrom: 2,
+		tasks: []types.RuntimeTaskInfo{{
+			ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal,
+			State: types.RuntimeTaskRetry, TenantID: 42, KnowledgeID: "knowledge-gone",
+			AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel},
+		}},
+	}
+	canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound}
+	handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller}
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Params = gin.Params{
+		{Key: "queue", Value: types.QueueMultimodal},
+		{Key: "task_id", Value: "task-orphan"},
+		{Key: "action", Value: string(types.RuntimeTaskActionCancel)},
+	}
+	ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil)
+
+	handler.MutateRuntimeTask(ctx)
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String())
+	}
+	if inspector.forceDeleted != "task-orphan" {
+		t.Fatalf("expected force delete attempt, got %q", inspector.forceDeleted)
+	}
+}