webber/webber-api/src/domains/agents/task/tools.py

"""
Tool registrations for the Task agent.

The Task agent has access to tools based on permission mode:
- Plan mode: Read-only tools only
- Default/auto_accept: All tools including write operations
"""
from pydantic_ai import Agent, RunContext

from src.domains.agents.base import AgentContext
from src.domains.tools.file.read import ReadFileTool
from src.domains.tools.file.glob import GlobFilesTool
from src.domains.tools.file.edit import EditFileTool
from src.domains.tools.file.write import WriteFileTool
from src.domains.tools.search.grep import GrepContentTool
from src.domains.tools.search.web import WebSearchTool
from src.domains.tools.shell.bash import BashReadOnlyTool
from src.domains.tools.shell.bash_full import BashTool


def _register_read_file(agent: Agent[AgentContext, str]) -> None:
    """Register read_file tool."""
    @agent.tool
    async def read_file(
        ctx: RunContext[AgentContext],
        file_path: str,
        offset: int = 0,
        limit: int = 2000
    ) -> str:
        """Read contents of a file with line numbers.

        Args:
            file_path: Absolute path to the file to read
            offset: Line number to start from (0-based, default: 0)
            limit: Maximum number of lines to read (default: 2000)

        Returns:
            File contents with line numbers, or error message.

        IMPORTANT: Always use absolute paths. Read files before editing them.
        """
        tool = ReadFileTool(allowed_paths=ctx.deps.allowed_paths)
        result = await tool.execute(
            file_path=file_path,
            offset=offset,
            limit=limit
        )
        return result.to_string()


def _register_glob_files(agent: Agent[AgentContext, str]) -> None:
    """Register glob_files tool."""
    @agent.tool
    async def glob_files(
        ctx: RunContext[AgentContext],
        pattern: str,
        path: str | None = None,
        limit: int = 100
    ) -> str:
        """Find files matching a glob pattern.

        Args:
            pattern: Glob pattern (e.g., "**/*.py", "src/**/*.ts", "*.md")
            path: Directory to search in (default: working directory)
            limit: Maximum number of files to return (default: 100)

        Returns:
            List of absolute file paths, sorted by modification time (newest first).

        Examples:
            - "**/*.py" finds all Python files
            - "src/**/*.ts" finds TypeScript files in src/
            - "**/test_*.py" finds all test files
        """
        tool = GlobFilesTool(allowed_paths=ctx.deps.allowed_paths)
        search_path = path or ctx.deps.working_dir
        result = await tool.execute(
            pattern=pattern,
            path=search_path,
            limit=limit
        )
        return result.to_string()


def _register_grep_content(agent: Agent[AgentContext, str]) -> None:
    """Register grep_content tool."""
    @agent.tool
    async def grep_content(
        ctx: RunContext[AgentContext],
        pattern: str,
        path: str | None = None,
        file_glob: str | None = None,
        context_lines: int = 0,
        case_sensitive: bool = True
    ) -> str:
        """Search file contents using regex pattern.

        Args:
            pattern: Regex pattern to search for (Python re syntax)
            path: Directory or file to search (default: working directory)
            file_glob: Filter files by glob (e.g., "*.py", "*.ts")
            context_lines: Lines of context before/after matches (default: 0)
            case_sensitive: Case-sensitive search (default: True)

        Returns:
            Matching lines with file paths and line numbers.
            Format: "filepath:line_num: content"
        """
        tool = GrepContentTool(allowed_paths=ctx.deps.allowed_paths)
        search_path = path or ctx.deps.working_dir
        result = await tool.execute(
            pattern=pattern,
            path=search_path,
            file_glob=file_glob,
            context_lines=context_lines,
            case_sensitive=case_sensitive
        )
        return result.to_string()


