audio_codebuddy/mcp-for-xiaozhi/self_test.py

"""
self_test.py — 不用 subprocess，直接用 stdio 当 MCP server 跟小智 AI 通信

xiaozhi.me 协议：
  1. xiaozhi 主动发 initialize request (id=0, method=initialize)
  2. 我们回 initialize response
  3. xiaozhi 发 notifications/initialized
  4. xiaozhi 发 tools/list, tools/call, ping
"""
import json
import logging
import os
import sys

import websockets.sync.client as ws_sync

logging.basicConfig(level="INFO", format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("self_test")

API_BASE = os.environ.get("AUDIOBOOK_API_BASE", "http://127.0.0.1:38381")
endpoint = os.environ["MCP_ENDPOINT"]

TOOLS = [
    {"name": "search_audiobooks", "description": "搜索有声书",
     "inputSchema": {"type": "object", "properties": {"keywords": {"type": "string"}}, "required": ["keywords"]}},
    {"name": "list_categories", "description": "分类列表",
     "inputSchema": {"type": "object", "properties": {}}},
    {"name": "list_books_by_category", "description": "分类下书籍",
     "inputSchema": {"type": "object", "properties": {"category_id": {"type": "integer"}, "page": {"type": "integer"}}, "required": ["category_id"]}},
    {"name": "get_book_details", "description": "书籍详情",
     "inputSchema": {"type": "object", "properties": {"book_id": {"type": "integer"}}, "required": ["book_id"]}},
    {"name": "get_chapter_audio_url", "description": "章节音频URL",
     "inputSchema": {"type": "object", "properties": {"chapter_id": {"type": "integer"}}, "required": ["chapter_id"]}},
    {"name": "get_chapter_text", "description": "章节文本",
     "inputSchema": {"type": "object", "properties": {"chapter_id": {"type": "integer"}, "max_chars": {"type": "integer"}}, "required": ["chapter_id"]}},
    {"name": "list_recent_history", "description": "历史",
     "inputSchema": {"type": "object", "properties": {"limit": {"type": "integer"}}}},
    {"name": "generate_audiobook", "description": "生成书",
     "inputSchema": {"type": "object", "properties": {"title": {"type": "string"}, "description": {"type": "string"}}, "required": ["title"]}},
]


def http_get(path: str):
    import urllib.request
    with urllib.request.urlopen(f"{API_BASE}{path}", timeout=10) as r:
        return json.loads(r.read().decode("utf-8"))


def handle_request(req: dict) -> dict:
    """处理 xiaozhi 发来的 MCP 请求，返回 response（notifications 返回 None）"""
    method = req.get("method")
    rid = req.get("id")
    params = req.get("params") or {}

    if method == "initialize":
        return {"jsonrpc": "2.0", "id": rid, "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {"listChanged": False}},
            "serverInfo": {"name": "AudioBook", "version": "1.0.0"},
        }}
    if method == "ping":
        return {"jsonrpc": "2.0", "id": rid, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": rid, "result": {"tools": TOOLS}}
    if method == "tools/call":
        name = params.get("name")
        args = params.get("arguments") or {}
        if name == "list_categories":
            data = http_get("/api/categories")
            items = data.get("data") or []
            result = {"count": len(items), "categories": [{"id": i.get("id"), "name": i.get("name") or i.get("title")} for i in items]}
            return {"jsonrpc": "2.0", "id": rid, "result": {
                "content": [{"type": "text", "text": json.dumps({"success": True, "result": json.dumps(result, ensure_ascii=False)}, ensure_ascii=False)}],
                "isError": False,
            }}
        if name == "get_chapter_audio_url":
            cid = args.get("chapter_id")
            data = http_get(f"/api/book-generator/books/chapters/{cid}")
            info = data.get("data") or {}
            audio_url = info.get("audioUrl") or ""
            text = json.dumps({"success": bool(audio_url and not audio_url.startswith("/uploads/")),
                              "chapter_id": cid, "title": info.get("title"),
                              "audio_url": audio_url, "book_id": info.get("bookId")}, ensure_ascii=False)
            return {"jsonrpc": "2.0", "id": rid, "result": {
                "content": [{"type": "text", "text": text}], "isError": False,
            }}
        if name == "get_book_details":
            bid = args.get("book_id")
            data = http_get(f"/api/book-generator/books/{bid}")
            info = data.get("data") or {}
            chapters = info.get("chapters") or []
            result = {"id": info.get("id"), "title": info.get("title"),
                     "chapters": [{"id": c.get("id"), "title": c.get("title"), "has_audio": bool(c.get("audioUrl"))} for c in chapters[:30]]}
            return {"jsonrpc": "2.0", "id": rid, "result": {
                "content": [{"type": "text", "text": json.dumps({"success": True, "result": json.dumps(result, ensure_ascii=False)}, ensure_ascii=False)}],
                "isError": False,
            }}
        if name == "search_audiobooks":
            kw = args.get("keywords", "")
            import urllib.parse
            with __import__("urllib.request").request.urlopen(f"{API_BASE}/api/search?q={urllib.parse.quote(kw)}&limit=8", timeout=10) as r:
                data = json.loads(r.read().decode("utf-8"))
            items = (data.get("data") or [])[:8]
            result = {"count": len(items), "items": [{"id": it.get("id") or it.get("bookId"), "title": it.get("title")} for it in items]}
            return {"jsonrpc": "2.0", "id": rid, "result": {
                "content": [{"type": "text", "text": json.dumps({"success": True, "result": json.dumps(result, ensure_ascii=False)}, ensure_ascii=False)}],
                "isError": False,
            }}
        return {"jsonrpc": "2.0", "id": rid, "error": {"code": -32602, "message": f"unknown tool: {name}"}}
    return {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": f"unknown method: {method}"}}


def main():
    logger.info("=== Self-test ===")
    ws = ws_sync.connect(endpoint, max_size=4 * 1024 * 1024)
    logger.info("WS connected, waiting for xiaozhi initialize...")

    while True:
        try:
            raw = ws.recv(timeout=60)
        except TimeoutError:
            logger.info("WS timeout 60s, exiting")
            break
        if not raw:
            break
        try:
            req = json.loads(raw)
        except Exception as e:
            logger.warning("parse err: %s", e)
            continue

        method = req.get("method")
        rid = req.get("id")
        logger.info("<- %s id=%s", method, rid)

        if rid is None and method and method.startswith("notifications/"):
            # notification, don't respond
            continue

        resp = handle_request(req)
        line = json.dumps(resp, ensure_ascii=False, separators=(",", ":")) + "\n"
        logger.info("-> response for %s id=%s, %d bytes", method, rid, len(line))
        ws.send(line)

    ws.close()


if __name__ == "__main__":
    main()