feat: tool node cancel

apps/application/workflow/nodes/tool_node/tool_node.py

@@ -17,7 +17,7 @@
 from rest_framework import serializers
 from rest_framework.utils.formatting import lazy_format
 
-from application.workflow.common import WorkflowType
+from application.workflow.common import CancelledException, WorkflowType
 from application.workflow.i_node import INode
 from application.workflow.message.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -26,7 +26,7 @@
 from common.field.common import ObjectField
 from common.utils.common import common_convert_value
 from common.utils.logger import maxkb_logger
-from common.utils.tool_code import ToolExecutor
+from common.utils.tool_code import ToolCancelledException, ToolExecution, ToolExecutor
 
 function_executor = ToolExecutor()
 
@@ -131,6 +131,16 @@ class ToolNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "tool-node"
 
+    def __init__(self, node, workflow_manage, get_node_parameters):
+        super().__init__(node, workflow_manage, get_node_parameters)
+        self._execution = ToolExecution()
+
+    def cancel(self):
+        super().cancel()
+        self._execution.cancelled.set()
+        if self._execution.proc:
+            self._execution.proc.terminate()
+
     def execute(self):
         node_params = self.get_parameters()
         input_field_list = node_params.get("input_field_list", [])
@@ -160,9 +170,13 @@ def execute(self):
                 all_params = init_params_default_value | params
         else:
             all_params = params
+        self.data["params"] = all_params
+        try:
+            result = function_executor.exec_code(code, all_params, execution=self._execution)
+        except ToolCancelledException:
+            # 引擎 run()(i_node.py:96) 捕获 CancelledException 并转 CANCELLED
+            raise CancelledException()
 
-        result = function_executor.exec_code(code, all_params)
-        self.write_context("params", all_params)
         self.write_context("result", result)
 
         if is_result:
@@ -175,7 +189,7 @@ def get_details(self, index: int = 0, position: dict = None, old_details: dict =
         details.update(
             {
                 "result": self.get_context("result"),
-                "params": self.get_context("params"),
+                "params": self.data.get("params"),
                 "enableException": self.node.properties.get("enableException"),
             }
         )

apps/common/utils/tool_code.py

@@ -12,6 +12,7 @@
 import subprocess
 import sys
 import tempfile
+import threading
 import time
 from contextlib import contextmanager, suppress
 from textwrap import dedent
@@ -104,7 +105,7 @@ def init_sandbox_dir():
     except Exception as e:
         maxkb_logger.error(f"Exception: {e}", exc_info=True)
 
-    def exec_code(self, code_str, keywords, function_name=None):
+    def exec_code(self, code_str, keywords, function_name=None, execution=None):
         _id = str(uuid.uuid7())
         action_function = (
             f"({function_name!a}, locals_v.get({function_name!a}))" if function_name else "locals_v.popitem()"
@@ -146,7 +147,7 @@ def exec_code(self, code_str, keywords, function_name=None):
             f.write(_exec_code)
             f.flush()
             with execution_timer(_id):
-                subprocess_result = self._exec(f.name, _id)
+                subprocess_result = self._exec(f.name, _id, execution=execution)
         if subprocess_result.returncode != 0:
             raise Exception(subprocess_result.stderr or subprocess_result.stdout or "Unknown exception occurred")
         lines = subprocess_result.stdout.splitlines()
@@ -387,7 +388,7 @@ def encode(data):
         maxkb_logger.warning("Chat files are too large to be passed to the nested agent, skipped")
         return None
 
-    def _exec(self, execute_file, _id):
+    def _exec(self, execute_file, _id, execution=None):
         kwargs = {
             "cwd": BASE_DIR,
             "env": {
@@ -404,23 +405,50 @@ def _set_resource_limit():
             with suppress(Exception):
                 os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), _process_limit_cpu_cores)))
 
+        proc = subprocess.Popen(
+            [sys.executable, execute_file],
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            **kwargs,
+            preexec_fn=_set_resource_limit,
+        )
+        if execution is not None:
+            execution.proc = proc
         try:
-            subprocess_result = subprocess.run(
-                [sys.executable, execute_file],
-                timeout=_process_limit_timeout_seconds,
-                text=True,
-                capture_output=True,
-                **kwargs,
-                preexec_fn=_set_resource_limit,
-            )
-            return subprocess_result
-        except subprocess.TimeoutExpired:
-            raise Exception(_("Process execution timed out after {} seconds.").format(_process_limit_timeout_seconds))
+            try:
+                stdout, stderr = proc.communicate(timeout=_process_limit_timeout_seconds)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                raise Exception(
+                    _("Process execution timed out after {} seconds.").format(_process_limit_timeout_seconds)
+                )
+            # 取消时外部线程已 terminate/kill 子进程,communicate 随即返回;据此区分取消与正常结束
+            if execution is not None and execution.cancelled.is_set():
+                raise ToolCancelledException("Tool execution cancelled")
+            return subprocess.CompletedProcess(proc.args, proc.returncode, stdout, stderr)
+        finally:
+            if execution is not None:
+                execution.proc = None
 
     def validate_mcp_transport(self, code_str):
         validate_mcp_servers(json.loads(code_str))
 
 
+class ToolCancelledException(Exception):
+    """工具子进程执行被外部取消(工作流中断)时抛出,用于与普通执行失败相区分。
+    取消异常向上传播到节点层,由节点转成 CANCELLED 状态,而非 FAIL。
+    """
+
+
+class ToolExecution:
+    """工具子进程存活句柄。由节点持有:外部线程置 cancelled 并 terminate/kill 子进程以中止执行。"""
+
+    def __init__(self):
+        self.cancelled = threading.Event()
+        self.proc = None
+
+
 @contextmanager
 def execution_timer(id=""):
     start = time.perf_counter()