python-sdk/pkg/hanzo-tools-shell/hanzo_tools/shell/dag_tool.py

"""DAG execution tool - directed acyclic graph for command execution.

Run commands/tools with proper dependency ordering using DAG semantics.
Supports serial (default), parallel, and complex mixed execution graphs.
"""

import os
import sys
import uuid
import shutil
import asyncio
from enum import Enum
from typing import Any, Dict, List, Union, Optional, Annotated, override
from pathlib import Path
from datetime import datetime
from dataclasses import field, dataclass

from pydantic import Field
from mcp.server import FastMCP
from mcp.server.fastmcp import Context as MCPContext

from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
from hanzo_tools.shell.base_process import get_shell_executor


class NodeStatus(Enum):
    """Execution status of a DAG node."""

    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class DagResult:
    """Result from a single DAG node execution."""

    node_id: str
    command: str
    stdout: str
    stderr: str
    status: NodeStatus
    exit_code: int
    duration_ms: int = 0
    node_type: str = "shell"


@dataclass
class DagNode:
    """A node in the execution DAG."""

    id: str
    command: Union[str, Dict[str, Any]]
    depends_on: List[str] = field(default_factory=list)
    status: NodeStatus = NodeStatus.PENDING
    result: Optional[DagResult] = None


Command = Union[str, Dict[str, Any], List[Any]]


