ws client hardening

Author: tmunzer-AIDE

README.md

@@ -32,7 +32,7 @@ A comprehensive Python package to interact with the Mist Cloud APIs, built from
     - [Examples](#examples)
 - [WebSocket Streaming](#websocket-streaming)
     - [Connection Parameters](#connection-parameters)
-    - [Callbacks](#callbacks)
+    - [Methods](#methods)
     - [Available Channels](#available-channels)
     - [Usage Patterns](#usage-patterns)
 - [Async Usage](#async-usage)
@@ -588,6 +588,7 @@ All channel classes accept the following optional keyword arguments:
 | `auto_reconnect` | `bool` | `False` | Automatically reconnect on transient failures using exponential backoff. |
 | `max_reconnect_attempts` | `int` | `5` | Maximum number of reconnect attempts before giving up. |
 | `reconnect_backoff` | `float` | `2.0` | Base backoff delay in seconds. Doubles after each failed attempt (2s, 4s, 8s, ...). Resets on successful reconnection. |
+| `queue_maxsize` | `int` | `0` | Maximum messages buffered in the internal queue for `receive()`. `0` means unbounded. When set, the receive thread blocks if the queue is full, providing backpressure for high-frequency streams. |
 
 ```python
 ws = mistapi.websockets.sites.DeviceStatsEvents(
@@ -600,15 +601,18 @@ ws = mistapi.websockets.sites.DeviceStatsEvents(
 ws.connect()
 ```
 
-### Callbacks
+### Methods
 
 | Method | Signature | Description |
 |--------|-----------|-------------|
-| `ws.on_open(cb)` | `cb()` | Called when the connection is established |
-| `ws.on_message(cb)` | `cb(data: dict)` | Called for every incoming message |
-| `ws.on_error(cb)` | `cb(error: Exception)` | Called on WebSocket errors |
-| `ws.on_close(cb)` | `cb(status_code: int, msg: str)` | Called when the connection closes |
-| `ws.ready()` | `-> bool \| None` | Returns `True` if the connection is open and ready |
+| `ws.on_open(cb)` | `cb()` | Register callback for connection established |
+| `ws.on_message(cb)` | `cb(data: dict)` | Register callback for incoming messages. Mutually exclusive with `receive()`. |
+| `ws.on_error(cb)` | `cb(error: Exception)` | Register callback for WebSocket errors |
+| `ws.on_close(cb)` | `cb(code: int, msg: str)` | Register callback for connection close. Safe to call `connect()` from within. |
+| `ws.connect(run_in_background)` | | Open the connection. `True` (default) runs in a daemon thread; `False` blocks. |
+| `ws.disconnect(wait, timeout)` | | Close the connection. `wait=True` blocks until the background thread finishes. |
+| `ws.receive()` | `-> Generator[dict]` | Blocking generator yielding messages. Mutually exclusive with `on_message`. |
+| `ws.ready()` | `-> bool` | Returns `True` if the connection is open and ready |
 
 ### Available Channels
 

src/mistapi/websockets/__ws_client.py

@@ -13,6 +13,7 @@
 """
 
 import json
+import logging
 import queue
 import ssl
 import threading
@@ -23,6 +24,10 @@
 
 from mistapi.__logger import logger
 
+# Prevent websocket-client from logging HTTP headers (which contain API
+# tokens) at DEBUG level.
+logging.getLogger("websocket").setLevel(logging.WARNING)
+
 if TYPE_CHECKING:
     from mistapi import APISession
 
@@ -48,6 +53,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         if max_reconnect_attempts < 0:
             raise ValueError("max_reconnect_attempts must be >= 0")
@@ -64,7 +70,7 @@ def __init__(
         self._lock = threading.Lock()
         self._ws: websocket.WebSocketApp | None = None
         self._thread: threading.Thread | None = None
-        self._queue: queue.Queue[dict | None] = queue.Queue()
+        self._queue: queue.Queue[dict | None] = queue.Queue(maxsize=queue_maxsize)
         self._connected = (
             threading.Event()
         )  # tracks whether the WebSocket connection is currently open

src/mistapi/websockets/location.py

@@ -39,7 +39,11 @@ class BleAssetsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
-
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -76,6 +80,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         channels = [f"/sites/{site_id}/stats/maps/{mid}/assets" for mid in map_ids]
         super().__init__(
@@ -86,6 +91,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -113,6 +119,11 @@ class ConnectedClientsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -149,6 +160,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         channels = [f"/sites/{site_id}/stats/maps/{mid}/clients" for mid in map_ids]
         super().__init__(
@@ -159,6 +171,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -186,6 +199,11 @@ class SdkClientsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -222,6 +240,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         channels = [f"/sites/{site_id}/stats/maps/{mid}/sdkclients" for mid in map_ids]
         super().__init__(
@@ -232,6 +251,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -259,6 +279,11 @@ class UnconnectedClientsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -295,6 +320,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         channels = [
             f"/sites/{site_id}/stats/maps/{mid}/unconnected_clients" for mid in map_ids
@@ -307,6 +333,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -334,6 +361,11 @@ class DiscoveredBleAssetsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -370,6 +402,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         channels = [
             f"/sites/{site_id}/stats/maps/{mid}/discovered_assets" for mid in map_ids
@@ -382,4 +415,5 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )

src/mistapi/websockets/orgs.py

@@ -37,6 +37,11 @@ class InsightsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -72,6 +77,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         super().__init__(
             mist_session,
@@ -81,6 +87,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -106,6 +113,11 @@ class MxEdgesStatsEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -141,6 +153,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         super().__init__(
             mist_session,
@@ -150,6 +163,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
 
@@ -175,6 +189,11 @@ class MxEdgesEvents(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -210,6 +229,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         super().__init__(
             mist_session,
@@ -219,4 +239,5 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )

src/mistapi/websockets/session.py

@@ -44,6 +44,11 @@ class SessionWithUrl(_MistWebsocket):
         Maximum number of reconnect attempts before giving up.
     reconnect_backoff : float, default 2.0
         Base backoff delay in seconds. Doubles after each failed attempt.
+    queue_maxsize : int, default 0
+        Maximum number of messages buffered in the internal queue for the
+        ``receive()`` generator. ``0`` means unbounded. When set, the
+        websocket-client receive thread blocks if the queue is full,
+        providing backpressure for high-frequency streams.
 
     EXAMPLE
     -----------
@@ -79,6 +84,7 @@ def __init__(
         auto_reconnect: bool = False,
         max_reconnect_attempts: int = 5,
         reconnect_backoff: float = 2.0,
+        queue_maxsize: int = 0,
     ) -> None:
         if not url.startswith("wss://"):
             raise ValueError("url must use the wss:// scheme")
@@ -91,6 +97,7 @@ def __init__(
             auto_reconnect=auto_reconnect,
             max_reconnect_attempts=max_reconnect_attempts,
             reconnect_backoff=reconnect_backoff,
+            queue_maxsize=queue_maxsize,
         )
 
     def _build_ws_url(self) -> str: