ForcomeBot 2.2.4__py3-none-any.whl

src/core/message_queue.py

@@ -0,0 +1,364 @@
+"""统一消息发送队列
+
+功能：
+- 所有消息发送请求统一排队
+- 随机延迟发送，避免并发过高
+- 支持优先级（高优先级插队）
+- 失败重试机制
+- 队列状态监控
+"""
+import asyncio
+import logging
+import random
+import time
+from dataclasses import dataclass, field
+from enum import IntEnum
+from typing import Any, Callable, Awaitable, Optional, Dict, List
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+
+class MessagePriority(IntEnum):
+    """消息优先级（数字越小优先级越高）"""
+    HIGH = 1      # 普通回复消息（用户触发）
+    NORMAL = 5    # 定时任务消息
+    LOW = 10      # 群发消息
+
+
+@dataclass(order=True)
+class QueuedMessage:
+    """队列中的消息"""
+    priority: int
+    timestamp: float = field(compare=False)
+    message_id: str = field(compare=False)
+    send_func: Callable[[], Awaitable[bool]] = field(compare=False)
+    message_type: str = field(compare=False, default="unknown")
+    target: str = field(compare=False, default="")
+    content_preview: str = field(compare=False, default="")
+    retry_count: int = field(compare=False, default=0)
+    max_retries: int = field(compare=False, default=2)
+    callback: Optional[Callable[[bool, str], Awaitable[None]]] = field(compare=False, default=None)
+
+
+class MessageQueue:
+    """统一消息发送队列"""
+    
+    def __init__(
+        self,
+        min_interval: float = 1.0,
+        max_interval: float = 3.0,
+        batch_min_interval: float = 2.0,
+        batch_max_interval: float = 5.0,
+        max_queue_size: int = 1000
+    ):
+        """初始化消息队列
+        
+        Args:
+            min_interval: 高优先级消息最小间隔（秒）
+            max_interval: 高优先级消息最大间隔（秒）
+            batch_min_interval: 低优先级消息最小间隔（秒）
+            batch_max_interval: 低优先级消息最大间隔（秒）
+            max_queue_size: 队列最大容量
+        """
+        self._queue: asyncio.PriorityQueue = asyncio.PriorityQueue(maxsize=max_queue_size)
+        self._min_interval = min_interval
+        self._max_interval = max_interval
+        self._batch_min_interval = batch_min_interval
+        self._batch_max_interval = batch_max_interval
+        self._max_queue_size = max_queue_size
+        
+        self._running = False
+        self._worker_task: Optional[asyncio.Task] = None
+        self._last_send_time: float = 0
+        self._message_counter: int = 0
+        
+        # 统计信息
+        self._stats = {
+            "total_sent": 0,
+            "total_failed": 0,
+            "total_retried": 0,
+            "started_at": None
+        }
+        
+        # 最近发送记录（用于监控）
+        self._recent_sends: List[Dict[str, Any]] = []
+        self._max_recent = 50
+    
+    def update_config(
+        self,
+        min_interval: float = None,
+        max_interval: float = None,
+        batch_min_interval: float = None,
+        batch_max_interval: float = None
+    ):
+        """更新配置"""
+        if min_interval is not None:
+            self._min_interval = min_interval
+        if max_interval is not None:
+            self._max_interval = max_interval
+        if batch_min_interval is not None:
+            self._batch_min_interval = batch_min_interval
+        if batch_max_interval is not None:
+            self._batch_max_interval = batch_max_interval
+        logger.info(f"消息队列配置已更新: interval={self._min_interval}-{self._max_interval}s, "
+                   f"batch={self._batch_min_interval}-{self._batch_max_interval}s")
+    
+    async def start(self):
+        """启动队列处理"""
+        if self._running:
+            return
+        
+        self._running = True
+        self._stats["started_at"] = datetime.now().isoformat()
+        self._worker_task = asyncio.create_task(self._worker())
+        logger.info("消息队列已启动")
+    
+    async def stop(self):
+        """停止队列处理"""
+        self._running = False
+        if self._worker_task:
+            self._worker_task.cancel()
+            try:
+                await self._worker_task
+            except asyncio.CancelledError:
+                pass
+        logger.info(f"消息队列已停止，统计: 发送={self._stats['total_sent']}, "
+                   f"失败={self._stats['total_failed']}, 重试={self._stats['total_retried']}")
+    
+    def _generate_message_id(self) -> str:
+        """生成消息ID"""
+        self._message_counter += 1
+        return f"msg_{int(time.time())}_{self._message_counter}"
+    
+    async def enqueue(
+        self,
+        send_func: Callable[[], Awaitable[bool]],
+        priority: MessagePriority = MessagePriority.NORMAL,
+        message_type: str = "unknown",
+        target: str = "",
+        content_preview: str = "",
+        max_retries: int = 2,
+        callback: Optional[Callable[[bool, str], Awaitable[None]]] = None
+    ) -> str:
+        """添加消息到队列
+        
+        Args:
+            send_func: 发送函数，返回是否成功
+            priority: 消息优先级
+            message_type: 消息类型（text, image, file 等）
+            target: 目标wxid
+            content_preview: 内容预览（用于日志）
+            max_retries: 最大重试次数
+            callback: 发送完成回调，参数为 (success, message_id)
+            
+        Returns:
+            消息ID
+        """
+        message_id = self._generate_message_id()
+        
+        msg = QueuedMessage(
+            priority=priority,
+            timestamp=time.time(),
+            message_id=message_id,
+            send_func=send_func,
+            message_type=message_type,
+            target=target,
+            content_preview=content_preview[:50] if content_preview else "",
+            max_retries=max_retries,
+            callback=callback
+        )
+        
+        try:
+            self._queue.put_nowait(msg)
+            logger.debug(f"消息入队: {message_id}, 优先级={priority}, 类型={message_type}, "
+                        f"目标={target}, 队列长度={self._queue.qsize()}")
+            return message_id
+        except asyncio.QueueFull:
+            logger.error(f"消息队列已满，丢弃消息: {message_id}")
+            if callback:
+                await callback(False, message_id)
+            return message_id
+    
+    async def enqueue_text(
+        self,
+        qianxun_client,
+        robot_wxid: str,
+        target: str,
+        message: str,
+        priority: MessagePriority = MessagePriority.NORMAL,
+        callback: Optional[Callable[[bool, str], Awaitable[None]]] = None
+    ) -> str:
+        """便捷方法：添加文本消息到队列"""
+        async def send():
+            return await qianxun_client.send_text(robot_wxid, target, message)
+        
+        return await self.enqueue(
+            send_func=send,
+            priority=priority,
+            message_type="text",
+            target=target,
+            content_preview=message,
+            callback=callback
+        )
+    
+    async def enqueue_image(
+        self,
+        qianxun_client,
+        robot_wxid: str,
+        target: str,
+        image_path: str,
+        file_name: str = "",
+        priority: MessagePriority = MessagePriority.NORMAL,
+        callback: Optional[Callable[[bool, str], Awaitable[None]]] = None
+    ) -> str:
+        """便捷方法：添加图片消息到队列"""
+        async def send():
+            return await qianxun_client.send_image(robot_wxid, target, image_path, file_name)
+        
+        return await self.enqueue(
+            send_func=send,
+            priority=priority,
+            message_type="image",
+            target=target,
+            content_preview=image_path,
+            callback=callback
+        )
+    
+    async def _worker(self):
+        """队列处理工作线程"""
+        logger.info("消息队列工作线程已启动")
+        
+        while self._running:
+            try:
+                # 等待消息，超时1秒检查运行状态
+                try:
+                    msg: QueuedMessage = await asyncio.wait_for(
+                        self._queue.get(), timeout=1.0
+                    )
+                except asyncio.TimeoutError:
+                    continue
+                
+                # 计算延迟
+                delay = self._calculate_delay(msg.priority)
+                
+                # 等待延迟时间
+                elapsed = time.time() - self._last_send_time
+                if elapsed < delay:
+                    wait_time = delay - elapsed
+                    logger.debug(f"等待 {wait_time:.2f}s 后发送: {msg.message_id}")
+                    await asyncio.sleep(wait_time)
+                
+                # 发送消息
+                success = await self._send_message(msg)
+                
+                # 更新最后发送时间
+                self._last_send_time = time.time()
+                
+                # 标记任务完成
+                self._queue.task_done()
+                
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"消息队列工作线程异常: {e}", exc_info=True)
+                await asyncio.sleep(1)
+        
+        logger.info("消息队列工作线程已停止")
+    
+    def _calculate_delay(self, priority: int) -> float:
+        """根据优先级计算延迟时间"""
+        if priority <= MessagePriority.HIGH:
+            # 高优先级：较短延迟
+            return random.uniform(self._min_interval, self._max_interval)
+        else:
+            # 普通/低优先级：较长延迟
+            return random.uniform(self._batch_min_interval, self._batch_max_interval)
+    
+    async def _send_message(self, msg: QueuedMessage) -> bool:
+        """发送单条消息"""
+        try:
+            success = await msg.send_func()
+            
+            # 记录发送结果
+            record = {
+                "message_id": msg.message_id,
+                "type": msg.message_type,
+                "target": msg.target,
+                "priority": msg.priority,
+                "success": success,
+                "retry_count": msg.retry_count,
+                "sent_at": datetime.now().isoformat()
+            }
+            self._recent_sends.append(record)
+            if len(self._recent_sends) > self._max_recent:
+                self._recent_sends.pop(0)
+            
+            if success:
+                self._stats["total_sent"] += 1
+                logger.info(f"消息发送成功: {msg.message_id} -> {msg.target}")
+                
+                if msg.callback:
+                    await msg.callback(True, msg.message_id)
+                return True
+            else:
+                # 发送失败，尝试重试
+                return await self._handle_failure(msg, "发送返回失败")
+                
+        except Exception as e:
+            logger.error(f"消息发送异常: {msg.message_id}, {e}")
+            return await self._handle_failure(msg, str(e))
+    
+    async def _handle_failure(self, msg: QueuedMessage, error: str) -> bool:
+        """处理发送失败"""
+        if msg.retry_count < msg.max_retries:
+            # 重试
+            msg.retry_count += 1
+            self._stats["total_retried"] += 1
+            logger.warning(f"消息发送失败，重试 {msg.retry_count}/{msg.max_retries}: "
+                          f"{msg.message_id}, 错误: {error}")
+            
+            # 重新入队（降低优先级避免阻塞其他消息）
+            msg.priority = max(msg.priority, MessagePriority.LOW)
+            try:
+                self._queue.put_nowait(msg)
+            except asyncio.QueueFull:
+                logger.error(f"重试入队失败，队列已满: {msg.message_id}")
+                self._stats["total_failed"] += 1
+                if msg.callback:
+                    await msg.callback(False, msg.message_id)
+            return False
+        else:
+            # 超过重试次数
+            self._stats["total_failed"] += 1
+            logger.error(f"消息发送最终失败: {msg.message_id} -> {msg.target}, 错误: {error}")
+            
+            if msg.callback:
+                await msg.callback(False, msg.message_id)
+            return False
+    
+    def get_status(self) -> Dict[str, Any]:
+        """获取队列状态"""
+        return {
+            "running": self._running,
+            "queue_size": self._queue.qsize(),
+            "max_queue_size": self._max_queue_size,
+            "config": {
+                "min_interval": self._min_interval,
+                "max_interval": self._max_interval,
+                "batch_min_interval": self._batch_min_interval,
+                "batch_max_interval": self._batch_max_interval
+            },
+            "stats": self._stats.copy(),
+            "recent_sends": self._recent_sends[-10:]  # 最近10条
+        }
+    
+    @property
+    def queue_size(self) -> int:
+        """当前队列长度"""
+        return self._queue.qsize()
+    
+    @property
+    def is_running(self) -> bool:
+        """是否运行中"""
+        return self._running

