ForcePilot/backend/package/yuxi/channel/runtime/trace.py

from __future__ import annotations

import time as _time
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any


class PhaseStatus(StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    OK = "ok"
    FAILED = "failed"
    SKIPPED = "skipped"


class GatewayReadyState(StrEnum):
    NOT_READY = "not_ready"
    READY = "ready"
    DEGRADED = "degraded"
    FAILED = "failed"


@dataclass
class PhaseMark:
    name: str
    status: PhaseStatus = PhaseStatus.PENDING
    started_at: float | None = None
    finished_at: float | None = None

    @property
    def elapsed_ms(self) -> float | None:
        if self.started_at and self.finished_at:
            return (self.finished_at - self.started_at) * 1000
        return None


@dataclass
class StartupTrace:
    """启动追踪器 — 记录系统启动各阶段耗时。"""

    started_at: float = 0.0
    completed_at: float = 0.0
    failed_at: float = 0.0
    error: str | None = None
    phases: list[dict[str, Any]] = field(default_factory=list)

    _phase_marks: list[PhaseMark] = field(default_factory=list, repr=False)
    _active_phase: str | None = field(default=None, repr=False)
    ready_state: GatewayReadyState = GatewayReadyState.NOT_READY
    trace_id: str = ""

    def start(self) -> None:
        self.started_at = _time.time()
        self.phases.clear()

    def add_phase(self, name: str, duration_ms: float, *, success: bool = True, error: str | None = None) -> None:
        self.phases.append(
            {
                "name": name,
                "duration_ms": round(duration_ms, 2),
                "success": success,
                "error": error,
            }
        )

    def complete(self) -> None:
        self.completed_at = _time.time()

    def fail(self, error: str) -> None:
        self.failed_at = _time.time()
        self.error = error

    @property
    def duration(self) -> float:
        if self.completed_at:
            return self.completed_at - self.started_at
        if self.failed_at:
            return self.failed_at - self.started_at
        return _time.time() - self.started_at

    @property
    def success(self) -> bool:
        return self.completed_at > 0 and self.failed_at == 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "trace_id": self.trace_id,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "failed_at": self.failed_at,
            "duration_seconds": round(self.duration, 3),
            "success": self.success,
            "ready_state": self.ready_state.value,
            "error": self.error,
            "phases": list(self.phases),
        }

    def start_phase(self, name: str) -> None:
        mark = PhaseMark(name=name, status=PhaseStatus.RUNNING, started_at=_time.time())
        self._upsert_mark(mark)
        self._active_phase = name

    def complete_phase(self, name: str, ok: bool = True, error: str | None = None) -> None:
        mark = self._find_mark(name)
        if mark is None:
            return
        mark.status = PhaseStatus.OK if ok else PhaseStatus.FAILED
        mark.finished_at = _time.time()
        self._active_phase = None
        duration_ms = mark.elapsed_ms or 0
        self.add_phase(name, duration_ms, success=ok, error=error)

    def skip_phase(self, name: str) -> None:
        mark = self._find_mark(name)
        if mark is not None:
            mark.status = PhaseStatus.SKIPPED
        else:
            mark = PhaseMark(name=name, status=PhaseStatus.SKIPPED)
            self._phase_marks.append(mark)
        self._active_phase = None

    def finalize(self, ready_state: GatewayReadyState) -> None:
        self.ready_state = ready_state
        if not self.completed_at and not self.failed_at:
            self.completed_at = _time.time()

    def summary(self) -> dict[str, Any]:
        return {
            "trace_id": self.trace_id,
            "ready_state": self.ready_state.value,
            "total_elapsed_ms": round(self.duration * 1000, 2) if self.duration else None,
            "phases": [
                {
                    "name": p.name,
                    "status": p.status.value,
                    "elapsed_ms": round(p.elapsed_ms, 2) if p.elapsed_ms else None,
                }
                for p in self._phase_marks
            ],
            "errors": [
                e["error"] for e in self.phases if not e.get("success") and e.get("error")
            ],
        }

    def _upsert_mark(self, mark: PhaseMark) -> None:
        existing = self._find_mark(mark.name)
        if existing:
            idx = self._phase_marks.index(existing)
            self._phase_marks[idx] = mark
        else:
            self._phase_marks.append(mark)

    def _find_mark(self, name: str) -> PhaseMark | None:
        for p in self._phase_marks:
            if p.name == name:
                return p
        return None