myagent-ai 1.0.0

package/core/config_broadcast.py

@@ -0,0 +1,276 @@
+"""
+core/config_broadcast.py - 配置热加载广播机制
+===============================================
+当配置发生变更（如添加/删除模型、修改LLM参数等）时，通知所有正在执行任务的 Agent，
+使它们能够安全地暂停、等待配置更新完成、然后恢复继续执行任务。
+
+核心设计:
+  - 基于 asyncio.Event 实现无轮询的高效等待
+  - Agent 在每次迭代循环中检查是否需要暂停
+  - 暂停时 Agent 保存当前上下文状态（checkpoint），等待配置更新
+  - 配置更新完成后 Agent 从 checkpoint 恢复继续执行
+"""
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Dict, Any, Optional, Set
+from dataclasses import dataclass, field
+
+from core.logger import get_logger
+
+logger = get_logger("myagent.broadcast")
+
+
+@dataclass
+class TaskCheckpoint:
+    """任务检查点 - 保存 Agent 暂停时的状态"""
+    task_id: str = ""
+    agent_name: str = ""
+    iteration: int = 0
+    paused_at: str = ""       # 暂停时间戳
+    resumed_at: str = ""      # 恢复时间戳
+    reload_version: int = 0   # 配置版本号（用于判断是否需要重新加载）
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class ConfigBroadcaster:
+    """
+    配置变更广播器。
+
+    使用方式:
+        # 初始化
+        broadcaster = ConfigBroadcaster()
+
+        # Agent 注册（开始任务时）
+        broadcaster.register(task_id="task_xxx", agent_name="main_agent")
+
+        # Agent 在循环中检查（每次迭代前）
+        if await broadcaster.check_and_wait(task_id="task_xxx"):
+            logger.info("配置已更新，继续执行")
+
+        # 配置变更时（热重载入口）
+        await broadcaster.request_reload()
+        # ... 更新配置 ...
+        await broadcaster.complete_reload()
+
+        # Agent 注销（任务完成时）
+        broadcaster.unregister(task_id="task_xxx")
+    """
+
+    def __init__(self):
+        # 事件: reload_requested -> reload_done
+        self._reload_requested: asyncio.Event = asyncio.Event()
+        self._reload_done: asyncio.Event = asyncio.Event()
+        self._lock: asyncio.Lock = asyncio.Lock()
+
+        # 活跃任务注册表: task_id -> TaskCheckpoint
+        self._active_tasks: Dict[str, TaskCheckpoint] = {}
+        # 已暂停的任务集合（用于等待所有任务暂停）
+        self._paused_tasks: Set[str] = set()
+
+        # 配置版本计数器（每次热重载递增）
+        self._reload_version: int = 0
+        # 是否正在重载中
+        self._reloading: bool = False
+
+        # 统计
+        self._total_reloads: int = 0
+        self._total_tasks_paused: int = 0
+
+    @property
+    def reload_version(self) -> int:
+        return self._reload_version
+
+    @property
+    def is_reloading(self) -> bool:
+        return self._reloading
+
+    @property
+    def active_task_count(self) -> int:
+        return len(self._active_tasks)
+
+    def register(self, task_id: str, agent_name: str = ""):
+        """
+        注册一个活跃任务。
+
+        Agent 在开始执行任务时调用此方法，以便在配置变更时能收到通知。
+
+        Args:
+            task_id: 任务唯一标识符
+            agent_name: Agent 名称（用于日志）
+        """
+        self._active_tasks[task_id] = TaskCheckpoint(
+            task_id=task_id,
+            agent_name=agent_name,
+            reload_version=self._reload_version,
+        )
+        logger.debug(f"任务已注册: {task_id} (agent={agent_name}, "
+                     f"活跃任务数={len(self._active_tasks)})")
+
+    def unregister(self, task_id: str):
+        """
+        注销一个任务。
+
+        Agent 在任务完成或失败时调用此方法。
+
+        Args:
+            task_id: 任务唯一标识符
+        """
+        removed = self._active_tasks.pop(task_id, None)
+        self._paused_tasks.discard(task_id)
+        if removed:
+            logger.debug(f"任务已注销: {task_id} (活跃任务数={len(self._active_tasks)})")
+
+    async def check_and_wait(self, task_id: str) -> bool:
+        """
+        检查是否需要暂停并等待配置重载完成。
+
+        Agent 应在其处理循环的每次迭代开始时调用此方法。
+        如果配置正在重载，此方法会阻塞（async）直到重载完成。
+
+        Returns:
+            bool: True 表示配置确实发生了重载（Agent 可能需要重新初始化某些状态），
+                  False 表示无配置变更。
+        """
+        if not self._reload_requested.is_set():
+            return False
+
+        checkpoint = self._active_tasks.get(task_id)
+        if not checkpoint:
+            return False
+
+        # 标记为暂停
+        self._paused_tasks.add(task_id)
+        checkpoint.paused_at = time.strftime("%Y-%m-%d %H:%M:%S")
+        self._total_tasks_paused += 1
+
+        logger.info(f"[{task_id}] ⏸️ 任务暂停，等待配置重载完成... "
+                     f"(agent={checkpoint.agent_name}, iteration={checkpoint.iteration})")
+
+        # 等待重载完成
+        try:
+            await asyncio.wait_for(self._reload_done.wait(), timeout=30.0)
+        except asyncio.TimeoutError:
+            logger.warning(f"[{task_id}] ⚠️ 等待配置重载超时，继续执行")
+            return False
+
+        checkpoint.resumed_at = time.strftime("%Y-%m-%d %H:%M:%S")
+        checkpoint.reload_version = self._reload_version
+        self._paused_tasks.discard(task_id)
+
+        logger.info(f"[{task_id}] ▶️ 任务恢复，配置已更新 "
+                     f"(version={self._reload_version})")
+
+        return True
+
+    async def request_reload(self) -> int:
+        """
+        请求配置重载 - 通知所有活跃任务暂停。
+
+        调用后会设置重载信号，等待所有活跃任务暂停后返回。
+        如果没有活跃任务，立即返回。
+
+        Returns:
+            int: 当前活跃任务数量
+        """
+        async with self._lock:
+            if self._reloading:
+                logger.warning("配置重载已在进行中，跳过")
+                return 0
+
+            self._reloading = True
+            active_count = len(self._active_tasks)
+
+            if active_count == 0:
+                logger.info("📡 配置重载广播: 无活跃任务，直接执行")
+                return 0
+
+            # 设置重载请求信号
+            self._reload_requested.set()
+            self._reload_done.clear()
+
+            logger.info(f"📡 配置重载广播: 已通知 {active_count} 个活跃任务暂停")
+
+            # 等待所有活跃任务暂停（最多 10 秒）
+            if active_count > 0:
+                try:
+                    await asyncio.wait_for(
+                        self._wait_all_paused(),
+                        timeout=10.0
+                    )
+                except asyncio.TimeoutError:
+                    paused_count = len(self._paused_tasks)
+                    logger.warning(
+                        f"⚠️ 等待任务暂停超时: {paused_count}/{active_count} 已暂停"
+                    )
+
+            return active_count
+
+    async def _wait_all_paused(self):
+        """等待所有活跃任务都暂停"""
+        while True:
+            all_paused = all(
+                tid in self._paused_tasks
+                for tid in self._active_tasks
+            )
+            if all_paused:
+                return
+            await asyncio.sleep(0.1)
+
+    async def complete_reload(self):
+        """
+        通知所有暂停的任务配置重载已完成。
+
+        在配置实际更新完成后调用此方法，所有暂停的任务将恢复执行。
+        """
+        self._reload_version += 1
+        self._total_reloads += 1
+
+        # 清除重载请求信号（防止重复触发）
+        self._reload_requested.clear()
+
+        # 设置重载完成信号（唤醒所有等待的 Agent）
+        self._reload_done.set()
+
+        # 给 Agent 一点时间恢复
+        await asyncio.sleep(0.3)
+
+        self._reloading = False
+
+        logger.info(f"✅ 配置重载广播完成 (version={self._reload_version}, "
+                     f"累计重载={self._total_reloads})")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """获取广播器统计信息"""
+        return {
+            "reload_version": self._reload_version,
+            "total_reloads": self._total_reloads,
+            "active_tasks": len(self._active_tasks),
+            "paused_tasks": len(self._paused_tasks),
+            "total_tasks_paused": self._total_tasks_paused,
+            "is_reloading": self._reloading,
+            "tasks": [
+                {
+                    "task_id": cp.task_id,
+                    "agent_name": cp.agent_name,
+                    "iteration": cp.iteration,
+                    "paused": cp.task_id in self._paused_tasks,
+                    "paused_at": cp.paused_at,
+                    "resumed_at": cp.resumed_at,
+                }
+                for cp in self._active_tasks.values()
+            ]
+        }
+
+    def force_reset(self):
+        """
+        强制重置广播器状态。
+
+        仅用于测试或异常恢复场景。
+        """
+        self._reload_requested.clear()
+        self._reload_done.clear()
+        self._reloading = False
+        self._paused_tasks.clear()
+        logger.warning("广播器状态已强制重置")