From 0236a0df9d79944514df1448053890c6ac448fbb Mon Sep 17 00:00:00 2001 From: hexiaonan-800 Date: Tue, 20 Oct 2026 13:38:24 +0800 Subject: [PATCH] feat: ai chat node with tools support cancellation --- .../workflow/nodes/ai_chat_node/agent.py | 11 ++-- .../nodes/ai_chat_node/ai_chat_node.py | 17 ++++++- .../nodes/ai_chat_node/tools/application.py | 14 ++++-- .../nodes/ai_chat_node/tools/tool/__init__.py | 8 +-- .../nodes/ai_chat_node/tools/tool/custom.py | 50 ++++++++++++++++--- .../nodes/ai_chat_node/tools/tool/workflow.py | 14 ++++-- 6 files changed, 89 insertions(+), 25 deletions(-) diff --git a/apps/application/workflow/nodes/ai_chat_node/agent.py b/apps/application/workflow/nodes/ai_chat_node/agent.py index 3d47543e500..f00b09166aa 100644 --- a/apps/application/workflow/nodes/ai_chat_node/agent.py +++ b/apps/application/workflow/nodes/ai_chat_node/agent.py @@ -149,6 +149,7 @@ def create_agent( chat_id=None, skill_tool_ids=None, extra_tools=None, + async_runner=None, ): # 创建临时文件夹 if chat_id: @@ -185,9 +186,13 @@ async def _run(): stream_mode="messages", ) - async for chunk in response: - msg = chunk[0] - call_back.on_next(msg) + async def consume(): + # runner 取消该 task 会让 CancelledError 注入在途的工具 await,从而打断正在运行的工具。 + async for chunk in response: + msg = chunk[0] + call_back.on_next(msg) + + await async_runner.await_or_cancel(consume()) def _classify_error(e): # 取消：原样保留（保持节点取消语义）；MCP TaskGroup 的 ExceptionGroup：展开取真实异常并包成 RuntimeError diff --git a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py index 72943ed5e28..fc1adee0c4b 100644 --- a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py +++ b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py @@ -543,8 +543,20 @@ def _handle_mcp( # 工具(workflow/custom) + 智能体(子应用) → LangChain tools； # MCP(自定义/库内) → mcp_servers 配置；技能 → 交给引擎侧 init_skills 初始化 tools = get_tool_tools( - source_type, source_id, tool_ids, workspace_id, self.get_workflow_parameters() - ) + get_application_tools(source_type, source_id, application_ids, workspace_id, self.get_workflow_parameters()) + source_type, + source_id, + tool_ids, + workspace_id, + self.async_runner.cancel_event, + self.get_workflow_parameters(), + ) + get_application_tools( + source_type, + source_id, + application_ids, + workspace_id, + self.get_workflow_parameters(), + cancel_event=self.async_runner.cancel_event, + ) mcp_servers_config = get_mcp_servers(mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids, self._handle_variables) ToolExecutor().validate_mcp_transport(json.dumps(mcp_servers_config)) @@ -663,6 +675,7 @@ def on_complete(error): chat_id, skill_tool_ids, tools, + async_runner=self.async_runner, ) return True diff --git a/apps/application/workflow/nodes/ai_chat_node/tools/application.py b/apps/application/workflow/nodes/ai_chat_node/tools/application.py index 272b1f70220..c3091a13293 100644 --- a/apps/application/workflow/nodes/ai_chat_node/tools/application.py +++ b/apps/application/workflow/nodes/ai_chat_node/tools/application.py @@ -36,7 +36,7 @@ def get_application_args(): ) -def get_application_func(source_type, source_id, application, workflow_params, workspace_id): +def get_application_func(source_type, source_id, application, workflow_params, workspace_id, cancel_event): """ 构建应用（Agent）工具的执行函数。 @@ -46,10 +46,11 @@ def get_application_func(source_type, source_id, application, workflow_params, w def inner(message: str = ""): from application.models import Application, ApplicationVersion, Chat, ChatRecord, ChatSourceChoices - from application.workflow.common import WorkflowType, new_instance + from application.workflow.common import CancelledException, WorkflowType, new_instance from application.workflow.content_type import ContentType from application.workflow.nodes import get_start_node from application.workflow.workflow_manage import CallBack, WorkflowManage + from common.utils.thread_pool import await_complete_or_cancel from chat.serializers.chat import get_work_flow from chat.serializers.chat_history import ChatHistory @@ -158,7 +159,10 @@ def get_start_node_fn(wf, wm): ) sub_manage.start_node.workflow_manage = sub_manage sub_manage.run() - done_event.wait() + + # 并发等子应用完成与取消信号;取消时级联取消子应用工作流,立即中断 + if await_complete_or_cancel(done_event, cancel_event, sub_manage.cancel): + raise CancelledException() if result_holder["error"]: raise result_holder["error"] @@ -170,7 +174,7 @@ def get_start_node_fn(wf, wm): return inner -def get_application_tools(source_type, source_id, application_ids, workspace_id, workflow_params): +def get_application_tools(source_type, source_id, application_ids, workspace_id, workflow_params, cancel_event): if not application_ids: return [] from application.models import Application @@ -178,7 +182,7 @@ def get_application_tools(source_type, source_id, application_ids, workspace_id, applications = QuerySet(Application).filter(id__in=application_ids, is_publish=True) results = [] for application in applications: - func = get_application_func(source_type, source_id, application, workflow_params, workspace_id) + func = get_application_func(source_type, source_id, application, workflow_params, workspace_id, cancel_event) args = get_application_args() structured_tool = StructuredTool.from_function( func=func, diff --git a/apps/application/workflow/nodes/ai_chat_node/tools/tool/__init__.py b/apps/application/workflow/nodes/ai_chat_node/tools/tool/__init__.py index 0138298dbd0..e44c2f9beff 100644 --- a/apps/application/workflow/nodes/ai_chat_node/tools/tool/__init__.py +++ b/apps/application/workflow/nodes/ai_chat_node/tools/tool/__init__.py @@ -13,7 +13,7 @@ __all__ = ["get_tool_tools", "get_workflow_tools", "get_custom_tools"] -def get_tool_tools(source_type, source_id, tool_ids, workspace_id, workflow_params=None): +def get_tool_tools(source_type, source_id, tool_ids, workspace_id, cancel_event, workflow_params=None): """ 构建工具（Tool）类工具：内部按 tool_type 拆分 workflow / custom，合并返回 LangChain tools。 @@ -21,6 +21,6 @@ def get_tool_tools(source_type, source_id, tool_ids, workspace_id, workflow_para """ if not tool_ids: return [] - return get_workflow_tools(source_type, source_id, tool_ids, workspace_id, workflow_params) + get_custom_tools( - source_type, source_id, tool_ids, workspace_id - ) + return get_workflow_tools( + source_type, source_id, tool_ids, workspace_id, workflow_params, cancel_event + ) + get_custom_tools(source_type, source_id, tool_ids, workspace_id, cancel_event) diff --git a/apps/application/workflow/nodes/ai_chat_node/tools/tool/custom.py b/apps/application/workflow/nodes/ai_chat_node/tools/tool/custom.py index a665f4b132a..737f6b36530 100644 --- a/apps/application/workflow/nodes/ai_chat_node/tools/tool/custom.py +++ b/apps/application/workflow/nodes/ai_chat_node/tools/tool/custom.py @@ -8,6 +8,7 @@ """ import json +import threading import time import uuid_utils.compat as uuid @@ -63,7 +64,33 @@ def _save_custom_tool_record(tool_id, workspace_id, source_type, source_id, inpu ).save() -def get_custom_func(source_type, source_id, tool, workspace_id): +def run_cancellable(fn, cancel_event, on_cancel): + """ + 在 worker 线程执行阻塞函数 fn,与取消信号并发;取消先到则调用 on_cancel() 中断并抛CancelledException。 + """ + from application.workflow.common import CancelledException + from common.utils.thread_pool import await_complete_or_cancel + + done_event = threading.Event() + result = {} + + def _run(): + try: + result["value"] = fn() + except BaseException as e: + result["error"] = e + finally: + done_event.set() + + threading.Thread(target=_run, daemon=True).start() + if await_complete_or_cancel(done_event, cancel_event, on_cancel): + raise CancelledException() + if result["error"]: + raise result["error"] + return result["value"] + + +def get_custom_func(source_type, source_id, tool, workspace_id, cancel_event): tool_id = tool.id code = tool.code init_field_list = tool.init_field_list or [] @@ -72,7 +99,7 @@ def get_custom_func(source_type, source_id, tool, workspace_id): def inner(**kwargs): # 在进程内直接跑沙箱代码（无 MCP 子进程），方式与工具调试执行 ToolExecutor.exec_code 一致。 from common.utils.rsa_util import rsa_long_decrypt - from common.utils.tool_code import ToolExecutor + from common.utils.tool_code import ToolExecution, ToolExecutor start_time = time.time() # 合并初始化参数（默认值 → 已保存的启动参数），服务端注入，模型不可见 @@ -83,10 +110,21 @@ def inner(**kwargs): init_params = init_params_default_value all_params = init_params | kwargs - error = None + # 持子进程句柄：取消时中断并 terminate 子进程，实现工具运行期间硬中断 + execution = ToolExecution() + + def _interrupt(): + execution.cancelled.set() + execution.proc.terminate() + result = None + error = None try: - result = ToolExecutor().exec_code(code, all_params) + result = run_cancellable( + lambda: ToolExecutor().exec_code(code, all_params, execution=execution), + cancel_event, + _interrupt, + ) except Exception as e: error = e finally: @@ -98,13 +136,13 @@ def inner(**kwargs): return inner -def get_custom_tools(source_type, source_id, tool_ids, workspace_id): +def get_custom_tools(source_type, source_id, tool_ids, workspace_id, cancel_event): if not tool_ids: return [] tools = QuerySet(Tool).filter(id__in=tool_ids, is_active=True, tool_type=ToolType.CUSTOM) results = [] for tool in tools: - func = get_custom_func(source_type, source_id, tool, workspace_id) + func = get_custom_func(source_type, source_id, tool, workspace_id, cancel_event) args = get_custom_args(tool) structured_tool = StructuredTool.from_function( func=func, diff --git a/apps/application/workflow/nodes/ai_chat_node/tools/tool/workflow.py b/apps/application/workflow/nodes/ai_chat_node/tools/tool/workflow.py index e70fc25e0fe..79b23e863ca 100644 --- a/apps/application/workflow/nodes/ai_chat_node/tools/tool/workflow.py +++ b/apps/application/workflow/nodes/ai_chat_node/tools/tool/workflow.py @@ -81,16 +81,17 @@ def _save_workflow_tool_record( ).save() -def get_workflow_func(source_type, source_id, tool, qv, workspace_id, workflow_params=None): +def get_workflow_func(source_type, source_id, tool, qv, workspace_id, cancel_event, workflow_params=None): from knowledge.services.retrieval_access import inherited_retrieval_context tool_id = tool.id def inner(**kwargs): # 使用新工作流引擎执行工具工作流，方式与 tool_workflow_lib_node 保持一致。 - from application.workflow.common import WorkflowType, new_instance + from application.workflow.common import CancelledException, WorkflowType, new_instance from application.workflow.nodes import get_node_class from application.workflow.workflow_manage import CallBack, WorkflowManage + from common.utils.thread_pool import await_complete_or_cancel tool_record_id = str(uuid.uuid7()) sub_workflow = new_instance(qv.work_flow, WorkflowType.TOOL) @@ -152,7 +153,10 @@ def get_start_node_fn(wf, wm): ) sub_manage.start_node.workflow_manage = sub_manage sub_manage.run() - done_event.wait() + + # 并发等子工作流完成与取消信号;取消时级联取消子工作流,立即中断 + if await_complete_or_cancel(done_event, cancel_event, sub_manage.cancel): + raise CancelledException() if result_holder["error"]: raise result_holder["error"] return result_holder["output"] @@ -160,7 +164,7 @@ def get_start_node_fn(wf, wm): return inner -def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id, workflow_params=None): +def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id, cancel_event, workflow_params=None): tools = QuerySet(Tool).filter( id__in=tool_workflow_ids, is_active=True, tool_type=ToolType.WORKFLOW, workspace_id=workspace_id ) @@ -173,7 +177,7 @@ def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id, results = [] for tool in tools: qv = qd.get(tool.id) - func = get_workflow_func(source_type, source_id, tool, qv, workspace_id, workflow_params) + func = get_workflow_func(source_type, source_id, tool, qv, workspace_id, cancel_event, workflow_params) args = get_workflow_args(tool, qv) tool = StructuredTool.from_function( func=func,