fix(ingestion): make document/KB counter application idempotent per run (#17995)

internal/ingestion/service/doc_state.go

@@ -33,7 +33,7 @@ import (
 type docStateSvc interface {
 	GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]any, error)
 	SetDocumentMetadata(ctx context.Context, docID string, meta map[string]any) error
-	IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error
+	ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error
 }
 
 // docStateUpdater applies a pipeline run's results to document state: it
@@ -61,8 +61,8 @@ func (u *docStateUpdater) apply(ctx context.Context, r *taskpkg.PipelineResult)
 			common.Warn(fmt.Sprintf("failed to update document metadata: %v", err))
 		}
 	}
-	if err := u.docSvc.IncrementChunkNum(ctx, r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil {
-		common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err))
+	if err := u.docSvc.ApplyDocCounts(ctx, r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil {
+		common.Warn(fmt.Sprintf("failed to apply doc counts: %v", err))
 	}
 }
 

internal/ingestion/service/doc_state_test.go

@@ -52,7 +52,7 @@ func (s *stubDocStateSvc) SetDocumentMetadata(ctx context.Context, docID string,
 	return nil
 }
 
-func (s *stubDocStateSvc) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error {
+func (s *stubDocStateSvc) ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error {
 	s.incrementCalled = true
 	s.gotDocID = docID
 	s.gotKbID = kbID

internal/ingestion/service/redelivery_counter_test.go

@@ -0,0 +1,131 @@
+//
+//  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 service
+
+import (
+	"context"
+	"testing"
+
+	"gorm.io/gorm"
+
+	"ragflow/internal/common"
+	"ragflow/internal/entity"
+	taskpkg "ragflow/internal/ingestion/task"
+	"ragflow/internal/ingestion/testutil"
+)
+
+// The two tests below guard counter idempotency across at-least-once redelivery.
+// runTask applies the pipeline's chunk/token counts to the document and its
+// knowledgebase (runDocumentTask -> docState.apply -> ApplyDocCounts) before it
+// marks the task complete, so a redelivered task re-runs the pipeline. Because
+// ApplyDocCounts sets the document's own counts and rolls only the delta into the
+// knowledge base aggregate, re-applying the same run contributes zero: the same
+// task, processed twice, applies its counters once. The chunk store is likewise
+// idempotent (upsert by deterministic id).
+
+const (
+	rcChunks int64 = 5
+	rcTokens int64 = 100
+)
+
+// applyResult stands in for defaultRunDocumentTask: a successful pipeline run
+// that applies its result to the doc + KB counters, as docState.apply does after
+// Execute returns.
+func applyResult(ingestor *Ingestor, docID, kbID string) func(context.Context, *entity.IngestionTask) error {
+	return func(ctx context.Context, _ *entity.IngestionTask) error {
+		ingestor.docState.apply(ctx, &taskpkg.PipelineResult{
+			DocID:            docID,
+			KbID:             kbID,
+			ChunkCount:       int(rcChunks),
+			TokenConsumption: int(rcTokens),
+			Duration:         1,
+		})
+		return nil
+	}
+}
+
+// assertCountersAppliedOnce fails when the doc/KB counters reflect more than a
+// single application of the pipeline result - i.e. a redelivery re-counted.
+func assertCountersAppliedOnce(t *testing.T, db *gorm.DB, kbID, docID string) {
+	t.Helper()
+	var kb entity.Knowledgebase
+	if err := db.First(&kb, "id = ?", kbID).Error; err != nil {
+		t.Fatalf("load kb: %v", err)
+	}
+	var doc entity.Document
+	if err := db.First(&doc, "id = ?", docID).Error; err != nil {
+		t.Fatalf("load doc: %v", err)
+	}
+	if kb.ChunkNum != rcChunks || kb.TokenNum != rcTokens {
+		t.Errorf("kb counters = (chunk %d, token %d), want (%d, %d) - redelivery double-counted", kb.ChunkNum, kb.TokenNum, rcChunks, rcTokens)
+	}
+	if doc.ChunkNum != rcChunks || doc.TokenNum != rcTokens {
+		t.Errorf("doc counters = (chunk %d, token %d), want (%d, %d) - redelivery double-counted", doc.ChunkNum, doc.TokenNum, rcChunks, rcTokens)
+	}
+}
+
+func rcMsg(taskID, docID, kbID string) *entity.IngestionTask {
+	return &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: kbID, Status: common.RUNNING}
+}
+
+// TestRunTask_RedeliveryOfCompletedTaskCountsOnce: the first delivery fully
+// completes (task -> COMPLETED), then the broker redelivers the same message
+// because its Ack was lost. runTask has no already-completed guard, so it
+// re-runs the pipeline; the counter application must stay idempotent.
+func TestRunTask_RedeliveryOfCompletedTaskCountsOnce(t *testing.T) {
+	db := testutil.SetupTestDB(t)
+	cleanup := testutil.ReplaceDBForTest(t, db)
+	defer cleanup()
+	_, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
+
+	ingestor := NewIngestor("test", 1, []string{"pdf"})
+	ingestor.runDocumentTask = applyResult(ingestor, docID, kbID)
+
+	// First delivery: parse succeeds, counters applied once, task -> COMPLETED.
+	if terminal := ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID)); !terminal {
+		t.Fatalf("expected the first delivery to complete (terminal=true)")
+	}
+	// Redelivery (Ack lost): the same message is processed again. Must NOT re-count.
+	ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID))
+
+	assertCountersAppliedOnce(t, db, kbID, docID)
+}
+
+// TestRunTask_RedeliveryAfterIncompleteRunCountsOnce: the crash/nack window. A
+// prior run applied the counters but died before MarkCompleted, so the task is
+// left RUNNING and the broker redelivers it; the redelivery re-runs and
+// completes, and the counters must not be applied twice.
+func TestRunTask_RedeliveryAfterIncompleteRunCountsOnce(t *testing.T) {
+	db := testutil.SetupTestDB(t)
+	cleanup := testutil.ReplaceDBForTest(t, db)
+	defer cleanup()
+	_, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
+
+	ingestor := NewIngestor("test", 1, []string{"pdf"})
+	ingestor.runDocumentTask = applyResult(ingestor, docID, kbID)
+
+	// Prior run: counters applied, but the task never completed (crash before
+	// MarkCompleted) - the task row is left RUNNING, so the broker redelivers.
+	applyResult(ingestor, docID, kbID)(context.Background(), nil)
+
+	// Redelivery of the still-RUNNING task: re-runs and completes.
+	if terminal := ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID)); !terminal {
+		t.Fatalf("expected the redelivery run to complete (terminal=true)")
+	}
+
+	assertCountersAppliedOnce(t, db, kbID, docID)
+}