class DagTool(BaseTool):
    """DAG-based execution with proper dependency ordering.

    Execute commands (shell or tools) with directed acyclic graph semantics.
    Supports serial, parallel, and complex mixed execution patterns.

    USAGE PATTERNS:

    1. Serial (default) - sequential execution:
       dag(["ls", "pwd", "git status"])

    2. Parallel - concurrent execution:
       dag(["npm install", "cargo build"], parallel=True)

    3. DAG - mixed serial and parallel blocks:
       dag([
           "mkdir -p dist",
           {"parallel": [
               "cp manifest.json dist/",
               "cp -rf assets/ dist/",
           ]},
           "zip -r package.zip dist/"
       ])

    4. Tool invocations - not just shell:
       dag([
           {"tool": "read", "input": {"file_path": "config.json"}},
           {"tool": "search", "input": {"pattern": "TODO"}},
       ])

    5. Named nodes with explicit dependencies:
       dag([
           {"id": "setup", "run": "mkdir -p dist"},
           {"id": "copy", "run": "cp *.txt dist/", "after": ["setup"]},
           {"id": "test", "run": "pytest", "after": ["setup"]},
           {"id": "package", "run": "tar -czf out.tar.gz dist/", "after": ["copy", "test"]},
       ])

    Uses zsh for shell execution.
    """

    name = "dag"

    def __init__(
        self, tools: Optional[Dict[str, BaseTool]] = None, default_shell: str = "zsh"
    ):
        """Initialize DAG execution tool."""
        super().__init__()
        self.tools = tools or {}
        self.default_shell = self._resolve_shell(default_shell)

    def _resolve_shell(self, preferred: str) -> str:
        """Resolve shell - prefer zsh, fallback to bash. On Windows, use pwsh/cmd."""
        if sys.platform == "win32":
            for shell in ("pwsh", "powershell", "cmd"):
                found = shutil.which(shell)
                if found:
                    return found
            return "cmd.exe"

        shell_priority = ["zsh", "bash"]
        search_paths = [
            "/opt/homebrew/bin",
            "/usr/local/bin",
            "/bin",
            "/usr/bin",
        ]

        for shell in shell_priority:
            for prefix in search_paths:
                full_path = f"{prefix}/{shell}"
                if os.path.isfile(full_path) and os.access(full_path, os.X_OK):
                    return full_path
            found = shutil.which(shell)
            if found:
                return found

        return "sh"

    @property
    @override
    def description(self) -> str:
        return """DAG execution - run commands with dependency ordering.

Execute shell commands or tools with DAG (directed acyclic graph) semantics.
Supports serial, parallel, and complex mixed execution patterns.

MODES:

Serial (default): dag(["ls", "pwd"])
  Commands run in sequence.

Parallel: dag(["npm install", "cargo build"], parallel=True)
  Commands run concurrently.

Mixed DAG:
  dag([
      "mkdir -p dist",
      {"parallel": ["cp a.txt dist/", "cp b.txt dist/"]},
      "zip -r out.zip dist/"
  ])

Tool invocations:
  dag([{"tool": "search", "input": {"pattern": "TODO"}}])

Named with dependencies:
  dag([
      {"id": "build", "run": "make build"},
      {"id": "test", "run": "make test", "after": ["build"]},
  ])

Uses zsh for shell execution.

AUTO-BACKGROUNDING: Commands that exceed timeout are automatically
backgrounded. Use ps tool to monitor: ps --logs <id>, ps --kill <id>"""

    async def _run_shell(
        self,
        cmd: str,
        shell: str,
        cwd: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        timeout: int = 30,
    ) -> DagResult:
        """Run a shell command with auto-backgrounding on timeout.

        Uses the shared ShellExecutor for consistent 45s auto-backgrounding.
        """
        start_time = datetime.now()
        node_id = f"shell_{id(cmd)}"
        executor = get_shell_executor()

        stdout, stderr, exit_code, was_backgrounded, process_id = (
            await executor.run_shell(
                command=cmd,
                shell=shell,
                cwd=cwd,
                env=env,
                timeout=float(timeout),
                tool_name="dag",
            )
        )

        duration = int((datetime.now() - start_time).total_seconds() * 1000)

        if was_backgrounded:
            return DagResult(
                node_id=node_id,
                command=cmd,
                stdout=stdout,
                stderr=stderr,
                status=NodeStatus.SUCCESS,  # Backgrounded = success from caller's POV
                exit_code=0,
                duration_ms=duration,
                node_type="shell",
            )

        return DagResult(
            node_id=node_id,
            command=cmd,
            stdout=stdout,
            stderr=stderr,
            status=NodeStatus.SUCCESS if exit_code == 0 else NodeStatus.FAILED,
            exit_code=exit_code,
            duration_ms=duration,
            node_type="shell",
        )

    async def _run_tool(
        self, tool_name: str, tool_input: Dict[str, Any], ctx: MCPContext
    ) -> DagResult:
        """Run an MCP tool invocation."""
        start_time = datetime.now()
        node_id = f"tool_{tool_name}"

        if tool_name not in self.tools:
            return DagResult(
                node_id=node_id,
                command=f"tool:{tool_name}",
                stdout="",
                stderr=f"Tool '{tool_name}' not found. Available: {list(self.tools.keys())}",
                status=NodeStatus.FAILED,
                exit_code=1,
                duration_ms=0,
                node_type="tool",
            )

        try:
            tool = self.tools[tool_name]
            result = await tool.call(ctx, **tool_input)
            duration = int((datetime.now() - start_time).total_seconds() * 1000)

            return DagResult(
                node_id=node_id,
                command=f"tool:{tool_name}",
                stdout=str(result) if result else "",
                stderr="",
                status=NodeStatus.SUCCESS,
                exit_code=0,
                duration_ms=duration,
                node_type="tool",
            )
        except Exception as e:
            duration = int((datetime.now() - start_time).total_seconds() * 1000)
            return DagResult(
                node_id=node_id,
                command=f"tool:{tool_name}",
                stdout="",
                stderr=str(e),
                status=NodeStatus.FAILED,
                exit_code=1,
                duration_ms=duration,
                node_type="tool",
            )

    async def _execute_node(
        self,
        cmd: Command,
        ctx: MCPContext,
        shell: str,
        cwd: Optional[str],
        env: Optional[Dict[str, str]],
        timeout: int,
    ) -> DagResult:
        """Execute a single DAG node."""

        if isinstance(cmd, str):
            return await self._run_shell(cmd, shell, cwd, env, timeout)

        # Nested array = auto-parallel (e.g., ["a", ["b", "c"], "d"] runs b,c in parallel)
        if isinstance(cmd, list):
            tasks = [self._execute_node(c, ctx, shell, cwd, env, timeout) for c in cmd]
            results = await asyncio.gather(*tasks, return_exceptions=True)

            combined_stdout = []
            combined_stderr = []
            all_success = True
            total_duration = 0

            for i, r in enumerate(results):
                if isinstance(r, Exception):
                    combined_stderr.append(f"[{i}] Error: {r}")
                    all_success = False
                else:
                    if r.stdout:
                        combined_stdout.append(f"[{i}] {r.stdout.rstrip()}")
                    if r.stderr:
                        combined_stderr.append(f"[{i}] {r.stderr.rstrip()}")
                    if r.status != NodeStatus.SUCCESS:
                        all_success = False
                    total_duration = max(total_duration, r.duration_ms)

            return DagResult(
                node_id=f"auto_parallel_{len(cmd)}",
                command=f"parallel[{len(cmd)} tasks]",
                stdout="\n".join(combined_stdout),
                stderr="\n".join(combined_stderr),
                status=NodeStatus.SUCCESS if all_success else NodeStatus.FAILED,
                exit_code=0 if all_success else 1,
                duration_ms=total_duration,
                node_type="parallel",
            )

        if isinstance(cmd, dict):
            if "tool" in cmd:
                return await self._run_tool(cmd["tool"], cmd.get("input", {}), ctx)

            if "parallel" in cmd:
                parallel_cmds = cmd["parallel"]
                tasks = [
                    self._execute_node(c, ctx, shell, cwd, env, timeout)
                    for c in parallel_cmds
                ]
                results = await asyncio.gather(*tasks, return_exceptions=True)

                combined_stdout = []
                combined_stderr = []
                all_success = True
                total_duration = 0

                for i, r in enumerate(results):
                    if isinstance(r, Exception):
                        combined_stderr.append(f"[{i}] Error: {r}")
                        all_success = False
                    else:
                        if r.stdout:
                            combined_stdout.append(f"[{i}] {r.stdout.rstrip()}")
                        if r.stderr:
                            combined_stderr.append(f"[{i}] {r.stderr.rstrip()}")
                        if r.status != NodeStatus.SUCCESS:
                            all_success = False
                        total_duration = max(total_duration, r.duration_ms)

                return DagResult(
                    node_id=f"parallel_{len(parallel_cmds)}",
                    command=f"parallel[{len(parallel_cmds)} tasks]",
                    stdout="\n".join(combined_stdout),
                    stderr="\n".join(combined_stderr),
                    status=NodeStatus.SUCCESS if all_success else NodeStatus.FAILED,
                    exit_code=0 if all_success else 1,
                    duration_ms=total_duration,
                    node_type="parallel",
                )

            if "run" in cmd:
                run_cmd = cmd["run"]
                return await self._execute_node(run_cmd, ctx, shell, cwd, env, timeout)

            return DagResult(
                node_id="unknown",
                command=str(cmd),
                stdout="",
                stderr=f"Unknown command format: {cmd}",
                status=NodeStatus.FAILED,
                exit_code=1,
                duration_ms=0,
                node_type="unknown",
            )

        return DagResult(
            node_id="unknown",
            command=str(cmd),
            stdout="",
            stderr=f"Unknown command type: {type(cmd).__name__}",
            status=NodeStatus.FAILED,
            exit_code=1,
            duration_ms=0,
            node_type="unknown",
        )

    @override
    @auto_timeout("dag")
    async def call(
        self,
        ctx: MCPContext,
        commands: List[Command],
        parallel: bool = False,
        shell: Optional[str] = None,
        cwd: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        timeout: int = 30,
        strict: bool = False,
        quiet: bool = False,
        **kwargs,
    ) -> str:
        """Execute commands with DAG semantics."""
        tool_ctx = create_tool_context(ctx)
        await tool_ctx.set_tool_info(self.name)

        shell = shell or self.default_shell
        results: List[DagResult] = []

        if parallel:
            tasks = [
                self._execute_node(cmd, ctx, shell, cwd, env, timeout)
                for cmd in commands
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            results = [
                (
                    r
                    if not isinstance(r, Exception)
                    else DagResult(
                        node_id=f"error_{i}",
                        command=str(commands[i]),
                        stdout="",
                        stderr=str(r),
                        status=NodeStatus.FAILED,
                        exit_code=1,
                        duration_ms=0,
                        node_type="error",
                    )
                )
                for i, r in enumerate(results)
            ]
        else:
            for cmd in commands:
                result = await self._execute_node(cmd, ctx, shell, cwd, env, timeout)
                results.append(result)

                if strict and result.status == NodeStatus.FAILED:
                    break

        return self._format_output(results, quiet)

    def _format_output(self, results: List[DagResult], quiet: bool = False) -> str:
        """Format DAG execution results."""
        output_parts = []
        total_duration = 0
        failed_count = 0

        for r in results:
            total_duration += r.duration_ms

            if r.status == NodeStatus.FAILED:
                failed_count += 1

            if r.stdout and not quiet:
                output_parts.append(r.stdout.rstrip())

            if r.stderr:
                output_parts.append(f"[stderr] {r.stderr.rstrip()}")

        if len(results) > 1:
            status = "✓" if failed_count == 0 else f"✗ ({failed_count} failed)"
            output_parts.append(
                f"\n[dag] {len(results)} nodes, {total_duration}ms, {status}"
            )

        return "\n".join(output_parts) if output_parts else "(no output)"

    @override
    def register(self, mcp_server: FastMCP) -> None:
        """Register DAG tool with MCP server."""
        tool_self = self

        @mcp_server.tool(name=self.name, description=self.description)
        async def dag_handler(
            commands: Annotated[
                List[Any],
                Field(
                    description="Commands to execute (strings, tool dicts, or parallel blocks)"
                ),
            ],
            parallel: Annotated[
                bool, Field(description="Run all commands in parallel", default=False)
            ] = False,
            shell: Annotated[
                Optional[str],
                Field(description="Shell to use (default: zsh)", default=None),
            ] = None,
            cwd: Annotated[
                Optional[str], Field(description="Working directory", default=None)
            ] = None,
            env: Annotated[
                Optional[Dict[str, str]],
                Field(description="Environment variables", default=None),
            ] = None,
            timeout: Annotated[
                int, Field(description="Timeout per command (seconds)", default=30)
            ] = 30,
            strict: Annotated[
                bool, Field(description="Stop on first error", default=False)
            ] = False,
            quiet: Annotated[
                bool, Field(description="Suppress stdout", default=False)
            ] = False,
            ctx: MCPContext = None,
        ) -> str:
            return await tool_self.call(
                ctx,
                commands=commands,
                parallel=parallel,
                shell=shell,
                cwd=cwd,
                env=env,
                timeout=timeout,
                strict=strict,
                quiet=quiet,
            )


def create_dag_tool(
    tools: Optional[Dict[str, BaseTool]] = None, default_shell: str = "zsh"
) -> DagTool:
    """Factory to create DAG execution tool."""
    return DagTool(tools, default_shell)