feat: workflow cancel

apps/application/workflow/common.py

@@ -7,6 +7,9 @@
 @desc:
 """
 
+import asyncio
+import contextlib
+import threading
 from enum import Enum
 from typing import List, Dict
 
@@ -246,3 +249,52 @@ def new_instance(flow_obj: Dict, workflow_type: WorkflowType = WorkflowType.APPL
     nodes = [Node(node.get("id"), node.get("type"), **node) for node in nodes]
     edges = [Edge(edge.get("id"), edge.get("type"), **edge) for edge in edges]
     return Workflow(nodes, edges)
+
+
+class CancelledException(Exception):
+    """工作流取消异常"""
+
+    pass
+
+
+class AsyncRunner:
+    """跨线程取消感知地运行异步协程;cancel() 置位事件使阻塞中的流式中断并抛 CancelledException。"""
+
+    def __init__(self):
+        self.cancel_event = threading.Event()
+
+    def cancel(self):
+        """请求取消(任意线程可调);只置旗标,真正中断发生在事件循环内。"""
+        self.cancel_event.set()
+
+    def run_async(self, async_factory):
+        """在独立事件循环里运行 async_factory() 的协程;被取消则抛 CancelledException。"""
+        return asyncio.run(self.await_or_cancel(async_factory()))
+
+    async def _wait_cancel(self):
+        # 轮询取消旗标,置位即结束,充当竞速对手
+        while not self.cancel_event.is_set():
+            await asyncio.sleep(0)
+
+    async def await_or_cancel(self, _async):
+        """竞速等待:流式完成返回其结果;取消先到则抛 CancelledException。"""
+
+        # 真正要跑的流式协程,包成 task 才能与 waiter 一起放进 asyncio.wait
+        stream = asyncio.create_task(_async)
+        # 把"是否已取消"伪装成一个可等待的 task,充当 stream 的竞速对手。
+        waiter = asyncio.create_task(self._wait_cancel())
+        # 同时等两个 task,谁先完成(返回)谁进入 done;FIRST_COMPLETED = 取先到者。
+        done, _ = await asyncio.wait({stream, waiter}, return_when=asyncio.FIRST_COMPLETED)
+        # stream_done 为真 -> 流式自己先跑完;为假 -> 是 waiter 先完成,即取消先到。
+        stream_done = stream in done
+        # 无论谁赢,都把两个 task cancel 并等它们退干净,防止"输掉的那一方"
+        # 还悬在事件循环里。await 吞掉 CancelledError(取消是本意,不算错误)。
+        for task in (stream, waiter):
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+        # 流式正常结束:把它的返回值交给调用方。
+        if stream_done:
+            return stream.result()
+        raise CancelledException()

apps/application/workflow/i_node.py

@@ -14,19 +14,13 @@
 
 from rest_framework import serializers
 
-from application.workflow.common import Node
+from application.workflow.common import CancelledException, Node
 from application.workflow.message.struct.content import Content, NodeInfo, Position
 from application.workflow.message.struct.progress_content import ProgressContent
 from application.workflow.status import Status
 from common.utils.logger import maxkb_logger
 
 
-class CancelledException(Exception):
-    """工作流取消异常"""
-
-    pass
-
-
 class Signal(str, Enum):
     BREAK = "BREAK"
     CONTINUE = "CONTINUE"

apps/application/workflow/nodes/ai_chat_node/agent.py

@@ -24,7 +24,7 @@
 from langgraph.checkpoint.memory import MemorySaver
 
 from application.workflow.backend.sandbox_shell import SandboxShellBackend
-from application.workflow.i_node import CancelledException
+from application.workflow.common import CancelledException
 from application.workflow.nodes.ai_chat_node.tools.skill import init_skills
 from maxkb.const import CONFIG
 

apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py

@@ -8,6 +8,7 @@
 """
 
 import base64
+import contextlib
 import json
 import re
 from functools import reduce