internal/service/document/apply_doc_counts_test.go

@@ -0,0 +1,126 @@
+//
+//  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 document
+
+import (
+	"context"
+	"testing"
+
+	"ragflow/internal/entity"
+)
+
+// TestApplyDocCounts_ReparseCarriesDelta checks that re-applying a document with
+// a different count - a re-parse - moves the knowledge base aggregate by the
+// delta to the new absolute value, rather than summing the two runs.
+func TestApplyDocCounts_ReparseCarriesDelta(t *testing.T) {
+	db := setupServiceTestDB(t)
+	pushServiceDB(t, db)
+	if err := db.Create(&entity.Knowledgebase{ID: "kb-1"}).Error; err != nil {
+		t.Fatalf("create kb: %v", err)
+	}
+	if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ParserConfig: entity.JSONMap{}}).Error; err != nil {
+		t.Fatalf("create doc: %v", err)
+	}
+	svc := testDocumentService(t)
+	ctx := context.Background()
+
+	// First parse produces 5 chunks / 100 tokens; a re-parse produces 7 / 140.
+	if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 1); err != nil {
+		t.Fatalf("first apply: %v", err)
+	}
+	if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 7, 140, 1); err != nil {
+		t.Fatalf("reparse apply: %v", err)
+	}
+
+	// The document holds the new absolute counts; the KB aggregate follows the
+	// delta to 7/140, not 5+7 / 100+140.
+	var kb entity.Knowledgebase
+	if err := db.First(&kb, "id = ?", "kb-1").Error; err != nil {
+		t.Fatalf("load kb: %v", err)
+	}
+	var doc entity.Document
+	if err := db.First(&doc, "id = ?", "doc-1").Error; err != nil {
+		t.Fatalf("load doc: %v", err)
+	}
+	if kb.ChunkNum != 7 || kb.TokenNum != 140 {
+		t.Errorf("kb = (chunk %d, token %d), want (7, 140)", kb.ChunkNum, kb.TokenNum)
+	}
+	if doc.ChunkNum != 7 || doc.TokenNum != 140 {
+		t.Errorf("doc = (chunk %d, token %d), want (7, 140)", doc.ChunkNum, doc.TokenNum)
+	}
+}
+
+// TestApplyDocCounts_KBAggregateClampsAtZero checks that an aggregate driven
+// below zero from an inconsistent starting state clamps to 0 instead of
+// underflowing.
+func TestApplyDocCounts_KBAggregateClampsAtZero(t *testing.T) {
+	db := setupServiceTestDB(t)
+	pushServiceDB(t, db)
+	// The KB aggregate is inconsistent: below this document's own contribution.
+	if err := db.Create(&entity.Knowledgebase{ID: "kb-1", ChunkNum: 2, TokenNum: 40}).Error; err != nil {
+		t.Fatalf("create kb: %v", err)
+	}
+	if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ChunkNum: 5, TokenNum: 100, ParserConfig: entity.JSONMap{}}).Error; err != nil {
+		t.Fatalf("create doc: %v", err)
+	}
+	svc := testDocumentService(t)
+
+	// A run that clears the document (0 chunks) drives the aggregate delta to -5;
+	// 2 - 5 is negative and must clamp to 0 rather than underflow.
+	if err := svc.ApplyDocCounts(context.Background(), "doc-1", "kb-1", 0, 0, 1); err != nil {
+		t.Fatalf("apply: %v", err)
+	}
+	var kb entity.Knowledgebase
+	if err := db.First(&kb, "id = ?", "kb-1").Error; err != nil {
+		t.Fatalf("load kb: %v", err)
+	}
+	if kb.ChunkNum != 0 || kb.TokenNum != 0 {
+		t.Errorf("kb = (chunk %d, token %d), want (0, 0) - aggregate must clamp, not underflow", kb.ChunkNum, kb.TokenNum)
+	}
+}
+
+// TestApplyDocCounts_ProcessDurationIsAbsolute checks that process_duration is
+// the last run's value, not a cumulative sum: a later run replaces the prior
+// value rather than adding to it.
+func TestApplyDocCounts_ProcessDurationIsAbsolute(t *testing.T) {
+	db := setupServiceTestDB(t)
+	pushServiceDB(t, db)
+	if err := db.Create(&entity.Knowledgebase{ID: "kb-1"}).Error; err != nil {
+		t.Fatalf("create kb: %v", err)
+	}
+	if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ParserConfig: entity.JSONMap{}}).Error; err != nil {
+		t.Fatalf("create doc: %v", err)
+	}
+	svc := testDocumentService(t)
+	ctx := context.Background()
+
+	// A run sets process_duration to its own value; a later run replaces it rather
+	// than accumulating, so the stored value is the last run's, not the 3.5+1.25 sum.
+	if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 3.5); err != nil {
+		t.Fatalf("first apply: %v", err)
+	}
+	if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 1.25); err != nil {
+		t.Fatalf("second apply: %v", err)
+	}
+	var doc entity.Document
+	if err := db.First(&doc, "id = ?", "doc-1").Error; err != nil {
+		t.Fatalf("load doc: %v", err)
+	}
+	if doc.ProcessDuration != 1.25 {
+		t.Errorf("process_duration = %v, want 1.25 (last run's value, not the 4.75 sum)", doc.ProcessDuration)
+	}
+}

