ForcePilot/backend/package/yuxi/channel/monitoring/diagnostic_heartbeat.py

import asyncio
import gc
import logging
import time
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

DEFAULT_HEARTBEAT_INTERVAL = 30
STUCK_SESSION_TIMEOUT = 600
MAX_CONSECUTIVE_FAILURES = 5

MB = 1024 * 1024
DEFAULT_RSS_WARNING_MB = 1536
DEFAULT_RSS_CRITICAL_MB = 3072
DEFAULT_DATA_WARNING_MB = 1024
DEFAULT_DATA_CRITICAL_MB = 2048
DEFAULT_GC_OBJECTS_WARNING = 500000
DEFAULT_GC_SAMPLE_INTERVAL = 5
DEFAULT_RSS_GROWTH_WARNING_MB = 512
DEFAULT_RSS_GROWTH_CRITICAL_MB = 1024
DEFAULT_GROWTH_WINDOW_S = 600
DEFAULT_PRESSURE_COOLDOWN_S = 300


@dataclass
class MemorySnapshot:
    rss_mb: float = 0.0
    vms_mb: float = 0.0
    data_mb: float = 0.0
    gc_objects: int = 0
    gc_g0_count: int = 0
    gc_g1_count: int = 0
    gc_g2_count: int = 0
    timestamp: float = 0.0


@dataclass
class MemoryPressureEvent:
    level: str  # "warning" | "critical"
    reason: str  # "rss_threshold" | "data_threshold" | "rss_growth"
    rss_mb: float = 0.0
    threshold_mb: float = 0.0
    rss_growth_mb: float = 0.0
    window_s: float = 0.0
    timestamp: float = 0.0


@dataclass
class HeartbeatReport:
    active_sessions: int = 0
    stuck_sessions: int = 0
    stuck_details: list[dict] = field(default_factory=list)
    tool_loop_sessions: int = 0
    tool_loop_details: list[dict] = field(default_factory=list)
    memory: MemorySnapshot = field(default_factory=MemorySnapshot)
    active_tasks: int = 0
    uptime_seconds: float = 0.0
    process_start_timestamp: float = 0.0
    timestamp: float = 0.0
    warnings: list[str] = field(default_factory=list)
    memory_pressure: MemoryPressureEvent | None = None

    def to_dict(self) -> dict:
        result = {
            "active_sessions": self.active_sessions,
            "stuck_sessions": self.stuck_sessions,
            "stuck_details": self.stuck_details,
            "tool_loop_sessions": self.tool_loop_sessions,
            "tool_loop_details": self.tool_loop_details,
            "memory": {
                "rss_mb": self.memory.rss_mb,
                "vms_mb": self.memory.vms_mb,
                "data_mb": self.memory.data_mb,
                "gc_objects": self.memory.gc_objects,
                "gc_g0_count": self.memory.gc_g0_count,
                "gc_g1_count": self.memory.gc_g1_count,
                "gc_g2_count": self.memory.gc_g2_count,
                "timestamp": self.memory.timestamp,
            },
            "active_tasks": self.active_tasks,
            "uptime_seconds": self.uptime_seconds,
            "process_start_timestamp": self.process_start_timestamp,
            "timestamp": self.timestamp,
            "warnings": self.warnings,
        }
        if self.memory_pressure:
            result["memory_pressure"] = {
                "level": self.memory_pressure.level,
                "reason": self.memory_pressure.reason,
                "rss_mb": self.memory_pressure.rss_mb,
                "threshold_mb": self.memory_pressure.threshold_mb,
                "rss_growth_mb": self.memory_pressure.rss_growth_mb,
                "window_s": self.memory_pressure.window_s,
                "timestamp": self.memory_pressure.timestamp,
            }
        return result