def _register_bash_readonly(agent: Agent[AgentContext, str]) -> None:
    """Register bash_readonly tool."""
    @agent.tool
    async def bash_readonly(
        ctx: RunContext[AgentContext],
        command: str,
        cwd: str | None = None,
        timeout: int = 30
    ) -> str:
        """Execute a read-only bash command.

        ALLOWED commands:
        - File inspection: ls, find, cat, head, tail, wc, file, stat, tree, du
        - Git (read-only): git status, git log, git diff, git show, git branch
        - Text processing: grep, awk, sed (read-only), sort, uniq
        - System info: pwd, whoami, hostname, which

        FORBIDDEN:
        - File modification (rm, mv, cp, mkdir, touch)
        - Redirects (>, >>)
        - Command chaining (&&, ||, ;)
        - Network (curl, wget)

        Args:
            command: The bash command to execute
            cwd: Working directory (default: agent working directory)
            timeout: Timeout in seconds (default: 30)
        """
        tool = BashReadOnlyTool(allowed_paths=ctx.deps.allowed_paths)
        working_dir = cwd or ctx.deps.working_dir
        result = await tool.execute(
            command=command,
            cwd=working_dir,
            timeout=min(timeout, ctx.deps.timeout_seconds)
        )
        return result.to_string()


def _register_spawn_agent(agent: Agent[AgentContext, str], readonly_only: bool = False) -> None:
    """Register spawn_agent tool."""
    @agent.tool
    async def spawn_agent(
        ctx: RunContext[AgentContext],
        agent_type: str,
        prompt: str,
        working_dir: str | None = None
    ) -> str:
        """Spawn a sub-agent to handle a focused task.

        Use this to offload work to specialized agents:
        - "explore": Fast codebase searches and analysis (read-only)
        - "plan": Design implementation strategies (read-only)

        Args:
            agent_type: Type of agent to spawn ("explore" or "plan")
            prompt: Task description for the sub-agent
            working_dir: Working directory for the sub-agent (default: current)

        Returns:
            Sub-agent's consolidated response.

        Examples:
            - spawn_agent(agent_type="explore", prompt="find all test files")
            - spawn_agent(agent_type="plan", prompt="design user auth feature")

        IMPORTANT:
        - Use sub-agents to keep context focused and efficient
        - Explore agent for research, Plan agent for design
        - Cannot spawn nested Task agents (recursion risk)
        """
        from src.domains.agents.base import get_agent

        # Validate agent type
        allowed_types = ["explore", "plan"]
        if agent_type not in allowed_types:
            if agent_type == "task":
                return "Error: Cannot spawn nested Task agents (recursion risk)"
            return f"Error: Unknown agent type '{agent_type}'. Allowed: {allowed_types}"

        sub_agent = get_agent(agent_type)
        if not sub_agent:
            return f"Error: Agent '{agent_type}' not found in registry"

        try:
            result = await sub_agent.run(
                prompt=prompt,
                working_dir=working_dir or ctx.deps.working_dir,
                allowed_paths=ctx.deps.allowed_paths,
            )
            return result
        except Exception as e:
            return f"Sub-agent error: {e}"


def register_readonly_tools(agent: Agent[AgentContext, str]) -> None:
    """
    Register read-only tools with the agent.

    Used in plan mode. Includes:
    - read_file, glob_files, grep_content, bash_readonly
    - spawn_agent (restricted to explore/plan)
    """
    _register_read_file(agent)
    _register_glob_files(agent)
    _register_grep_content(agent)
    _register_bash_readonly(agent)
    _register_spawn_agent(agent, readonly_only=True)


