312 lines
11 KiB
Python
312 lines
11 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import base64
|
|
from typing import Optional, List
|
|
|
|
NAPCT_WS_URL = 'ws://127.0.0.1:3001'
|
|
RECEIVE_TIMEOUT = 30
|
|
|
|
_ws = None
|
|
_msg_queue: List[dict] = []
|
|
|
|
|
|
def is_image_file(path: str) -> bool:
|
|
ext = os.path.splitext(path)[1].lower()
|
|
return ext in ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp')
|
|
|
|
|
|
def is_video_file(path: str) -> bool:
|
|
ext = os.path.splitext(path)[1].lower()
|
|
return ext in ('.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm')
|
|
|
|
|
|
def build_message_segments(text: str, file_paths: Optional[List[str]] = None) -> list:
|
|
"""构建 NapCat 消息段数组"""
|
|
segments = []
|
|
if text:
|
|
segments.append({"type": "text", "data": {"text": text}})
|
|
if file_paths:
|
|
for path in file_paths:
|
|
abs_path = os.path.abspath(path)
|
|
if not os.path.exists(abs_path):
|
|
raise Exception(f"文件不存在: {abs_path}")
|
|
if is_image_file(abs_path):
|
|
segments.append({"type": "image", "data": {"file": f"file://{abs_path}"}})
|
|
elif is_video_file(abs_path):
|
|
segments.append({"type": "video", "data": {"file": f"file://{abs_path}"}})
|
|
else:
|
|
segments.append({"type": "file", "data": {"file": f"file://{abs_path}"}})
|
|
if not segments:
|
|
segments.append({"type": "text", "data": {"text": ""}})
|
|
return segments
|
|
|
|
|
|
async def ensure_ws():
|
|
global _ws
|
|
if _ws is None or _ws.closed:
|
|
import websockets
|
|
_ws = await websockets.connect(NAPCT_WS_URL)
|
|
return _ws
|
|
|
|
|
|
async def ws_send_action(action: str, params: dict) -> dict:
|
|
"""发送 action 到 NapCat 并等待响应"""
|
|
ws = await ensure_ws()
|
|
echo = f"mcp_{int(time.time() * 1000)}"
|
|
payload = {"action": action, "params": params, "echo": echo}
|
|
await ws.send(json.dumps(payload))
|
|
# 等待对应的响应
|
|
timeout_ts = time.time() + 15
|
|
while time.time() < timeout_ts:
|
|
try:
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=2)
|
|
resp = json.loads(raw)
|
|
if resp.get("echo") == echo:
|
|
return resp
|
|
# 不是我们的响应，放入消息队列
|
|
_msg_queue.append(resp)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
raise Exception("发送消息超时")
|
|
|
|
|
|
async def send_private_message(user_id: int, text: str = "", file_paths: Optional[List[str]] = None) -> dict:
|
|
"""发送私聊消息，可附带图片/视频"""
|
|
segments = build_message_segments(text, file_paths)
|
|
return await ws_send_action("send_private_msg", {
|
|
"user_id": user_id,
|
|
"message": segments
|
|
})
|
|
|
|
|
|
async def send_group_message(group_id: int, text: str = "", file_paths: Optional[List[str]] = None) -> dict:
|
|
"""发送群聊消息，可附带图片/视频"""
|
|
segments = build_message_segments(text, file_paths)
|
|
return await ws_send_action("send_group_msg", {
|
|
"group_id": group_id,
|
|
"message": segments
|
|
})
|
|
|
|
|
|
async def receive_messages(count: int = 1, user_id: Optional[int] = None,
|
|
group_id: Optional[int] = None, timeout: int = RECEIVE_TIMEOUT) -> list:
|
|
"""接收 N 条消息，可按 user_id 或 group_id 过滤"""
|
|
ws = await ensure_ws()
|
|
results = []
|
|
deadline = time.time() + timeout
|
|
while len(results) < count and time.time() < deadline:
|
|
remain = deadline - time.time()
|
|
if remain <= 0:
|
|
break
|
|
try:
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=min(remain, 2))
|
|
msg = json.loads(raw)
|
|
if msg.get("post_type") != "message":
|
|
continue
|
|
msg_user_id = msg.get("user_id")
|
|
msg_group_id = msg.get("group_id") if msg.get("message_type") == "group" else None
|
|
if user_id is not None and msg_user_id != user_id:
|
|
continue
|
|
if group_id is not None and msg_group_id != group_id:
|
|
continue
|
|
|
|
# 提取纯文本和媒体信息
|
|
segments = msg.get("message", [])
|
|
text_parts = []
|
|
media_info = []
|
|
for seg in segments:
|
|
if seg.get("type") == "text":
|
|
text_parts.append(seg.get("data", {}).get("text", ""))
|
|
elif seg.get("type") in ("image", "video"):
|
|
media_info.append({
|
|
"type": seg.get("type"),
|
|
"url": seg.get("data", {}).get("url", ""),
|
|
"file": seg.get("data", {}).get("file", "")
|
|
})
|
|
|
|
results.append({
|
|
"message_id": msg.get("message_id"),
|
|
"user_id": msg_user_id,
|
|
"group_id": msg_group_id,
|
|
"message_type": msg.get("message_type"),
|
|
"text": "".join(text_parts).strip(),
|
|
"media": media_info,
|
|
"raw": msg
|
|
})
|
|
_msg_queue.clear()
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
return results
|
|
|
|
|
|
# ============================================================
|
|
# 同步包装器（供 stdio MCP 调用）
|
|
# ============================================================
|
|
|
|
_loop = None
|
|
|
|
|
|
def get_loop():
|
|
global _loop
|
|
if _loop is None or _loop.is_closed():
|
|
_loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(_loop)
|
|
return _loop
|
|
|
|
|
|
def sync_send_private_message(user_id: int, text: str = "", file_paths: list = None):
|
|
loop = get_loop()
|
|
result = loop.run_until_complete(send_private_message(user_id, text, file_paths or []))
|
|
loop.run_until_complete(_ws.close())
|
|
return result
|
|
|
|
|
|
def sync_send_group_message(group_id: int, text: str = "", file_paths: list = None):
|
|
loop = get_loop()
|
|
result = loop.run_until_complete(send_group_message(group_id, text, file_paths or []))
|
|
loop.run_until_complete(_ws.close())
|
|
return result
|
|
|
|
|
|
def sync_receive_messages(count: int = 1, user_id: int = None,
|
|
group_id: int = None, timeout: int = RECEIVE_TIMEOUT):
|
|
loop = get_loop()
|
|
result = loop.run_until_complete(receive_messages(count, user_id, group_id, timeout))
|
|
loop.run_until_complete(_ws.close())
|
|
return result
|
|
|
|
|
|
# ============================================================
|
|
# MCP JSON-RPC 处理
|
|
# ============================================================
|
|
|
|
TOOLS = [
|
|
{
|
|
"name": "send_private_message",
|
|
"description": "向指定 QQ 用户发送私聊消息，可附带任意类型的文件（图片/视频/文档等）。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"user_id": {"type": "number", "description": "目标 QQ 号"},
|
|
"text": {"type": "string", "description": "要发送的文字内容（可为空，但需要至少有文字或媒体文件之一）"},
|
|
"file_paths": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "文件的本地绝对路径列表，图片用 image 类型，视频用 video 类型，其余后缀一律用 file 类型发送"
|
|
}
|
|
},
|
|
"required": ["user_id"]
|
|
}
|
|
},
|
|
{
|
|
"name": "send_group_message",
|
|
"description": "向指定 QQ 群发送消息，可附带任意类型的文件（图片/视频/文档等）。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"group_id": {"type": "number", "description": "目标群号"},
|
|
"text": {"type": "string", "description": "要发送的文字内容（可为空，但需要至少有文字或媒体文件之一）"},
|
|
"file_paths": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "文件的本地绝对路径列表，图片用 image 类型，视频用 video 类型，其余后缀一律用 file 类型发送"
|
|
}
|
|
},
|
|
"required": ["group_id"]
|
|
}
|
|
},
|
|
{
|
|
"name": "receive_messages",
|
|
"description": "接收来自 NapCat 的消息，可按 QQ 号或群号过滤，接收指定条数后返回（有超时限制）。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"count": {"type": "number", "description": "要接收的消息条数，默认 1"},
|
|
"user_id": {"type": "number", "description": "按 QQ 号过滤（可选，不传则不过滤）"},
|
|
"group_id": {"type": "number", "description": "按群号过滤（可选，不传则不过滤）"},
|
|
"timeout": {"type": "number", "description": "超时秒数，默认 30"}
|
|
},
|
|
"required": []
|
|
}
|
|
}
|
|
]
|
|
|
|
TOOL_MAP = {
|
|
"send_private_message": sync_send_private_message,
|
|
"send_group_message": sync_send_group_message,
|
|
"receive_messages": sync_receive_messages
|
|
}
|
|
|
|
|
|
def handle_request(req):
|
|
method = req.get("method", "")
|
|
req_id = req.get("id")
|
|
|
|
if method == "initialize":
|
|
return {
|
|
"jsonrpc": "2.0",
|
|
"id": req_id,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {
|
|
"tools": {}
|
|
},
|
|
"serverInfo": {
|
|
"name": "ros-qq-messages",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
}
|
|
|
|
if method == "ping":
|
|
return {"jsonrpc": "2.0", "id": req_id, "result": {}}
|
|
|
|
if method == "notifications/initialized":
|
|
return None
|
|
|
|
if method == "tools/list":
|
|
return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}
|
|
|
|
if method == "tools/call":
|
|
tool_name = req["params"]["name"]
|
|
arguments = req["params"].get("arguments", {})
|
|
tool_func = TOOL_MAP.get(tool_name)
|
|
if not tool_func:
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"error": {"code": -32601, "message": f"未知工具: {tool_name}"}}
|
|
try:
|
|
result = tool_func(**arguments)
|
|
text = json.dumps(result, ensure_ascii=False, indent=2)
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"result": {"content": [{"type": "text", "text": text}]}}
|
|
except Exception as e:
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"error": {"code": -32000, "message": str(e)}}
|
|
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"error": {"code": -32601, "message": f"未知方法: {method}"}}
|
|
|
|
|
|
def main():
|
|
for line in sys.stdin:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
req = json.loads(line)
|
|
resp = handle_request(req)
|
|
if resp is None:
|
|
continue
|
|
sys.stdout.write(json.dumps(resp, ensure_ascii=False) + '\n')
|
|
sys.stdout.flush()
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|