feat: switch telemetry and events transport from HTTP to WebSocket

- Telemetry and events are now sent over a persistent WebSocket connection
  to /p/ws instead of individual HTTP POSTs to /p/telemetry and /p/events
- Commands from Connect arrive asynchronously via server push (J-prefix for
  JSON commands, G-prefix for GCode, F-prefix for forced GCode) instead of
  as HTTP responses to telemetry
- Registration (/p/register) and camera handling remain on HTTP unchanged
- WS connection failure sets INTERNET error; send failure sets HTTP error

Author: ondratu

prusa/connect/printer/__init__.py

@@ -3,6 +3,7 @@
     Copyright (C) 2024 PrusaResearch
 """
 import configparser
+import json
 import os
 import re
 from logging import getLogger
@@ -11,7 +12,7 @@
 from typing import Any, Callable, Dict, List, Optional
 
 from gcode_metadata import get_metadata
-from requests import RequestException, Response, Session  # type: ignore
+from requests import RequestException, Session  # type: ignore
 from requests.exceptions import (
     ConnectionError as RequestsConnectionError,  # type: ignore
 )
@@ -34,6 +35,7 @@
     Telemetry,
 )
 from .util import RetryingSession, get_timestamp
+from .ws import PrinterWS
 
 __version__ = "0.8.2"
 __date__ = "5 May 2025"  # version date
@@ -96,6 +98,7 @@ def __init__(self,
                  max_retries: int = 1,
                  mmu_supported: bool = True):
         # pylint: disable=too-many-positional-arguments
+        # pylint: disable=too-many-statements
         self.__type = type_
         self.__sn = sn
         self.__fingerprint = fingerprint
@@ -172,6 +175,8 @@ def __init__(self,
                                         self.download_finished_cb)
         self.camera_controller = CameraController(self.conn, self.server,
                                                   self.send_cb)
+        self.ws = PrinterWS(self.handle_ws_message)
+        self.__ws_token: Optional[str] = None
         self.__running_loop = False
 
     @staticmethod
@@ -644,65 +649,106 @@ def wrapper(handler: Callable[[Command], Dict[str, Any]]):
 
         return wrapper
 
-    def parse_command(self, res: Response):
-        """Parse telemetry response.
+    def ws_url(self) -> str:
+        """Return the WebSocket URL derived from self.server."""
+        assert self.server
+        url = self.server.replace("https://", "wss://", 1)
+        return url.replace("http://", "ws://", 1) + "/p/ws"
+
+    def ensure_ws_connected(self) -> None:
+        """Open (or reopen) the WS connection when credentials change."""
+        if not self.server or not self.token:
+            return
+        if self.ws.connected and self.__ws_token == self.token:
+            return
+        self.ws.connect(self.ws_url(), self.make_headers())
+        self.__ws_token = self.token
+
+    def handle_ws_message(self, message: str) -> None:
+        """Process a text frame pushed by the server over the WebSocket.
 
-        When response from connect is command (HTTP Status: 200 OK), it
-        will set a command object, if the printer is initialized properly.
+        Wire format (BuddyEncoder):
+          J{08x command_id}{JSON body}  — high-level JSON command
+          G{08x command_id}{gcode text} — low-level GCode command
+          F{08x command_id}{gcode text} — forced GCode command
+          T…                            — file transfer block (ignored)
         """
-        if res.status_code == 200:
-            command_id: Optional[int] = None
-            command_id_string: str
+        if not message or len(message) < 9:
+            log.error("WS: message too short: %r", message)
+            return
+
+        msg_type = message[0]
+        try:
+            command_id = int(message[1:9], 16)
+        except ValueError:
+            log.error("WS: invalid command_id in: %r", message[:9])
+            self.event_cb(const.Event.REJECTED,
+                          const.Source.CONNECT,
+                          reason="Invalid command_id")
+            return
+
+        body = message[9:]
+
+        if msg_type == 'T':
+            log.debug("WS: file transfer message, ignoring")
+            return
+
+        if not self.is_initialised():
+            self.event_cb(const.Event.REJECTED,
+                          const.Source.WUI,
+                          command_id=command_id,
+                          reason=self.NOT_INITIALISED_MSG)
+            return
+
+        if msg_type in ('G', 'F'):
+            # Low-level GCode command; F prefix means force=True
+            force = msg_type == 'F'
+            command_name = const.Command.GCODE.value
+            log.debug("WS GCode command: id=%s force=%s", command_id, force)
             try:
