175 lines
5.3 KiB
Python
175 lines
5.3 KiB
Python
"""
|
|
SSE client for workflow events.
|
|
|
|
Connects to the /workflow/sse endpoint and emits signals for:
|
|
- Runtime state changes
|
|
- Control state changes
|
|
- Individual workflow events
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from PySide6.QtCore import QObject, Signal, Slot, QTimer
|
|
from PySide6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply
|
|
from PySide6.QtCore import QUrl, QByteArray
|
|
|
|
from aare.common.automation_models import RuntimeState, ControlState, WorkflowEvent
|
|
from aare.common.logger_config import setup_logger
|
|
|
|
logger = setup_logger("aareGUI")
|
|
|
|
|
|
class WorkflowSSEClient(QObject):
|
|
"""
|
|
SSE client that subscribes to workflow events.
|
|
|
|
Emits signals when state changes are received.
|
|
"""
|
|
|
|
# Signals
|
|
runtime_changed = Signal(object) # RuntimeState
|
|
control_changed = Signal(object) # ControlState
|
|
workflow_event = Signal(object) # WorkflowEvent
|
|
connected = Signal()
|
|
disconnected = Signal()
|
|
error = Signal(str)
|
|
|
|
def __init__(
|
|
self,
|
|
base_url: str,
|
|
token: str,
|
|
parent: QObject | None = None,
|
|
):
|
|
super().__init__(parent)
|
|
|
|
self._base_url = base_url
|
|
self._token = token
|
|
self._manager = QNetworkAccessManager(self)
|
|
self._reply: QNetworkReply | None = None
|
|
self._buffer = ""
|
|
|
|
# Reconnection
|
|
self._reconnect_timer = QTimer(self)
|
|
self._reconnect_timer.setInterval(5000) # 5 seconds
|
|
self._reconnect_timer.timeout.connect(self.connect)
|
|
self._should_reconnect = False
|
|
|
|
def connect(self) -> None:
|
|
"""Start SSE connection."""
|
|
if self._reply is not None:
|
|
return # Already connected
|
|
|
|
self._should_reconnect = True
|
|
|
|
url = QUrl(f"{self._base_url}/workflow/sse")
|
|
|
|
request = QNetworkRequest(url)
|
|
request.setRawHeader(b"Authorization", f"Bearer {self._token}".encode())
|
|
request.setRawHeader(b"Accept", b"text/event-stream")
|
|
request.setRawHeader(b"Cache-Control", b"no-cache")
|
|
|
|
self._reply = self._manager.get(request)
|
|
self._reply.readyRead.connect(self._on_data_ready)
|
|
self._reply.finished.connect(self._on_finished)
|
|
self._reply.errorOccurred.connect(self._on_error)
|
|
|
|
self._reconnect_timer.stop()
|
|
logger.debug("Workflow SSE: connecting...")
|
|
|
|
def disconnect(self) -> None:
|
|
"""Stop SSE connection."""
|
|
self._should_reconnect = False
|
|
self._reconnect_timer.stop()
|
|
|
|
if self._reply is not None:
|
|
try:
|
|
self._reply.abort()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self._reply.deleteLater()
|
|
except Exception:
|
|
pass
|
|
self._reply = None
|
|
|
|
self.disconnected.emit()
|
|
|
|
@Slot()
|
|
def _on_data_ready(self) -> None:
|
|
"""Handle incoming SSE data."""
|
|
if self._reply is None:
|
|
return
|
|
|
|
try:
|
|
data = self._reply.readAll().data().decode("utf-8")
|
|
self._buffer += data
|
|
|
|
# Process complete events (separated by double newlines)
|
|
while "\n\n" in self._buffer:
|
|
event_data, self._buffer = self._buffer.split("\n\n", 1)
|
|
self._parse_event(event_data)
|
|
except Exception as e:
|
|
logger.warning(f"Workflow SSE data read error: {e}")
|
|
|
|
def _parse_event(self, event_data: str) -> None:
|
|
"""Parse a single SSE event."""
|
|
event_type = "message"
|
|
data_lines = []
|
|
|
|
for line in event_data.split("\n"):
|
|
if line.startswith("event:"):
|
|
event_type = line[6:].strip()
|
|
elif line.startswith("data:"):
|
|
data_lines.append(line[5:].strip())
|
|
|
|
if not data_lines:
|
|
return
|
|
|
|
data_str = "\n".join(data_lines)
|
|
|
|
try:
|
|
if event_type == "runtime":
|
|
runtime = RuntimeState.model_validate_json(data_str)
|
|
self.runtime_changed.emit(runtime)
|
|
elif event_type == "control":
|
|
control = ControlState.model_validate_json(data_str)
|
|
self.control_changed.emit(control)
|
|
elif event_type == "workflow_event":
|
|
event = WorkflowEvent.model_validate_json(data_str)
|
|
self.workflow_event.emit(event)
|
|
except Exception as e:
|
|
logger.warning(f"Workflow SSE: failed to parse {event_type}: {e}")
|
|
|
|
@Slot()
|
|
def _on_finished(self) -> None:
|
|
"""Handle connection finished."""
|
|
if self._reply is not None:
|
|
try:
|
|
self._reply.deleteLater()
|
|
except Exception:
|
|
pass
|
|
self._reply = None
|
|
|
|
self._buffer = ""
|
|
self.disconnected.emit()
|
|
|
|
# Reconnect if desired
|
|
if self._should_reconnect:
|
|
logger.debug("Workflow SSE: disconnected, will reconnect...")
|
|
self._reconnect_timer.start()
|
|
|
|
@Slot(QNetworkReply.NetworkError)
|
|
def _on_error(self, error: QNetworkReply.NetworkError) -> None:
|
|
"""Handle connection error."""
|
|
error_msg = ""
|
|
if self._reply is not None:
|
|
try:
|
|
error_msg = self._reply.errorString()
|
|
except Exception:
|
|
error_msg = str(error)
|
|
else:
|
|
error_msg = str(error)
|
|
logger.warning(f"Workflow SSE error: {error_msg}")
|
|
self.error.emit(error_msg) |