ai-agent-admin/backend-fastapi/online_dev/workflow/engine/handlers/parallel_handler.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
并行分支节点处理器
处理并行分支的执行和汇聚
"""
import copy
import logging
from typing import Dict, TYPE_CHECKING

from sqlalchemy.orm.attributes import flag_modified

from online_dev.workflow.engine.handlers.base_handler import BaseNodeHandler

if TYPE_CHECKING:
    from online_dev.workflow.engine.base import ExecutionContext

logger = logging.getLogger(__name__)


class ParallelHandler(BaseNodeHandler):
    """
    并行分支节点处理器
    
    同时执行所有分支，所有分支完成后汇聚继续
    """

    async def execute(self, context: 'ExecutionContext', node: Dict) -> None:
        """
        执行并行节点：同时启动所有分支
        """
        node_id = self.get_node_id(node)
        node_name = self.get_node_name(node)
        branches = node.get('branches', [])

        logger.info(f"处理并行分支 - 节点: {node_id}, 分支数: {len(branches)}")

        if not branches:
            logger.warning(f"并行节点 {node_id} 没有分支，自动跳过")
            await self.advance_to_next(context, node)
            return

        # 更新实例当前节点
        await self.update_instance_node(context, node)

        # 初始化并行分支状态（深拷贝确保 SQLAlchemy 检测到变更）
        parallel_status = copy.deepcopy(context.instance.parallel_branch_status or {})
        parallel_status[node_id] = {
            branch.get('id', ''): 'pending' for branch in branches
        }
        context.instance.parallel_branch_status = parallel_status
        flag_modified(context.instance, 'parallel_branch_status')
        context.db.add(context.instance)
        await context.db.flush()

        # 记录进入并行分支
        await self.create_log(
            context, node, 'parallel_start',
            comment=f'进入并行分支，共 {len(branches)} 个分支',
            extra_data={
                'branch_count': len(branches),
                'branch_ids': [b.get('id') for b in branches],
            },
        )

        # 同时执行所有分支
        for branch in branches:
            branch_children = branch.get('children')
            if branch_children:
                await self._execute_branch(context, node, branch, branch_children)
            else:
                # 空分支，标记为完成
                await self.mark_branch_complete(context, node, branch)

        # 检查是否所有分支都已完成
        await self.check_completion(context, node)

    async def _execute_branch(
            self,
            context: 'ExecutionContext',
            parallel_node: Dict,
            branch: Dict,
            first_node: Dict,
    ) -> None:
        """
        执行并行分支中的第一个节点。
        统一通过引擎 handler 分发，支持所有节点类型（approval, handle, copy, notify,
        condition, delay, service, subflow, data_update, parallel 等）。
        非阻塞节点执行后会自动调用 advance_to_next 推进；阻塞节点（approval/handle）
        在任务完成时由引擎推进。当分支走到末尾时，引擎的 _advance_to_next 会通过
        find_parallel_branch_for_node 找到所属并行分支并 mark_branch_complete。
        """
        node_type = first_node.get('type')
        handler = self.engine._get_handler(node_type)

        if handler:
            await handler.execute(context, first_node)
        else:
            logger.warning(f"并行分支内未知节点类型: {node_type}，标记分支完成")
            await self.mark_branch_complete(context, parallel_node, branch)

    async def mark_branch_complete(
            self,
            context: 'ExecutionContext',
            parallel_node: Dict,
            branch: Dict,
    ) -> None:
        """
        标记并行分支完成
        """
        node_id = self.get_node_id(parallel_node)
        branch_id = branch.get('id', '')
        branch_name = branch.get('name', '')

        # 更新分支状态（深拷贝确保 SQLAlchemy 检测到 JSON 字段变更）
        await context.db.refresh(context.instance)
        parallel_status = copy.deepcopy(context.instance.parallel_branch_status or {})

        if node_id in parallel_status:
            parallel_status[node_id][branch_id] = 'completed'
            context.instance.parallel_branch_status = parallel_status
            flag_modified(context.instance, 'parallel_branch_status')
            context.db.add(context.instance)
            await context.db.flush()

        logger.info(f"并行分支完成 - 节点: {node_id}, 分支: {branch_id}")

        # 记录分支完成
        await self.create_log(
            context, parallel_node, 'parallel_branch_complete',
            comment=f'并行分支 {branch_name} 完成',
            extra_data={'branch_id': branch_id},
        )

        # 检查是否所有分支都已完成
        await self.check_completion(context, parallel_node)

    async def check_completion(self, context: 'ExecutionContext', parallel_node: Dict) -> None:
        """
        检查并行分支是否全部完成
        """
        node_id = self.get_node_id(parallel_node)

        # 刷新实例获取最新状态
        await context.db.refresh(context.instance)
        parallel_status = copy.deepcopy(context.instance.parallel_branch_status or {})

        branch_statuses = parallel_status.get(node_id, {})

        if not branch_statuses:
            logger.warning(f"并行节点 {node_id} 没有分支状态记录")
            return

        # 检查是否所有分支都已完成
        all_completed = all(status == 'completed' for status in branch_statuses.values())
        pending_branches = [bid for bid, status in branch_statuses.items() if status != 'completed']

        logger.info(f"并行节点 {node_id} 状态检查: 全部完成={all_completed}, 待完成={pending_branches}")

        if all_completed:
            logger.info(f"并行节点 {node_id} 所有分支完成，继续推进")

            # 清理已完成的并行节点状态（深拷贝确保变更检测）
            parallel_status = copy.deepcopy(parallel_status)
            del parallel_status[node_id]
            context.instance.parallel_branch_status = parallel_status
            flag_modified(context.instance, 'parallel_branch_status')
            context.db.add(context.instance)
            await context.db.flush()

            # 记录日志
            await self.create_log(
                context, parallel_node, 'parallel_complete',
                comment='所有并行分支完成',
            )

            # 推进到下一节点
            await self.advance_to_next(context, parallel_node)