AstrBot/astrbot/core/computer/computer_tool_provider.py

"""ComputerToolProvider — decoupled tool injection for computer-use runtimes.

Encapsulates all sandbox / local tool injection logic previously hardcoded in
``astr_main_agent.py``.  The main agent now calls
``provider.get_tools(ctx)`` / ``provider.get_system_prompt_addon(ctx)``
without knowing about specific tool classes.

Tool lists are delegated to booter subclasses via ``get_default_tools()``
and ``get_tools()`` (see ``booters/base.py``), so adding a new booter type
does not require changes here.
"""

from __future__ import annotations

import os
import platform
from functools import cache
from typing import TYPE_CHECKING

from astrbot.api import logger
from astrbot.core.tool_provider import ToolProviderContext

if TYPE_CHECKING:
    from astrbot.core.agent.tool import FunctionTool


# ---------------------------------------------------------------------------
# Lazy local-mode tool cache
# ---------------------------------------------------------------------------

_LOCAL_TOOLS_CACHE: list[FunctionTool] | None = None


def _get_local_tools() -> list[FunctionTool]:
    global _LOCAL_TOOLS_CACHE
    if _LOCAL_TOOLS_CACHE is None:
        from astrbot.core.computer.tools import ExecuteShellTool, LocalPythonTool

        _LOCAL_TOOLS_CACHE = [
            ExecuteShellTool(is_local=True),
            LocalPythonTool(),
        ]
    return list(_LOCAL_TOOLS_CACHE)


# ---------------------------------------------------------------------------
# System-prompt helpers
# ---------------------------------------------------------------------------

SANDBOX_MODE_PROMPT = (
    "You have access to a sandboxed environment and can execute "
    "shell commands and Python code securely."
)


def _build_local_mode_prompt() -> str:
    system_name = platform.system() or "Unknown"
    shell_hint = (
        "The runtime shell is Windows Command Prompt (cmd.exe). "
        "Use cmd-compatible commands and do not assume Unix commands like cat/ls/grep are available."
        if system_name.lower() == "windows"
        else "The runtime shell is Unix-like. Use POSIX-compatible shell commands."
    )
    return (
        "You have access to the host local environment and can execute shell commands and Python code. "
        f"Current operating system: {system_name}. "
        f"{shell_hint}"
    )


# ---------------------------------------------------------------------------
# ComputerToolProvider
# ---------------------------------------------------------------------------


