beaver_project/app-instance/backend/beaver/team_workflows/executor.py

"""Runtime bridge for local team workflow MCP tools."""

from __future__ import annotations

import json
from typing import Any, Callable

from beaver.coordinator.models import ExecutionGraph, TeamRunResult
from beaver.tools.base import ToolContext, ToolResult

from . import agent_rearrange, concurrent, graph, mixture_of_agents, sequential

GraphBuilder = Callable[..., ExecutionGraph]


class TeamWorkflowExecutor:
    """Execute workflow MCP calls inside the current Beaver runtime."""

    _BUILDERS: dict[str, GraphBuilder] = {
        "SequentialWorkflow": sequential.build_graph,
        "ConcurrentWorkflow": concurrent.build_graph,
        "MixtureOfAgents": mixture_of_agents.build_graph,
        "AgentRearrange": agent_rearrange.build_graph,
        "GraphWorkflow": graph.build_graph,
    }

    async def execute(
        self,
        workflow_name: str,
        arguments: dict[str, Any],
        context: ToolContext,
        *,
        tool_name: str | None = None,
    ) -> ToolResult:
        exposed_name = tool_name or workflow_name
        try:
            if str(context.metadata.get("source") or "").startswith("team:"):
                raise ValueError("nested_team_workflow_not_allowed")
            builder = self._BUILDERS.get(workflow_name)
            if builder is None:
                raise ValueError(f"unknown team workflow tool: {workflow_name}")
            graph = builder(**dict(arguments or {}))
            parent_task_id = _task_id(context)
            parent_session_id = _session_id(context)
            result = await self._run_team(
                context=context,
                graph=graph,
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
            )
            payload = _success_payload(
                workflow_name=workflow_name,
                graph=graph,
                result=result,
            )
            return ToolResult(
                success=True,
                content=json.dumps(payload, ensure_ascii=False),
                tool_name=exposed_name,
                raw_output=payload,
            )
        except Exception as exc:
            payload = {
                "success": False,
                "workflow": workflow_name,
                "error": str(exc),
            }
            return ToolResult(
                success=False,
                content=json.dumps(payload, ensure_ascii=False),
                tool_name=exposed_name,
                error=str(exc),
                raw_output=payload,
            )

    async def _run_team(
        self,
        *,
        context: ToolContext,
        graph: ExecutionGraph,
        parent_task_id: str,
        parent_session_id: str,
    ) -> TeamRunResult:
        runner = context.services.get("agent_team_runner")
        parent_run_id = _run_id(context)
        if runner is not None:
            return await runner(
                graph,
                parent_task_id=parent_task_id,
                parent_session_id=parent_session_id,
                parent_run_id=parent_run_id,
            )

        agent_loop = context.services.get("agent_loop")
        if agent_loop is None:
            raise ValueError("team workflow execution requires agent_loop or agent_team_runner")
        provider_bundle = context.services.get("provider_bundle")

        def provider_bundle_factory(_node: Any) -> Any:
            return provider_bundle

        from beaver.engine import AgentLoop
        from beaver.services.team_service import TeamService

        loaded = context.services.get("loaded")
        team_loop = AgentLoop(profile=agent_loop.profile, loader=agent_loop.loader)
        team_loop.loaded = loaded
        return await TeamService(team_loop).run_team(
            graph,
            parent_task_id=parent_task_id,
            parent_session_id=parent_session_id,
            parent_run_id=parent_run_id,
            provider_bundle_factory=provider_bundle_factory if provider_bundle is not None else None,
            allow_candidate_generation=False,
        )


def _task_id(context: ToolContext) -> str:
    value = str(context.services.get("task_id") or context.metadata.get("task_id") or "").strip()
    if not value:
        raise ValueError("team workflow execution requires task_id")
    return value


def _session_id(context: ToolContext) -> str:
    value = str(context.session_id or context.services.get("session_id") or "").strip()
    if not value:
        raise ValueError("team workflow execution requires session_id")
    return value


def _run_id(context: ToolContext) -> str | None:
    return str(context.services.get("run_id") or context.metadata.get("run_id") or "").strip() or None


def _success_payload(
    *,
    workflow_name: str,
    graph: ExecutionGraph,
    result: TeamRunResult,
) -> dict[str, Any]:
    return {
        "success": result.success,
        "workflow": workflow_name,
        "summary": result.summary,
        "run_ids": list(result.run_ids),
        "session_ids": list(result.session_ids),
        "node_results": [item.to_dict() for item in result.node_results],
        "graph": _graph_to_dict(graph),
    }


def _graph_to_dict(graph: ExecutionGraph) -> dict[str, Any]:
    return {
        "strategy": graph.strategy,
        "nodes": [
            {
                "node_id": node.node_id,
                "task": node.task,
                "depends_on": list(node.depends_on),
                "allowed_tool_names": (
                    None if node.allowed_tool_names is None else list(node.allowed_tool_names)
                ),
                "required_evidence": list(node.required_evidence),
                "evidence_contract": dict(node.evidence_contract),
                "validation_rules": list(node.validation_rules),
                "required_for_completion": node.required_for_completion,
                "block_downstream_on_partial": node.block_downstream_on_partial,
                "max_tool_iterations": node.max_tool_iterations,
                "metadata": dict(node.agent.metadata),
            }
            for node in graph.nodes
        ],
    }