From 919673ca16beedebe82179be7036338ac7b37510 Mon Sep 17 00:00:00 2001 From: hexiaonan-800 Date: Thu, 15 Oct 2026 18:33:52 +0800 Subject: [PATCH] feat: tool node cancel --- .../workflow/nodes/tool_node/tool_node.py | 24 ++++++-- apps/common/utils/tool_code.py | 56 ++++++++++++++----- 2 files changed, 61 insertions(+), 19 deletions(-) diff --git a/apps/application/workflow/nodes/tool_node/tool_node.py b/apps/application/workflow/nodes/tool_node/tool_node.py index a43d561042c..02c276d1447 100644 --- a/apps/application/workflow/nodes/tool_node/tool_node.py +++ b/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"), } ) diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index b95612c24e2..61ed8b7c8c7 100644 --- a/apps/common/utils/tool_code.py +++ b/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()