refactor(ingestion): use worker-driven task pull dispatcher (#19314)

Author: jay77721

cmd/ragflow_server.go

@@ -593,7 +593,7 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg
 	// Search and UpsertDoc can embed queries/summaries automatically.
 	nav.SetNavService(nlp.NewNavService(service.NewNavEmbedder(service.NewModelProviderService(), "")))
 	// Memory extraction runs on the Ingestor's shared NATS consumer + worker
-	// pool (task_type="memory" dispatched by processMessage -> executeMemoryTask),
+	// pool (task_type="memory" dispatched by handleAndExecute -> executeMemoryTask),
 	// so there is no longer a dedicated Redis memory consumer to start.
 	ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService()))
 

internal/admin/handler.go

@@ -17,6 +17,7 @@
 package admin
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -923,7 +924,7 @@ func (h *Handler) PublishMessageToQueue(c *gin.Context) {
 	}
 
 	msgQueueEngine := engine.GetMessageQueueEngine()
-	err = msgQueueEngine.PublishTask("tasks.RAGFLOW", taskMessageStr)
+	err = msgQueueEngine.PublishTask(common.TaskSubject, taskMessageStr)
 	if err != nil {
 		common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
 		return
@@ -933,7 +934,7 @@ func (h *Handler) PublishMessageToQueue(c *gin.Context) {
 }
 
 type PullMessageFromQueueRequest struct {
-	MessageCount int    `json:"message_count" binding:"required"`
+	MessageCount int    `json:"message_count" binding:"required,gt=0"`
 	AckPolicy    string `json:"ack_policy" binding:"required"`
 }
 
@@ -943,14 +944,20 @@ func (h *Handler) PullMessageFromQueue(c *gin.Context) {
 		common.ErrorWithCode(c, common.CodeBadRequest, fmt.Sprintf("Message count error: %s", err.Error()))
 		return
 	}
+	if req.MessageCount > common.MaxManualPullMessages {
+		common.ErrorWithCode(c, common.CodeBadRequest,
+			fmt.Sprintf("message count must be between 1 and %d", common.MaxManualPullMessages))
+		return
+	}
 
 	msgQueueEngine := engine.GetMessageQueueEngine()
-	err := msgQueueEngine.InitConsumer("tasks.RAGFLOW")
-	if err != nil {
+	if err := msgQueueEngine.InitConsumer(common.TaskSubject); err != nil {
 		common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
 		return
 	}
-	messages, err := msgQueueEngine.GetMessages(req.MessageCount)
+	pullCtx, cancel := context.WithTimeout(c.Request.Context(), time.Second)
+	defer cancel()
+	messages, err := msgQueueEngine.PullMessages(pullCtx, req.MessageCount)
 	if err != nil {
 		common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
 		return
@@ -978,6 +985,7 @@ func (h *Handler) PullMessageFromQueue(c *gin.Context) {
 				"id":   taskMessage.TaskID,
 				"type": taskMessage.TaskType,
 			}
+			err = message.Nack()
 			if err == nil {
 				resultMessage["nack"] = "true"
 			} else {

internal/admin/handler_queue_test.go

@@ -0,0 +1,118 @@
+//
+// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package admin
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"ragflow/internal/common"
+	"ragflow/internal/engine"
+	"ragflow/internal/ingestion/testutil"
+
+	"github.com/gin-gonic/gin"
+)
+
+// TestPullMessageFromQueueInitializesSharedConsumer preserves the admin
+// endpoint's ability to pull messages before the ingestor starts. The endpoint
+// must initialize the existing durable consumer rather than creating a
+// dispatcher-specific consumer.
+func TestPullMessageFromQueueInitializesSharedConsumer(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	queue := testutil.SetupNatsEngine(t)
+	previousQueue := engine.GetMessageQueueEngine()
+	engine.SetMessageQueueEngine(queue)
+	t.Cleanup(func() { engine.SetMessageQueueEngine(previousQueue) })
+
+	payload, err := json.Marshal(common.TaskMessage{
+		TaskID:   "admin-pull-before-ingestor",
+		TaskType: common.TaskTypeIngestionTask,
+	})
+	if err != nil {
+		t.Fatalf("marshal task: %v", err)
+	}
+	if err := queue.PublishTask(common.TaskSubject, payload); err != nil {
+		t.Fatalf("publish task: %v", err)
+	}
+
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = httptest.NewRequest(
+		http.MethodPost,
+		"/",
+		bytes.NewBufferString(`{"message_count":1,"ack_policy":"ACK"}`),
+	)
+	ctx.Request.Header.Set("Content-Type", "application/json")
+
+	(&Handler{}).PullMessageFromQueue(ctx)
+
+	var response struct {
+		Code int `json:"code"`
+		Data []struct {
+			ID  string `json:"id"`
+			Ack string `json:"ack"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
+		t.Fatalf("decode response: %v", err)
+	}
+	if response.Code != int(common.CodeSuccess) {
+		t.Fatalf("response code = %d, want %d; body = %s", response.Code, common.CodeSuccess, recorder.Body.String())
+	}
+	if len(response.Data) != 1 || response.Data[0].ID != "admin-pull-before-ingestor" || response.Data[0].Ack != "true" {
+		t.Fatalf("pull response = %+v, want acked admin-pull-before-ingestor", response.Data)
+	}
+}
+
+func TestPullMessageFromQueueRejectsOutOfRangeMessageCount(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	for _, testCase := range []struct {
+		name         string
+		messageCount int
+	}{
+		{name: "negative", messageCount: -1},
+		{name: "zero", messageCount: 0},
+		{name: "above limit", messageCount: 101},
+	} {
+		t.Run(testCase.name, func(t *testing.T) {
+			recorder := httptest.NewRecorder()
+			ctx, _ := gin.CreateTestContext(recorder)
+			ctx.Request = httptest.NewRequest(
+				http.MethodPost,
+				"/",
+				bytes.NewBufferString(fmt.Sprintf(`{"message_count":%d,"ack_policy":"ACK"}`, testCase.messageCount)),
+			)
+			ctx.Request.Header.Set("Content-Type", "application/json")
+
+			(&Handler{}).PullMessageFromQueue(ctx)
+
+			var response struct {
+				Code int `json:"code"`
+			}
+			if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
+				t.Fatalf("decode response: %v", err)
+			}
+			if response.Code != int(common.CodeBadRequest) {
+				t.Fatalf("response code = %d, want %d; body = %s", response.Code, common.CodeBadRequest, recorder.Body.String())
+			}
+		})
+	}
+}

internal/cli/admin_parser.go

@@ -2714,8 +2714,8 @@ func (p *Parser) parseMessageQueueCommand() (*Command, error) {
 			p.nextToken() // consume NUMBER
 		}
 
-		if messageCount <= 0 || messageCount > 100 {
-			return nil, fmt.Errorf("message count cannot be less than 0 or greater than 100")
+		if messageCount <= 0 || messageCount > common.MaxManualPullMessages {
+			return nil, fmt.Errorf("message count must be between 1 and %d", common.MaxManualPullMessages)
 		}
 
 		cmd = NewCommand("user_pull_message_command")

internal/common/task.go

@@ -24,14 +24,17 @@ const (
 	// symbol so the routing contract cannot diverge (mirrors the RAGFLOW_TASKS
 	// JetStream subject in internal/engine/nats).
 	TaskSubject = "tasks.RAGFLOW"
+	// MaxManualPullMessages is the largest task batch the administrative queue
+	// pull endpoint accepts.
+	MaxManualPullMessages = 100
 
 	TaskTypeIngestionTask = "ingestion_task"
 	TaskTypeIngestionTest = "ingestion_test"
 	// TaskTypeSyncer is the NATS wake-up message type for datasource sync_logs tasks.
 	TaskTypeSyncer = "syncer"
 	// TaskTypeMemory is the async memory-extraction task type. Memory tasks
 	// share the tasks.RAGFLOW subject and the Ingestor's consumer + worker
-	// pool with ingestion tasks; processMessage dispatches them by TaskType.
+	// pool with ingestion tasks; handleAndExecute dispatches them by TaskType.
 	// The memory-specific payload (message_dict/memory_id/source_id) is
 	// carried in TaskMessage.Payload.
 	TaskTypeMemory = "memory"

internal/engine/engine.go

@@ -113,7 +113,8 @@ type MessageQueue interface {
 	Type() string
 	InitConsumer(subject string) error
 	PublishTask(subject string, payload []byte) error
-	GetMessages(messageCount int) ([]common.TaskHandle, error)
+	PullMessages(ctx context.Context, messageCount int) ([]common.TaskHandle, error)
+	PullMessage(ctx context.Context) (common.TaskHandle, error)
 	ListMessages(messageType string, pending bool) ([]map[string]string, error)
 	ShowMessageQueue() (map[string]string, error)
 	CheckStatus() string

internal/engine/global.go

@@ -111,7 +111,10 @@ func InitMessageQueue() error {
 	switch messageQueueType {
 	case "nats":
 		natsConfig := globalConfig.GetNATSConfig()
-		messageQueueEngine = nats.NewNatsEngine(natsConfig.Host, natsConfig.Port)
+		messageQueueEngine = nats.NewNatsEngine(
+			natsConfig.Host,
+			natsConfig.Port,
+		)
 		err := messageQueueEngine.Init()
 		if err != nil {
 			return err

internal/engine/nats/nats.go

@@ -121,7 +121,7 @@ func (n *NatsEngine) PublishTask(subject string, payload []byte) error {
 	// Duplicate delivery is instead made safe at the consumer level:
 	// StartRunning's CREATED/SCHEDULED→RUNNING CAS plus the in-process claim guard
 	// prevent a second copy from executing while the first owner is active (see
-	// Ingestor.processMessage).
+	// Ingestor.handleAndExecute).
 	ack, err := n.jetStream.Publish(ctx, subject, payload)
 	if err != nil {
 		return err
@@ -256,22 +256,54 @@ func (n *NatsEngine) InitConsumer(subject string) error {
 	}
 	return nil
 }
-func (n *NatsEngine) GetMessages(messageCount int) ([]common.TaskHandle, error) {
+
+// PullMessages fetches up to messageCount messages before ctx expires.
+func (n *NatsEngine) PullMessages(ctx context.Context, messageCount int) ([]common.TaskHandle, error) {
+	if messageCount < 1 || messageCount > common.MaxManualPullMessages {
+		return nil, fmt.Errorf("message count must be between 1 and %d", common.MaxManualPullMessages)
+	}
 	if n.consumer == nil {
 		return nil, errors.New("NATS consumer is nil, engine not properly initialized")
 	}
+	if _, ok := ctx.Deadline(); !ok {
+		return nil, errors.New("pull messages context must have a deadline")
+	}
 
-	resultMessages := make([]common.TaskHandle, 0)
-	messages, err := n.consumer.Fetch(messageCount, jetstream.FetchMaxWait(1*time.Second))
+	resultMessages := make([]common.TaskHandle, 0, messageCount)
+	messages, err := n.consumer.Fetch(messageCount, jetstream.FetchContext(ctx))
 	if err != nil {
 		return nil, fmt.Errorf("failed to fetch messages: %w", err)
 	}
-	for msg := range messages.Messages() {
-		resultMessages = append(resultMessages, NewNatsMessageHandle(msg))
+	for message := range messages.Messages() {
+		resultMessages = append(resultMessages, NewNatsMessageHandle(message))
+	}
+	if batchErr := messages.Error(); batchErr != nil {
+		if errors.Is(batchErr, context.DeadlineExceeded) {
+			return resultMessages, nil
+		}
+		for _, message := range resultMessages {
+			if nackErr := message.Nack(); nackErr != nil {
+				common.Error("nack message after failed pull", nackErr)
+			}
+		}
+		return nil, fmt.Errorf("failed to fetch messages: %w", batchErr)
 	}
 	return resultMessages, nil
 }
 
+// PullMessage returns one task handle from PullMessages. A nil handle
+// with a nil error means the pull expired without an available task.
+func (n *NatsEngine) PullMessage(ctx context.Context) (common.TaskHandle, error) {
+	messages, err := n.PullMessages(ctx, 1)
+	if err != nil {
+		return nil, err
+	}
+	if len(messages) == 0 {
+		return nil, nil
+	}
+	return messages[0], nil
+}
+
 func (n *NatsEngine) CheckStatus() string {
 	if n.nc == nil {
 		return "NATS connection is nil, engine not properly initialized"

internal/engine/nats/nats_uninitialized_test.go

@@ -19,8 +19,8 @@ func TestUninitializedEngineMethodsReturnErrors(t *testing.T) {
 		t.Fatalf("PublishTask on uninitialized engine: err = %v, want 'not properly initialized'", err)
 	}
 
-	if _, err := e.GetMessages(1); err == nil || !strings.Contains(err.Error(), "not properly initialized") {
-		t.Fatalf("GetMessages on uninitialized engine: err = %v, want 'not properly initialized'", err)
+	if _, err := e.PullMessages(t.Context(), 1); err == nil || !strings.Contains(err.Error(), "not properly initialized") {
+		t.Fatalf("PullMessages on uninitialized engine: err = %v, want 'not properly initialized'", err)
 	}
 
 	if _, err := e.ShowMessageQueue(); err == nil || !strings.Contains(err.Error(), "not properly initialized") {