AstrBot/astrbot/core/tools/message_tools.py

import json
import os
import shlex
import uuid
from pathlib import Path
from typing import TypedDict

from pydantic import Field
from pydantic.dataclasses import dataclass

import astrbot.core.message.components as Comp
from astrbot.api import logger
from astrbot.core.agent.run_context import ContextWrapper
from astrbot.core.agent.tool import FunctionTool, ToolExecResult
from astrbot.core.astr_agent_context import AstrAgentContext
from astrbot.core.computer.computer_client import get_booter
from astrbot.core.message.message_event_result import MessageChain
from astrbot.core.platform.message_session import MessageSession
from astrbot.core.tools.computer_tools.fs import _remote_basename
from astrbot.core.tools.computer_tools.util import (
    check_admin_permission,
    is_local_runtime,
    workspace_root,
    workspace_root_for_context,
)
from astrbot.core.tools.registry import builtin_tool
from astrbot.core.utils.astrbot_path import (
    get_astrbot_system_tmp_path,
    get_astrbot_temp_path,
)


class MessageComponentPayload(TypedDict, total=False):
    type: str
    text: str
    path: str
    url: str
    mention_user_id: str


def _file_send_allowed_roots(
    umo: str | None,
    current_workspace_root: Path | None = None,
) -> tuple[Path, ...]:
    roots = []
    if umo:
        roots.append(current_workspace_root or workspace_root(umo))
    roots.extend(
        [
            Path(get_astrbot_temp_path()).resolve(strict=False),
            Path(get_astrbot_system_tmp_path()).resolve(strict=False),
        ]
    )
    return tuple(roots)


def _is_path_within(path: Path, roots: tuple[Path, ...]) -> bool:
    return any(path == root or path.is_relative_to(root) for root in roots)


def _normalize_message_component(raw_msg: object) -> MessageComponentPayload | None:
    if not isinstance(raw_msg, dict):
        return None

    normalized: MessageComponentPayload = {}
    for key, value in raw_msg.items():
        if not isinstance(key, str):
            continue
        if key == "type" and isinstance(value, str):
            normalized["type"] = value
        elif key == "text" and isinstance(value, str):
            normalized["text"] = value
        elif key == "path" and isinstance(value, str):
            normalized["path"] = value
        elif key == "url" and isinstance(value, str):
            normalized["url"] = value
        elif key == "mention_user_id" and isinstance(value, str):
            normalized["mention_user_id"] = value
    return normalized


def _is_restricted_local_env(context: ContextWrapper[AstrAgentContext]) -> bool:
    if not is_local_runtime(context):
        return False
    cfg = context.context.context.get_config(
        umo=context.context.event.unified_msg_origin
    )
    provider_settings = cfg.get("provider_settings", {})
    require_admin = provider_settings.get("computer_use_require_admin", True)
    return require_admin and context.context.event.role != "admin"


def _can_send_local_file(
    context: ContextWrapper[AstrAgentContext],
    local_path: Path,
    current_workspace_root: Path | None = None,
) -> bool:
    umo = context.context.event.unified_msg_origin
    allowed_roots = _file_send_allowed_roots(umo, current_workspace_root)
    if _is_path_within(local_path, allowed_roots):
        return True
    return is_local_runtime(context) and not _is_restricted_local_env(context)


