tools_service.py

from __future__ import annotations

import traceback
from typing import Any

from astrbot.core import logger, sp
from astrbot.core.agent.mcp_client import MCPTool, validate_mcp_stdio_config
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
from astrbot.core.star import star_map
from astrbot.core.tools.registry import get_builtin_tool_config_statuses


class ToolsServiceError(Exception):
    pass


class EmptyMcpServersError(ValueError):
    pass


def extract_mcp_server_config(mcp_servers_value: object) -> dict:
    if not isinstance(mcp_servers_value, dict):
        raise ValueError("mcpServers must be a JSON object")
    if not mcp_servers_value:
        raise EmptyMcpServersError("mcpServers configuration cannot be empty")
    key_0 = next(iter(mcp_servers_value))
    extracted = mcp_servers_value[key_0]
    if not isinstance(extracted, dict):
        raise ValueError(
            "Invalid mcpServers format. Ensure each key in mcpServers is a server name, "
            "and each value is an object containing fields like command/url."
        )
    return extracted


class ToolsService:
    def __init__(self, core_lifecycle: AstrBotCoreLifecycle) -> None:
        self.core_lifecycle = core_lifecycle
        self.tool_mgr = core_lifecycle.provider_manager.llm_tools

    def rollback_mcp_server(self, name: str) -> bool:
        try:
            rollback_config = self.tool_mgr.load_mcp_config()
            if name in rollback_config["mcpServers"]:
                rollback_config["mcpServers"].pop(name)
                return self.tool_mgr.save_mcp_config(rollback_config)
            return True
        except Exception:
            logger.error(traceback.format_exc())
            return False

    def get_mcp_servers(self) -> list[dict]:
        try:
            config = self.tool_mgr.load_mcp_config()
            servers = []
            mcp_servers = config.get("mcpServers", {})

            if not isinstance(mcp_servers, dict):
                logger.warning(
                    f"Invalid MCP server config type: {type(mcp_servers).__name__}. Expected object/dict; skipped all MCP servers."
                )
                mcp_servers = {}

            for name, server_config in mcp_servers.items():
                if not isinstance(server_config, dict):
                    logger.warning(
                        f"Invalid config for MCP server '{name}' (type: {type(server_config).__name__}); skipped."
                    )
                    continue

                server_info = {
                    "name": name,
                    "active": server_config.get("active", True),
                }
                for key, value in server_config.items():
                    if key != "active":
                        server_info[key] = value

                for name_key, runtime in self.tool_mgr.mcp_server_runtime_view.items():
                    if name_key == name:
                        mcp_client = runtime.client
                        server_info["tools"] = [tool.name for tool in mcp_client.tools]
                        server_info["errlogs"] = mcp_client.server_errlogs
                        break
                else:
                    server_info["tools"] = []

                servers.append(server_info)

            return servers
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to get MCP server list: {exc!s}") from exc

    def get_mcp_server_config(self, name: str) -> dict | None:
        config = self.tool_mgr.load_mcp_config()
        mcp_servers = config.get("mcpServers", {})
        if not isinstance(mcp_servers, dict):
            return None

        server_config = mcp_servers.get(name)
        if not isinstance(server_config, dict):
            return None
        return dict(server_config)

    async def add_mcp_server(self, server_data: Any) -> str:
        try:
            name = server_data.get("name", "")
            if not name:
                raise ToolsServiceError("Server name cannot be empty")

            has_valid_config, server_config = self._build_server_config(server_data)
            if not has_valid_config:
                raise ToolsServiceError("A valid server configuration is required")

            self._validate_server_config(server_config)

            config = self.tool_mgr.load_mcp_config()
            if name in config["mcpServers"]:
                raise ToolsServiceError(f"Server {name} already exists")

            try:
                await self.tool_mgr.test_mcp_server_connection(server_config)
            except Exception as exc:
                logger.error(traceback.format_exc())
                raise ToolsServiceError(f"MCP connection test failed: {exc!s}") from exc

            config["mcpServers"][name] = server_config

            if self.tool_mgr.save_mcp_config(config):
                await self._enable_added_server(name, server_config)
                return f"Successfully added MCP server {name}"
            raise ToolsServiceError("Failed to save configuration")
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to add MCP server: {exc!s}") from exc

    async def update_mcp_server(self, server_data: Any) -> str:
        try:
            name = server_data.get("name", "")
            old_name = server_data.get("oldName") or name

            if not name:
                raise ToolsServiceError("Server name cannot be empty")

            config = self.tool_mgr.load_mcp_config()

            if old_name not in config["mcpServers"]:
                raise ToolsServiceError(f"Server {old_name} does not exist")

            is_rename = name != old_name
            if name in config["mcpServers"] and is_rename:
                raise ToolsServiceError(f"Server {name} already exists")

            old_config = config["mcpServers"][old_name]
            old_active = (
                old_config.get("active", True) if isinstance(old_config, dict) else True
            )
            active = server_data.get("active", old_active)

            only_update_active, server_config = self._build_updated_server_config(
                server_data,
                old_config,
                active,
            )
            self._validate_server_config(server_config)

            if is_rename:
                config["mcpServers"].pop(old_name)
                config["mcpServers"][name] = server_config
            else:
                config["mcpServers"][name] = server_config

            if self.tool_mgr.save_mcp_config(config):
                await self._sync_updated_server_runtime(
                    name=name,
                    old_name=old_name,
                    active=active,
                    is_rename=is_rename,
                    only_update_active=only_update_active,
                    server_config=config["mcpServers"][name],
                )
                return f"Successfully updated MCP server {name}"
            raise ToolsServiceError("Failed to save configuration")
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to update MCP server: {exc!s}") from exc

    async def delete_mcp_server(self, server_data: Any) -> str:
        try:
            name = server_data.get("name", "")

            if not name:
                raise ToolsServiceError("Server name cannot be empty")

            config = self.tool_mgr.load_mcp_config()

            if name not in config["mcpServers"]:
                raise ToolsServiceError(f"Server {name} does not exist")

            del config["mcpServers"][name]

            if self.tool_mgr.save_mcp_config(config):
                if name in self.tool_mgr.mcp_server_runtime_view:
                    await self._disable_server(name)
                return f"Successfully deleted MCP server {name}"
            raise ToolsServiceError("Failed to save configuration")
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to delete MCP server: {exc!s}") from exc

    async def test_mcp_connection(self, server_data: Any) -> list:
        try:
            config = server_data.get("mcp_server_config", None)

            if not isinstance(config, dict) or not config:
                raise ToolsServiceError("Invalid MCP server configuration")

            if "mcpServers" in config:
                mcp_servers = config["mcpServers"]
                if isinstance(mcp_servers, dict) and len(mcp_servers) > 1:
                    raise ToolsServiceError(
                        "Only one MCP server configuration can be tested at a time"
                    )
                try:
                    config = extract_mcp_server_config(mcp_servers)
                except EmptyMcpServersError as exc:
                    raise ToolsServiceError(
                        "MCP server configuration cannot be empty"
                    ) from exc
                except ValueError as exc:
                    raise ToolsServiceError(f"{exc!s}") from exc
            elif not config:
                raise ToolsServiceError("MCP server configuration cannot be empty")

            self._validate_server_config(config)
            return await self.tool_mgr.test_mcp_server_connection(config)
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to test MCP connection: {exc!s}") from exc

    def get_tool_list(self) -> list[dict]:
        try:
            tools = list(self.tool_mgr.func_list)
            existing_names = {tool.name for tool in tools}
            for tool in self.tool_mgr.iter_builtin_tools():
                if tool.name not in existing_names:
                    tools.append(tool)

            config_entries = self._get_config_entries()
            tools_dict = []
            for tool in tools:
                tools_dict.append(self._serialize_tool(tool, config_entries))
            return tools_dict
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to get tool list: {exc!s}") from exc

    def update_tool_permission(self, data: Any) -> str:
        """Set a tool permission level.

        Args:
            data: Legacy dashboard payload with ``name`` and ``permission``.

        Returns:
            A success message for the response body.

        Raises:
            ToolsServiceError: If the payload is invalid, the tool is unknown,
                or permission storage cannot be updated.
        """
        try:
            tool_name = data.get("name") if isinstance(data, dict) else None
            permission = data.get("permission") if isinstance(data, dict) else None

            if not tool_name or permission not in ("admin", "member"):
                raise ToolsServiceError(
                    "name and permission (admin or member) are required"
                )

            if self.tool_mgr.is_builtin_tool(tool_name):
                raise ToolsServiceError(
                    "Builtin tools do not support per-tool permission configuration."
                )

            if not any(t.name == tool_name for t in self.tool_mgr.func_list):
                raise ToolsServiceError(f"Tool '{tool_name}' not found")

            perms_store = sp.get(
                "tool_permissions",
                {},
                scope="global",
                scope_id="global",
            )
            if not isinstance(perms_store, dict):
                perms_store = {}
            defaults = perms_store.get("_default", {})
            if not isinstance(defaults, dict):
                defaults = {}
            defaults[tool_name] = permission
            perms_store["_default"] = defaults
            sp.put(
                "tool_permissions",
                perms_store,
                scope="global",
                scope_id="global",
            )

            return f"Tool '{tool_name}' permission set to {permission}"
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(
                f"Failed to update tool permission: {exc!s}"
            ) from exc

    def toggle_tool(self, data: Any) -> str:
        try:
            tool_name = data.get("name")
            action = data.get("activate")

            if not tool_name or action is None:
                raise ToolsServiceError("Missing required parameters: name or activate")

            if self.tool_mgr.is_builtin_tool(tool_name):
                raise ToolsServiceError(
                    "Builtin tools are read-only and cannot be toggled."
                )

            if action:
                try:
                    ok = self.tool_mgr.activate_llm_tool(tool_name, star_map=star_map)
                except ValueError as exc:
                    raise ToolsServiceError(
                        f"Failed to activate tool: {exc!s}"
                    ) from exc
            else:
                ok = self.tool_mgr.deactivate_llm_tool(tool_name)

            if ok:
                return "Operation successful."
            raise ToolsServiceError(
                f"Tool {tool_name} does not exist or the operation failed."
            )
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Failed to operate tool: {exc!s}") from exc

    async def sync_provider(self, data: Any) -> str:
        try:
            provider_name = data.get("name")
            match provider_name:
                case "modelscope":
                    access_token = data.get("access_token", "")
                    await self.tool_mgr.sync_modelscope_mcp_servers(access_token)
                case _:
                    raise ToolsServiceError(f"Unknown provider: {provider_name}")

            return "Sync completed"
        except ToolsServiceError:
            raise
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(f"Sync failed: {exc!s}") from exc

    @staticmethod
    def _build_server_config(server_data: dict) -> tuple[bool, dict]:
        has_valid_config = False
        server_config = {"active": server_data.get("active", True)}

        for key, value in server_data.items():
            if key in ["name", "active", "tools", "errlogs"]:
                continue
            if key == "mcpServers":
                try:
                    server_config = extract_mcp_server_config(server_data["mcpServers"])
                except ValueError as exc:
                    raise ToolsServiceError(f"{exc!s}") from exc
            else:
                server_config[key] = value
            has_valid_config = True

        return has_valid_config, server_config

    @staticmethod
    def _build_updated_server_config(
        server_data: dict,
        old_config: object,
        active: bool,
    ) -> tuple[bool, dict]:
        server_config = {"active": active}
        only_update_active = True

        for key, value in server_data.items():
            if key in ["name", "active", "tools", "errlogs", "oldName"]:
                continue
            if key == "mcpServers":
                try:
                    server_config = extract_mcp_server_config(server_data["mcpServers"])
                except ValueError as exc:
                    raise ToolsServiceError(f"{exc!s}") from exc
            else:
                server_config[key] = value
            only_update_active = False

        if only_update_active and isinstance(old_config, dict):
            for key, value in old_config.items():
                if key != "active":
                    server_config[key] = value

        return only_update_active, server_config

    @staticmethod
    def _validate_server_config(server_config: dict) -> None:
        try:
            validate_mcp_stdio_config(server_config)
        except ValueError as exc:
            raise ToolsServiceError(f"{exc!s}") from exc

    async def _enable_added_server(self, name: str, server_config: dict) -> None:
        try:
            await self.tool_mgr.enable_mcp_server(name, server_config, timeout=30)
        except TimeoutError as exc:
            rollback_ok = self.rollback_mcp_server(name)
            err_msg = f"Timed out while enabling MCP server {name}."
            if not rollback_ok:
                err_msg += (
                    " Configuration rollback failed. Please check the config manually."
                )
            raise ToolsServiceError(err_msg) from exc
        except Exception as exc:
            logger.error(traceback.format_exc())
            rollback_ok = self.rollback_mcp_server(name)
            err_msg = f"Failed to enable MCP server {name}: {exc!s}"
            if not rollback_ok:
                err_msg += (
                    " Configuration rollback failed. Please check the config manually."
                )
            raise ToolsServiceError(err_msg) from exc

    async def _sync_updated_server_runtime(
        self,
        *,
        name: str,
        old_name: str,
        active: bool,
        is_rename: bool,
        only_update_active: bool,
        server_config: dict,
    ) -> None:
        if active:
            if (
                old_name in self.tool_mgr.mcp_server_runtime_view
                or not only_update_active
                or is_rename
            ):
                await self._disable_server_before_enable(old_name)
            await self._enable_updated_server(name, server_config)
        elif old_name in self.tool_mgr.mcp_server_runtime_view:
            await self._disable_server(old_name)

    async def _disable_server_before_enable(self, old_name: str) -> None:
        try:
            await self.tool_mgr.disable_mcp_server(old_name, timeout=10)
        except TimeoutError as exc:
            raise ToolsServiceError(
                f"Timed out while disabling MCP server {old_name} before enabling: {exc!s}"
            ) from exc
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(
                f"Failed to disable MCP server {old_name} before enabling: {exc!s}"
            ) from exc

    async def _enable_updated_server(self, name: str, server_config: dict) -> None:
        try:
            await self.tool_mgr.enable_mcp_server(name, server_config, timeout=30)
        except TimeoutError as exc:
            raise ToolsServiceError(
                f"Timed out while enabling MCP server {name}."
            ) from exc
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(
                f"Failed to enable MCP server {name}: {exc!s}"
            ) from exc

    async def _disable_server(self, name: str) -> None:
        try:
            await self.tool_mgr.disable_mcp_server(name, timeout=10)
        except TimeoutError as exc:
            raise ToolsServiceError(
                f"Timed out while disabling MCP server {name}."
            ) from exc
        except Exception as exc:
            logger.error(traceback.format_exc())
            raise ToolsServiceError(
                f"Failed to disable MCP server {name}: {exc!s}"
            ) from exc

    def _get_config_entries(self) -> list[dict]:
        conf_list = self.core_lifecycle.astrbot_config_mgr.get_conf_list()
        conf_name_map = {conf["id"]: conf["name"] for conf in conf_list}
        config_entries = []
        for conf_id, conf in self.core_lifecycle.astrbot_config_mgr.confs.items():
            config_entries.append(
                {
                    "conf_id": conf_id,
                    "conf_name": conf_name_map.get(conf_id, conf_id),
                    "config": conf,
                }
            )
        return config_entries

    def _serialize_tool(self, tool, config_entries: list[dict]) -> dict:
        readonly = False
        builtin_config_statuses = []
        builtin_config_tags = []
        if self.tool_mgr.is_builtin_tool(tool.name):
            origin = "builtin"
            origin_name = "AstrBot Core"
            readonly = True
            builtin_config_statuses = get_builtin_tool_config_statuses(
                tool.name,
                config_entries,
            )
            builtin_config_tags = [
                status for status in builtin_config_statuses if status["enabled"]
            ]
        elif isinstance(tool, MCPTool):
            origin = "mcp"
            origin_name = tool.mcp_server_name
        elif tool.handler_module_path and star_map.get(tool.handler_module_path):
            star = star_map[tool.handler_module_path]
            origin = "plugin"
            origin_name = star.name
        else:
            origin = "unknown"
            origin_name = "unknown"

        tool_info = {
            "name": tool.name,
            "description": tool.description,
            "parameters": tool.parameters,
            "active": tool.active,
            "origin": origin,
            "origin_name": origin_name,
            "readonly": readonly,
            "builtin_config_statuses": builtin_config_statuses,
            "builtin_config_tags": builtin_config_tags,
        }
        if not readonly:
            perms_store = sp.get(
                "tool_permissions",
                {},
                scope="global",
                scope_id="global",
            )
            defaults = (
                perms_store.get("_default", {}) if isinstance(perms_store, dict) else {}
            )
            configured = tool.name in defaults
            permission = (
                defaults[tool.name]
                if configured
                else self.tool_mgr._default_permission(tool.name)
            )
            tool_info["permission"] = permission
            tool_info["permission_configured"] = configured
            tool_info["declared_permission_type"] = getattr(
                tool, "declared_permission_type", None
            )
        return tool_info