@@ -19,7 +20,7 @@
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage, AIMessageChunk
 from rest_framework import serializers
 
-from application.workflow.common import WorkflowType
+from application.workflow.common import AsyncRunner, WorkflowType
 from application.workflow.i_node import INode
 from application.workflow.message.aggregator import AggregationManager
 from application.workflow.message.struct.content import NodeInfo, Position, Content
@@ -218,6 +219,14 @@ class AIChatNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "ai-chat-node"
 
+    def __init__(self, node, workflow_manage, get_node_parameters):
+        super().__init__(node, workflow_manage, get_node_parameters)
+        self.async_runner = AsyncRunner()
+
+    def cancel(self):
+        self.async_runner.cancel()
+        super().cancel()
+
     def write(self, message: Content):
         super().write(message)
         if not self.data.get("messages"):
@@ -345,9 +354,15 @@ def execute(self):
             message_list_with_system = [SystemMessage(system)] + message_list
 
             if stream:
-                r = chat_model.stream(message_list_with_system)
-                self._stream_response(
-                    r, chat_model, message_list_with_system, question.content, reasoning_content_id, text_content_id
+                self.async_runner.run_async(
+                    lambda: self._stream_response(
+                        chat_model.astream(message_list_with_system),
+                        chat_model,
+                        message_list_with_system,
+                        question.content,
+                        reasoning_content_id,
+                        text_content_id,
+                    )
                 )
             else:
                 r = chat_model.invoke(message_list_with_system)
@@ -370,40 +385,71 @@ def _generate_prompt_question(self, question_str, model, vision, image_list, vid
                 videos = _process_videos(video, model)
         return HumanMessage(content=[*videos, *images, {"type": "text", "text": question_str}])
 
-    def _stream_response(self, response, chat_model, message_list, question, reasoning_content_id, text_content_id):
-        node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.RUNNING)
-        model_setting = self.get_context("model_setting") or {}
-        reasoning = Reasoning(
-            model_setting.get("reasoning_content_start", "<think>"),
-            model_setting.get("reasoning_content_end", "</think>"),
-        )
-        answer = ""
-        reasoning_content = ""
-        response_reasoning_content = False
-
-        for chunk in response:
-            self._check_cancelled()
-            reasoning_chunk = reasoning.get_reasoning_content(chunk)
-            content_chunk = reasoning_chunk.get("content")
-            if "reasoning_content" in chunk.additional_kwargs:
-                response_reasoning_content = True
-                reasoning_content_chunk = chunk.additional_kwargs.get("reasoning_content", "")
-            else:
-                reasoning_content_chunk = reasoning_chunk.get("reasoning_content")
-            answer += content_chunk
-            if reasoning_content_chunk is None:
-                reasoning_content_chunk = ""
-            reasoning_content += reasoning_content_chunk
-            reasoning_end = False
-            if content_chunk:
-                if not reasoning_end:
+    async def _stream_response(
+        self, response, chat_model, message_list, question, reasoning_content_id, text_content_id
+    ):
+        try:
+            node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.RUNNING)
+            model_setting = self.get_context("model_setting") or {}
+            reasoning = Reasoning(
+                model_setting.get("reasoning_content_start", "<think>"),
+                model_setting.get("reasoning_content_end", "</think>"),
+            )
+            answer = ""
+            reasoning_content = ""
+            response_reasoning_content = False
+
+            async for chunk in response:
+                self._check_cancelled()
+                reasoning_chunk = reasoning.get_reasoning_content(chunk)
+                content_chunk = reasoning_chunk.get("content")
+                if "reasoning_content" in chunk.additional_kwargs:
+                    response_reasoning_content = True
+                    reasoning_content_chunk = chunk.additional_kwargs.get("reasoning_content", "")
+                else:
+                    reasoning_content_chunk = reasoning_chunk.get("reasoning_content")
+                answer += content_chunk
+                if reasoning_content_chunk is None:
+                    reasoning_content_chunk = ""
+                reasoning_content += reasoning_content_chunk
+                reasoning_end = False
+                if content_chunk:
+                    if not reasoning_end:
+                        self.write(
+                            ReasoningContent(
+                                reasoning_content_id, "", Status.SUCCESS, node_info, Position(self.get_node_id())
+                            )
+                        )
+                    self.write(
+                        TextContent(
+                            text_content_id, content_chunk, Status.RUNNING, node_info, Position(self.get_node_id())
+                        )
+                    )
+                if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
                     self.write(
                         ReasoningContent(
-                            reasoning_content_id, "", Status.SUCCESS, node_info, Position(self.get_node_id())
+                            reasoning_content_id,
+                            reasoning_content_chunk,
+                            Status.RUNNING,
+                            node_info,
+                            Position(self.get_node_id()),
                         )
                     )