class DiagnosticHeartbeat:
    """诊断心跳：30s 间隔检测卡住 session、内存快照 + 内存压力。

    集成到 FastAPI lifespan 后启动，可被 health endpoint 消费。
    """

    def __init__(
        self,
        interval: int = DEFAULT_HEARTBEAT_INTERVAL,
        stuck_timeout: int = STUCK_SESSION_TIMEOUT,
        rss_warning_mb: float = DEFAULT_RSS_WARNING_MB,
        rss_critical_mb: float = DEFAULT_RSS_CRITICAL_MB,
        data_warning_mb: float = DEFAULT_DATA_WARNING_MB,
        data_critical_mb: float = DEFAULT_DATA_CRITICAL_MB,
        rss_growth_warning_mb: float = DEFAULT_RSS_GROWTH_WARNING_MB,
        rss_growth_critical_mb: float = DEFAULT_RSS_GROWTH_CRITICAL_MB,
        growth_window_s: float = DEFAULT_GROWTH_WINDOW_S,
        pressure_cooldown_s: float = DEFAULT_PRESSURE_COOLDOWN_S,
        gc_objects_warning: int = DEFAULT_GC_OBJECTS_WARNING,
        gc_sample_interval: int = DEFAULT_GC_SAMPLE_INTERVAL,
    ):
        self._interval = interval
        self._stuck_timeout = stuck_timeout
        self._rss_warning_mb = rss_warning_mb
        self._rss_critical_mb = rss_critical_mb
        self._data_warning_mb = data_warning_mb
        self._data_critical_mb = data_critical_mb
        self._rss_growth_warning_mb = rss_growth_warning_mb
        self._rss_growth_critical_mb = rss_growth_critical_mb
        self._growth_window_s = growth_window_s
        self._pressure_cooldown_s = pressure_cooldown_s
        self._gc_objects_warning = gc_objects_warning
        self._gc_sample_interval = gc_sample_interval
        self._running = False
        self._task: asyncio.Task | None = None
        self._last_report: HeartbeatReport | None = None
        self._last_memory: MemorySnapshot | None = None
        self._last_pressure_at: dict[str, float] = {}
        self._start_time: float | None = None
        self._start_monotonic: float | None = None
        self._consecutive_failures = 0
        self._cycle_count = 0

    @property
    def last_report(self) -> HeartbeatReport | None:
        return self._last_report

    @staticmethod
    def _resolve_start_time() -> float:
        try:
            import psutil

            return psutil.Process().create_time()
        except Exception:
            return time.time()

    async def start(self):
        self._running = True
        self._start_time = self._resolve_start_time()
        self._start_monotonic = time.monotonic()
        self._task = asyncio.create_task(self._loop())
        logger.info(
            "DiagnosticHeartbeat started (interval=%ds, stuck_timeout=%ds)",
            self._interval,
            self._stuck_timeout,
        )

    async def stop(self):
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        logger.info("DiagnosticHeartbeat stopped")

    async def _loop(self):
        while self._running:
            await asyncio.sleep(self._interval)
            try:
                self._last_report = await self._collect()
                self._consecutive_failures = 0
                await self._push_metrics(self._last_report)
                if self._last_report.warnings:
                    for w in self._last_report.warnings:
                        logger.warning("DiagnosticHeartbeat: %s", w)
                if self._last_report.memory_pressure:
                    mp = self._last_report.memory_pressure
                    logger.warning(
                        "DiagnosticHeartbeat: 内存压力 %s (reason=%s, rss=%.0fMB, threshold=%.0fMB)",
                        mp.level,
                        mp.reason,
                        mp.rss_mb,
                        mp.threshold_mb,
                    )
            except Exception:
                self._consecutive_failures += 1
                logger.exception(
                    "DiagnosticHeartbeat: collect failed (consecutive=%d/%d)",
                    self._consecutive_failures,
                    MAX_CONSECUTIVE_FAILURES,
                )
                if self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
                    logger.critical(
                        "DiagnosticHeartbeat: %d consecutive collect failures, heartbeat may be broken",
                        self._consecutive_failures,
                    )

    async def _collect(self) -> HeartbeatReport:
        now = time.monotonic()
        self._cycle_count += 1
        memory = await self._capture_memory(now)
        self._last_memory = memory

        snapshots = _get_all_snapshots()
        stuck_details, stuck_count = await self._detect_stuck_sessions(now, snapshots)
        active_tasks = len(asyncio.all_tasks())
        active_sessions = _count_active_sessions_from(snapshots)

        uptime_seconds = now - self._start_monotonic if self._start_monotonic else 0.0
        process_start_timestamp = self._start_time or 0.0

        memory_pressure = self._detect_memory_pressure(memory, now)

        warnings: list[str] = []
        if stuck_count > 0:
            warnings.append(f"检测到 {stuck_count} 个卡住 session")
        if memory.rss_mb > self._rss_warning_mb:
            warnings.append(f"内存 RSS {memory.rss_mb:.0f}MB > {self._rss_warning_mb:.0f}MB")
        if memory.data_mb > self._data_warning_mb:
            warnings.append(f"内存 DATA {memory.data_mb:.0f}MB > {self._data_warning_mb:.0f}MB")
        if memory.gc_objects > self._gc_objects_warning:
            warnings.append(f"GC 对象数 {memory.gc_objects} > {self._gc_objects_warning}")
        if memory_pressure:
            level_label = "严重" if memory_pressure.level == "critical" else ""
            warnings.append(f"内存压力{level_label}: {memory_pressure.reason} (RSS={memory_pressure.rss_mb:.0f}MB)")

        return HeartbeatReport(
            active_sessions=active_sessions,
            stuck_sessions=stuck_count,
            stuck_details=stuck_details,
            tool_loop_sessions=0,
            tool_loop_details=[],
            memory=memory,
            active_tasks=active_tasks,
            uptime_seconds=uptime_seconds,
            process_start_timestamp=process_start_timestamp,
            timestamp=now,
            warnings=warnings,
            memory_pressure=memory_pressure,
        )

    async def _push_metrics(self, report: HeartbeatReport):
        try:
            from yuxi.channel.monitoring.metrics import (
                active_sessions,
                active_tasks,
                memory_gc_objects,
                memory_rss_mb,
                stuck_sessions,
            )

            active_sessions.set(report.active_sessions)
            stuck_sessions.set(report.stuck_sessions)
            memory_rss_mb.set(report.memory.rss_mb)
            memory_gc_objects.set(report.memory.gc_objects)
            active_tasks.set(report.active_tasks)
        except Exception:
            pass

    def _detect_memory_pressure(self, current: MemorySnapshot, now: float) -> MemoryPressureEvent | None:
        pressure = _check_threshold_pressure(
            current,
            rss_warning_mb=self._rss_warning_mb,
            rss_critical_mb=self._rss_critical_mb,
        )
        if pressure is None:
            pressure = _check_data_threshold_pressure(
                current,
                data_warning_mb=self._data_warning_mb,
                data_critical_mb=self._data_critical_mb,
            )
        if pressure is None:
            pressure = _check_growth_pressure(
                self._last_memory,
                current,
                now,
                growth_window_s=self._growth_window_s,
                rss_growth_warning_mb=self._rss_growth_warning_mb,
                rss_growth_critical_mb=self._rss_growth_critical_mb,
            )

        if pressure is None:
            return None

        key = pressure.reason
        last_at = self._last_pressure_at.get(key, 0)
        if now - last_at < self._pressure_cooldown_s:
            return None

        self._last_pressure_at[key] = now
        return pressure

    async def _detect_stuck_sessions(self, now: float, snapshots: dict) -> tuple[list[dict], int]:
        stuck: list[dict] = []
        for key, snap in snapshots.items():
            if snap.last_event_at and (now - snap.last_event_at) > self._stuck_timeout:
                state_value = _get_state_value(snap.state)
                if snap.connected and state_value in ("running", "retrying"):
                    stuck.append(
                        {
                            "key": key,
                            "state": state_value,
                            "last_event_at": snap.last_event_at,
                            "idle_seconds": round(now - snap.last_event_at, 1),
                        }
                    )
        return stuck, len(stuck)

    async def _capture_memory(self, now: float) -> MemorySnapshot:
        snapshot = MemorySnapshot(timestamp=now)
        try:
            import psutil

            proc = psutil.Process()
            mem = proc.memory_info()
            snapshot.rss_mb = round(mem.rss / 1024 / 1024, 2)
            snapshot.vms_mb = round(mem.vms / 1024 / 1024, 2)
            snapshot.data_mb = round(getattr(mem, "data", 0) / 1024 / 1024, 2)
        except Exception:
            pass

        gc_counts = gc.get_count()
        snapshot.gc_g0_count = gc_counts[0]
        snapshot.gc_g1_count = gc_counts[1]
        snapshot.gc_g2_count = gc_counts[2]

        if self._cycle_count % self._gc_sample_interval == 0:
            snapshot.gc_objects = len(gc.get_objects())

        return snapshot

    def get_memory_trend(self) -> dict:
        if self._last_memory:
            return {
                "rss_mb": self._last_memory.rss_mb,
                "vms_mb": self._last_memory.vms_mb,
                "data_mb": self._last_memory.data_mb,
                "gc_objects": self._last_memory.gc_objects,
                "timestamp": self._last_memory.timestamp,
            }
        return {}

    def reset(self):
        self._last_memory = None
        self._last_pressure_at.clear()
        self._last_report = None
        self._cycle_count = 0