src/core/state_store.py

@@ -0,0 +1,242 @@
+"""状态存储器 - 管理内存状态和持久化"""
+import asyncio
+import json
+import logging
+import time
+from collections import OrderedDict
+from pathlib import Path
+from typing import Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class StateStore:
+    """状态存储器"""
+
+    def __init__(self, data_dir: str = "data"):
+        self._data_dir = Path(data_dir)
+        self._data_dir.mkdir(parents=True, exist_ok=True)
+
+        # 数据存储
+        self._processed_msg_ids: OrderedDict[str, float] = OrderedDict()
+        self._image_cache: OrderedDict[tuple, Dict[str, Any]] = OrderedDict()
+        self._group_id_mapping: OrderedDict[str, str] = OrderedDict()
+        self._message_user_mapping: OrderedDict[str, str] = OrderedDict()
+
+        # 清理任务
+        self._cleanup_task: Optional[asyncio.Task] = None
+        self._running = False
+
+        # 配置限制
+        self.max_msg_ids = 10000
+        self.max_image_cache = 1000
+        self.max_group_mapping = 10000
+        self.max_message_mapping = 10000
+        self.msg_id_ttl = 60
+        self.image_cache_ttl = 120
+
+        # 持久化
+        self._state_file = self._data_dir / "state.json"
+        self._persist_lock = asyncio.Lock()
+        self._dirty = False
+
+    def _enforce_limit(self, collection: OrderedDict, max_size: int):
+        """通用的集合大小限制执行"""
+        while len(collection) > max_size:
+            collection.popitem(last=False)
+
+    def _set_mapping(self, collection: OrderedDict, key: str, value: Any, max_size: int):
+        """通用的映射设置方法"""
+        collection[key] = value
+        self._dirty = True
+        collection.move_to_end(key)
+        self._enforce_limit(collection, max_size)
+
+    async def start(self):
+        """启动状态存储器"""
+        logger.info("启动状态存储器...")
+        await self._load_persisted()
+        self._running = True
+        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+        logger.info("状态存储器已启动")
+
+    async def stop(self):
+        """停止状态存储器"""
+        logger.info("停止状态存储器...")
+        self._running = False
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+        await self._persist()
+        logger.info("状态存储器已停止")
+
+    def is_duplicate_message(self, msg_id: str) -> bool:
+        """检查消息是否重复"""
+        if msg_id in self._processed_msg_ids:
+            return True
+        self._processed_msg_ids[msg_id] = time.time()
+        self._dirty = True
+        self._enforce_limit(self._processed_msg_ids, self.max_msg_ids)
+        return False
+
+    def cache_image(self, group_id: str, user_id: str, image_url: str):
+        """缓存图片"""
+        cache_key = (group_id, user_id)
+        self._image_cache[cache_key] = {"image_url": image_url, "timestamp": time.time()}
+        self._dirty = True
+        self._image_cache.move_to_end(cache_key)
+        self._enforce_limit(self._image_cache, self.max_image_cache)
+        logger.debug(f"缓存图片: [{group_id}][{user_id}] -> {image_url}")
+
+    def get_cached_image(self, group_id: str, user_id: str) -> Optional[str]:
+        """获取缓存的图片"""
+        cache_key = (group_id, user_id)
+        cached = self._image_cache.get(cache_key)
+
+        if cached:
+            age = time.time() - cached["timestamp"]
+            if age < self.image_cache_ttl:
+                del self._image_cache[cache_key]
+                self._dirty = True
+                logger.info(f"命中图片缓存: [{group_id}][{user_id}], 年龄: {age:.1f}秒")
+                return cached["image_url"]
+            else:
+                logger.debug(f"图片缓存已过期: {age:.1f}秒 > {self.image_cache_ttl}秒")
+                del self._image_cache[cache_key]
+                self._dirty = True
+        return None
+
+    def set_group_mapping(self, numeric_id: str, original_id: str):
+        """设置群ID映射"""
+        self._set_mapping(self._group_id_mapping, numeric_id, original_id, self.max_group_mapping)
+
+    def get_original_group_id(self, numeric_id: str) -> str:
+        """获取原始群ID"""
+        return self._group_id_mapping.get(numeric_id, numeric_id)
+
+    def set_message_user(self, message_id: str, user_id: str):
+        """设置消息-用户映射"""
+        self._set_mapping(self._message_user_mapping, message_id, user_id, self.max_message_mapping)
+
+    def get_user_by_message(self, message_id: str) -> Optional[str]:
+        """根据消息ID获取用户ID"""
+        return self._message_user_mapping.get(message_id)
+
+    async def _cleanup_loop(self):
+        """定期清理过期数据（每60秒）"""
+        while self._running:
+            try:
+                await asyncio.sleep(60)
+                if not self._running:
+                    break
+                self._cleanup_expired()
+                if self._dirty:
+                    await self._persist()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"清理任务异常: {e}", exc_info=True)
+
+    def _cleanup_expired(self):
+        """清理过期数据"""
+        now = time.time()
+
+        # 清理过期的消息ID
+        expired_msg_ids = [k for k, v in self._processed_msg_ids.items() if now - v > self.msg_id_ttl]
+        for k in expired_msg_ids:
+            del self._processed_msg_ids[k]
+
+        # 清理过期的图片缓存
+        expired_images = [k for k, v in self._image_cache.items() if now - v["timestamp"] > self.image_cache_ttl]
+        for k in expired_images:
+            del self._image_cache[k]
+
+        if expired_msg_ids or expired_images:
+            self._dirty = True
+            logger.debug(f"清理过期数据: 消息ID={len(expired_msg_ids)}, 图片缓存={len(expired_images)}")
+
+    async def _persist(self):
+        """异步持久化数据"""
+        async with self._persist_lock:
+            try:
+                data = {
+                    "processed_msg_ids": dict(self._processed_msg_ids),
+                    "group_id_mapping": dict(self._group_id_mapping),
+                    "message_user_mapping": dict(self._message_user_mapping),
+                }
+
+                def write_file():
+                    with open(self._state_file, 'w', encoding='utf-8') as f:
+                        json.dump(data, f, ensure_ascii=False, indent=2)
+
+                await asyncio.get_event_loop().run_in_executor(None, write_file)
+                self._dirty = False
+                logger.debug(f"状态已持久化到 {self._state_file}")
+            except Exception as e:
+                logger.error(f"持久化状态失败: {e}", exc_info=True)
+
+    async def _load_persisted(self):
+        """加载持久化数据"""
+        if not self._state_file.exists():
+            logger.info("没有找到持久化文件，使用空状态启动")
+            return
+
+        try:
+            def read_file():
+                with open(self._state_file, 'r', encoding='utf-8') as f:
+                    return json.load(f)
+
+            data = await asyncio.get_event_loop().run_in_executor(None, read_file)
+            now = time.time()
+
+            # 恢复消息ID记录（过滤过期的）
+            if "processed_msg_ids" in data:
+                self._processed_msg_ids = OrderedDict(
+                    (k, v) for k, v in data["processed_msg_ids"].items()
+                    if now - v <= self.msg_id_ttl
+                )
+
+            # 恢复群ID映射
+            if "group_id_mapping" in data:
+                self._group_id_mapping = OrderedDict(data["group_id_mapping"])
+
+            # 恢复消息-用户映射
+            if "message_user_mapping" in data:
+                self._message_user_mapping = OrderedDict(data["message_user_mapping"])
+
+            logger.info(f"已从 {self._state_file} 恢复状态: "
+                       f"消息ID={len(self._processed_msg_ids)}, "
+                       f"群映射={len(self._group_id_mapping)}, "
+                       f"消息映射={len(self._message_user_mapping)}")
+
+        except json.JSONDecodeError as e:
+            logger.warning(f"持久化文件格式错误，使用空状态启动: {e}")
+        except Exception as e:
+            logger.error(f"加载持久化数据失败: {e}", exc_info=True)
+
+    def get_stats(self) -> Dict[str, int]:
+        """获取状态统计信息"""
+        return {
+            "processed_msg_ids": len(self._processed_msg_ids),
+            "image_cache": len(self._image_cache),
+            "group_id_mapping": len(self._group_id_mapping),
+            "message_user_mapping": len(self._message_user_mapping),
+        }
+
+    def clear_image_cache(self):
+        """清空图片缓存"""
+        self._image_cache.clear()
+        self._dirty = True
+        logger.info("图片缓存已清空")
+
+    def clear_all(self):
+        """清空所有缓存"""
+        self._processed_msg_ids.clear()
+        self._image_cache.clear()
+        self._group_id_mapping.clear()
+        self._message_user_mapping.clear()
+        self._dirty = True
+        logger.info("所有缓存已清空")

src/handlers/__init__.py

@@ -0,0 +1,8 @@
+# Handlers layer
+"""Business logic handlers for message processing and task scheduling."""
+
+from .message_parser import MessageParser, ParsedMessage
+from .message_handler import MessageHandler
+from .scheduler import TaskScheduler
+
+__all__ = ["MessageParser", "ParsedMessage", "MessageHandler", "TaskScheduler"]