-                command_id_string = res.headers.get("Command-Id", "")
-                command_id = int(command_id_string)
-            except (TypeError, ValueError):
-                log.error("Invalid Command-Id header. Headers: %s",
-                          res.headers)
+                if self.command.check_state(command_id, command_name):
+                    self.command.accept(command_id,
+                                        command_name=command_name,
+                                        args=[body],
+                                        kwargs={"gcode": body},
+                                        force=force)
+            except Exception as exc:  # pylint: disable=broad-except
+                log.exception("")
                 self.event_cb(const.Event.REJECTED,
                               const.Source.CONNECT,
-                              reason="Invalid Command-Id header")
-                return res
-            if not self.is_initialised():
+                              command_id=command_id,
+                              reason=str(exc))
+
+        elif msg_type == 'J':
+            # High-level JSON command — body is raw Connect HTTP response body
+            try:
+                data = json.loads(body)
+            except json.JSONDecodeError:
+                log.error("WS: invalid JSON body: %r", body)
                 self.event_cb(const.Event.REJECTED,
-                              const.Source.WUI,
+                              const.Source.CONNECT,
                               command_id=command_id,
-                              reason=self.NOT_INITIALISED_MSG)
-                return res
-            content_type = res.headers.get("content-type", "")
-            log.debug("parse_command res: %s", res.text)
+                              reason="Invalid JSON body")
+                return
+
+            command_name = data.get("command", "")
+            log.debug("WS JSON command: id=%s cmd=%s", command_id,
+                      command_name)
             try:
-                if content_type.startswith("application/json"):
-                    data = res.json()
-                    command_name = data.get("command", "")
-                    if self.command.check_state(command_id, command_name):
-                        self.command.accept(command_id,
-                                            command_name=command_name,
-                                            args=data.get("args"),
-                                            kwargs=data.get('kwargs'))
-                elif content_type == "text/x.gcode":
-                    command_name = const.Command.GCODE.value
-                    if self.command.check_state(command_id, command_name):
-                        force = ("Force" in res.headers
-                                 and res.headers["Force"] == "1")
-                        self.command.accept(command_id,
-                                            command_name, [res.text],
-                                            {"gcode": res.text},
-                                            force=force)
-                else:
-                    raise ValueError("Invalid command content type")
-            except Exception as e:  # pylint: disable=broad-except
+                if self.command.check_state(command_id, command_name):
+                    self.command.accept(command_id,
+                                        command_name=command_name,
+                                        args=data.get("args"),
+                                        kwargs=data.get("kwargs"))
+            except Exception as exc:  # pylint: disable=broad-except
                 log.exception("")
                 self.event_cb(const.Event.REJECTED,
                               const.Source.CONNECT,
                               command_id=command_id,
-                              reason=str(e))
-        elif res.status_code == 204:  # no cmd in telemetry
-            pass
+                              reason=str(exc))
+
         else:
-            log.info("Got unexpected telemetry response (%s): %s",
-                     res.status_code, res.text)
-        return res
+            log.warning("WS: unknown message type: %r", msg_type)
 
     def register(self):
         """Register the printer with Connect and return a registration
@@ -759,21 +805,19 @@ def loop(self):
 
     def loop_step(self):
         """