def _check_threshold_pressure(
    memory: MemorySnapshot,
    *,
    rss_warning_mb: float,
    rss_critical_mb: float,
) -> MemoryPressureEvent | None:
    if memory.rss_mb >= rss_critical_mb:
        return MemoryPressureEvent(
            level="critical",
            reason="rss_threshold",
            rss_mb=memory.rss_mb,
            threshold_mb=rss_critical_mb,
            timestamp=memory.timestamp,
        )
    if memory.rss_mb >= rss_warning_mb:
        return MemoryPressureEvent(
            level="warning",
            reason="rss_threshold",
            rss_mb=memory.rss_mb,
            threshold_mb=rss_warning_mb,
            timestamp=memory.timestamp,
        )
    return None


def _check_data_threshold_pressure(
    memory: MemorySnapshot,
    *,
    data_warning_mb: float,
    data_critical_mb: float,
) -> MemoryPressureEvent | None:
    if memory.data_mb >= data_critical_mb:
        return MemoryPressureEvent(
            level="critical",
            reason="data_threshold",
            rss_mb=memory.rss_mb,
            threshold_mb=data_critical_mb,
            timestamp=memory.timestamp,
        )
    if memory.data_mb >= data_warning_mb:
        return MemoryPressureEvent(
            level="warning",
            reason="data_threshold",
            rss_mb=memory.rss_mb,
            threshold_mb=data_warning_mb,
            timestamp=memory.timestamp,
        )
    return None