internal/service/document/document_crud.go

@@ -13,6 +13,7 @@ import (
 	"ragflow/internal/storage"
 
 	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
 )
 
 // Accessible reports whether docID belongs to a knowledge base
@@ -142,31 +143,49 @@ func (s *DocumentService) UpdateDocument(ctx context.Context, id string, req *Up
 	return s.documentDAO.Update(ctx, dao.DB, document)
 }
 
-// IncrementChunkNum atomically increments chunk/token counters on the document and its knowledge base in a transaction
-func (s *DocumentService) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error {
+// ApplyDocCounts records a pipeline run's chunk/token/duration counts on the
+// document and rolls the change into its knowledge base aggregate. It is
+// idempotent per document: the document row holds this document's own counts, so
+// re-applying the same run - e.g. an at-least-once task redelivery that re-runs
+// the pipeline - sets the same document values and adjusts the knowledge base by
+// a zero delta. It also carries a changed count (a re-parse producing a
+// different result) correctly, since only the delta reaches the aggregate. The
+// document row is locked so the read-modify-write of its current counts is
+// atomic against a concurrent apply, and the aggregate is clamped at zero.
+func (s *DocumentService) ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error {
 	return dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		// Update document
+		var doc entity.Document
+		if err := tx.WithContext(ctx).Clauses(clause.Locking{Strength: "UPDATE"}).
+			Where("id = ? AND kb_id = ?", docID, kbID).
+			First(&doc).Error; err != nil {
+			return err
+		}
+
+		chunkDelta := int64(chunkNum) - doc.ChunkNum
+		tokenDelta := int64(tokenNum) - doc.TokenNum
+
+		// Set the document to this run's absolute counts.
 		if err := tx.WithContext(ctx).Model(&entity.Document{}).
 			Where("id = ? AND kb_id = ?", docID, kbID).
 			Updates(map[string]interface{}{
-				"chunk_num":        gorm.Expr("chunk_num + ?", int64(chunkNum)),
-				"token_num":        gorm.Expr("token_num + ?", int64(tokenNum)),
-				"process_duration": gorm.Expr("process_duration + ?", duration),
+				"chunk_num":        int64(chunkNum),
+				"token_num":        int64(tokenNum),
+				"process_duration": duration,
 			}).Error; err != nil {
 			return err
 		}
 
-		// Update knowledgebase
-		if err := tx.WithContext(ctx).Model(&entity.Knowledgebase{}).
+		// Roll only the delta into the knowledge base aggregate; a re-applied run
+		// contributes zero. Clamp at zero so a stale aggregate cannot go negative.
+		if chunkDelta == 0 && tokenDelta == 0 {
+			return nil
+		}
+		return tx.WithContext(ctx).Model(&entity.Knowledgebase{}).
 			Where("id = ?", kbID).
 			Updates(map[string]interface{}{
-				"chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)),
-				"token_num": gorm.Expr("token_num + ?", int64(tokenNum)),
-			}).Error; err != nil {
-			return err
-		}
-
-		return nil
+				"chunk_num": gorm.Expr("CASE WHEN chunk_num + ? >= 0 THEN chunk_num + ? ELSE 0 END", chunkDelta, chunkDelta),
+				"token_num": gorm.Expr("CASE WHEN token_num + ? >= 0 THEN token_num + ? ELSE 0 END", tokenDelta, tokenDelta),
+			}).Error
 	})
 }