zk-data-agent/src/agent_tool_core.py

from __future__ import annotations

import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Union

from .agent_types import AgentPermissions, AgentRuntimeConfig, ToolExecutionResult

if TYPE_CHECKING:
    from .account_runtime import AccountRuntime
    from .ask_user_runtime import AskUserRuntime
    from .config_runtime import ConfigRuntime
    from .lsp_runtime import LSPRuntime
    from .mcp_runtime import MCPRuntime
    from .plan_runtime import PlanRuntime
    from .remote_runtime import RemoteRuntime
    from .remote_trigger_runtime import RemoteTriggerRuntime
    from .search_runtime import SearchRuntime
    from .task_runtime import TaskRuntime
    from .team_runtime import TeamRuntime
    from .workflow_runtime import WorkflowRuntime
    from .worktree_runtime import WorktreeRuntime


class ToolPermissionError(RuntimeError):
    """Raised when the runtime configuration does not allow a tool action."""


class ToolExecutionError(RuntimeError):
    """Raised when a tool cannot complete because of invalid input or state."""


@dataclass(frozen=True)
class ToolExecutionContext:
    root: Path
    command_timeout_seconds: float
    max_output_chars: int
    permissions: AgentPermissions
    scratchpad_directory: Path | None = None
    python_env_dir: Path | None = None
    extra_env: dict[str, str] = field(default_factory=dict)
    cancel_event: Any | None = None
    process_registry: Any | None = None
    tool_registry: dict[str, 'AgentTool'] | None = None
    search_runtime: 'SearchRuntime | None' = None
    account_runtime: 'AccountRuntime | None' = None
    ask_user_runtime: 'AskUserRuntime | None' = None
    config_runtime: 'ConfigRuntime | None' = None
    lsp_runtime: 'LSPRuntime | None' = None
    mcp_runtime: 'MCPRuntime | None' = None
    remote_runtime: 'RemoteRuntime | None' = None
    remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None
    plan_runtime: 'PlanRuntime | None' = None
    task_runtime: 'TaskRuntime | None' = None
    team_runtime: 'TeamRuntime | None' = None
    workflow_runtime: 'WorkflowRuntime | None' = None
    worktree_runtime: 'WorktreeRuntime | None' = None


ToolHandler = Callable[
    [dict[str, Any], ToolExecutionContext],
    Union[str, tuple[str, dict[str, Any]]],
]


@dataclass(frozen=True)
class AgentTool:
    name: str
    description: str
    parameters: dict[str, Any]
    handler: ToolHandler

    def to_openai_tool(self) -> dict[str, object]:
        return {
            'type': 'function',
            'function': {
                'name': self.name,
                'description': self.description,
                'parameters': self.parameters,
            },
        }

    def execute(
        self,
        arguments: dict[str, Any],
        context: ToolExecutionContext,
    ) -> ToolExecutionResult:
        try:
            result = self.handler(arguments, context)
            if isinstance(result, tuple):
                content, metadata = result
            else:
                content, metadata = result, {}
            return ToolExecutionResult(
                name=self.name,
                ok=True,
                content=content,
                metadata=metadata,
            )
        except ToolPermissionError as exc:
            return ToolExecutionResult(
                name=self.name,
                ok=False,
                content=str(exc),
                metadata={'error_kind': 'permission_denied'},
            )
        except (ToolExecutionError, OSError, subprocess.SubprocessError) as exc:
            return ToolExecutionResult(
                name=self.name,
                ok=False,
                content=str(exc),
                metadata={'error_kind': 'tool_execution_error'},
            )


@dataclass(frozen=True)
class ToolStreamUpdate:
    kind: str
    content: str = ''
    stream: str | None = None
    result: ToolExecutionResult | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


def build_tool_context(
    config: AgentRuntimeConfig,
    *,
    scratchpad_directory: Path | None = None,
    python_env_dir: Path | None = None,
    extra_env: dict[str, str] | None = None,
    cancel_event: Any | None = None,
    process_registry: Any | None = None,
    tool_registry: dict[str, AgentTool] | None = None,
    search_runtime: 'SearchRuntime | None' = None,
    account_runtime: 'AccountRuntime | None' = None,
    ask_user_runtime: 'AskUserRuntime | None' = None,
    config_runtime: 'ConfigRuntime | None' = None,
    lsp_runtime: 'LSPRuntime | None' = None,
    mcp_runtime: 'MCPRuntime | None' = None,
    remote_runtime: 'RemoteRuntime | None' = None,
    remote_trigger_runtime: 'RemoteTriggerRuntime | None' = None,
    plan_runtime: 'PlanRuntime | None' = None,
    task_runtime: 'TaskRuntime | None' = None,
    team_runtime: 'TeamRuntime | None' = None,
    workflow_runtime: 'WorkflowRuntime | None' = None,
    worktree_runtime: 'WorktreeRuntime | None' = None,
) -> ToolExecutionContext:
    return ToolExecutionContext(
        root=config.cwd.resolve(),
        command_timeout_seconds=config.command_timeout_seconds,
        max_output_chars=config.max_output_chars,
        permissions=config.permissions,
        scratchpad_directory=scratchpad_directory.resolve() if scratchpad_directory else None,
        python_env_dir=(
            python_env_dir.resolve()
            if python_env_dir
            else config.python_env_dir.resolve()
            if config.python_env_dir
            else None
        ),
        extra_env=dict(extra_env or {}),
        cancel_event=cancel_event,
        process_registry=process_registry,
        tool_registry=tool_registry,
        search_runtime=search_runtime,
        account_runtime=account_runtime,
        ask_user_runtime=ask_user_runtime,
        config_runtime=config_runtime,
        lsp_runtime=lsp_runtime,
        mcp_runtime=mcp_runtime,
        remote_runtime=remote_runtime,
        remote_trigger_runtime=remote_trigger_runtime,
        plan_runtime=plan_runtime,
        task_runtime=task_runtime,
        team_runtime=team_runtime,
        workflow_runtime=workflow_runtime,
        worktree_runtime=worktree_runtime,
    )