-        Gets an item LoopObject from queue, sends it and handles the response
-        The LoopObject is either an Event - in which case it's just sent,
-        a Telemetry, in which case the response might contain a command to
-        execute, a Register object in which case the response contains the
-        credentials for further communication.
+        Gets a LoopObject from the queue and sends it.
+
+        Telemetry and Event travel over the persistent WebSocket (/p/ws).
+        Commands from Connect arrive asynchronously via handle_ws_message.
+        Register (/p/register) and CameraRegister (/p/camera) remain HTTP.
         """
         # pylint: disable=too-many-branches
         # pylint: disable=too-many-statements
         try:
-            # Get the item to send
             item = self.queue.get(timeout=const.TIMESTAMP_PRECISION)
         except Empty:
             return
 
-        # Make sure we're able to send it
         if not self.server:
             log.warning("Server is not set, skipping item: %s", item)
             return
@@ -786,7 +830,26 @@ def loop_step(self):
             log.warning("No token, skipping item: %s", item)
             return
 
-        # Send it
+        # --- WebSocket path: Telemetry and Event ---
+        if isinstance(item, (Telemetry, Event)):
+            self.ensure_ws_connected()
+            if not self.ws.connected:
+                # Can't reach server at all — network/internet issue
+                errors.INTERNET.ok = False
+                INTERNET.state = CondState.NOK
+                log.warning("WS not connected, dropping item: %s", item)
+                return
+            sent = self.ws.send(item.to_payload())
+            if sent:
+                errors.API.ok = True
+                API.state = CondState.OK
+            else:
+                # WS connected but send failed — HTTP/WS layer issue
+                errors.HTTP.ok = False
+                HTTP.state = CondState.NOK
+            return
+
+        # --- HTTP path: Register and CameraRegister ---
         headers = self.make_headers(item.timestamp)
         try:
             res = item.send(self.conn, self.server, headers)
@@ -807,10 +870,7 @@ def loop_step(self):
             INTERNET.state = CondState.NOK
             log.exception('Unhandled error')
         else:
-            # Handle the response
-            if isinstance(item, Telemetry):
-                self.parse_command(res)
-            elif isinstance(item, Register):
+            if isinstance(item, Register):
                 if res.status_code == 200:
                     self.token = res.headers["Token"]
                     errors.TOKEN.ok = True
@@ -823,7 +883,6 @@ def loop_step(self):
                     sleep(1)
             elif isinstance(item, CameraRegister):
                 camera = item.camera
-                # pylint: disable=unused-argument
                 if res.status_code == 200:
                     camera_token = res.headers["Token"]
                     camera.set_token(camera_token)

prusa/connect/printer/ws.py

@@ -0,0 +1,112 @@
+"""WebSocket transport for Connect SDK — telemetry and events."""
+import json
+from logging import getLogger
+from threading import Event, Lock, Thread
+from typing import Callable, Optional
+
+import websocket  # type: ignore
+
+from . import const
+
+log = getLogger("connect-printer")
+
+
+class PrinterWS:
+    """Persistent WebSocket connection to Connect.
+
+    Runs WebSocketApp.run_forever() in a background daemon thread so that
+    Printer.loop (synchronous/threaded) can call send() without blocking.
+    Incoming frames are forwarded to the on_message callback supplied at
+    construction (Printer.handle_ws_message).
+    """
+
+    def __init__(self, on_message: Callable[[str], None]):
+        self._on_message = on_message
+        self._ws: Optional[websocket.WebSocketApp] = None
+        self._thread: Optional[Thread] = None
+        self._open_event = Event()
+        self._send_lock = Lock()
+
+    @property
+    def connected(self) -> bool:
+        """True when the WebSocket handshake completed and socket is open."""
+        return (self._ws is not None and self._ws.sock is not None
+                and self._ws.sock.connected)
+
+    def connect(self, url: str, headers: dict) -> None:
+        """(Re)connect to *url* using *headers* for the WS handshake.
+
+        Blocks until on_open fires or CONNECTION_TIMEOUT elapses.
+        Safe to call when already connected — disconnects first.
+        """
+        self.disconnect()
+        self._open_event.clear()
+
+        header_list = [
+            f"{key}: {value}" for key, value in headers.items()
+            if value is not None
+        ]
+
+        self._ws = websocket.WebSocketApp(
+            url,
+            header=header_list,
+            on_open=self._handle_open,
+            on_message=self._handle_message,
+            on_error=self._handle_error,
+            on_close=self._handle_close,
+        )
+
+        self._thread = Thread(
+            target=self._ws.run_forever,
+            kwargs={
+                "ping_interval": 30,
+                "ping_timeout": 10,
+            },
+            daemon=True,
+        )
+        self._thread.start()
+
+        if not self._open_event.wait(timeout=const.CONNECTION_TIMEOUT):
+            log.warning("WS connect timed out: %s", url)
+
+    def disconnect(self) -> None:
+        """Close the WebSocket and wait for the background thread to stop."""
+        if self._ws is not None:
+            self._ws.close()
+        if self._thread is not None and self._thread.is_alive():
+            self._thread.join(timeout=const.CONNECTION_TIMEOUT)
+        self._ws = None
+        self._thread = None
+
+    def send(self, payload: dict) -> bool:
+        """Serialize *payload* to JSON and send as a text frame.
+
+        Returns True on success, False when not connected or on any error.
+        Thread-safe (one send at a time via lock).
+        """
+        if not self.connected:
+            log.warning("WS not connected, dropping: %s", payload)
+            return False
+        try:
+            with self._send_lock:
+                self._ws.send(json.dumps(payload))  # type: ignore[union-attr]
+            return True
+        except Exception:  # pylint: disable=broad-except
+            log.exception("WS send failed")
+            return False
+
+    # --- internal WebSocketApp callbacks ---
+
+    def _handle_open(self, _ws) -> None:
+        log.debug("WS connection opened")
+        self._open_event.set()
+
+    def _handle_message(self, _ws, message: str) -> None:
+        self._on_message(message)
+
+    def _handle_error(self, _ws, error) -> None:
+        log.error("WS error: %s", error)
+        self._open_event.set()  # unblock connect() if still waiting
+
+    def _handle_close(self, _ws, close_code, close_msg) -> None:
+        log.debug("WS closed: %s %s", close_code, close_msg)

requirements.txt

@@ -3,3 +3,4 @@ requests>=2.32.3
 inotify_simple~=1.3.5
 mypy-extensions~=1.0.0
 urllib3>=1.21.1,<3
+websocket-client>=1.6.0