ADR-041: WebSocket-Based CML Worker Monitoring — Implementation Plan¶

class CmlWebSocketMonitor:
    """Persistent WebSocket monitor for a single CML worker.

    Connects to wss://<host>/ws/ui and receives real-time events:
    - system_stats: System resource utilization (CPU, memory, disk, dominfo)
    - lab_stats: Per-lab node/link statistics
    - state_change: Node/interface lifecycle transitions
    - lab_event: Lab-level state changes

    Lifecycle managed by WorkerReconciler:
    - start() when worker transitions to RUNNING
    - stop() when worker transitions to STOPPING/STOPPED/TERMINATED
    """

    def __init__(
        self,
        worker_id: str,
        host: str,
        username: str,
        password: str,
        *,
        verify_ssl: bool = False,
        metrics_report_interval: int = 10,  # Max frequency for metrics reporting (seconds)
        on_system_stats: Callable[[str, CmlSystemStats], Awaitable[None]] | None = None,
        on_activity_event: Callable[[str, dict], Awaitable[None]] | None = None,
        on_lab_stats: Callable[[str, str, dict], Awaitable[None]] | None = None,
        on_lab_state_change: Callable[[str, dict], Awaitable[None]] | None = None,
        on_connection_change: Callable[[str, bool, str | None], Awaitable[None]] | None = None,
    ):
        ...

    async def start(self) -> None:
        """Start WebSocket connection. Non-blocking; runs read loop as background task."""
        ...

    async def stop(self) -> None:
        """Gracefully close WebSocket connection."""
        ...

    @property
    def is_connected(self) -> bool:
        """Whether the WebSocket is currently connected and receiving."""
        ...

    @property
    def last_message_at(self) -> datetime | None:
        """Timestamp of the last received message (for health monitoring)."""
        ...

    @property
    def latest_system_stats(self) -> CmlSystemStats | None:
        """Most recent system_stats received via WebSocket."""
        ...

    @property
    def recent_activity_events(self) -> list[dict]:
        """Activity events received since last drain (for idle detection)."""
        ...

    def drain_activity_events(self) -> list[dict]:
        """Return and clear accumulated activity events."""
        ...

    async def _connect(self) -> None:
        """Establish WebSocket connection with authentication."""
        # 1. Get bearer token via CmlSystemSpiClient._authenticate()
        # 2. Connect to wss://<host>/ws/ui?token=<token> (or header, per validation)
        # 3. Start _read_loop() as asyncio.Task

    async def _read_loop(self) -> None:
        """Main event loop: read messages, parse, and dispatch."""
        # Runs until stop() is called or connection drops
        # On connection loss → _schedule_reconnect()

    async def _handle_message(self, raw: str) -> None:
        """Parse JSON message and route to appropriate handler."""
        # Switch on event_type: system_stats, lab_stats, state_change, lab_event

    async def _handle_system_stats(self, data: dict) -> None:
        """Process system_stats event. Throttle outbound reporting."""
        # Parse into CmlSystemStats DTO (reuse existing from_api_response())
        # Store as latest_system_stats
        # If metrics_report_interval elapsed → invoke on_system_stats callback

    async def _handle_lab_stats(self, lab_id: str, data: dict) -> None:
        """Process lab_stats event."""
        # Invoke on_lab_stats callback with parsed data

    async def _handle_state_change(self, msg: dict) -> None:
        """Process state_change event. Track as activity if relevant."""
        # Classify: is this a user-activity signal?
        # Activity categories (matches existing telemetry filter):
        #   - element_type == "node" AND event in ("QUEUED", "STARTED")
        # If activity → append to recent_activity_events
        # Invoke on_lab_state_change callback

    async def _handle_lab_event(self, msg: dict) -> None:
        """Process lab_event event. Track as activity if relevant."""
        # Activity categories:
        #   - event == "state" AND data.state in ("STARTED", "STOPPED")
        # If activity → append to recent_activity_events
        # Invoke on_activity_event callback

    async def _schedule_reconnect(self, reason: str) -> None:
        """Schedule reconnection with exponential backoff."""
        # Backoff: 1s, 2s, 4s, 8s, 16s, max 30s (matches SSEClient frontend pattern)
        # On reconnect → re-authenticate (token may have expired)
        # Invoke on_connection_change(worker_id, False, reason)

    async def _on_reconnected(self) -> None:
        """Handle successful reconnection."""
        # Reset backoff counter
        # Invoke on_connection_change(worker_id, True, None)

ACTIVITY_CATEGORIES = {
    "lab_event": {"STARTED", "STOPPED"},  # lab_event with data.state in these
    "state_change_node": {"QUEUED", "STARTED"},  # state_change on element_type=="node"
}

