From ce816edb5ff0a2b5a2f3ca7912ba9dc881c7ca16 Mon Sep 17 00:00:00 2001
From: Stephen Hu <stephenhu@seismic.com>
Date: Thu, 22 May 2025 09:28:08 +0800
Subject: [PATCH] Fix: improve task cancel lag (#7765)

### What problem does this PR solve?

https://github.com/infiniflow/ragflow/issues/7761

but it may be difficult to achieve 0 delay (which need to pass the
cancel token to all parts)

Another solution is just 0 delay effect at UI.
And task will stop latter

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 rag/svr/task_executor.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py
index 6375be6dd..4e9eb3d2a 100644
--- a/rag/svr/task_executor.py
+++ b/rag/svr/task_executor.py
@@ -368,6 +368,10 @@ async def build_chunks(task, progress_callback):
 
         docs_to_tag = []
         for d in docs:
+            task_canceled = TaskService.do_cancel(task["id"])
+            if task_canceled:
+                progress_callback(-1, msg="Task has been canceled.")
+                return
             if settings.retrievaler.tag_content(tenant_id, kb_ids, d, all_tags, topn_tags=topn_tags, S=S) and len(d[TAG_FLD]) > 0:
                 examples.append({"content": d["content_with_weight"], TAG_FLD: d[TAG_FLD]})
             else:
@@ -589,6 +593,10 @@ async def do_handle_task(task):
 
     for b in range(0, len(chunks), es_bulk_size):
         doc_store_result = await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert(chunks[b:b + es_bulk_size], search.index_name(task_tenant_id), task_dataset_id))
+        task_canceled = TaskService.do_cancel(task_id)
+        if task_canceled:
+            progress_callback(-1, msg="Task has been canceled.")
+            return
         if b % 128 == 0:
             progress_callback(prog=0.8 + 0.1 * (b + 1) / len(chunks), msg="")
         if doc_store_result: