astrbotmcp 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

astrbot_mcp/astrbot_client.py

@@ -98,7 +98,16 @@ class AstrBotClient:
             pwd = hashlib.md5(pwd.encode("utf-8")).hexdigest()
 
         url = f"{self.base_url}/api/auth/login"
-        client_kwargs = {"timeout": self.timeout}
+        # SSE endpoints can legitimately stay quiet for a long time while work is happening.
+        # Use an infinite read timeout, while keeping connect/write/pool bounded.
+        client_kwargs = {
+            "timeout": httpx.Timeout(
+                connect=self.timeout,
+                read=None,
+                write=self.timeout,
+                pool=self.timeout,
+            )
+        }
         if self.settings.disable_proxy:
             client_kwargs["trust_env"] = False  # 禁用代理，忽略环境变量设置
 
@@ -175,9 +184,12 @@ class AstrBotClient:
         max_events: Optional[int] = None,
     ) -> List[Dict[str, Any]]:
         """
-        Consume a simple SSE endpoint and return parsed JSON payloads.
+        Consume an SSE endpoint and return parsed event payloads.
 
-        AstrBot's SSE endpoints use `data: {...}\\n\\n` format per event.
+        AstrBot's SSE endpoints typically use `data: {...}\\n\\n` format per event.
+        This parser is tolerant:
+          - Supports multi-line `data:` frames per SSE spec (joined with `\\n`).
+          - If `data:` is not valid JSON, returns it as `{\"type\":\"raw\",\"data\":...}`.
 
         `max_seconds` is a soft upper bound for how long we wait:
         - 如果持续有事件流入，最多等待约 `max_seconds` 秒；
@@ -223,31 +235,63 @@ class AstrBotClient:
                     )
 
                 async def consume() -> None:
-                    async for line in response.aiter_lines():
-                        if not line:
-                            # Heartbeats / blank lines
-                            continue
+                    current_event: str | None = None
+                    data_lines: List[str] = []
 
-                        if not line.startswith("data:"):
-                            continue
-
-                        _, data_str = line.split("data:", 1)
-                        data_str = data_str.strip()
+                    def flush() -> None:
+                        nonlocal current_event, data_lines
+                        if not data_lines:
+                            current_event = None
+                            return
+                        data_str = "\n".join(data_lines).strip()
+                        data_lines = []
 
                         if not data_str:
-                            continue
+                            current_event = None
+                            return
 
                         try:
                             payload = json.loads(data_str)
                         except json.JSONDecodeError:
-                            continue
+                            payload = None
 
                         if isinstance(payload, dict):
+                            if current_event and "event" not in payload:
+                                payload = {**payload, "event": current_event}
                             events.append(payload)
+                        else:
+                            events.append(
+                                {
+                                    "type": "raw",
+                                    "event": current_event,
+                                    "data": data_str,
+                                }
+                            )
+                        current_event = None
+
+                    async for line in response.aiter_lines():
+                        # Blank line terminates an SSE event.
+                        if line == "":
+                            flush()
+                            if max_events is not None and len(events) >= max_events:
+                                break
+                            continue
+
+                        # Comments / heartbeats
+                        if line.startswith(":"):
+                            continue
+
+                        if line.startswith("event:"):
+                            current_event = line.split("event:", 1)[1].strip() or None
+                            continue
+
+                        if line.startswith("data:"):
+                            data_lines.append(line.split("data:", 1)[1].lstrip())
+                            continue
+
+                        continue
 
-                        if max_events is not None and len(events) >= max_events:
-                            # Enough events collected; stop consuming.
-                            break
+                    flush()
 
                 if max_seconds is not None and max_seconds > 0:
                     try:

astrbot_mcp/tools/__init__.py

@@ -18,7 +18,7 @@ All functions are re-exported from this module for convenience.
 # 导入所有工具函数，保持向后兼容
 from .control_tools import restart_astrbot
 from .log_tools import get_astrbot_logs
-from .message_tools import (
+from .message import (
     send_platform_message,
     send_platform_message_direct,
 )

astrbot_mcp/tools/message/__init__.py

@@ -0,0 +1,4 @@
+from .direct import send_platform_message_direct
+from .webchat import send_platform_message
+
+__all__ = ["send_platform_message", "send_platform_message_direct"]

astrbot_mcp/tools/message/cache.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Dict, Tuple
+
+from ...astrbot_client import AstrBotClient
+
+_SESSION_CACHE_LOCK = asyncio.Lock()
+_SESSION_CACHE: Dict[Tuple[str, str, str], str] = {}
+
+_LAST_SAVED_MESSAGE_ID_LOCK = asyncio.Lock()
+_LAST_SAVED_MESSAGE_ID_BY_SESSION: Dict[Tuple[str, str, str], str] = {}
+
+_LAST_USER_MESSAGE_ID_LOCK = asyncio.Lock()
+_LAST_USER_MESSAGE_ID_BY_SESSION: Dict[Tuple[str, str, str], str] = {}
+
+
+def _session_cache_key(client: AstrBotClient, platform_id: str) -> Tuple[str, str, str]:
+    return (client.base_url, client.settings.username or "", platform_id)
+
+
+def _last_saved_key(client: AstrBotClient, session_id: str) -> Tuple[str, str, str]:
+    return (client.base_url, client.settings.username or "", session_id)

astrbot_mcp/tools/message/direct.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List, Literal, Optional
+
+from ...astrbot_client import AstrBotClient
+from ..helpers import (
+    _as_file_uri,
+    _attachment_download_url,
+    _direct_media_mode,
+    _httpx_error_detail,
+    _resolve_local_file_path,
+)
+from ..types import MessagePart
+
+
+async def send_platform_message_direct(
+    platform_id: str,
+    target_id: str,
+    message_chain: Optional[List[MessagePart]] = None,
+    message: Optional[str] = None,
+    images: Optional[List[str]] = None,
+    files: Optional[List[str]] = None,
+    videos: Optional[List[str]] = None,
+    records: Optional[List[str]] = None,
+    message_type: Literal["GroupMessage", "FriendMessage"] = "GroupMessage",
+) -> Dict[str, Any]:
+    """
+    Directly send a message chain to a platform group/user (bypass LLM).
+
+    This calls AstrBot dashboard endpoint: POST /api/platform/send_message
+
+    Notes:
+      - This is for sending to a real platform target (group/user), not WebChat.
+      - Media parts:
+        - If `file_path` is a local path, this tool will upload it to AstrBot first, then send it as an AstrBot-hosted URL.
+        - If `file_path`/`url` is an http(s) URL, it will be forwarded as-is.
+    """
+    client = AstrBotClient.from_env()
+    onebot_like = platform_id.strip().lower() in {
+        "napcat",
+        "onebot",
+        "cqhttp",
+        "gocqhttp",
+        "llonebot",
+    }
+
+    if message_chain is None:
+        message_chain = []
+        if message:
+            message_chain.append({"type": "plain", "text": message})
+        for src in images or []:
+            message_chain.append({"type": "image", "file_path": src})
+        for src in files or []:
+            message_chain.append({"type": "file", "file_path": src})
+        for src in records or []:
+            message_chain.append({"type": "record", "file_path": src})
+        for src in videos or []:
+            message_chain.append({"type": "video", "file_path": src})
+
+    async def build_chain(mode: str) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+        normalized_chain: List[Dict[str, Any]] = []
+        uploaded_attachments: List[Dict[str, Any]] = []
+
+        for part in message_chain or []:
+            p_type = part.get("type")
+            if p_type in ("image", "file", "record", "video"):
+                file_path = part.get("file_path")
+                url = part.get("url")
+                file_name = part.get("file_name")
+                mime_type = part.get("mime_type")
+                src = url or file_path
+                if not src:
+                    continue
+
+                normalized = dict(part)
+                if not isinstance(src, str):
+                    raise ValueError(f"Invalid media source (expected str): {src!r}")
+
+                if src.startswith(("http://", "https://")):
+                    normalized["file_path"] = src
+                    if onebot_like:
+                        normalized.setdefault("file", src)
+                    normalized.pop("url", None)
+                    normalized_chain.append(normalized)
+                    continue
+
+                try:
+                    local_path = _resolve_local_file_path(client, src)
+                except ValueError as e:
+                    raise ValueError(str(e)) from e
+                except FileNotFoundError as e:
+                    raise FileNotFoundError(f"Local file_path does not exist: {src!r}") from e
+
+                if mode == "local":
+                    normalized["file_path"] = local_path
+                    normalized.pop("url", None)
+                    if onebot_like:
+                        uri = _as_file_uri(local_path)
+                        normalized.setdefault("file", uri or local_path)
+                    normalized_chain.append(normalized)
+                    continue
+
+                if mode != "upload":
+                    raise ValueError(f"Unknown direct media mode: {mode!r}")
+
+                if not file_name:
+                    file_name = os.path.basename(local_path) or None
+
+                attach_resp = await client.post_attachment_file(
+                    local_path,
+                    file_name=file_name,
+                    mime_type=mime_type,
+                )
+
+                if attach_resp.get("status") != "ok":
+                    raise RuntimeError(attach_resp.get("message") or "Attachment upload failed")
+
+                attach_data = attach_resp.get("data") or {}
+                attachment_id = attach_data.get("attachment_id")
+                if not attachment_id:
+                    raise RuntimeError(
+                        "Attachment upload succeeded but attachment_id is missing"
+                    )
+
+                download_url = _attachment_download_url(client, str(attachment_id))
+                normalized["file_path"] = download_url
+                if onebot_like:
+                    normalized.setdefault("file", download_url)
+                normalized.pop("url", None)
+                normalized.pop("file_name", None)
+                normalized.pop("mime_type", None)
+                uploaded_attachments.append(attach_data)
+                normalized_chain.append(normalized)
+            else:
+                normalized_chain.append(dict(part))
+
+        return normalized_chain, uploaded_attachments
+
+    # Prefer local paths (more compatible with Napcat / Windows), but keep an upload fallback.
+    try:
+        mode = _direct_media_mode(client)
+    except ValueError as e:
+        return {
+            "status": "error",
+            "message": str(e),
+            "platform_id": platform_id,
+            "session_id": str(target_id),
+            "message_type": message_type,
+        }
+    modes_to_try = ["local", "upload"] if mode == "auto" else [mode]
+    last_error: Dict[str, Any] | None = None
+
+    for attempt_mode in modes_to_try:
+        try:
+            normalized_chain, uploaded_attachments = await build_chain(attempt_mode)
+        except FileNotFoundError as e:
+            return {
+                "status": "error",
+                "message": str(e),
+                "platform_id": platform_id,
+                "session_id": str(target_id),
+                "message_type": message_type,
+                "hint": "If you passed a relative path, set ASTRBOTMCP_FILE_ROOT (or run the server in the correct working directory).",
+            }
+        except ValueError as e:
+            return {
+                "status": "error",
+                "message": str(e),
+                "platform_id": platform_id,
+                "session_id": str(target_id),
+                "message_type": message_type,
+                "hint": "Set ASTRBOTMCP_FILE_ROOT to control how relative paths are resolved.",
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e),
+                "platform_id": platform_id,
+                "session_id": str(target_id),
+                "message_type": message_type,
+                "attempt_mode": attempt_mode,
+            }
+
+        if not normalized_chain:
+            return {
+                "status": "error",
+                "message": "message_chain did not produce any valid message parts",
+                "platform_id": platform_id,
+                "session_id": str(target_id),
+                "message_type": message_type,
+            }
+
+        try:
+            direct_resp = await client.send_platform_message_direct(
+                platform_id=platform_id,
+                message_type=message_type,
+                session_id=str(target_id),
+                message_chain=normalized_chain,
+            )
+        except Exception as e:
+            status_code = getattr(getattr(e, "response", None), "status_code", None)
+            hint = "Ensure AstrBot includes /api/platform/send_message and you are authenticated."
+            if status_code in (404, 405):
+                hint = (
+                    "Your AstrBot may not expose /api/platform/send_message (some versions only provide "
+                    "/api/platform/stats and /api/platform/webhook). Upgrade AstrBot or add an HTTP route for sending."
+                )
+            return {
+                "status": "error",
+                "message": (
+                    f"AstrBot API error: HTTP {status_code}"
+                    if status_code is not None
+                    else f"AstrBot API error: {e}"
+                ),
+                "platform_id": platform_id,
+                "session_id": str(target_id),
+                "message_type": message_type,
+                "attempt_mode": attempt_mode,
+                "detail": _httpx_error_detail(e),
+                "hint": hint,
+            }
+
+        status = direct_resp.get("status")
+        if status == "ok":
+            data = direct_resp.get("data") or {}
+            return {
+                "status": "ok",
+                "platform_id": data.get("platform_id", platform_id),
+                "session_id": data.get("session_id", str(target_id)),
+                "message_type": data.get("message_type", message_type),
+                "attempt_mode": attempt_mode,
+                "uploaded_attachments": uploaded_attachments,
+            }
+
+        last_error = {
+            "status": status,
+            "platform_id": platform_id,
+            "session_id": str(target_id),
+            "message_type": message_type,
+            "attempt_mode": attempt_mode,
+            "message": direct_resp.get("message"),
+            "raw": direct_resp,
+        }
+
+    return last_error or {
+        "status": "error",
+        "message": "Failed to send message",
+        "platform_id": platform_id,
+        "session_id": str(target_id),
+        "message_type": message_type,
+    }

astrbot_mcp/tools/message/quote.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Tuple
+
+from ...astrbot_client import AstrBotClient
+from .utils import _extract_plain_text_from_history_item, _format_quote_block
+
+
+async def _resolve_webchat_quotes(
+    client: AstrBotClient, *, session_id: str, reply_ids: List[str]
+) -> Tuple[str, Dict[str, Any]]:
+    """
+    Resolve WebChat `message_saved.id` -> quoted text by calling /api/chat/get_session.
+    Best-effort: returns a quote prefix text and debug info.
+    """
+    cleaned: List[str] = []
+    for rid in reply_ids:
+        s = str(rid).strip()
+        if s:
+            cleaned.append(s)
+    if not cleaned:
+        return "", {"resolved": {}, "missing": []}
+
+    try:
+        resp = await client.get_platform_session(session_id=session_id)
+    except Exception as e:
+        return "", {"error": str(e), "resolved": {}, "missing": cleaned}
+
+    if resp.get("status") != "ok":
+        return "", {"status": resp.get("status"), "message": resp.get("message"), "raw": resp}
+
+    data = resp.get("data") or {}
+    history = data.get("history") or []
+    if not isinstance(history, list):
+        return "", {"resolved": {}, "missing": cleaned, "raw_history_type": str(type(history))}
+
+    index: Dict[str, Dict[str, Any]] = {}
+    for item in history:
+        if not isinstance(item, dict):
+            continue
+        mid = item.get("id")
+        if mid is None:
+            continue
+        index[str(mid)] = item
+
+    resolved: Dict[str, str] = {}
+    missing: List[str] = []
+    blocks: List[str] = []
+    for rid in cleaned:
+        item = index.get(str(rid))
+        if not item:
+            missing.append(rid)
+            blocks.append(
+                _format_quote_block(
+                    message_id=str(rid),
+                    sender="missing",
+                    text="<not found in /api/chat/get_session history>",
+                )
+            )
+            continue
+        sender = (
+            item.get("sender_name")
+            or item.get("sender_id")
+            or "unknown"
+        )
+        txt = _extract_plain_text_from_history_item(item)
+        block = _format_quote_block(message_id=str(rid), sender=str(sender), text=txt)
+        resolved[str(rid)] = block
+        blocks.append(block)
+
+    return "".join(blocks), {"resolved": resolved, "missing": missing}

astrbot_mcp/tools/message/utils.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import textwrap
+from typing import Any, Dict, List
+
+
+def _extract_plain_text_from_history_item(item: Dict[str, Any]) -> str:
+    content = item.get("content") or {}
+    if not isinstance(content, dict):
+        return str(content)
+    message = content.get("message") or []
+    if not isinstance(message, list):
+        return str(message)
+
+    chunks: List[str] = []
+    for part in message:
+        if not isinstance(part, dict):
+            continue
+        p_type = part.get("type")
+        if p_type == "plain":
+            txt = part.get("text")
+            if isinstance(txt, str) and txt:
+                chunks.append(txt)
+        elif p_type in ("image", "file", "record", "video"):
+            name = part.get("filename") or part.get("attachment_id") or ""
+            if name:
+                chunks.append(f"[{p_type}:{name}]")
+            else:
+                chunks.append(f"[{p_type}]")
+        else:
+            if p_type:
+                chunks.append(f"[{p_type}]")
+    return "".join(chunks).strip()
+
+
+def _format_quote_block(*, message_id: str, sender: str, text: str) -> str:
+    sender = (sender or "unknown").strip() or "unknown"
+    text = (text or "").strip()
+    if not text:
+        text = "<empty>"
+    text = textwrap.shorten(text, width=800, placeholder="…")
+    return f"[引用消息 {message_id} | {sender}] {text}\n"
+
+
+def _normalize_history_message_id(value: Any) -> Any:
+    """
+    AstrBot WebChat reply expects `message_id` to be the history record primary key (usually int).
+    Keep original value if it cannot be safely converted.
+    """
+    if value is None:
+        return None
+    if isinstance(value, int):
+        return value
+    s = str(value).strip()
+    if not s:
+        return value
+    if s.isdigit():
+        try:
+            return int(s)
+        except Exception:
+            return value
+    return value