class CmlWebSocketMonitorRegistry:
    """Registry of active CmlWebSocketMonitor instances.

    The WorkerReconciler calls ensure_monitoring(worker_id, host) for each
    RUNNING worker, and stop_monitoring(worker_id) when workers leave RUNNING.
    """

    def __init__(self, settings: Settings, cml_spi: CmlSystemSpiClient):
        self._monitors: dict[str, CmlWebSocketMonitor] = {}
        self._settings = settings
        self._cml_spi = cml_spi

    async def ensure_monitoring(
        self,
        worker_id: str,
        host: str,
        *,
        on_system_stats: ...,
        on_activity_event: ...,
        on_lab_stats: ...,
        on_lab_state_change: ...,
        on_connection_change: ...,
    ) -> CmlWebSocketMonitor:
        """Ensure a monitor exists and is connected for the given worker."""
        ...

    async def stop_monitoring(self, worker_id: str) -> None:
        """Stop and remove the monitor for a worker."""
        ...

    async def stop_all(self) -> None:
        """Stop all monitors (shutdown)."""
        ...

    def get_monitor(self, worker_id: str) -> CmlWebSocketMonitor | None:
        """Get the monitor for a worker (if exists)."""
        ...

    @property
    def active_count(self) -> int:
        """Number of active monitors."""
        ...

# ============================================================================
# CML WebSocket Monitoring (ADR-041)
# ============================================================================
cml_websocket_enabled: bool = True  # Master flag to enable WS monitoring
cml_websocket_metrics_report_interval: int = 10  # Seconds between metrics reports to CPA
cml_websocket_reconnect_max_interval: int = 30  # Max reconnect backoff (seconds)
cml_websocket_health_timeout: int = 60  # Consider disconnected if no message in N seconds

# Register CmlWebSocketMonitorRegistry as singleton
builder.services.add_singleton(CmlWebSocketMonitorRegistry)

# Step: Ensure CML WebSocket monitor is connected (ADR-041)
if self._settings.cml_websocket_enabled:
    monitor = await self._ws_registry.ensure_monitoring(
        worker_id=worker.id,
        host=worker.public_ip or worker.private_ip,
        on_system_stats=self._on_ws_system_stats,
        on_activity_event=self._on_ws_activity_event,
        on_lab_stats=self._on_ws_lab_stats,
        on_lab_state_change=self._on_ws_lab_state_change,
        on_connection_change=self._on_ws_connection_change,
    )

# Disconnect WebSocket monitor when worker leaves RUNNING
if self._settings.cml_websocket_enabled:
    await self._ws_registry.stop_monitoring(worker_id)

async def _on_ws_system_stats(self, worker_id: str, stats: CmlSystemStats) -> None:
    """Handle system_stats from WebSocket. Report to CPA."""
    metrics = {
        "collected_at": datetime.utcnow().isoformat(),
        "cml": {
            "cpu_percent": stats.cpu.percent,
            "memory_total": stats.memory.total,
            "memory_used": stats.memory.used,
            "memory_free": stats.memory.free,
            "disk_total": stats.disk.total,
            "disk_used": stats.disk.used,
            "disk_free": stats.disk.free,
            # Per-compute dominfo for capacity tracking
            "computes": [
                {
                    "compute_id": c.compute_id,
                    "hostname": c.hostname,
                    "allocated_cpus": c.stats.dominfo.allocated_cpus,
                    "allocated_memory": c.stats.dominfo.allocated_memory,
                    "running_nodes": c.stats.dominfo.running_nodes,
                    "total_nodes": c.stats.dominfo.total_nodes,
                }
                for c in stats.computes
            ],
        },
        "source": "websocket",  # Distinguish from polling-collected metrics
    }
    await self._cpa_client.report_worker_metrics(worker_id, metrics)

async def _on_ws_activity_event(self, worker_id: str, event: dict) -> None:
    """Handle activity event from WebSocket. Report to CPA for idle detection."""
    await self._cpa_client.report_worker_activity(
        worker_id,
        activity_events=[event],
    )

async def _on_ws_lab_stats(self, worker_id: str, lab_id: str, data: dict) -> None:
    """Handle lab_stats from WebSocket. Report to CPA for per-lab metrics."""
    await self._cpa_client.report_worker_lab_stats(worker_id, lab_id, data)

async def _on_ws_lab_state_change(self, worker_id: str, event: dict) -> None:
    """Handle state_change from WebSocket. Report to CPA."""
    await self._cpa_client.report_worker_lab_state_change(worker_id, event)

async def _on_ws_connection_change(self, worker_id: str, connected: bool, reason: str | None) -> None:
    """Handle WebSocket connection state change."""
    await self._cpa_client.report_worker_ws_status(worker_id, connected, reason)

# If WebSocket is connected for this worker, skip polling CML system_stats
monitor = self._ws_registry.get_monitor(worker_id)
if monitor and monitor.is_connected:
    # Use latest WS-collected stats if available
    if monitor.latest_system_stats:
        # Already reported via callback — skip redundant poll
        logger.debug(f"Skipping CML system_stats poll for {worker_id} (WS connected)")
    else:
        # WS connected but no stats yet — poll as fallback
        stats = await self._cml_spi.get_system_stats(host, username, password)
        ...
