diff --git a/internal/dao/ingestion.go b/internal/dao/ingestion.go index a41d741ad1..82f6630924 100644 --- a/internal/dao/ingestion.go +++ b/internal/dao/ingestion.go @@ -194,6 +194,21 @@ func (dao *IngestionTaskDAO) GetByDocumentID(documentId string) (*entity.Ingesti return tasks[0], nil } +// DeleteIfTerminal deletes ingestion tasks for a document that are in a +// terminal state (COMPLETED, STOPPED, FAILED) or still queued (CREATED). +// RUNNING and STOPPING tasks are NOT deleted because an in-flight worker +// would keep writing chunks and corrupt a new run's results. +// Returns the number of rows deleted. +func (dao *IngestionTaskDAO) DeleteIfTerminal(documentID string) (int64, error) { + result := DB.Where("document_id = ? AND status NOT IN (?, ?)", + documentID, common.RUNNING, common.STOPPING). + Delete(&entity.IngestionTask{}) + if result.Error != nil { + return 0, result.Error + } + return result.RowsAffected, nil +} + type IngestionTaskLogDAO struct{} func NewIngestionTaskLogDAO() *IngestionTaskLogDAO { diff --git a/internal/dao/ingestion_task_test.go b/internal/dao/ingestion_task_test.go index 8378f2cf71..6899ddb6b8 100644 --- a/internal/dao/ingestion_task_test.go +++ b/internal/dao/ingestion_task_test.go @@ -4,6 +4,7 @@ import ( "errors" "testing" + "fmt" "ragflow/internal/common" "ragflow/internal/entity" @@ -133,3 +134,61 @@ func TestIngestionTaskDAOUpdateStatusIfCurrentRejectsMismatchedStatus(t *testing t.Fatalf("status = %q, want %q", reloaded.Status, common.STOPPING) } } + +func TestIngestionTaskDAODeleteIfTerminal_RemovesOnlyTerminal(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + // Create tasks in different statuses, each with a unique docID. + statuses := []string{common.CREATED, common.RUNNING, common.STOPPING, common.COMPLETED, common.STOPPED, common.FAILED} + for i, status := range statuses { + docID := fmt.Sprintf("doc-%d", i) + task := &entity.IngestionTask{ + ID: fmt.Sprintf("task-%d", i), + UserID: "user-1", + DocumentID: docID, + DatasetID: "kb-1", + Status: status, + } + if err := db.Create(task).Error; err != nil { + t.Fatalf("create task %s: %v", status, err) + } + } + + // DeleteIfTerminal deletes everything except RUNNING and STOPPING. + // CREATED is safe to delete (no worker has claimed it yet); + // COMPLETED/STOPPED/FAILED are terminal. + // Call it for every doc and verify the negative cases survived. + for i := 0; i < len(statuses); i++ { + docID := fmt.Sprintf("doc-%d", i) + _, err := NewIngestionTaskDAO().DeleteIfTerminal(docID) + if err != nil { + t.Fatalf("DeleteIfTerminal(doc-%d): %v", i, err) + } + } + + // RUNNING and STOPPING must survive. + for _, i := range []int{1, 2} { + docID := fmt.Sprintf("doc-%d", i) + task, err := NewIngestionTaskDAO().GetByDocumentID(docID) + if err != nil { + t.Fatalf("GetByDocumentID %s: %v", docID, err) + } + if task == nil { + t.Fatalf("%s task (doc=%d) must not be deleted", statuses[i], i) + } + } + // CREATED, COMPLETED, STOPPED, FAILED must be gone. + for _, i := range []int{0, 3, 4, 5} { + docID := fmt.Sprintf("doc-%d", i) + task, err := NewIngestionTaskDAO().GetByDocumentID(docID) + if err != nil { + t.Fatalf("GetByDocumentID %s: %v", docID, err) + } + if task != nil { + t.Fatalf("%s task (doc=%d) should be deleted, still present", statuses[i], i) + } + } +} diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index fd0a2f4fab..d1da9bb83d 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -49,6 +49,25 @@ func (m *mockCanvasStage) Invoke(_ context.Context, inputs map[string]any) (map[ func (m *mockCanvasStage) Inputs() map[string]string { return map[string]string{"name": "string"} } func (m *mockCanvasStage) Outputs() map[string]string { return map[string]string{"output": "any"} } +// oneShotErrStage errors on the first Invoke (simulating a component crash +// mid-run), then delegates to the embedded mock on subsequent calls. Used to +// test that a second pipeline Run on the same taskID resumes past non-terminal +// checkpoints instead of re-executing completed components. +type oneShotErrStage struct { + mockCanvasStage + n int +} + +func (s *oneShotErrStage) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { + s.n++ + if s.n == 1 { + return nil, errors.New("simulated crash") + } + return s.mockCanvasStage.Invoke(ctx, inputs) +} +func (s *oneShotErrStage) Inputs() map[string]string { return s.mockCanvasStage.Inputs() } +func (s *oneShotErrStage) Outputs() map[string]string { return s.mockCanvasStage.Outputs() } + func TestPipelineRunHappyPath(t *testing.T) { stageA := &mockCanvasStage{output: map[string]any{"a": 1}} stageB := &mockCanvasStage{output: map[string]any{"b": 2}} @@ -361,8 +380,81 @@ func TestPipelineRunResumableAutoResumes(t *testing.T) { } } -// TestPipelineRun_RequireResumeRejectsWithoutStore verifies M4 (plan §6.a -// 方案 A): with WithRequireResume set and no checkpoint store resolvable (no +// TestPipelineRunResumableCrossRunResume validates crash-recovery resume +// across two Run calls: when the terminal component errors mid-run (simulated +// crash), non-terminal checkpoints + interrupt state persist. The second Run +// on the same taskID resumes past completed non-terminal components instead of +// re-executing them. A non-terminal stage (A) runs exactly once across both runs; the +// terminal stage (B, a oneShotErrStage) errors on run 1 and succeeds on run 2. +func TestPipelineRunResumableCrossRunResume(t *testing.T) { + mockA := &mockCanvasStage{output: map[string]any{"a": 1}} + mockB := &mockCanvasStage{output: map[string]any{"b": 2}} + termStage := &oneShotErrStage{mockCanvasStage: *mockB} + + const ( + nameA = "p.XRunStageA" + nameB = "p.XRunStageB" + ) + runtime.MustRegister(nameA, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return mockA, nil }, + runtime.Metadata{Version: "1.0.0"}) + runtime.MustRegister(nameB, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return termStage, nil }, + runtime.Metadata{Version: "1.0.0"}) + + store := newMemCheckpointStore() + mr := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { client.Close() }) + tracker := canvas.NewRunTrackerWithClient(client, time.Hour) + + const taskID = "task-cross-run-resume" + pipe, err := NewPipelineFromDSL([]byte(`{ + "dsl": { + "components": { + "begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]}, + "a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]}, + "b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]} + }, + "path": ["begin", "a", "b"], + "graph": {"nodes": []} + } + }`), taskID, WithCheckPointStore(store), WithRunTracker(tracker)) + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + + // Run 1: terminal B errors (oneShotErrStage n=1), simulating a crash. + // Non-terminal A's checkpoint + interrupt persist because the error path + // does not call ClearInterruptID or store.Delete. + _, err = pipe.Run(context.Background(), map[string]any{"name": "doc-cross-run"}) + if err == nil { + t.Fatal("Run 1: expected error from simulated crash, got nil") + } + if mockA.calls != 1 { + t.Fatalf("Run 1: expected A to run once, got %d", mockA.calls) + } + // oneShotErrStage did not delegate to its embedded mock on the first call. + if termStage.calls != 0 { + t.Fatalf("Run 1: expected B (embedded mock) calls=0 (error before delegate), got %d", termStage.calls) + } + + // Run 2: resume from after A via tracker.GetInterruptID. A is skipped; + // B's oneShotErrStage (n=2) delegates to its embedded mock successfully. + _, err = pipe.Run(context.Background(), map[string]any{"name": "doc-cross-run"}) + if err != nil { + t.Fatalf("Run 2: expected recovery success, got error: %v", err) + } + if mockA.calls != 1 { + t.Fatalf("Run 2: expected A to still have calls=1 (was skipped by resume), got %d", mockA.calls) + } + if termStage.calls != 1 { + t.Fatalf("Run 2: expected B (embedded mock) calls=1 (delegated once), got %d", termStage.calls) + } +} + +// TestPipelineRun_RequireResumeRejectsWithoutStore verifies that with +// WithRequireResume set and no checkpoint store resolvable (no // injected store, no global Redis in unit scope), Run must refuse to start // and return ErrResumeUnavailable — a clear, distinguishable signal — rather // than silently degrading to a non-resumable runPlain. The reject fires diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go index 44ec2d7e47..7bb93fc2c0 100644 --- a/internal/ingestion/service/process_message_test.go +++ b/internal/ingestion/service/process_message_test.go @@ -50,12 +50,26 @@ func TestProcessMessage_TaskNotFoundAcks(t *testing.T) { } // TestProcessMessage_AlreadyCompletedAcks: a task already in a terminal state -// (COMPLETED) is acked and skipped — no enqueue, no status change. +// (COMPLETED) is acked and skipped — no enqueue, and the document is NOT +// resurrected to RUNNING. A redelivered terminal task must not reset a +// finished document's run status or counters. func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Mark the document as a finished parse: a non-RUNNING run status and + // non-zero counters that StartRunning would clobber to "1"/0. + finishedRun := "3" + if err := db.Model(&entity.Document{}).Where("id = ?", docID). + Updates(map[string]interface{}{ + "run": finishedRun, + "chunk_num": 42, + "token_num": 99, + }).Error; err != nil { + t.Fatalf("set document finished sentinel: %v", err) + } // Set the task COMPLETED so the status switch skips it. if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID). @@ -73,6 +87,18 @@ func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) { if len(ingestor.taskChan) != 0 { t.Fatal("expected no task enqueued for completed task") } + + // The finished document must be untouched - no resurrection to RUNNING. + var doc entity.Document + if err := db.Where("id = ?", docID).First(&doc).Error; err != nil { + t.Fatalf("reload document: %v", err) + } + if doc.Run == nil || *doc.Run != finishedRun { + t.Fatalf("document.run = %v, want %q (must not be resurrected to RUNNING)", doc.Run, finishedRun) + } + if doc.ChunkNum != 42 || doc.TokenNum != 99 { + t.Fatalf("document counters changed: chunk_num=%d token_num=%d, want 42/99", doc.ChunkNum, doc.TokenNum) + } } // TestProcessMessage_ClaimFailsAcks: a RUNNING task whose claim fails diff --git a/internal/parser/parser/xlsx_parse_method_test.go b/internal/parser/parser/xlsx_parse_method_test.go new file mode 100644 index 0000000000..f50f3086b2 --- /dev/null +++ b/internal/parser/parser/xlsx_parse_method_test.go @@ -0,0 +1,121 @@ +package parser + +import ( + "strings" + "testing" + + "github.com/xuri/excelize/v2" +) + +// TestNormalizeXLSXParseMethod verifies the parse_method normalization +// shared by the XLSX/XLS/CSV parsers. "deepdoc" is the default +// spreadsheet parse_method (see schema.ParserParam.Defaults and the +// matching Python ParserParam), and the DSL templates ship "DeepDOC". +// Both must normalize to "" so the default Excelize/CSV path is taken, +// matching rag/flow/parser/parser.py:_spreadsheet which only special-cases +// "tcadp parser" and routes everything else to the default parser. +func TestNormalizeXLSXParseMethod(t *testing.T) { + cases := []struct { + in, want string + }{ + {"", ""}, + {"deepdoc", ""}, + {"DeepDOC", ""}, // canonical casing used by DSL templates + {"DEEPDOC", ""}, + {" deepdoc ", ""}, + {"deepdoc parser", ""}, + {"DeepDOC Parser", ""}, + {"tcadp parser", "tcadp"}, + {"TCADP Parser", "tcadp"}, + {"excelize", "excelize"}, + {"csv", "csv"}, + {"unknown", "unknown"}, // preserved so the switch rejects it + } + for _, c := range cases { + if got := normalizeXLSXParseMethod(c.in); got != c.want { + t.Errorf("normalizeXLSXParseMethod(%q) = %q, want %q", c.in, got, c.want) + } + } +} + +// TestXLSXParser_DeepDocParseMethod verifies that both the lowercase "deepdoc" +// and the uppercase "DeepDOC" (as shipped by the ingestion pipeline DSL templates) +// parse_method values produce the default HTML table output. +func TestXLSXParser_DeepDocParseMethod(t *testing.T) { + cases := []struct { + name string + method string + cellValue string + }{ + {"Lowercase", "deepdoc", "hello"}, + {"Uppercase", "DeepDOC", "world"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + f := excelize.NewFile() + defer f.Close() + if err := f.SetCellValue("Sheet1", "A1", tc.cellValue); err != nil { + t.Fatalf("SetCellValue: %v", err) + } + buf, err := f.WriteToBuffer() + if err != nil { + t.Fatalf("WriteToBuffer: %v", err) + } + + p, err := NewXLSXParser("") + if err != nil { + t.Fatalf("NewXLSXParser: %v", err) + } + p.ConfigureFromSetup(map[string]any{"parse_method": tc.method}) + + res := p.ParseWithResult("test.xlsx", buf.Bytes()) + if res.Err != nil { + t.Fatalf("ParseWithResult(%s): %v", tc.method, res.Err) + } + if got, want := res.OutputFormat, "html"; got != want { + t.Fatalf("OutputFormat = %q, want %q", got, want) + } + if !strings.Contains(res.HTML, tc.cellValue) { + t.Fatalf("HTML = %q, want it to contain cell content %q", res.HTML, tc.cellValue) + } + }) + } +} + +// TestCSVParser_DeepDocParseMethod asserts the CSV parser accepts the +// default "deepdoc" parse_method and renders the default HTML table. +func TestCSVParser_DeepDocParseMethod(t *testing.T) { + p := NewCSVParser() + p.ConfigureFromSetup(map[string]any{"parse_method": "deepdoc"}) + + res := p.ParseWithResult("test.csv", []byte("a,b\n1,2")) + if res.Err != nil { + t.Fatalf("ParseWithResult(deepdoc): %v", res.Err) + } + if got, want := res.OutputFormat, "html"; got != want { + t.Fatalf("OutputFormat = %q, want %q", got, want) + } + if !strings.Contains(res.HTML, "") { + t.Fatalf("HTML = %q, want a rendered
", res.HTML) + } +} + +// TestXLSParser_DeepDocParseMethod_NoUnsupportedError asserts the XLS +// parser no longer rejects "deepdoc". A real .xls blob is hard to +// synthesize in a unit test, so we only assert that the error is not the +// "unsupported XLS parse method" rejection; an open/parse failure from +// the fake blob is acceptable. The shared normalization is already +// covered by TestNormalizeXLSXParseMethod. +func TestXLSParser_DeepDocParseMethod_NoUnsupportedError(t *testing.T) { + p, err := NewXLSParser("") + if err != nil { + t.Fatalf("NewXLSParser: %v", err) + } + p.ConfigureFromSetup(map[string]any{"parse_method": "deepdoc"}) + + res := p.ParseWithResult("test.xls", []byte("not a real xls")) + if res.Err != nil && strings.Contains(res.Err.Error(), "unsupported XLS parse method") { + t.Fatalf("deepdoc must not be rejected as unsupported: %v", res.Err) + } +} diff --git a/internal/parser/parser/xlsx_parser.go b/internal/parser/parser/xlsx_parser.go index a88cccbd07..a61d4d30ad 100644 --- a/internal/parser/parser/xlsx_parser.go +++ b/internal/parser/parser/xlsx_parser.go @@ -78,6 +78,15 @@ func normalizeXLSXParseMethod(raw string) string { if method == "tcadp parser" { return "tcadp" } + // "deepdoc" / "deepdoc parser" is the default spreadsheet parse_method + // (see schema.ParserParam.Defaults and the matching Python ParserParam), + // and the DSL templates ship "DeepDOC". Normalize to "" so the default + // Excelize/CSV path is taken by all three spreadsheet parsers, mirroring + // rag/flow/parser/parser.py:_spreadsheet which only special-cases + // "tcadp parser" and routes every other value to the default parser. + if method == "deepdoc" || method == "deepdoc parser" { + return "" + } return method } diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index 679496b8c9..dcabe1517b 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -17,28 +17,18 @@ package chunk import ( - "archive/zip" "bytes" "context" "encoding/base64" - "encoding/csv" - "encoding/json" - "encoding/xml" "fmt" "image" "image/color" "image/draw" "image/jpeg" - "io" "math" - "math/rand" - "path/filepath" "ragflow/internal/common" - "ragflow/internal/engine/redis" "ragflow/internal/entity" "ragflow/internal/entity/models" - "regexp" - "sort" "strconv" "strings" "sync" @@ -61,11 +51,6 @@ import ( "ragflow/internal/utility" ) -const ( - maximumPageNumber = 100000 - maximumTaskPageNumber = maximumPageNumber * 1000 -) - var chunkImageMergeLocks = struct { sync.Mutex locks map[string]*chunkImageMergeLock @@ -97,20 +82,24 @@ type ChunkService struct { taskDAO *dao.TaskDAO searchService *service.SearchService - accessibleFunc func(string, string) bool - getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error) - getDocumentsByIDsFunc func([]string) ([]*entity.Document, error) - getDocumentStorageAddressFunc func(*entity.Document) (string, string, error) - queueParseTasksFunc func(*entity.Document, string, string, int64) error - beginParseDocumentFunc func(string) error - deleteTasksByDocIDsFunc func([]string) (int64, error) - getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error) - incrementChunkStatsFunc func(string, string, int64, int64, float64) error - decrementChunkStatsFunc func(string, string, int64, int64, float64) error - storeChunkImageFunc func(string, string, []byte) error - tokenizeFunc func(string) (string, error) - fineGrainedTokenizeFunc func(string) (string, error) - numTokensFunc func(string) int + accessibleFunc func(string, string) bool + getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error) + getDocumentsByIDsFunc func([]string) ([]*entity.Document, error) + // startParseDocumentsFunc overrides the DSL start-parse flow. Production + // uses service.DocumentService.StartParseDocuments; tests inject a fake + // to avoid the MQ publisher. + startParseDocumentsFunc func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts service.StartParseOptions) error + // cancelIngestionTaskFunc overrides the document-parsing cancellation. + // Production uses service.DocumentService.CancelDocParse; tests inject + // a fake to avoid the MQ publisher. + cancelIngestionTaskFunc func(doc *entity.Document) error + getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error) + incrementChunkStatsFunc func(string, string, int64, int64, float64) error + decrementChunkStatsFunc func(string, string, int64, int64, float64) error + storeChunkImageFunc func(string, string, []byte) error + tokenizeFunc func(string) (string, error) + fineGrainedTokenizeFunc func(string) (string, error) + numTokensFunc func(string) int } // NewChunkService creates chunk service @@ -611,26 +600,12 @@ const ( docStopParsingInvalidStateErrorCode = "DOC_STOP_PARSING_INVALID_STATE" ) -func (s *ChunkService) cancelAllTasksOfDoc(docID string) error { - tasks, err := s.taskDAO.GetByDocID(docID) - if err != nil { - return fmt.Errorf("failed to get tasks for document %s: %w", docID, err) +func (s *ChunkService) cancelAllTasksOfDoc(doc *entity.Document) error { + cancel := s.cancelIngestionTaskFunc + if cancel == nil { + cancel = service.NewDocumentService().CancelDocParse } - - redisClient := redis.Get() - if redisClient == nil { - common.Warn(fmt.Sprintf("Redis unavailable; cannot cancel tasks for document %s", docID)) - return nil - } - - for _, task := range tasks { - if task == nil { - continue - } - redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 0) - } - - return nil + return cancel(doc) } func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) { @@ -642,15 +617,13 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required") } - kb, err := s.kbDAO.GetByID(datasetID) + _, err := s.kbDAO.GetByID(datasetID) if err != nil { return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset %s", datasetID) } docIDs, duplicateMessages := service.CheckDuplicateIDs(req.DocumentIDs, "document") successCount := 0 - ctx := context.Background() - indexName := service.IndexName(kb.TenantID) for _, docID := range docIDs { doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(docID, datasetID) @@ -658,40 +631,21 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar return nil, common.CodeDataError, fmt.Errorf("You don't own the document %s", docID) } - if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) { + task, _ := dao.NewIngestionTaskDAO().GetByDocumentID(docID) + if task == nil || task.Status == common.COMPLETED || + task.Status == common.STOPPED || task.Status == common.FAILED { return &service.StopParsingResponse{ Data: map[string]interface{}{"error_code": docStopParsingInvalidStateErrorCode}, }, common.CodeDataError, fmt.Errorf("%s", docStopParsingInvalidStateMessage) } - if err := s.cancelAllTasksOfDoc(docID); err != nil { + if err := s.cancelAllTasksOfDoc(doc); err != nil { return nil, common.CodeServerError, err } - - updates := map[string]interface{}{ - "run": string(entity.TaskStatusCancel), - "progress": 0, - "chunk_num": 0, - } - if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil { - return nil, common.CodeServerError, fmt.Errorf("failed to update document %s: %w", doc.ID, err) - } - - if s.docEngine != nil { - exists, err := s.docEngine.ChunkStoreExists(ctx, indexName, datasetID) - if err != nil { - return nil, common.CodeServerError, fmt.Errorf("failed to check chunk store %s/%s: %w", indexName, datasetID, err) - } - if exists { - if _, err := s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": doc.ID}, indexName, datasetID); err != nil { - return nil, common.CodeServerError, fmt.Errorf("failed to delete chunks for document %s: %w", doc.ID, err) - } - } else { - common.Info(fmt.Sprintf("Skipping chunk delete during stop_parsing for doc %s: index %s/%s does not exist", doc.ID, indexName, datasetID)) - } - } else { - common.Info(fmt.Sprintf("Skipping chunk delete during stop_parsing for doc %s: index %s/%s does not exist", doc.ID, indexName, datasetID)) - } + // CancelDocParse (inside cancelAllTasksOfDoc) already issues + // RequestStop (STOPPING) and updates doc.run=CANCEL. Defer + // destruction (chunk deletion, counter reset) until the worker + // detects STOPPING and reaches a terminal state. successCount++ } @@ -729,184 +683,6 @@ func checkDuplicateIDs(documentIDs []string, idTypes string) ([]string, []string return uniqueDocIDs, duplicateMessages } -func (s *ChunkService) queueParseTasks(doc *entity.Document, bucket, objectName string, priority int64) error { - if s.queueParseTasksFunc != nil { - return s.queueParseTasksFunc(doc, bucket, objectName, priority) - } - tasks, err := s.buildParseTasks(doc, bucket, objectName, priority) - if err != nil { - return err - } - if len(tasks) == 0 { - return nil - } - if err := dao.NewTaskDAO().CreateMany(tasks); err != nil { - return err - } - - queueName := s.parseQueueName(doc, priority) - for _, task := range tasks { - if task.Progress >= 1 { - continue - } - message := parseTaskMessage(task) - if ok := redis.Get().QueueProduct(queueName, message); !ok { - if _, err := dao.NewTaskDAO().DeleteByDocIDs([]string{doc.ID}); err != nil { - common.Warn("Failed to clean parse tasks after Redis enqueue failure", - zap.String("docID", doc.ID), - zap.Error(err)) - } - return fmt.Errorf("Can't access Redis. Please check the Redis' status.") - } - } - return nil -} - -func (s *ChunkService) buildParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) { - now := time.Now() - ranges, err := s.parseTaskRanges(doc, bucket, objectName) - if err != nil { - return nil, err - } - tasks := make([]*entity.Task, 0, len(ranges)) - for _, pageRange := range ranges { - taskID := utility.GenerateUUID() - progressMsg := "" - digest := s.parseTaskDigest(doc, pageRange.from, pageRange.to) - chunkIDs := "" - tasks = append(tasks, &entity.Task{ - ID: taskID, - DocID: doc.ID, - FromPage: pageRange.from, - ToPage: pageRange.to, - TaskType: "", - Priority: priority, - BeginAt: &now, - Progress: 0, - ProgressMsg: &progressMsg, - Digest: &digest, - ChunkIDs: &chunkIDs, - }) - } - return tasks, nil -} - -type parsePageRange struct { - from int64 - to int64 -} - -func (s *ChunkService) parseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { - if doc.Type == "pdf" { - return s.pdfParseTaskRanges(doc, bucket, objectName) - } - if doc.ParserID == string(entity.ParserTypeTable) { - return s.tableParseTaskRanges(doc, bucket, objectName) - } - return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil -} - -func (s *ChunkService) pdfParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { - binary, err := s.getStorageBinary(bucket, objectName) - if err != nil { - return nil, err - } - pages := estimatePDFPageCount(binary) - pageSize := int64(parserConfigInt(doc.ParserConfig, "task_page_size", 12)) - if doc.ParserID == string(entity.ParserTypePaper) { - pageSize = int64(parserConfigInt(doc.ParserConfig, "task_page_size", 22)) - } - if doc.ParserID == string(entity.ParserTypeOne) || - parserConfigString(doc.ParserConfig, "layout_recognize", "DeepDOC") != "DeepDOC" || - parserConfigBool(doc.ParserConfig, "toc_extraction", false) { - pageSize = maximumTaskPageNumber - } - if pageSize <= 0 { - pageSize = 12 - } - - pageRanges := parserConfigPageRanges(doc.ParserConfig) - ranges := make([]parsePageRange, 0) - for _, configuredRange := range pageRanges { - start := configuredRange.from - 1 - if start < 0 { - start = 0 - } - end := configuredRange.to - 1 - if pages >= 0 && end > pages { - end = pages - } - for page := start; page < end; page += pageSize { - to := page + pageSize - if to > end { - to = end - } - ranges = append(ranges, parsePageRange{from: page, to: to}) - } - } - if len(ranges) == 0 { - ranges = append(ranges, parsePageRange{from: 0, to: maximumTaskPageNumber}) - } - return ranges, nil -} - -func (s *ChunkService) tableParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) { - binary, err := s.getStorageBinary(bucket, objectName) - if err != nil { - return nil, err - } - rows := estimateTableRowCount(docName(doc), binary) - if rows <= 0 { - return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil - } - ranges := make([]parsePageRange, 0, (rows+2999)/3000) - for row := int64(0); row < int64(rows); row += 3000 { - to := row + 3000 - if to > int64(rows) { - to = int64(rows) - } - ranges = append(ranges, parsePageRange{from: row, to: to}) - } - return ranges, nil -} - -func (s *ChunkService) getStorageBinary(bucket, objectName string) ([]byte, error) { - storageImpl := storage.GetStorageFactory().GetStorage() - if storageImpl == nil { - return nil, fmt.Errorf("storage not initialized") - } - return storageImpl.Get(bucket, objectName) -} - -func (s *ChunkService) beginParseDocument(docID string) error { - if s.beginParseDocumentFunc != nil { - return s.beginParseDocumentFunc(docID) - } - now := time.Now() - return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{ - "progress_msg": "Task is queued...", - "process_begin_at": now, - "progress": rand.Float64() * 0.01, - "run": string(entity.TaskStatusRunning), - "chunk_num": 0, - "token_num": 0, - }).Error -} - -func (s *ChunkService) getDocumentStorageAddress(doc *entity.Document) (string, string, error) { - if s.getDocumentStorageAddressFunc != nil { - return s.getDocumentStorageAddressFunc(doc) - } - return service.NewDocumentService().GetDocumentStorageAddress(doc) -} - -func (s *ChunkService) deleteTasksByDocIDs(docIDs []string) (int64, error) { - if s.deleteTasksByDocIDsFunc != nil { - return s.deleteTasksByDocIDsFunc(docIDs) - } - return dao.NewTaskDAO().DeleteByDocIDs(docIDs) -} - func (s *ChunkService) accessible(datasetID, userID string) bool { if s.accessibleFunc != nil { return s.accessibleFunc(datasetID, userID) @@ -928,282 +704,6 @@ func (s *ChunkService) getDocumentsByIDs(docIDs []string) ([]*entity.Document, e return s.documentDAO.GetByIDs(docIDs) } -func (s *ChunkService) parseQueueName(doc *entity.Document, priority int64) string { - suffix := "common" - if doc.ParserID == string(entity.ParserTypeResume) { - suffix = "resume" - } - return fmt.Sprintf("te.%d.%s", priority, suffix) -} - -func (s *ChunkService) parseTaskDigest(doc *entity.Document, fromPage, toPage int64) string { - hasher := xxhash.New() - config := chunkingConfigForDigest(doc) - keys := make([]string, 0, len(config)) - for key := range config { - keys = append(keys, key) - } - sort.Strings(keys) - for _, key := range keys { - hasher.WriteString(stableString(config[key])) - } - hasher.WriteString(doc.ID) - hasher.WriteString(strconv.FormatInt(fromPage, 10)) - hasher.WriteString(strconv.FormatInt(toPage, 10)) - return fmt.Sprintf("%x", hasher.Sum64()) -} - -func parseTaskMessage(task *entity.Task) map[string]interface{} { - beginAt := "" - if task.BeginAt != nil { - beginAt = task.BeginAt.Format("2006-01-02 15:04:05") - } - digest := "" - if task.Digest != nil { - digest = *task.Digest - } - return map[string]interface{}{ - "id": task.ID, - "doc_id": task.DocID, - "from_page": task.FromPage, - "to_page": task.ToPage, - "progress": task.Progress, - "priority": task.Priority, - "begin_at": beginAt, - "digest": digest, - } -} - -func chunkingConfigForDigest(doc *entity.Document) map[string]interface{} { - return map[string]interface{}{ - "doc_id": doc.ID, - "kb_id": doc.KbID, - "parser_id": doc.ParserID, - "parser_config": copyParserConfigForDigest(doc.ParserConfig), - } -} - -func copyParserConfigForDigest(config map[string]interface{}) map[string]interface{} { - copied := make(map[string]interface{}, len(config)) - for key, value := range config { - if key == "raptor" || key == "graphrag" { - continue - } - copied[key] = value - } - return copied -} - -func stableString(value interface{}) string { - binary, err := json.Marshal(value) - if err != nil { - return fmt.Sprint(value) - } - return string(binary) -} - -func parserConfigInt(config map[string]interface{}, key string, fallback int) int { - value, ok := config[key] - if !ok || value == nil { - return fallback - } - switch typedValue := value.(type) { - case int: - return typedValue - case int64: - return int(typedValue) - case float64: - return int(typedValue) - case json.Number: - if intValue, err := typedValue.Int64(); err == nil { - return int(intValue) - } - case string: - if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil { - return intValue - } - } - return fallback -} - -func parserConfigString(config map[string]interface{}, key, fallback string) string { - value, ok := config[key] - if !ok || value == nil { - return fallback - } - if stringValue, ok := value.(string); ok { - return stringValue - } - return fmt.Sprint(value) -} - -func parserConfigBool(config map[string]interface{}, key string, fallback bool) bool { - value, ok := config[key] - if !ok || value == nil { - return fallback - } - switch typedValue := value.(type) { - case bool: - return typedValue - case string: - switch strings.ToLower(strings.TrimSpace(typedValue)) { - case "true", "1", "yes", "on": - return true - case "false", "0", "no", "off": - return false - } - } - return fallback -} - -func parserConfigPageRanges(config map[string]interface{}) []parsePageRange { - defaultRanges := []parsePageRange{{from: 1, to: maximumPageNumber}} - raw, ok := config["pages"] - if !ok || raw == nil { - return defaultRanges - } - rawRanges, ok := raw.([]interface{}) - if !ok || len(rawRanges) == 0 { - return defaultRanges - } - - ranges := make([]parsePageRange, 0, len(rawRanges)) - for _, rawRange := range rawRanges { - rangeValues, ok := rawRange.([]interface{}) - if !ok || len(rangeValues) < 2 { - continue - } - from, okFrom := toInt64(rangeValues[0]) - to, okTo := toInt64(rangeValues[1]) - if okFrom && okTo && to > from { - ranges = append(ranges, parsePageRange{from: from, to: to}) - } - } - if len(ranges) == 0 { - return defaultRanges - } - return ranges -} - -func toInt64(value interface{}) (int64, bool) { - switch typedValue := value.(type) { - case int: - return int64(typedValue), true - case int64: - return typedValue, true - case float64: - return int64(typedValue), true - case json.Number: - intValue, err := typedValue.Int64() - return intValue, err == nil - case string: - intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64) - return intValue, err == nil - default: - return 0, false - } -} - -var pdfPagePattern = regexp.MustCompile(`/Type\s*/Page\b`) - -func estimatePDFPageCount(binary []byte) int64 { - if len(binary) == 0 { - return 0 - } - return int64(len(pdfPagePattern.FindAll(binary, -1))) -} - -func estimateTableRowCount(name string, binary []byte) int { - switch strings.ToLower(filepath.Ext(name)) { - case ".xlsx": - if rows, err := countXLSXRows(binary); err == nil { - return rows - } - case ".csv", ".tsv", ".txt": - return countDelimitedRows(name, binary) - } - return 0 -} - -func countDelimitedRows(name string, binary []byte) int { - reader := csv.NewReader(bytes.NewReader(binary)) - reader.FieldsPerRecord = -1 - reader.ReuseRecord = true - if strings.EqualFold(filepath.Ext(name), ".tsv") { - reader.Comma = '\t' - } - rows := 0 - for { - _, err := reader.Read() - if err == nil { - rows++ - continue - } - if err == io.EOF { - break - } - rows += bytes.Count(binary, []byte{'\n'}) - if len(binary) > 0 && binary[len(binary)-1] != '\n' { - rows++ - } - break - } - return rows -} - -func countXLSXRows(binary []byte) (int, error) { - zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary))) - if err != nil { - return 0, err - } - maxRows := 0 - for _, file := range zipReader.File { - if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") { - continue - } - rows, err := countWorksheetRows(file) - if err != nil { - return 0, err - } - if rows > maxRows { - maxRows = rows - } - } - return maxRows, nil -} - -func countWorksheetRows(file *zip.File) (int, error) { - reader, err := file.Open() - if err != nil { - return 0, err - } - defer reader.Close() - - decoder := xml.NewDecoder(reader) - rows := 0 - for { - token, err := decoder.Token() - if err == io.EOF { - break - } - if err != nil { - return 0, err - } - start, ok := token.(xml.StartElement) - if ok && start.Name.Local == "row" { - rows++ - } - } - return rows, nil -} - -func docName(doc *entity.Document) string { - if doc.Name == nil { - return "" - } - return *doc.Name -} - func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) { if !s.accessible(datasetID, userID) { return nil, common.CodeOperatingError, fmt.Errorf("You don't own the dataset %s.", datasetID) @@ -1244,34 +744,23 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq } } + // Batch pre-check: refuse the whole request if any document's ingestion + // task is non-terminal (RUNNING/STOPPING), so we never partially clean. + if err := (service.NewDocumentService().AssertIngestionTasksTerminal(docIDs)); err != nil { + return nil, common.CodeDataError, err + } + + startParse := s.startParseDocumentsFunc + if startParse == nil { + docSvc := service.NewDocumentService() + startParse = docSvc.StartParseDocuments + } + successCount := 0 for _, docID := range docIDs { doc := docByID[docID] - - if s.docEngine != nil { - indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": docID}, indexName, datasetID); err != nil { - return nil, common.CodeServerError, err - } - } - if _, err := s.deleteTasksByDocIDs([]string{docID}); err != nil { - return nil, common.CodeServerError, err - } - - bucket, objectName, err := s.getDocumentStorageAddress(doc) - if err != nil { - return nil, common.CodeServerError, err - } - if err := s.queueParseTasks(doc, bucket, objectName, 0); err != nil { - return nil, common.CodeServerError, err - } - if err := s.beginParseDocument(doc.ID); err != nil { - if _, delErr := s.deleteTasksByDocIDs([]string{doc.ID}); delErr != nil { - common.Warn("Failed to clean parse tasks after document state update failure", - zap.String("docID", doc.ID), - zap.Error(delErr)) - } + if err := startParse(doc, kb, userID, service.StartParseOptions{RerunWithDelete: true}); err != nil { return nil, common.CodeServerError, err } successCount++ diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index e2bdc1e2ed..97f5509eb8 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -274,200 +274,6 @@ func TestParseRejectsRunningDocument(t *testing.T) { } } -func TestParseReturnsServerErrorWhenDeleteChunksFails(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - deleteErr := errors.New("delete chunks failed") - engine := &parseTestDocEngine{deleteChunksErr: deleteErr} - svc := newParseTestService(t) - svc.docEngine = engine - - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if !errors.Is(err, deleteErr) { - t.Fatalf("expected delete chunks error, got %v", err) - } - if code != common.CodeServerError { - t.Fatalf("expected CodeServerError, got %v", code) - } - if engine.deleteChunksCalls != 1 { - t.Fatalf("expected DeleteChunks once, got %d", engine.deleteChunksCalls) - } -} - -func TestParseReturnsServerErrorWhenDeleteTasksFails(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - deleteErr := errors.New("delete tasks failed") - svc := newParseTestService(t) - svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) { return 0, deleteErr } - - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if !errors.Is(err, deleteErr) { - t.Fatalf("expected delete tasks error, got %v", err) - } - if code != common.CodeServerError { - t.Fatalf("expected CodeServerError, got %v", code) - } -} - -func TestParseReturnsServerErrorWhenStorageAddressFails(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - storageErr := errors.New("storage address failed") - svc := newParseTestService(t) - svc.getDocumentStorageAddressFunc = func(*entity.Document) (string, string, error) { - return "", "", storageErr - } - - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if !errors.Is(err, storageErr) { - t.Fatalf("expected storage error, got %v", err) - } - if code != common.CodeServerError { - t.Fatalf("expected CodeServerError, got %v", code) - } -} - -func TestParseReturnsServerErrorWhenQueueFails(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - queueErr := errors.New("queue failed") - svc := newParseTestService(t) - svc.queueParseTasksFunc = func(*entity.Document, string, string, int64) error { - return queueErr - } - - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if !errors.Is(err, queueErr) { - t.Fatalf("expected queue error, got %v", err) - } - if code != common.CodeServerError { - t.Fatalf("expected CodeServerError, got %v", code) - } -} - -func TestParseCleansTasksWhenBeginParseFails(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - beginErr := errors.New("begin parse failed") - deleteCalls := 0 - svc := newParseTestService(t) - svc.beginParseDocumentFunc = func(string) error { return beginErr } - svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) { - deleteCalls++ - return 1, nil - } - - _, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if !errors.Is(err, beginErr) { - t.Fatalf("expected begin error, got %v", err) - } - if code != common.CodeServerError { - t.Fatalf("expected CodeServerError, got %v", code) - } - if deleteCalls != 2 { - t.Fatalf("expected initial delete and cleanup delete, got %d calls", deleteCalls) - } -} - -func TestParseReturnsPartialSuccessForDuplicateDocumentIDs(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - - svc := newParseTestService(t) - result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ - DocumentIDs: []string{"doc-1", "doc-1"}, - }) - if err == nil { - t.Fatal("expected duplicate warning error") - } - if code != common.CodeSuccess { - t.Fatalf("expected CodeSuccess, got %v", code) - } - if result["success_count"] != 1 { - t.Fatalf("expected one successful parse, got %v", result["success_count"]) - } - errorsValue, ok := result["errors"].([]string) - if !ok || len(errorsValue) != 1 || !strings.Contains(errorsValue[0], "Duplicate document ids: doc-1") { - t.Fatalf("unexpected duplicate errors: %#v", result["errors"]) - } -} - -func TestParseQueuesAndBeginsDocument(t *testing.T) { - db := setupChunkTestDB(t) - pushChunkTestDB(t, db) - userID, datasetID := "user-1", "kb-1" - insertChunkTestKB(t, datasetID, userID) - insertChunkTestDoc(t, "doc-1", datasetID) - insertChunkTestTask(t, "task-1", "doc-1") - - queueCalls := 0 - svc := newParseTestService(t) - svc.queueParseTasksFunc = func(doc *entity.Document, bucket, objectName string, priority int64) error { - queueCalls++ - if doc.ID != "doc-1" || bucket != datasetID || objectName != "doc-1.txt" || priority != 0 { - t.Fatalf("unexpected queue args: doc=%s bucket=%s object=%s priority=%d", doc.ID, bucket, objectName, priority) - } - return nil - } - - result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) - if err != nil { - t.Fatalf("expected parse success, got %v", err) - } - if code != common.CodeSuccess { - t.Fatalf("expected CodeSuccess, got %v", code) - } - if result != nil { - t.Fatalf("expected nil result, got %#v", result) - } - if queueCalls != 1 { - t.Fatalf("expected queue once, got %d", queueCalls) - } - - var taskCount int64 - if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { - t.Fatalf("count tasks: %v", err) - } - if taskCount != 0 { - t.Fatalf("expected old tasks to be deleted, got %d", taskCount) - } - - doc, err := dao.NewDocumentDAO().GetByID("doc-1") - if err != nil { - t.Fatalf("get doc: %v", err) - } - if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) { - t.Fatalf("expected document to be running, got %v", doc.Run) - } - if doc.ChunkNum != 0 { - t.Fatalf("expected chunk_num reset to 0, got %d", doc.ChunkNum) - } -} - func TestAddChunkSuccess(t *testing.T) { db := setupChunkTestDB(t) pushChunkTestDB(t, db) @@ -961,6 +767,7 @@ func setupChunkTestDB(t *testing.T) *gorm.DB { &entity.Knowledgebase{}, &entity.Document{}, &entity.Task{}, + &entity.IngestionTask{}, ); err != nil { t.Fatalf("failed to migrate: %v", err) } @@ -984,15 +791,6 @@ func newParseTestService(t *testing.T) *ChunkService { docEngine: nil, kbDAO: dao.NewKnowledgebaseDAO(), documentDAO: dao.NewDocumentDAO(), - getDocumentStorageAddressFunc: func(doc *entity.Document) (string, string, error) { - if doc.Location == nil { - return doc.KbID, "", nil - } - return doc.KbID, *doc.Location, nil - }, - queueParseTasksFunc: func(*entity.Document, string, string, int64) error { - return nil - }, } } @@ -1057,6 +855,21 @@ func insertChunkTestDoc(t *testing.T, id, kbID string) { } } +func insertChunkTestIngestionTask(t *testing.T, id, userID, docID, datasetID, status string) { + t.Helper() + task := &entity.IngestionTask{ + ID: id, + UserID: userID, + DocumentID: docID, + DatasetID: datasetID, + Schema: entity.JSONMap{}, + Status: status, + } + if err := dao.DB.Create(task).Error; err != nil { + t.Fatalf("insert ingestion task: %v", err) + } +} + func setChunkTestStats(t *testing.T, docID, kbID string, docTokenNum, docChunkNum, kbTokenNum, kbChunkNum int64) { t.Helper() @@ -1093,6 +906,7 @@ type parseTestDocEngine struct { deleteChunksCondition map[string]interface{} deleteIndexName string deleteDatasetID string + chunkStoreExists bool // if true, ChunkStoreExists returns true } func (e *parseTestDocEngine) CreateChunkStore(context.Context, string, string, int, string) error { @@ -1128,7 +942,7 @@ func (e *parseTestDocEngine) DropChunkStore(context.Context, string, string) err } func (e *parseTestDocEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { - return false, nil + return e.chunkStoreExists, nil } func (e *parseTestDocEngine) CreateMetadataStore(context.Context, string) error { @@ -1517,3 +1331,175 @@ func copyMap(in map[string]interface{}) map[string]interface{} { } return out } + +func TestChunkServiceParse_RejectsBatchWithRunningIngestionTask(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestUserTenant(t, "user-1", "tenant-1") + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + insertChunkTestDoc(t, "doc-2", "kb-1") + insertChunkTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING) + + svc := newParseTestService(t) + svc.accessibleFunc = func(string, string) bool { return true } + _, _, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}}) + if err == nil { + t.Fatal("expected error for RUNNING ingestion task, got nil") + } +} + +func TestChunkServiceParse_CallsStartParseDocumentsWithRerunWithDelete(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestUserTenant(t, "user-1", "tenant-1") + insertChunkTestKB(t, "kb-1", "tenant-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + + svc := newParseTestService(t) + svc.accessibleFunc = func(string, string) bool { return true } + var calledOpts service.StartParseOptions + var calledDocID string + svc.startParseDocumentsFunc = func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts service.StartParseOptions) error { + calledOpts = opts + calledDocID = doc.ID + return nil + } + + _, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}}) + if err != nil { + t.Fatalf("Parse: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("code = %d, want success", code) + } + if calledDocID != "doc-1" { + t.Fatalf("startParseDocuments called with doc %s, want doc-1", calledDocID) + } + if !calledOpts.RerunWithDelete { + t.Fatal("expected RerunWithDelete=true") + } +} + +func TestChunkServiceParse_ReturnsPartialSuccessForDuplicateDocumentIDs(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + userID, datasetID := "user-1", "kb-1" + insertChunkTestKB(t, datasetID, userID) + insertChunkTestDoc(t, "doc-1", datasetID) + + svc := newParseTestService(t) + svc.startParseDocumentsFunc = func(*entity.Document, *entity.Knowledgebase, string, service.StartParseOptions) error { + return nil + } + + result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{ + DocumentIDs: []string{"doc-1", "doc-1"}, + }) + if err == nil { + t.Fatal("expected duplicate warning error") + } + if code != common.CodeSuccess { + t.Fatalf("expected CodeSuccess, got %v", code) + } + if result["success_count"] != 1 { + t.Fatalf("expected one successful parse, got %v", result["success_count"]) + } + errorsValue, ok := result["errors"].([]string) + if !ok || len(errorsValue) != 1 || !strings.Contains(errorsValue[0], "Duplicate document ids: doc-1") { + t.Fatalf("unexpected duplicate errors: %#v", result["errors"]) + } +} + +func TestStopParsing_CallsCancelIngestionTask(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestKB(t, "kb-1", "user-1") // owner = user-1, Accessible passes + insertChunkTestDoc(t, "doc-1", "kb-1") + // StopParsing now checks the IngestionTask status (not doc.Run). + insertChunkTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING) + + svc := newParseTestService(t) + var calledDocID string + svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { + calledDocID = doc.ID + return nil + } + + _, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + if err != nil { + t.Fatalf("StopParsing: %v", err) + } + if calledDocID != "doc-1" { + t.Fatalf("cancelIngestionTaskFunc called with doc %s, want doc-1", calledDocID) + } +} + +func TestStopParsing_RejectsDocumentWithoutIngestionTask(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestKB(t, "kb-1", "user-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + // No IngestionTask — should be rejected (nothing to stop). + + svc := newParseTestService(t) + var called bool + svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { called = true; return nil } + + _, code, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + if err == nil { + t.Fatal("expected error for document without ingestion task") + } + if code != common.CodeDataError { + t.Fatalf("expected CodeDataError, got %v", code) + } + if called { + t.Fatal("cancel should not be called for doc without ingestion task") + } +} + +func TestStopParsing_DoesNotDeleteChunksOrResetCountersAfterCancel(t *testing.T) { + db := setupChunkTestDB(t) + pushChunkTestDB(t, db) + insertChunkTestKB(t, "kb-1", "user-1") + insertChunkTestDoc(t, "doc-1", "kb-1") + // Pre-set counters on the doc so we can verify they are NOT reset. + if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Updates(map[string]interface{}{ + "token_num": int64(10), + "chunk_num": int64(5), + }).Error; err != nil { + t.Fatalf("set counters: %v", err) + } + insertChunkTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING) + + svc := newParseTestService(t) + svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { + // Simulate CancelDocParse: set doc.run=CANCEL. + return dao.DB.Model(&entity.Document{}).Where("id = ?", doc.ID). + Update("run", string(entity.TaskStatusCancel)).Error + } + engine := &parseTestDocEngine{chunkStoreExists: true} + svc.docEngine = engine + + _, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}}) + if err != nil { + t.Fatalf("StopParsing: %v", err) + } + + // Counters must NOT be reset. + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("reload doc: %v", err) + } + if doc.TokenNum != 10 { + t.Fatalf("token_num = %d, want 10 (not reset)", doc.TokenNum) + } + if doc.ChunkNum != 5 { + t.Fatalf("chunk_num = %d, want 5 (not reset)", doc.ChunkNum) + } + + // Chunks must NOT be deleted during cancel. + if engine.deleteChunksCalls != 0 { + t.Fatalf("DeleteChunks called %d times, want 0", engine.deleteChunksCalls) + } +} diff --git a/internal/service/document.go b/internal/service/document.go index 5dab489041..d3343bf8da 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -17,17 +17,12 @@ package service import ( - "archive/zip" - "bytes" "context" - "encoding/csv" "encoding/json" - "encoding/xml" "errors" "fmt" "io" "math" - "math/rand" "mime/multipart" "net/http" "path/filepath" @@ -40,16 +35,13 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" - "ragflow/internal/deepdoc/parser/pdf/pdfoxide" "ragflow/internal/engine" - "ragflow/internal/engine/redis" enginetypes "ragflow/internal/engine/types" "ragflow/internal/entity" "ragflow/internal/storage" "ragflow/internal/tokenizer" "ragflow/internal/utility" - "github.com/cespare/xxhash/v2" "go.uber.org/zap" "gorm.io/gorm" "gorm.io/gorm/clause" @@ -721,7 +713,7 @@ func (s *DocumentService) deleteDocumentFull(docID string) error { // Delete tasks from DB ingestionTask, err := s.ingestionTaskDAO.GetByDocumentID(docID) if err != nil { - common.Error(fmt.Sprintf("failed to get ingestion task by doc:%s", docID), err) + common.Error(fmt.Sprintf("failed to get ingestion task by doc:%s", doc.ID), err) return err } if ingestionTask != nil { @@ -1200,9 +1192,63 @@ type IngestDocumentRequest struct { ApplyKB bool `json:"apply_kb"` } -type documentParsePageRange struct { - from int64 - to int64 +// StartParseOptions controls StartParseDocuments behavior. +type StartParseOptions struct { + // ApplyKB merges the knowledgebase's parser_config (llm_id, metadata) + // into the document before parsing. + ApplyKB bool + // RerunWithDelete clears prior chunks/tasks/counters before re-parsing. + RerunWithDelete bool +} + +// StartParseDocuments starts parsing a document via the DSL ingestion +// pipeline. It optionally clears prior results (RerunWithDelete), applies +// KB config (ApplyKB), validates storage, and enqueues an ingestion task. +// The document run status is NOT set here; IngestionTaskService.StartRunning +// sets it to RUNNING when the worker picks up the task and transitions it from +// CREATED. Extracted from Ingest so other entry points (e.g. ChunkService.Parse) +// can reuse the same start-parse flow. +func (s *DocumentService) StartParseDocuments(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts StartParseOptions) error { + // Validate storage first so we don't clear prior results and then fail + // because the document can't be read, leaving the document with neither + // old nor new parse results. + if _, _, err := s.GetDocumentStorageAddress(doc); err != nil { + return err + } + + if opts.RerunWithDelete { + if err := s.clearDocumentParseResults(doc, kb.TenantID); err != nil { + return err + } + } + + if opts.ApplyKB { + if doc.ParserConfig == nil { + doc.ParserConfig = entity.JSONMap{} + } + config := map[string]interface{}{ + "llm_id": kb.ParserConfig["llm_id"], + "enable_metadata": false, + "metadata": map[string]interface{}{}, + } + if value, ok := kb.ParserConfig["enable_metadata"]; ok { + config["enable_metadata"] = value + } + if value, ok := kb.ParserConfig["metadata"]; ok { + config["metadata"] = value + } + if err := s.updateDocumentParserConfig(doc.ID, config); err != nil { + return err + } + for key, value := range config { + doc.ParserConfig[key] = value + } + } + + if _, err := s.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil { + return err + } + return nil } func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (common.ErrorCode, error) { @@ -1220,150 +1266,148 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com } } - tableDoneCountByKB := make(map[string]int64) - + // First pass: validate every document exists and is accessible before + // mutating any state, so a single invalid doc rejects the whole request. + type validatedDoc struct { + doc *entity.Document + kb *entity.Knowledgebase + } + validated := make([]validatedDoc, 0, len(req.DocIDs)) + validatedIDs := make([]string, 0, len(req.DocIDs)) for _, docID := range req.DocIDs { doc := docsByID[docID] if doc == nil { return common.CodeDataError, fmt.Errorf("document not found") } - kb, err := s.kbDAO.GetByID(doc.KbID) if err != nil { return common.CodeDataError, fmt.Errorf("dataset not found") } - if !s.kbDAO.Accessible(kb.ID, userID) { return common.CodeAuthenticationError, fmt.Errorf("no authorization") } + validated = append(validated, validatedDoc{doc, kb}) + validatedIDs = append(validatedIDs, docID) + } - updates := map[string]interface{}{ - "run": run, - "progress": 0, + // Batch pre-check for re-parse with delete: use the validated doc IDs + // so we don't silently skip non-existent or unauthorized documents. + if run == string(entity.TaskStatusRunning) && req.Delete { + if err := s.AssertIngestionTasksTerminal(validatedIDs); err != nil { + return common.CodeDataError, err } + } - rerunWithDelete := run == string(entity.TaskStatusRunning) && req.Delete - if rerunWithDelete { - updates["progress_msg"] = "" - updates["chunk_num"] = 0 - updates["token_num"] = 0 - } + for _, vd := range validated { + doc := vd.doc + kb := vd.kb - if run == string(entity.TaskStatusCancel) { - if err := s.cancelDocParse(doc); err != nil { - common.Error(fmt.Sprintf("go side, start to process %s, run is cancel", docID), err) - return common.CodeDataError, err - } - } - - if rerunWithDelete { - if err := s.prepareDocumentRerunWithDelete(doc, kb.TenantID); err != nil { - common.Error(fmt.Sprintf("go side, start to process %s, error when rerun with delete", docID), err) + // Start parsing: delegates to the shared start-parse flow. The + // document run status is set by IngestionTaskService.StartRunning + // when the task transitions from CREATED, not here. + if run == string(entity.TaskStatusRunning) { + if err := s.StartParseDocuments(doc, kb, userID, StartParseOptions{ + ApplyKB: req.ApplyKB, + RerunWithDelete: req.Delete, + }); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, start parse", doc.ID), err) return common.CodeExceptionError, err } + continue } - if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil { - common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", docID), err) + // Cancel: RequestStop (STOPPING) and update doc state. Do NOT + // delete the ingestion task or chunks here — deletion races with + // the worker's async markStopped/settleToTerminal flow. Once the + // worker detects STOPPING and transitions to STOPPED, the task + // is terminal and can be safely cleaned up. + if run == string(entity.TaskStatusCancel) { + if err := s.CancelDocParse(doc); err != nil { + common.Error(fmt.Sprintf("go side, start to process %s, run is cancel", doc.ID), err) + return common.CodeDataError, err + } + if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{ + "run": string(entity.TaskStatusCancel), + "progress": 0, + }); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", doc.ID), err) + return common.CodeExceptionError, err + } + continue + } + + // Delete-only: user asked to remove prior parse results without + // starting a new parse. RUNNING already continued above. + if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{ + "run": run, + "progress": 0, + }); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", doc.ID), err) return common.CodeExceptionError, err } - if req.Delete && !rerunWithDelete { + if req.Delete { _, _ = s.taskDAO.DeleteIngestionTasksByDocIDs([]string{doc.ID}) indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) if s.docEngine != nil { exists, err := s.docEngine.ChunkStoreExists(context.Background(), indexName, doc.KbID) if err != nil { - common.Error(fmt.Sprintf("go side, doc %s, ChunkStoreExists failed", docID), err) + common.Error(fmt.Sprintf("go side, doc %s, ChunkStoreExists failed", doc.ID), err) return common.CodeExceptionError, err } if exists { if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil { - common.Error(fmt.Sprintf("go side, doc %s, DeleteChunks failed", docID), err) + common.Error(fmt.Sprintf("go side, doc %s, DeleteChunks failed", doc.ID), err) return common.CodeExceptionError, err } } } } - - if run == string(entity.TaskStatusRunning) { - if req.ApplyKB { - if doc.ParserConfig == nil { - doc.ParserConfig = entity.JSONMap{} - } - config := map[string]interface{}{ - "llm_id": kb.ParserConfig["llm_id"], - "enable_metadata": false, - "metadata": map[string]interface{}{}, - } - if value, ok := kb.ParserConfig["enable_metadata"]; ok { - config["enable_metadata"] = value - } - if value, ok := kb.ParserConfig["metadata"]; ok { - config["metadata"] = value - } - if err := s.updateDocumentParserConfig(doc.ID, config); err != nil { - return common.CodeExceptionError, err - } - for key, value := range config { - doc.ParserConfig[key] = value - } - } - if doc.PipelineID != nil && strings.TrimSpace(*doc.PipelineID) != "" { - if err := s.queueDocumentDataflowTask(kb, doc, userID); err != nil { - return common.CodeExceptionError, err - } - continue - } - if doc.ParserID == string(entity.ParserTypeTable) { - doneCount, ok := tableDoneCountByKB[doc.KbID] - if !ok { - count, err := s.countDoneDocuments(doc.KbID) - if err != nil { - return common.CodeExceptionError, err - } - doneCount = count - tableDoneCountByKB[doc.KbID] = doneCount - if doneCount <= 0 { - if err := s.kbDAO.DeleteFieldMap(doc.KbID); err != nil && !dao.IsNotFoundErr(err) { - return common.CodeExceptionError, err - } - } - } - } - if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { - common.Error(fmt.Sprintf("go side, doc %s, DeleteByDocIDs", docID), err) - return common.CodeExceptionError, err - } - _, _, err := s.GetDocumentStorageAddress(doc) - if err != nil { - common.Error(fmt.Sprintf("go side, doc %s, GetDocumentStorageAddress", docID), err) - return common.CodeExceptionError, err - } - - if _, err := s.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil { - common.Error(fmt.Sprintf("go side, doc %s, IngestDocuments", docID), err) - return common.CodeExceptionError, err - } - - if err := s.beginDocumentParse(doc.ID); err != nil { - common.Error(fmt.Sprintf("go side, doc %s, beginDocumentParse", docID), err) - return common.CodeExceptionError, err - } - } } return common.CodeSuccess, nil } -func (s *DocumentService) prepareDocumentRerunWithDelete(doc *entity.Document, tenantID string) error { +// AssertIngestionTasksTerminal verifies none of the documents has an +// in-flight (RUNNING/STOPPING) ingestion task. Used as a batch pre-check +// before re-parsing so a single non-terminal doc rejects the whole request +// up front instead of partially cleaning some docs then failing. +func (s *DocumentService) AssertIngestionTasksTerminal(docIDs []string) error { + for _, docID := range docIDs { + task, err := s.ingestionTaskDAO.GetByDocumentID(docID) + if err != nil { + return fmt.Errorf("check ingestion task for %s: %w", docID, err) + } + if task == nil { + continue + } + if task.Status == common.RUNNING || task.Status == common.STOPPING { + return fmt.Errorf("document %s ingestion task is %s; stop it and wait for a terminal state before re-parsing", docID, task.Status) + } + } + return nil +} + +func (s *DocumentService) clearDocumentParseResults(doc *entity.Document, tenantID string) error { if doc == nil { return fmt.Errorf("document is nil") } - s.cancelExistingParseTasksBestEffort(doc.ID) + // Refuse to clear a non-terminal ingestion task. An in-flight worker + // (RUNNING) or one mid-stop (STOPPING) would keep writing chunks and + // corrupt the new run's results. The caller must stop the task first + // and wait for a terminal state (COMPLETED/STOPPED/FAILED) or CREATED. + if task, _ := s.ingestionTaskDAO.GetByDocumentID(doc.ID); task != nil { + if task.Status == common.RUNNING || task.Status == common.STOPPING { + return fmt.Errorf("document %s ingestion task is %s; stop it and wait for a terminal state before re-parsing", doc.ID, task.Status) + } + } - if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { + // Delete terminal and CREATED ingestion tasks atomically, leaving + // RUNNING/STOPPING tasks untouched so the check-then-delete window + // between GetByDocumentID and the delete above cannot delete a task + // that just transitioned to RUNNING. + if _, err := s.ingestionTaskDAO.DeleteIfTerminal(doc.ID); err != nil { return err } @@ -1389,24 +1433,6 @@ func (s *DocumentService) prepareDocumentRerunWithDelete(doc *entity.Document, t return nil } -func (s *DocumentService) cancelExistingParseTasksBestEffort(docID string) { - tasks, err := s.taskDAO.GetByDocID(docID) - if err != nil { - common.Logger.Warn(fmt.Sprintf("cancelExistingParseTasksBestEffort: failed to get tasks for %s: %v", docID, err)) - return - } - redisClient := redis.Get() - if redisClient == nil { - return - } - for _, task := range tasks { - if task == nil { - continue - } - redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 24*time.Hour) - } -} - func (s *DocumentService) clearDocumentAndKBCountersForRerun(docID, kbID string) error { return dao.DB.Transaction(func(tx *gorm.DB) error { var current entity.Document @@ -1458,368 +1484,6 @@ func (s *DocumentService) countDoneDocuments(datasetID string) (int64, error) { return count, err } -func (s *DocumentService) queueDocumentDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, userID string) error { - if err := s.beginDocumentParse(doc.ID); err != nil { - return err - } - _, err := s.ingestionTaskSvc.CreateAndEnqueue(&entity.IngestionTask{ - DocumentID: doc.ID, - UserID: userID, - DatasetID: kb.ID, - Schema: nil, - Status: common.CREATED, - }) - return err -} - -func (s *DocumentService) newDocumentParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) { - ranges, err := documentParseTaskRanges(doc, bucket, objectName) - if err != nil { - return nil, err - } - tasks := make([]*entity.Task, 0, len(ranges)) - for _, pageRange := range ranges { - tasks = append(tasks, s.newDocumentParseTask(doc, pageRange.from, pageRange.to, priority)) - } - return tasks, nil -} - -func (s *DocumentService) newDocumentParseTask(doc *entity.Document, fromPage, toPage, priority int64) *entity.Task { - now := time.Now() - progressMsg := "" - digest := documentParseTaskDigest(doc, fromPage, toPage) - chunkIDs := "" - return &entity.Task{ - ID: utility.GenerateUUID(), - DocID: doc.ID, - FromPage: fromPage, - ToPage: toPage, - TaskType: "", - Priority: priority, - BeginAt: &now, - Progress: 0, - ProgressMsg: &progressMsg, - Digest: &digest, - ChunkIDs: &chunkIDs, - } -} - -func documentParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]documentParsePageRange, error) { - if doc.Type == "pdf" { - binary, err := documentStorageBinary(bucket, objectName) - if err != nil { - return nil, err - } - pages := documentEstimatePDFPageCount(binary) - pageSize := int64(documentParserConfigInt(doc.ParserConfig, "task_page_size", 12)) - if doc.ParserID == string(entity.ParserTypePaper) { - pageSize = int64(documentParserConfigInt(doc.ParserConfig, "task_page_size", 22)) - } - if doc.ParserID == string(entity.ParserTypeOne) || - documentParserConfigBool(doc.ParserConfig, "toc_extraction", false) { - pageSize = maximumTaskPageNumber - } - if pageSize <= 0 { - pageSize = 12 - } - ranges := make([]documentParsePageRange, 0) - for _, configuredRange := range documentParserConfigPageRanges(doc.ParserConfig) { - start := configuredRange.from - 1 - if start < 0 { - start = 0 - } - end := configuredRange.to - 1 - if pages >= 0 && end > pages { - end = pages - } - for page := start; page < end; page += pageSize { - to := page + pageSize - if to > end { - to = end - } - ranges = append(ranges, documentParsePageRange{from: page, to: to}) - } - } - if len(ranges) == 0 { - // pages == 0 means page count detection failed (e.g. compressed - // PDF where both regex and pdfoxide fallbacks failed). Fall back - // to maximumTaskPageNumber so the Python parser processes all - // pages via slicing (Python gracefully caps at actual page count). - ranges = append(ranges, documentParsePageRange{from: 0, to: maximumTaskPageNumber}) - } - return ranges, nil - } - if doc.ParserID == string(entity.ParserTypeTable) { - binary, err := documentStorageBinary(bucket, objectName) - if err != nil { - return nil, err - } - rows := documentEstimateTableRowCount(documentName(doc), binary) - if rows <= 0 { - return []documentParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil - } - ranges := make([]documentParsePageRange, 0, (rows+2999)/3000) - for row := int64(0); row < int64(rows); row += 3000 { - to := row + 3000 - if to > int64(rows) { - to = int64(rows) - } - ranges = append(ranges, documentParsePageRange{from: row, to: to}) - } - return ranges, nil - } - return []documentParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil -} - -func documentStorageBinary(bucket, objectName string) ([]byte, error) { - storageImpl := storage.GetStorageFactory().GetStorage() - if storageImpl == nil { - return nil, fmt.Errorf("storage not initialized") - } - return storageImpl.Get(bucket, objectName) -} - -func documentName(doc *entity.Document) string { - if doc == nil || doc.Name == nil { - return "" - } - return *doc.Name -} - -func documentParserConfigInt(config map[string]interface{}, key string, fallback int) int { - value, ok := config[key] - if !ok || value == nil { - return fallback - } - switch typedValue := value.(type) { - case int: - return typedValue - case int64: - return int(typedValue) - case float64: - return int(typedValue) - case json.Number: - if intValue, err := typedValue.Int64(); err == nil { - return int(intValue) - } - case string: - if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil { - return intValue - } - } - return fallback -} - -func documentParserConfigBool(config map[string]interface{}, key string, fallback bool) bool { - value, ok := config[key] - if !ok || value == nil { - return fallback - } - switch typedValue := value.(type) { - case bool: - return typedValue - case string: - switch strings.ToLower(strings.TrimSpace(typedValue)) { - case "true", "1", "yes", "on": - return true - case "false", "0", "no", "off": - return false - } - } - return fallback -} - -func documentParserConfigPageRanges(config map[string]interface{}) []documentParsePageRange { - defaultRanges := []documentParsePageRange{{from: 1, to: 100000}} - raw, ok := config["pages"] - if !ok || raw == nil { - return defaultRanges - } - rawRanges, ok := raw.([]interface{}) - if !ok || len(rawRanges) == 0 { - return defaultRanges - } - ranges := make([]documentParsePageRange, 0, len(rawRanges)) - for _, rawRange := range rawRanges { - rangeValues, ok := rawRange.([]interface{}) - if !ok || len(rangeValues) < 2 { - continue - } - from, okFrom := documentToInt64(rangeValues[0]) - to, okTo := documentToInt64(rangeValues[1]) - if okFrom && okTo && to > from { - ranges = append(ranges, documentParsePageRange{from: from, to: to}) - } - } - if len(ranges) == 0 { - return defaultRanges - } - return ranges -} - -func documentToInt64(value interface{}) (int64, bool) { - switch typedValue := value.(type) { - case int: - return int64(typedValue), true - case int64: - return typedValue, true - case float64: - return int64(typedValue), true - case json.Number: - intValue, err := typedValue.Int64() - return intValue, err == nil - case string: - intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64) - return intValue, err == nil - default: - return 0, false - } -} - -var documentPDFPagePattern = regexp.MustCompile(`/Type\s*/Page\b`) - -func documentEstimatePDFPageCount(binary []byte) int64 { - if len(binary) == 0 { - return 0 - } - // Fast path: regex works for uncompressed PDFs. - count := int64(len(documentPDFPagePattern.FindAll(binary, -1))) - if count > 0 { - return count - } - // Fallback for compressed PDFs where /Type /Page is inside a - // compressed object stream: use pdf_oxide to get the real page count. - if doc, err := pdfoxide.OpenBytes(binary); err == nil { - defer doc.Close() - if pages, err := doc.PageCount(); err == nil { - return int64(pages) - } - } - return 0 -} - -func documentEstimateTableRowCount(name string, binary []byte) int { - switch strings.ToLower(filepath.Ext(name)) { - case ".xlsx": - if rows, err := documentCountXLSXRows(binary); err == nil { - return rows - } - case ".csv", ".tsv", ".txt": - return documentCountDelimitedRows(name, binary) - } - return 0 -} - -func documentCountDelimitedRows(name string, binary []byte) int { - reader := csv.NewReader(bytes.NewReader(binary)) - reader.FieldsPerRecord = -1 - reader.ReuseRecord = true - if strings.EqualFold(filepath.Ext(name), ".tsv") { - reader.Comma = '\t' - } - rows := 0 - for { - _, err := reader.Read() - if err == nil { - rows++ - continue - } - if err == io.EOF { - break - } - rows += bytes.Count(binary, []byte{'\n'}) - if len(binary) > 0 && binary[len(binary)-1] != '\n' { - rows++ - } - break - } - return rows -} - -func documentCountXLSXRows(binary []byte) (int, error) { - zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary))) - if err != nil { - return 0, err - } - maxRows := 0 - for _, file := range zipReader.File { - if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") { - continue - } - rows, err := documentCountWorksheetRows(file) - if err != nil { - return 0, err - } - if rows > maxRows { - maxRows = rows - } - } - return maxRows, nil -} - -func documentCountWorksheetRows(file *zip.File) (int, error) { - reader, err := file.Open() - if err != nil { - return 0, err - } - defer reader.Close() - decoder := xml.NewDecoder(reader) - rows := 0 - for { - token, err := decoder.Token() - if err == io.EOF { - break - } - if err != nil { - return 0, err - } - start, ok := token.(xml.StartElement) - if ok && start.Name.Local == "row" { - rows++ - } - } - return rows, nil -} - -func (s *DocumentService) beginDocumentParse(docID string) error { - now := time.Now() - return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{ - "progress_msg": "Task is queued...", - "process_begin_at": now, - "progress": rand.Float64() * 0.01, - "run": string(entity.TaskStatusRunning), - "chunk_num": 0, - "token_num": 0, - }).Error -} - -func documentParseTaskDigest(doc *entity.Document, fromPage, toPage int64) string { - hasher := xxhash.New() - config := map[string]interface{}{ - "doc_id": doc.ID, - "kb_id": doc.KbID, - "parser_id": doc.ParserID, - "parser_config": doc.ParserConfig, - } - keys := make([]string, 0, len(config)) - for key := range config { - keys = append(keys, key) - } - sort.Strings(keys) - for _, key := range keys { - b, err := json.Marshal(config[key]) - if err != nil { - hasher.WriteString(fmt.Sprint(config[key])) - } else { - hasher.Write(b) - } - } - hasher.WriteString(doc.ID) - hasher.WriteString(strconv.FormatInt(fromPage, 10)) - hasher.WriteString(strconv.FormatInt(toPage, 10)) - return fmt.Sprintf("%x", hasher.Sum64()) -} - func (s *DocumentService) clearKBChunkNumWhenRerun(doc *entity.Document) error { if doc == nil { return fmt.Errorf("document is nil") @@ -1915,7 +1579,7 @@ func (s *DocumentService) StopParseDocuments(datasetID string, docIDs []string) var errors []string successCount := 0 for _, doc := range docs { - if cancelErr := s.cancelDocParse(doc); cancelErr != nil { + if cancelErr := s.CancelDocParse(doc); cancelErr != nil { errors = append(errors, cancelErr.Error()) continue } @@ -1952,10 +1616,9 @@ func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID strin return docs, nil } -// cancelDocParse stops the ingestion task for the document by calling -// RequestStop (which sets Redis {taskID}-cancel + DB STOPPING), then marks -// the document run status as CANCEL. -func (s *DocumentService) cancelDocParse(doc *entity.Document) error { +// CancelDocParse stops the ingestion task for the document by calling +// RequestStop (STOPPING), then marks the document run status as CANCEL. +func (s *DocumentService) CancelDocParse(doc *entity.Document) error { task, err := s.ingestionTaskDAO.GetByDocumentID(doc.ID) if err != nil { return fmt.Errorf("failed to get ingestion task for %s: %v", doc.ID, err) diff --git a/internal/service/document_test.go b/internal/service/document_test.go index 09ebcdaad5..e6f6d7b297 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -517,6 +517,11 @@ func insertTestTask(t *testing.T, id, docID string) { } func insertTestIngestionTask(t *testing.T, id, userID, docID, datasetID string) { + t.Helper() + insertTestIngestionTaskWithStatus(t, id, userID, docID, datasetID, common.CREATED) +} + +func insertTestIngestionTaskWithStatus(t *testing.T, id, userID, docID, datasetID, status string) { t.Helper() task := &entity.IngestionTask{ ID: id, @@ -524,7 +529,7 @@ func insertTestIngestionTask(t *testing.T, id, userID, docID, datasetID string) DocumentID: docID, DatasetID: datasetID, Schema: entity.JSONMap{}, - Status: common.CREATED, + Status: status, } if err := dao.DB.Create(task).Error; err != nil { t.Fatalf("insert test ingestion task: %v", err) @@ -1077,12 +1082,16 @@ func insertTestTaskWithProgress(t *testing.T, id, docID string, progress float64 } } -func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) { +func TestStartParseDocuments_EnqueuesIngestionTask(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) insertTestDoc(t, "doc-1", "kb-1", 9, 4) + // GetDocumentStorageAddress requires a non-empty document location. + if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("location", "loc-1").Error; err != nil { + t.Fatalf("set location: %v", err) + } publisher := &recordingTaskPublisher{} svc := testDocumentService(t) @@ -1097,8 +1106,8 @@ func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) { t.Fatalf("load doc: %v", err) } - if err := svc.queueDocumentDataflowTask(kb, doc, "user-1"); err != nil { - t.Fatalf("queueDocumentDataflowTask: %v", err) + if err := svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{}); err != nil { + t.Fatalf("StartParseDocuments: %v", err) } if publisher.subject != "tasks.RAGFLOW" { @@ -1805,12 +1814,12 @@ func TestResetDocumentForReparseSkipsSecondCounterDecrement(t *testing.T) { } } -func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T) { +func TestClearDocumentParseResultsClearsCountersTasksAndChunks(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) - insertTestTask(t, "task-1", "doc-1") + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) doc, err := dao.NewDocumentDAO().GetByID("doc-1") if err != nil { @@ -1821,8 +1830,8 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T svc := testDocumentService(t) svc.docEngine = engine - if err := svc.prepareDocumentRerunWithDelete(doc, "tenant-1"); err != nil { - t.Fatalf("prepareDocumentRerunWithDelete failed: %v", err) + if err := svc.clearDocumentParseResults(doc, "tenant-1"); err != nil { + t.Fatalf("clearDocumentParseResults failed: %v", err) } updatedDoc, _ := dao.NewDocumentDAO().GetByID("doc-1") @@ -1833,12 +1842,10 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T if kb.TokenNum != 0 || kb.ChunkNum != 0 { t.Fatalf("kb counters = token:%d chunk:%d, want zero", kb.TokenNum, kb.ChunkNum) } - var taskCount int64 - if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { - t.Fatalf("count tasks: %v", err) - } - if taskCount != 0 { - t.Fatalf("task count = %d, want zero", taskCount) + // The completed ingestion task must be deleted so the new run can proceed. + remainingTask, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if remainingTask != nil { + t.Fatalf("ingestion task should be deleted, status was %q", remainingTask.Status) } if engine.deleteCalls != 1 { t.Fatalf("deleteCalls = %d, want 1", engine.deleteCalls) @@ -1848,7 +1855,7 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T } } -func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testing.T) { +func TestClearDocumentParseResultsIsIdempotentForStaleDocSnapshot(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) @@ -1860,11 +1867,11 @@ func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testin } svc := testDocumentService(t) - if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil { - t.Fatalf("first prepareDocumentRerunWithDelete failed: %v", err) + if err := svc.clearDocumentParseResults(staleDoc, "tenant-1"); err != nil { + t.Fatalf("first clearDocumentParseResults failed: %v", err) } - if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil { - t.Fatalf("second prepareDocumentRerunWithDelete failed: %v", err) + if err := svc.clearDocumentParseResults(staleDoc, "tenant-1"); err != nil { + t.Fatalf("second clearDocumentParseResults failed: %v", err) } doc, _ := dao.NewDocumentDAO().GetByID("doc-1") @@ -1877,6 +1884,126 @@ func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testin } } +// TestClearDocumentParseResults_RejectsNonTerminalIngestionTask verifies +// that re-parsing is refused while a document's ingestion task is still +// RUNNING or STOPPING. Deleting a non-terminal task would let the in-flight +// worker keep writing chunks and corrupt the new run's results. +func TestClearDocumentParseResults_RejectsNonTerminalIngestionTask(t *testing.T) { + for _, status := range []string{common.RUNNING, common.STOPPING} { + t.Run(status, func(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) + insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusRunning), 10, 5) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status) + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + svc := testDocumentService(t) + + if err := svc.clearDocumentParseResults(doc, "tenant-1"); err == nil { + t.Fatalf("expected error for %s ingestion task, got nil", status) + } + // The non-terminal task must NOT be deleted. + task, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if task == nil { + t.Fatalf("%s ingestion task must not be deleted", status) + } + }) + } +} + +// TestClearDocumentParseResults_DeletesTerminalIngestionTask verifies +// that a terminal ingestion task (STOPPED/COMPLETED/FAILED) and a CREATED +// (queued) task are deleted so the new run can proceed. +func TestClearDocumentParseResults_DeletesTerminalIngestionTask(t *testing.T) { + for _, status := range []string{common.CREATED, common.COMPLETED, common.STOPPED, common.FAILED} { + t.Run(status, func(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) + insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status) + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + svc := testDocumentService(t) + + if err := svc.clearDocumentParseResults(doc, "tenant-1"); err != nil { + t.Fatalf("clearDocumentParseResults for %s task: %v", status, err) + } + task, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if task != nil { + t.Fatalf("%s ingestion task should be deleted, still present", status) + } + }) + } +} + +func TestAssertIngestionTasksTerminal_RejectsNonTerminal(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 0, 0) + insertTestDoc(t, "doc-2", "kb-1", 0, 0) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) + insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING) + + svc := testDocumentService(t) + if err := svc.AssertIngestionTasksTerminal([]string{"doc-1", "doc-2"}); err == nil { + t.Fatal("expected error for RUNNING task, got nil") + } +} + +func TestAssertIngestionTasksTerminal_AcceptsAllTerminal(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 0, 0) + insertTestDoc(t, "doc-2", "kb-1", 0, 0) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) + insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.STOPPED) + + svc := testDocumentService(t) + if err := svc.AssertIngestionTasksTerminal([]string{"doc-1", "doc-2"}); err != nil { + t.Fatalf("expected nil for all-terminal batch, got %v", err) + } +} + +// TestIngest_RerunWithDelete_RejectsBatchWithRunningTask verifies the batch +// pre-check: when re-parsing with delete, if any document's ingestion task is +// non-terminal, the whole request is rejected up front so no document is +// partially cleaned. +func TestIngest_RerunWithDelete_RejectsBatchWithRunningTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertUserTenantForAccessCheck(t, "user-1", "tenant-1") + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 0, 0) + insertTestDoc(t, "doc-2", "kb-1", 0, 0) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) + insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING) + + svc := testDocumentService(t) + _, err := svc.Ingest("user-1", &IngestDocumentRequest{ + DocIDs: []string{"doc-1", "doc-2"}, + Run: string(entity.TaskStatusRunning), + Delete: true, + }) + if err == nil { + t.Fatal("expected error for batch with RUNNING task, got nil") + } + // doc-1 (terminal) task must NOT be deleted - the whole batch was rejected. + task1, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if task1 == nil { + t.Fatal("doc-1 terminal task should not be deleted when batch is rejected") + } +} + func TestUpdateDatasetDocumentPropagatesMetadataDeleteFailure(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) @@ -2506,3 +2633,69 @@ func TestGetThumbnails_AlignsWithPythonFormatting(t *testing.T) { t.Fatalf("did not expect other tenant doc in result: %#v", got) } } + +// TestStartParseDocuments_FailsBeforeClearing verifies that GetDocumentStorageAddress +// runs before clearDocumentParseResults: when storage is invalid, we fail without +// deleting existing parse results. +func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 0, 10, 5) + insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED) + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get kb: %v", err) + } + + // Force GetDocumentStorageAddress to fail by clearing the document location. + if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("location", "").Error; err != nil { + t.Fatalf("clear location: %v", err) + } + + svc := testDocumentService(t) + err = svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{RerunWithDelete: true}) + if err == nil { + t.Fatal("expected error from GetDocumentStorageAddress, got nil") + } + + // The old ingestion task must still exist — we failed before clearing. + remaining, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if remaining == nil { + t.Fatal("ingestion task should NOT be deleted when storage validation fails") + } +} + +// TestIngest_CancelDoesNotDeleteIngestionTask verifies that cancel+delete +// no longer races with the worker by deleting the ingestion task while the +// worker is still stopping. Cancel only issues RequestStop (STOPPING); +// deletion is deferred until the task reaches a terminal state. +func TestIngest_CancelDoesNotDeleteIngestionTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertUserTenantForAccessCheck(t, "user-1", "tenant-1") + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 10, 5) + insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING) + + svc := testDocumentService(t) + _, err := svc.Ingest("user-1", &IngestDocumentRequest{ + DocIDs: []string{"doc-1"}, + Run: string(entity.TaskStatusCancel), + Delete: true, + }) + if err != nil { + t.Fatalf("Ingest(cancel+delete): %v", err) + } + + // The ingestion task must NOT be deleted — cancel only stops it. + remaining, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1") + if remaining == nil { + t.Fatal("ingestion task must NOT be deleted by cancel") + } +} diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index da08e68ef5..7da77bb64b 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -205,7 +205,27 @@ func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTas } switch task.Status { case common.CREATED: - return s.transition(taskID, common.RUNNING) + task, err := s.transition(taskID, common.RUNNING) + if err != nil { + return nil, err + } + // The task just started running: mirror it to the document so its + // run status and progress counters reflect real processing, not + // just API acceptance. Best-effort - a DB blip here must not fail + // the task transition and trigger a redelivery loop. run uses the + // document's numeric TaskStatus enum ("1"), not the task's string + // status label. + if err := s.documentDAO.UpdateByID(task.DocumentID, map[string]interface{}{ + "run": string(entity.TaskStatusRunning), + "progress": float64(0), + "chunk_num": int64(0), + "token_num": int64(0), + "process_begin_at": time.Now(), + "progress_msg": "", + }); err != nil { + common.Warn(fmt.Sprintf("StartRunning: mark document %s running for task %s: %v", task.DocumentID, taskID, err)) + } + return task, nil case common.STOPPING: return s.transition(taskID, common.STOPPED) case common.RUNNING, common.COMPLETED, common.STOPPED, common.FAILED: diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index 3f0714207c..c4d1b88563 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -213,6 +213,102 @@ func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) { } } +// TestStartRunningMarksDocumentRunning locks in that starting a CREATED task +// mirrors the transition to its document: run=RUNNING and progress counters +// reset, with a fresh process_begin_at. The document bookkeeping is owned by +// the task-lifecycle transition, not the ingestion worker's execution path. +func TestStartRunningMarksDocumentRunning(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 100, 10) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + // Seed the document as a partially-processed, non-RUNNING state that the + // start transition must clobber. + if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1"). + Updates(map[string]interface{}{ + "run": string(entity.TaskStatusDone), + "progress": float64(0.5), + "progress_msg": "partial", + }).Error; err != nil { + t.Fatalf("seed document: %v", err) + } + + svc := NewIngestionTaskService() + if _, err := svc.StartRunning("task-1"); err != nil { + t.Fatalf("StartRunning failed: %v", err) + } + + var doc entity.Document + if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil { + t.Fatalf("reload document: %v", err) + } + if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) { + t.Fatalf("run = %v, want RUNNING(%q)", doc.Run, string(entity.TaskStatusRunning)) + } + if doc.Progress != 0 { + t.Fatalf("progress = %f, want 0", doc.Progress) + } + if doc.ChunkNum != 0 { + t.Fatalf("chunk_num = %d, want 0", doc.ChunkNum) + } + if doc.TokenNum != 0 { + t.Fatalf("token_num = %d, want 0", doc.TokenNum) + } + if doc.ProgressMsg != nil && *doc.ProgressMsg != "" { + t.Fatalf("progress_msg = %q, want empty", *doc.ProgressMsg) + } + if doc.ProcessBeginAt == nil || doc.ProcessBeginAt.IsZero() { + t.Fatal("process_begin_at not set") + } +} + +// TestStartRunningLeavesTerminalDocumentUntouched locks in the no-resurrection +// invariant: a task already in a terminal status is returned as-is by +// StartRunning, and its document's finished run status/counters are not reset. +func TestStartRunningLeavesTerminalDocumentUntouched(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 100, 10) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + finishedRun := string(entity.TaskStatusDone) + if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1"). + Updates(map[string]interface{}{ + "run": finishedRun, + "progress": float64(1.0), + "progress_msg": "done", + }).Error; err != nil { + t.Fatalf("seed document: %v", err) + } + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", "task-1"). + Update("status", common.COMPLETED).Error; err != nil { + t.Fatalf("set COMPLETED: %v", err) + } + + svc := NewIngestionTaskService() + task, err := svc.StartRunning("task-1") + if err != nil { + t.Fatalf("StartRunning failed: %v", err) + } + if task.Status != common.COMPLETED { + t.Fatalf("status = %q, want %q (terminal must be preserved)", task.Status, common.COMPLETED) + } + + var doc entity.Document + if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil { + t.Fatalf("reload document: %v", err) + } + if doc.Run == nil || *doc.Run != finishedRun { + t.Fatalf("run = %v, want %q (terminal document must not be resurrected)", doc.Run, finishedRun) + } + if doc.ChunkNum != 10 || doc.TokenNum != 100 { + t.Fatalf("counters changed: chunk_num=%d token_num=%d, want 10/100", doc.ChunkNum, doc.TokenNum) + } +} + func TestIngestionTaskServiceRequestStopTransitionsCreatedTaskToStopped(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db)