From 3bca4b3485dcdd35ea04f724cf073da58abf75eb Mon Sep 17 00:00:00 2001
From: Stephen Zhou <38493346+hyoban@users.noreply.github.com>
Date: Thu, 30 Jul 2026 18:26:04 +0800
Subject: [PATCH] fix(dataset): retry failed New RAG document processing

---
 .../new-rag/__tests__/documents-page.spec.tsx | 59 +++++++++++++++++++
 web/features/new-rag/documents-page.tsx       | 29 +++++++++
 2 files changed, 88 insertions(+)

diff --git a/web/features/new-rag/__tests__/documents-page.spec.tsx b/web/features/new-rag/__tests__/documents-page.spec.tsx
index 671aab7f279..a291187e93e 100644
--- a/web/features/new-rag/__tests__/documents-page.spec.tsx
+++ b/web/features/new-rag/__tests__/documents-page.spec.tsx
@@ -780,6 +780,65 @@ describe('DocumentsPage', () => {
     )
   })
 
+  it('retries failed initial processing from a document row action', async () => {
+    const user = userEvent.setup()
+    documentsQuery.data = {
+      pages: [
+        {
+          items: [
+            document({
+              active: null,
+              activeRevision: undefined,
+              id: 'failed-document',
+              status: 'failed',
+              title: 'Failed.pdf',
+            }),
+          ],
+        },
+      ],
+    }
+    tasksQuery.data = {
+      pages: [
+        {
+          items: [
+            task({
+              documentId: 'failed-document',
+              documentRevision: 1,
+              id: 'failed-task',
+              state: 'failed',
+            }),
+          ],
+        },
+      ],
+    }
+
+    render(<DocumentsPage knowledgeSpaceId="space-1" />)
+    await user.click(
+      screen.getByRole('button', {
+        name: /dataset\.newKnowledge\.documentActions/,
+      }),
+    )
+    await user.click(
+      await screen.findByRole('menuitem', {
+        name: 'dataset.newKnowledge.reindexDocument',
+      }),
+    )
+
+    expect(retryMutation.mutateAsync).toHaveBeenCalledWith({
+      params: {
+        control_space_id: 'space-1',
+        task_id: 'failed-task',
+        task_kind: 'document',
+      },
+    })
+    expect(reindexMutation.mutateAsync).not.toHaveBeenCalled()
+    await waitFor(() =>
+      expect(toastMock.success).toHaveBeenCalledWith(
+        'dataset.newKnowledge.documentsReindexStarted',
+      ),
+    )
+  })
+
   it('opens the upload form and consumes the one-shot URL request', async () => {
     const user = userEvent.setup()
     const { onUrlUpdate } = render(<DocumentsPage knowledgeSpaceId="space-1" />, {
diff --git a/web/features/new-rag/documents-page.tsx b/web/features/new-rag/documents-page.tsx
index af2b4b953ef..f3a74a7c37e 100644
--- a/web/features/new-rag/documents-page.tsx
+++ b/web/features/new-rag/documents-page.tsx
@@ -39,6 +39,7 @@ import {
   documentDisplayStatus,
   newestTaskByDocument,
   sourceName,
+  taskCanRetry,
   taskIsActive,
   taskNeedsAttention,
   taskVersionIsAfter,
@@ -279,6 +280,23 @@ export function DocumentsPage({ knowledgeSpaceId }: { knowledgeSpaceId: string }
   const { mutateAsync: reindexDocuments } = useMutation(
     consoleQuery.knowledgeFs.spaces.byControlSpaceId.documents.reindex.post.mutationOptions(),
   )
+  const { mutateAsync: retryProcessingTask } = useMutation({
+    mutationFn: async (task: DocumentProcessingTask) => {
+      const updated = documentTaskFromApi(
+        await consoleClient.knowledgeFs.spaces.byControlSpaceId.backgroundTasks.byTaskKind.byTaskId.retry.post(
+          {
+            params: {
+              control_space_id: knowledgeSpaceId,
+              task_id: task.id,
+              task_kind: task.taskKind ?? 'document',
+            },
+          },
+        ),
+      )
+      if (!updated) throw new Error('KnowledgeFS returned a non-document task')
+      return updated
+    },
+  })
   const {
     configureModelSetup,
     ensureModelSetupReady,
@@ -1622,6 +1640,14 @@ export function DocumentsPage({ knowledgeSpaceId }: { knowledgeSpaceId: string }
       setReindexing(true)
       try {
         if (!(await ensureModelSetupReady())) return
+        const document = documents.find((candidate) => candidate.id === documentId)
+        const retryableTask = document?.active ? undefined : taskByDocument.get(documentId)
+        if (retryableTask && taskCanRetry(retryableTask)) {
+          await retryProcessingTask(retryableTask)
+          toast.success(t(($) => $['newKnowledge.documentsReindexStarted']))
+          refreshDocumentsAndTasks()
+          return
+        }
         const result = await reindexDocuments({
           body: { documentIds: [documentId] },
           params: { control_space_id: knowledgeSpaceId },
@@ -1645,12 +1671,15 @@ export function DocumentsPage({ knowledgeSpaceId }: { knowledgeSpaceId: string }
     },
     [
       canWrite,
+      documents,
       ensureModelSetupReady,
       handleWritePermissionDenied,
       knowledgeSpaceId,
       refreshDocumentsAndTasks,
       reindexDocuments,
+      retryProcessingTask,
       t,
+      taskByDocument,
     ],
   )