else:
    # Fallback: poll as before
    stats = await self._cml_spi.get_system_stats(host, username, password)
    ...

monitor = self._ws_registry.get_monitor(worker_id) if self._settings.cml_websocket_enabled else None

if monitor and monitor.is_connected:
    # Drain accumulated activity events from WebSocket
    activity_events = monitor.drain_activity_events()
    if activity_events:
        # Report directly to CPA (already in processed form)
        result = await self._cpa_client.detect_worker_idle(
            worker_id,
            force_check=False,
            raw_telemetry_events=activity_events,
        )
    else:
        # No new activity events — still call detect_idle for time-based evaluation
        result = await self._cpa_client.detect_worker_idle(
            worker_id,
            force_check=False,
            raw_telemetry_events=[],  # Empty = no new activity
        )
else:
    # Fallback: fetch telemetry via REST (existing behavior)
    raw_events = await self._cml_spi.get_telemetry_events(host, username, password)
    result = await self._cpa_client.detect_worker_idle(
        worker_id,
        force_check=False,
        raw_telemetry_events=raw_events,
    )

if self._ws_registry:
    await self._ws_registry.stop_all()

@router.post("/api/internal/workers/{worker_id}/activity")
async def report_worker_activity(worker_id: str, request: ReportActivityRequest):
    """Report real-time activity events from WebSocket stream (ADR-041)."""
    ...

class ReportActivityRequest(BaseModel):
    activity_events: list[dict[str, Any]] = Field(default_factory=list)
    source: str = "websocket"  # "websocket" or "telemetry_poll"

@router.post("/api/internal/workers/{worker_id}/labs/{lab_id}/stats")
async def report_worker_lab_stats(worker_id: str, lab_id: str, request: LabStatsRequest):
    """Report per-lab metrics from WebSocket lab_stats events (ADR-041)."""
    ...

class LabStatsRequest(BaseModel):
    nodes: dict[str, Any] = Field(default_factory=dict)  # node_id -> {cpu_usage, ram_usage, disk_usage, ...}
    links: dict[str, Any] = Field(default_factory=dict)  # link_id -> {readbytes, writebytes, ...}
    collected_at: str | None = None

@router.post("/api/internal/workers/{worker_id}/labs/{lab_id}/state-change")
async def report_worker_lab_state_change(worker_id: str, lab_id: str, request: LabStateChangeRequest):
    """Report real-time node/interface state change from WebSocket (ADR-041)."""
    ...

class LabStateChangeRequest(BaseModel):
    event: str  # QUEUED, STARTED, BOOTED, STOPPED, etc.
    element_type: str  # "node", "interface", "lab"
    element_id: str | None = None
    data: dict[str, Any] = Field(default_factory=dict)

@router.post("/api/internal/workers/{worker_id}/ws-status")
async def report_worker_ws_status(worker_id: str, request: WsStatusRequest):
    """Report WebSocket connection status change (ADR-041)."""
    ...

class WsStatusRequest(BaseModel):
    connected: bool
    reason: str | None = None
    connected_at: str | None = None
    disconnected_at: str | None = None

@dataclass
class UpdateWorkerLabStatsCommand(Command[OperationResult[dict[str, Any]]]):
    worker_id: str
    lab_id: str
    nodes: dict[str, Any] = field(default_factory=dict)
    links: dict[str, Any] = field(default_factory=dict)
    collected_at: str | None = None

async def report_worker_activity(self, worker_id: str, activity_events: list[dict]) -> dict:
    """Report push-based activity events (ADR-041)."""
    ...

async def report_worker_lab_stats(self, worker_id: str, lab_id: str, data: dict) -> dict:
    """Report per-lab resource metrics (ADR-041)."""
    ...

async def report_worker_lab_state_change(self, worker_id: str, event: dict) -> dict:
    """Report lab/node state change (ADR-041)."""
    ...

async def report_worker_ws_status(self, worker_id: str, connected: bool, reason: str | None = None) -> dict:
    """Report WebSocket connection status (ADR-041)."""
    ...

BATCHABLE_EVENT_TYPES = {
    "worker.metrics.updated",
    "worker.lab.stats_updated",  # ADR-041: Per-lab metrics (high frequency)
}

# WebSocket monitoring metrics (ADR-041)
lcm_worker_controller_ws_connections_active = Gauge(
    "lcm_worker_controller_ws_connections_active",
    "Number of active WebSocket connections to CML workers",
)
lcm_worker_controller_ws_messages_total = Counter(
    "lcm_worker_controller_ws_messages_total",
    "Total WebSocket messages received",
    ["worker_id", "event_type"],
)
lcm_worker_controller_ws_reconnections_total = Counter(
    "lcm_worker_controller_ws_reconnections_total",
    "Total WebSocket reconnection attempts",
    ["worker_id", "reason"],
)
lcm_worker_controller_ws_message_latency_seconds = Histogram(
    "lcm_worker_controller_ws_message_latency_seconds",
    "Time from WS message receipt to CPA report",
    ["event_type"],
)