@builtin_tool
@dataclass
class SendMessageToUserTool(FunctionTool[AstrAgentContext]):
    name: str = "send_message_to_user"
    description: str = (
        "Send message to the user. "
        "Supports various message types including `plain`, `image`, `record`, `video`, `file`, and `mention_user`. "
        "Use this tool to send media files (`image`, `record`, `video`, `file`), "
        "or when you need to proactively message the user(such as cron job). For other normal text replies, you can output directly and no need to use this tool."
    )
    parameters: dict = Field(
        default_factory=lambda: {
            "type": "object",
            "properties": {
                "messages": {
                    "type": "array",
                    "description": "An ordered list of message components to send. `mention_user` type can be used to mention the user.",
                    "items": {
                        "type": "object",
                        "properties": {
                            "type": {
                                "type": "string",
                                "description": (
                                    "Component type. One of: "
                                    "plain, image, record, video, file, mention_user. Record is voice message."
                                ),
                            },
                            "text": {
                                "type": "string",
                                "description": "Text content for `plain` type.",
                            },
                            "path": {
                                "type": "string",
                                "description": "File path for `image`, `record`, `video`, or `file` types. Both local path and sandbox path are supported.",
                            },
                            "url": {
                                "type": "string",
                                "description": "URL for `image`, `record`, `video`, or `file` types.",
                            },
                            "mention_user_id": {
                                "type": "string",
                                "description": "User ID to mention for `mention_user` type.",
                            },
                        },
                        "required": ["type"],
                    },
                },
                "session": {
                    "type": "string",
                    "description": (
                        "Optional. Leave empty for the current session. "
                        "Use 'platform_id:message_type:session_id' to target another session."
                    ),
                },
            },
            "required": ["messages"],
        },
    )

    async def _resolve_path_from_sandbox(
        self,
        context: ContextWrapper[AstrAgentContext],
        path: str,
        *,
        component_type: str = "file",
    ) -> tuple[str, bool]:
        path = str(path).strip()
        if not path:
            raise FileNotFoundError(f"{component_type} path is empty")

        current_workspace_root = (
            await workspace_root_for_context(context)
            if is_local_runtime(context)
            else None
        )

        # Relative host paths are resolved only inside the user's workspace.
        if not os.path.isabs(path):
            unified_msg_origin = context.context.event.unified_msg_origin
            if unified_msg_origin:
                ws_path = current_workspace_root or workspace_root(unified_msg_origin)
                try:
                    ws_candidate = (ws_path / path).resolve(strict=False)
                    if ws_candidate.is_file() and ws_candidate.is_relative_to(ws_path):
                        return str(ws_candidate), False
                except Exception:
                    pass
        else:
            local_candidate = Path(path).expanduser().resolve(strict=False)
            if local_candidate.is_file():
                if _can_send_local_file(
                    context, local_candidate, current_workspace_root
                ):
                    return str(local_candidate), False
                if is_local_runtime(context):
                    allowed = ", ".join(
                        str(root)
                        for root in _file_send_allowed_roots(
                            context.context.event.unified_msg_origin,
                            current_workspace_root,
                        )
                    )
                    raise PermissionError(
                        "Local file send is restricted for this user. "
                        f"Allowed directories: {allowed}. "
                        f"Blocked path: {local_candidate}."
                    )

        try:
            sb = await get_booter(
                context.context.context,
                context.context.event.unified_msg_origin,
            )
            quoted_path = shlex.quote(path)
            result = await sb.shell.exec(f"test -f {quoted_path} && echo '_&exists_'")
            if "_&exists_" in json.dumps(result):
                name = _remote_basename(path) or os.path.basename(path)
                local_path = os.path.join(
                    get_astrbot_temp_path(),
                    f"sandbox_{uuid.uuid4().hex[:4]}_{name}",
                )
                await sb.download_file(path, local_path)
                logger.info(f"Downloaded file from sandbox: {path} -> {local_path}")
                return local_path, True
        except Exception as exc:
            logger.warning(f"Failed to check/download file from sandbox: {exc}")
            raise

        raise FileNotFoundError(f"{component_type} path does not exist: {path}")

    async def call(
        self,
        context: ContextWrapper[AstrAgentContext],
        **kwargs,
    ) -> ToolExecResult:
        # Security: only AstrBot admins can send messages to other sessions.
        # Non-admin users are always restricted to their own session.
        # See https://github.com/AstrBotDevs/AstrBot/issues/7822
        current_session = context.context.event.unified_msg_origin
        session = kwargs.get("session") or current_session
        if session != current_session:
            if permission_error := check_admin_permission(
                context,
                "Send message to another session",
            ):
                return permission_error
        messages = kwargs.get("messages")
        if not isinstance(messages, list) or not messages:
            return "error: messages parameter is empty or invalid."

        components: list[Comp.BaseMessageComponent] = []
        for idx, msg in enumerate(messages):
            normalized_msg = _normalize_message_component(msg)
            if normalized_msg is None:
                return f"error: messages[{idx}] should be an object."

            msg_type = normalized_msg.get("type", "").lower()
            if not msg_type:
                return f"error: messages[{idx}].type is required."

            try:
                if msg_type == "plain":
                    text = normalized_msg.get("text", "").strip()
                    if not text:
                        return f"error: messages[{idx}].text is required for plain component."
                    components.append(Comp.Plain(text=text))
                elif msg_type == "image":
                    path = normalized_msg.get("path")
                    url = normalized_msg.get("url")
                    if path:
                        local_path, _ = await self._resolve_path_from_sandbox(
                            context,
                            path,
                            component_type="image",
                        )
                        components.append(Comp.Image.fromFileSystem(path=local_path))
                    elif url:
                        components.append(Comp.Image.fromURL(url=url))
                    else:
                        return f"error: messages[{idx}] must include path or url for image component."
                elif msg_type == "record":
                    path = normalized_msg.get("path")
                    url = normalized_msg.get("url")
                    if path:
                        local_path, _ = await self._resolve_path_from_sandbox(
                            context,
                            path,
                            component_type="record",
                        )
                        components.append(Comp.Record.fromFileSystem(path=local_path))
                    elif url:
                        components.append(Comp.Record.fromURL(url=url))
                    else:
                        return f"error: messages[{idx}] must include path or url for record component."
                elif msg_type == "video":
                    path = normalized_msg.get("path")
                    url = normalized_msg.get("url")
                    if path:
                        local_path, _ = await self._resolve_path_from_sandbox(
                            context,
                            path,
                            component_type="video",
                        )
                        components.append(Comp.Video.fromFileSystem(path=local_path))
                    elif url:
                        components.append(Comp.Video.fromURL(url=url))
                    else:
                        return f"error: messages[{idx}] must include path or url for video component."
                elif msg_type == "file":
                    path = normalized_msg.get("path")
                    url = normalized_msg.get("url")
                    name = (
                        normalized_msg.get("text")
                        or (_remote_basename(path) if path else "")
                        or (os.path.basename(url) if url else "")
                        or "file"
                    )
                    if path:
                        local_path, _ = await self._resolve_path_from_sandbox(
                            context,
                            path,
                            component_type="file",
                        )
                        components.append(Comp.File(name=name, file=local_path))
                    elif url:
                        components.append(Comp.File(name=name, url=url))
                    else:
                        return f"error: messages[{idx}] must include path or url for file component."
                elif msg_type == "mention_user":
                    mention_user_id = normalized_msg.get("mention_user_id")
                    if not mention_user_id:
                        return f"error: messages[{idx}].mention_user_id is required for mention_user component."
                    components.append(Comp.At(qq=mention_user_id))
                else:
                    return (
                        f"error: unsupported message type '{msg_type}' at index {idx}."
                    )
            except FileNotFoundError as exc:
                return f"error: {exc}"
            except PermissionError as exc:
                return f"error: {exc}"
            except Exception as exc:
                return f"error: failed to build messages[{idx}] component: {exc}"

        try:
            target_session = (
                MessageSession.from_str(session)
                if isinstance(session, str)
                else session
            )
        except Exception:
            # LLM 在 cron 等主动场景下可能只传 session_id（如 oc_xxx），
            # 而不是完整的三段式 platform_id:message_type:session_id。
            # 此时用 current_session 的前两段补全。
            # 注意：这里的session是传入的session参数，实际上是用户输入的session_id
            # current_session才是完整的三段式session字符串。
            # 仅当传入字符串不含 ':'（明显是裸 session_id）时才用 current_session 补全，
            # 避免 LLM 传了带 ':' 但格式错误的目标 session 被错误修复。
            # issue: https://github.com/AstrBotDevs/AstrBot/issues/7907
            if isinstance(session, str) and current_session and ":" not in session:
                try:
                    cur = MessageSession.from_str(current_session)
                    target_session = MessageSession(
                        platform_name=cur.platform_id,
                        message_type=cur.message_type,
                        session_id=session,
                    )
                except Exception:
                    return f"error: invalid session: {session}"
            else:
                return f"error: invalid session: {session}"

        message_chain = MessageChain(chain=components)
        await context.context.context.send_message(target_session, message_chain)
        if str(target_session) == current_session:
            context.context.event._has_send_oper = True
            sent_plain_text = message_chain.get_plain_text().strip()
            if sent_plain_text:
                sent_plain_texts = context.context.event.get_extra(
                    "_send_message_to_user_current_session_plain_texts",
                    [],
                )
                if not isinstance(sent_plain_texts, list):
                    sent_plain_texts = []
                sent_plain_texts.append(sent_plain_text)
                context.context.event.set_extra(
                    "_send_message_to_user_current_session_plain_texts",
                    sent_plain_texts,
                )
        return f"Message sent to session {target_session}"


__all__ = [
    "SendMessageToUserTool",
]