Improve data skill routing and tool organization

src/agent_tool_core.py

@@ -0,0 +1,173 @@
+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)
+    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,
+    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 {}),
+        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,
+    )