def _check_growth_pressure(
    previous: MemorySnapshot | None,
    current: MemorySnapshot,
    now: float,
    *,
    growth_window_s: float,
    rss_growth_warning_mb: float,
    rss_growth_critical_mb: float,
) -> MemoryPressureEvent | None:
    if previous is None or previous.rss_mb <= 0:
        return None

    window_s = now - previous.timestamp
    if window_s <= 0 or window_s > growth_window_s:
        return None

    rss_growth_mb = current.rss_mb - previous.rss_mb
    if rss_growth_mb <= 0:
        return None

    if rss_growth_mb >= rss_growth_critical_mb:
        return MemoryPressureEvent(
            level="critical",
            reason="rss_growth",
            rss_mb=current.rss_mb,
            threshold_mb=rss_growth_critical_mb,
            rss_growth_mb=round(rss_growth_mb, 2),
            window_s=round(window_s, 1),
            timestamp=current.timestamp,
        )
    if rss_growth_mb >= rss_growth_warning_mb:
        return MemoryPressureEvent(
            level="warning",
            reason="rss_growth",
            rss_mb=current.rss_mb,
            threshold_mb=rss_growth_warning_mb,
            rss_growth_mb=round(rss_growth_mb, 2),
            window_s=round(window_s, 1),
            timestamp=current.timestamp,
        )
    return None


def _get_state_value(state) -> str:
    return state.value if hasattr(state, "value") else str(state)


def _get_all_snapshots() -> dict:
    try:
        from yuxi.channel.runtime.manager import gateway  # 延迟导入以避免循环依赖；Python import 缓存保证无重复加载开销

        return gateway.get_all_snapshots()
    except Exception:
        logger.warning("DiagnosticHeartbeat: get_all_snapshots failed", exc_info=True)
        return {}


def _count_active_sessions_from(snapshots: dict) -> int:
    return sum(1 for s in snapshots.values() if _get_state_value(s.state) == "running")


diagnostic_heartbeat = DiagnosticHeartbeat()