def register_task_tools(agent: Agent[AgentContext, str]) -> None:
    """
    Register all tools with the Task agent.

    Includes:
    - Read-only tools: read_file, glob_files, grep_content, bash_readonly
    - Write tools: edit_file, write_file, bash
    - External: web_search
    - Orchestration: spawn_agent
    """
    # Register read-only tools via helpers
    _register_read_file(agent)
    _register_glob_files(agent)
    _register_grep_content(agent)
    _register_bash_readonly(agent)

    # === Write tools ===

    @agent.tool
    async def edit_file(
        ctx: RunContext[AgentContext],
        file_path: str,
        old_string: str,
        new_string: str,
        replace_all: bool = False
    ) -> str:
        """Make targeted edits to a file using find-and-replace.

        Args:
            file_path: Absolute path to the file to edit
            old_string: The exact text to find and replace (must exist in file)
            new_string: The replacement text
            replace_all: If True, replace all occurrences. If False (default),
                         old_string must be unique (appear exactly once).

        Returns:
            Success message with diff preview, or error.

        IMPORTANT:
        - old_string must exactly match file content (including whitespace)
        - By default, old_string must appear exactly once (for safety)
        - Always read the file first to verify exact content before editing
        """
        tool = EditFileTool(allowed_paths=ctx.deps.allowed_paths)
        result = await tool.execute(
            file_path=file_path,
            old_string=old_string,
            new_string=new_string,
            replace_all=replace_all
        )
        return result.to_string()

    @agent.tool
    async def write_file(
        ctx: RunContext[AgentContext],
        file_path: str,
        content: str
    ) -> str:
        """Create a new file or overwrite an existing file.

        Args:
            file_path: Absolute path to the file to create/write
            content: The content to write to the file

        Returns:
            Success message with file path and size.

        IMPORTANT:
        - Parent directory must exist (use bash mkdir first if needed)
        - For editing existing files, prefer edit_file instead
        - Will overwrite existing files without confirmation
        """
        tool = WriteFileTool(allowed_paths=ctx.deps.allowed_paths)
        result = await tool.execute(
            file_path=file_path,
            content=content
        )
        return result.to_string()

    @agent.tool
    async def bash(
        ctx: RunContext[AgentContext],
        command: str,
        cwd: str | None = None,
        timeout: int = 60
    ) -> str:
        """Execute a bash command with write capabilities.

        ALLOWED:
        - File operations: ls, find, mkdir, touch, cp, mv, rm (single files)
        - Git (full): git add, git commit, git checkout, git merge, git pull
        - Python: python, pip install, pytest, mypy, ruff
        - Text processing: grep, awk, sed, sort
        - Command chaining: && and || are allowed

        FORBIDDEN:
        - sudo, su (privilege escalation)
        - Network: curl, wget, ssh, scp, rsync
        - Dangerous: rm -rf, chmod 777, dd, mkfs

        Args:
            command: The bash command to execute
            cwd: Working directory (default: agent working directory)
            timeout: Timeout in seconds (default: 60)

        Examples:
            - "mkdir -p src/utils" creates directory
            - "git add . && git commit -m 'fix: bug'" commits changes
            - "pytest tests/ -v" runs tests
        """
        tool = BashTool(allowed_paths=ctx.deps.allowed_paths)
        working_dir = cwd or ctx.deps.working_dir
        result = await tool.execute(
            command=command,
            cwd=working_dir,
            timeout=min(timeout, ctx.deps.timeout_seconds)
        )
        return result.to_string()

    # === External tools ===

    @agent.tool
    async def web_search(
        ctx: RunContext[AgentContext],
        query: str,
        num_results: int = 5,
        categories: str | None = None
    ) -> str:
        """Search the web for current information.

        Args:
            query: Search query (e.g., "Python 3.12 new features")
            num_results: Number of results to return (1-10, default: 5)
            categories: Optional category filter ("general", "it", "news", "science")

        Returns:
            Search results with titles, URLs, and snippets.

        Use this for:
        - Current events or recent information
        - Documentation updates
        - Technical references with URLs
        """
        tool = WebSearchTool()
        result = await tool.execute(
            query=query,
            num_results=num_results,
            categories=categories
        )
        return result.to_string()

    # === Orchestration tools ===
    _register_spawn_agent(agent, readonly_only=False)