+
+            reasoning_end = reasoning.get_end_reasoning_content()
+            answer += reasoning_end.get("content")
+            reasoning_content_chunk = ""
+            if not response_reasoning_content:
+                reasoning_content_chunk = reasoning_end.get("reasoning_content")
+            if reasoning_end.get("content"):
                 self.write(
-                    TextContent(text_content_id, content_chunk, Status.RUNNING, node_info, Position(self.get_node_id()))
+                    TextContent(
+                        text_content_id,
+                        reasoning_end.get("content"),
+                        Status.RUNNING,
+                        node_info,
+                        Position(self.get_node_id()),
+                    )
                 )
             if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
                 self.write(
@@ -416,33 +462,11 @@ def _stream_response(self, response, chat_model, message_list, question, reasoni
                     )
                 )
 
-        reasoning_end = reasoning.get_end_reasoning_content()
-        answer += reasoning_end.get("content")
-        reasoning_content_chunk = ""
-        if not response_reasoning_content:
-            reasoning_content_chunk = reasoning_end.get("reasoning_content")
-        if reasoning_end.get("content"):
-            self.write(
-                TextContent(
-                    text_content_id,
-                    reasoning_end.get("content"),
-                    Status.RUNNING,
-                    node_info,
-                    Position(self.get_node_id()),
-                )
-            )
-        if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
-            self.write(
-                ReasoningContent(
-                    reasoning_content_id,
-                    reasoning_content_chunk,
-                    Status.RUNNING,
-                    node_info,
-                    Position(self.get_node_id()),
-                )
-            )
-
-        self._write_final_context(chat_model, message_list, question, answer, reasoning_content)
+            self._write_final_context(chat_model, message_list, question, answer, reasoning_content)
+        # 无论正常结束、抛异常还是被取消中断,都关闭 LLM 流式响应底层连接
+        finally:
+            with contextlib.suppress(Exception):
+                await response.aclose()
 
     def _invoke_response(self, response, chat_model, message_list, question, is_result=False, text_content_id=None):
         model_setting = self.get_context("model_setting") or {}

apps/application/workflow/nodes/data_source_web_node/data_source_web_node.py

@@ -12,8 +12,8 @@
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
 
-from application.workflow.common import WorkflowType
-from application.workflow.i_node import CancelledException, INode
+from application.workflow.common import CancelledException, WorkflowType
+from application.workflow.i_node import INode
 from common.utils.fork import ChildLink, Fork, ForkManage
 from common.utils.logger import maxkb_logger
 

apps/application/workflow/workflow_manage.py

@@ -273,6 +273,10 @@ def from_context(cls, get_context, workflow, parameters, workflow_type, call_bac
             return None
 
     def cancel(self):
-        self.signal = Signal.CANCELLED
-        for node in self.nodes:
+        with self._lock:
+            if self.signal == Signal.CANCELLED:
+                return
+            self.signal = Signal.CANCELLED
+            nodes = list(self.nodes)
+        for node in nodes:
             node.cancel()