multiAgentTry/backend/app/services/heartbeat.py

"""
心跳服务 - 管理 Agent 心跳和超时检测
用于监控 Agent 活跃状态和检测掉线 Agent
"""

import asyncio
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from dataclasses import dataclass, asdict

from .storage import get_storage


@dataclass
class HeartbeatInfo:
    """心跳信息"""
    agent_id: str
    last_heartbeat: str       # 最后心跳时间 (ISO format)
    status: str               # Agent 状态：working, waiting, idle, error
    current_task: str = ""    # 当前任务描述
    progress: int = 0         # 任务进度 0-100

    @property
    def elapsed_seconds(self) -> int:
        """距最后心跳的秒数"""
        last_time = datetime.fromisoformat(self.last_heartbeat)
        return int((datetime.now() - last_time).total_seconds())

    def is_timeout(self, timeout_seconds: int = 60) -> bool:
        """是否超时"""
        return self.elapsed_seconds > timeout_seconds

    @property
    def elapsed_display(self) -> str:
        """格式化的时间差"""
        seconds = self.elapsed_seconds
        if seconds < 10:
            return f"{seconds}s ago"
        elif seconds < 60:
            return f"{seconds}s"
        minutes = seconds // 60
        secs = seconds % 60
        return f"{minutes}m {secs:02d}s"


class HeartbeatService:
    """
    心跳服务

    管理所有 Agent 的心跳记录，检测超时 Agent
    """

    HEARTBEATS_FILE = "cache/heartbeats.json"
    DEFAULT_TIMEOUT = 60  # 默认超时时间（秒）

    def __init__(self):
        self._storage = get_storage()
        self._lock = asyncio.Lock()

    async def _load_heartbeats(self) -> Dict[str, Dict]:
        """加载心跳记录"""
        return await self._storage.read_json(self.HEARTBEATS_FILE)

    async def _save_heartbeats(self, heartbeats: Dict[str, Dict]) -> None:
        """保存心跳记录"""
        await self._storage.write_json(self.HEARTBEATS_FILE, heartbeats)

    async def update_heartbeat(
        self,
        agent_id: str,
        status: str,
        current_task: str = "",
        progress: int = 0
    ) -> None:
        """
        更新 Agent 心跳

        Args:
            agent_id: Agent ID
            status: 状态 (working, waiting, idle, error)
            current_task: 当前任务
            progress: 进度 0-100
        """
        async with self._lock:
            heartbeats = await self._load_heartbeats()

            heartbeat_info = HeartbeatInfo(
                agent_id=agent_id,
                last_heartbeat=datetime.now().isoformat(),
                status=status,
                current_task=current_task,
                progress=progress
            )

            heartbeats[agent_id] = asdict(heartbeat_info)
            await self._save_heartbeats(heartbeats)

    async def get_heartbeat(self, agent_id: str) -> Optional[HeartbeatInfo]:
        """
        获取指定 Agent 的心跳信息

        Args:
            agent_id: Agent ID

        Returns:
            心跳信息，如果不存在返回 None
        """
        heartbeats = await self._load_heartbeats()
        data = heartbeats.get(agent_id)
        if data:
            return HeartbeatInfo(**data)
        return None

    async def get_all_heartbeats(self) -> Dict[str, HeartbeatInfo]:
        """
        获取所有 Agent 的心跳信息

        Returns:
            Agent ID -> 心跳信息 的字典
        """
        heartbeats = await self._load_heartbeats()
        result = {}
        for agent_id, data in heartbeats.items():
            result[agent_id] = HeartbeatInfo(**data)
        return result

    async def check_timeout(self, timeout_seconds: int = None) -> List[str]:
        """
        检查超时的 Agent

        Args:
            timeout_seconds: 超时秒数，默认使用 DEFAULT_TIMEOUT

        Returns:
            超时的 Agent ID 列表
        """
        if timeout_seconds is None:
            timeout_seconds = self.DEFAULT_TIMEOUT

        all_heartbeats = await self.get_all_heartbeats()
        timeout_agents = []

        for agent_id, heartbeat in all_heartbeats.items():
            if heartbeat.is_timeout(timeout_seconds):
                timeout_agents.append(agent_id)

        return timeout_agents

    async def remove_heartbeat(self, agent_id: str) -> bool:
        """
        移除 Agent 的心跳记录

        Args:
            agent_id: Agent ID

        Returns:
            是否成功移除
        """
        async with self._lock:
            heartbeats = await self._load_heartbeats()
            if agent_id in heartbeats:
                del heartbeats[agent_id]
                await self._save_heartbeats(heartbeats)
                return True
            return False

    async def get_active_agents(self, within_seconds: int = 120) -> List[str]:
        """
        获取活跃的 Agent 列表

        Args:
            within_seconds: 活跃判定时间窗口（秒）

        Returns:
            活跃 Agent ID 列表
        """
        all_heartbeats = await self.get_all_heartbeats()
        active_agents = []

        for agent_id, heartbeat in all_heartbeats.items():
            if heartbeat.elapsed_seconds <= within_seconds:
                active_agents.append(agent_id)

        return active_agents

    async def get_agents_by_status(self, status: str) -> List[HeartbeatInfo]:
        """
        获取指定状态的所有 Agent

        Args:
            status: 状态 (working, waiting, idle, error)

        Returns:
            符合条件的 Agent 心跳信息列表
        """
        all_heartbeats = await self.get_all_heartbeats()
        return [
            hb for hb in all_heartbeats.values()
            if hb.status == status
        ]


# 全局单例
_heartbeat_service_instance: Optional[HeartbeatService] = None


def get_heartbeat_service() -> HeartbeatService:
    """获取心跳服务单例"""
    global _heartbeat_service_instance
    if _heartbeat_service_instance is None:
        _heartbeat_service_instance = HeartbeatService()
    return _heartbeat_service_instance