class ComputerToolProvider:
    """Provides computer-use tools (local / sandbox) based on session context.

    Sandbox tool lists are delegated to booter subclasses so that each booter
    declares its own capabilities.  ``get_tools`` prefers the precise
    post-boot tool list from a running session; when the sandbox has not yet
    been booted it falls back to the conservative pre-boot default.
    """

    @staticmethod
    def get_all_tools() -> list[FunctionTool]:
        """Return ALL computer-use tools across all runtimes for registration.

        Creates **fresh instances** separate from the runtime caches so that
        setting ``active=False`` on them does not affect runtime behaviour.
        These registration-only instances let the WebUI display and assign
        tools without injecting them into actual LLM requests.

        At request time, ``get_tools(ctx)`` provides the real, active
        instances filtered by runtime.
        """
        from astrbot.core.computer.tools import (
            AnnotateExecutionTool,
            BrowserBatchExecTool,
            BrowserExecTool,
            CreateSkillCandidateTool,
            CreateSkillPayloadTool,
            EvaluateSkillCandidateTool,
            ExecuteShellTool,
            FileDownloadTool,
            FileUploadTool,
            GetExecutionHistoryTool,
            GetSkillPayloadTool,
            ListSkillCandidatesTool,
            ListSkillReleasesTool,
            LocalPythonTool,
            PromoteSkillCandidateTool,
            PythonTool,
            RollbackSkillReleaseTool,
            RunBrowserSkillTool,
            SyncSkillReleaseTool,
        )

        all_tools: list[FunctionTool] = [
            ExecuteShellTool(),
            PythonTool(),
            FileUploadTool(),
            FileDownloadTool(),
            LocalPythonTool(),
            BrowserExecTool(),
            BrowserBatchExecTool(),
            RunBrowserSkillTool(),
            GetExecutionHistoryTool(),
            AnnotateExecutionTool(),
            CreateSkillPayloadTool(),
            GetSkillPayloadTool(),
            CreateSkillCandidateTool(),
            ListSkillCandidatesTool(),
            EvaluateSkillCandidateTool(),
            PromoteSkillCandidateTool(),
            ListSkillReleasesTool(),
            RollbackSkillReleaseTool(),
            SyncSkillReleaseTool(),
        ]

        # De-duplicate by name and mark inactive so they are visible
        # in WebUI but never sent to the LLM via func_list.
        seen: set[str] = set()
        result: list[FunctionTool] = []
        for tool in all_tools:
            if tool.name not in seen:
                tool.active = False
                result.append(tool)
                seen.add(tool.name)
        return result

    def get_tools(self, ctx: ToolProviderContext) -> list[FunctionTool]:
        runtime = ctx.computer_use_runtime
        if runtime == "none":
            return []

        if runtime == "local":
            return _get_local_tools()

        if runtime == "sandbox":
            return self._sandbox_tools(ctx)

        logger.warning("[ComputerToolProvider] Unknown runtime: %s", runtime)
        return []

    def get_system_prompt_addon(self, ctx: ToolProviderContext) -> str:
        runtime = ctx.computer_use_runtime
        if runtime == "none":
            return ""

        if runtime == "local":
            return f"\n{_build_local_mode_prompt()}\n"

        if runtime == "sandbox":
            return self._sandbox_prompt_addon(ctx)

        return ""

    # -- sandbox helpers ----------------------------------------------------

    def _sandbox_tools(self, ctx: ToolProviderContext) -> list[FunctionTool]:
        """Collect tools for sandbox mode.

        Prefers the precise post-boot tool list from the running session's
        booter instance (``get_sandbox_tools``).  When the sandbox has not
        yet been booted, falls back to the conservative pre-boot default
        declared by the booter class (``get_default_sandbox_tools``).
        """
        from astrbot.core.computer.computer_client import (
            get_default_sandbox_tools,
            get_sandbox_tools,
        )

        booter_type = ctx.sandbox_cfg.get("booter", "shipyard_neo")

        # Validate shipyard (non-neo) config
        if booter_type == "shipyard":
            ep = ctx.sandbox_cfg.get("shipyard_endpoint", "")
            at = ctx.sandbox_cfg.get("shipyard_access_token", "")
            if not ep or not at:
                logger.error("Shipyard sandbox configuration is incomplete.")
                return []
            os.environ["SHIPYARD_ENDPOINT"] = ep
            os.environ["SHIPYARD_ACCESS_TOKEN"] = at

        # Prefer precise post-boot tools from the running session
        booted_tools = get_sandbox_tools(ctx.session_id)
        if booted_tools:
            return booted_tools

        # Pre-boot: conservative default from booter class
        return get_default_sandbox_tools(ctx.sandbox_cfg)

    def _sandbox_prompt_addon(self, ctx: ToolProviderContext) -> str:
        """Build system-prompt addon for sandbox mode."""
        from astrbot.core.computer.computer_client import get_sandbox_prompt_parts

        parts = get_sandbox_prompt_parts(ctx.sandbox_cfg)
        parts.append(f"\n{SANDBOX_MODE_PROMPT}\n")
        return "".join(parts)


def get_all_tools() -> list[FunctionTool]:
    """Module-level entry point for ``FunctionToolManager.register_internal_tools()``.

    Delegates to ``ComputerToolProvider.get_all_tools()`` which collects
    tools from all runtimes (local, sandbox, browser, neo).
    """
    return ComputerToolProvider.get_all_tools()