"""WebSocket (Socket.IO) client for BragerOne realtime events."""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Awaitable, Callable
from contextlib import suppress
from typing import (
Any,
Protocol,
TypedDict,
runtime_checkable,
)
import socketio
from ..models.events import MODULE_CONNECTION_STATUS_CHANGED, MODULE_MEMORY_UPDATED
from ..utils import spawn
from .client import format_expected_failure_reason, is_expected_upstream_unavailable
from .constants import IO_BASE, ONE_BASE, SOCK_PATH, WS_NAMESPACE
log = logging.getLogger(__name__)
sio_log = logging.getLogger(__name__ + ".sio")
eio_log = logging.getLogger(__name__ + ".eio")
def _is_expected_ws_reconnect_failure(err: BaseException) -> bool:
"""Return whether a reconnect failure is an expected upstream/transient error."""
if is_expected_upstream_unavailable(err):
return True
if isinstance(err, TimeoutError):
return True
current: BaseException | None = err
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
module = getattr(type(current), "__module__", "") or ""
if isinstance(current, ConnectionError) and (module.startswith("socketio") or module.startswith("engineio")):
return True
current = current.__cause__ or (current.__context__ if current.__context__ is not current.__cause__ else None)
return False
# Signature for a generic event handler used by the gateway.
EventHandler = Callable[[str, Any], None]
# Signature for a callback invoked on connection establishment.
ConnectedCb = Callable[[], Awaitable[None] | None]
[docs]
@runtime_checkable
class EventDispatcher(Protocol):
"""Protocol for an event dispatcher used by the realtime manager."""
def __call__(self, event_name: str, payload: Any) -> Awaitable[None] | None:
"""Handle an event with the given name and payload.
Args:
event_name: The event name.
payload: The event payload (varies by event).
"""
...
class _SubPayload(TypedDict, total=False):
"""Payload for subscription events."""
modules: list[str]
devids: list[str]
group_id: int
[docs]
class RealtimeManager:
"""Thin Socket.IO wrapper for BragerOne realtime channel.
The manager keeps a single AsyncClient connection, exposes the Engine.IO
SID and the namespace SID, and forwards selected events to a user-provided
callback (``EventHandler``). It does **not** interpret payloads; that is the
gateway's responsibility.
Notes:
- Authentication is provided **only** via HTTP headers (Bearer token).
- We always connect to the `:data:`~.constants.WS_NAMESPACE`` namespace.
- We listen to: ``snapshot`` and the various ``*:parameters:change`` events.
- Subscriptions are emitted in a few payload variants (``modules`` /
``devids``) and optionally include ``group_id``.
"""
def __init__(
self,
token: str,
*,
origin: str = ONE_BASE,
referer: str = f"{ONE_BASE}/",
io_base: str = IO_BASE,
socket_path: str = SOCK_PATH,
namespace: str = WS_NAMESPACE,
token_provider: Callable[[], Awaitable[str]] | None = None,
connect_timeout_s: float = 20.0,
) -> None:
"""Initialize the realtime manager.
Args:
token: Bearer token used for the initial Socket.IO HTTP upgrade.
origin: HTTP ``Origin`` header value (default: :data:`~.constants.ONE_BASE`).
referer: HTTP ``Referer`` header value (default: :data:`~.constants.ONE_BASE` + ``/``).
io_base: Base URL of the Engine.IO/Socket.IO server (default: :data:`~.constants.IO_BASE`).
socket_path: Socket.IO path on the server (default: :data:`~.constants.SOCK_PATH`).
namespace: The namespace to join (default: :data:`~.constants.WS_NAMESPACE`).
token_provider: Optional async callable returning a fresh access token (the
``Bearer `` prefix is added by the manager). When set, every (re)connect
attempt resolves the token through it, so long outages that outlive the
token TTL can still recover. Falls back to the static ``token`` when unset
or on error.
connect_timeout_s: Maximum seconds a single connect attempt may take before it is aborted
and retried by the supervisor (default: 20). Guards against hung TCP/DNS handshakes.
"""
self._token = token
self._token_provider = token_provider
self._connect_timeout_s = connect_timeout_s
self._origin = origin
self._referer = referer
self._io_base = io_base.rstrip("/")
self._socket_path = socket_path
self._namespace = namespace
self._connected = asyncio.Event()
self._on_connected: list[ConnectedCb] = []
self._on_disconnected: list[ConnectedCb] = []
self._on_event: EventDispatcher | None = None
self._modules: list[str] = []
self._group_id: int | None = None
self._connect_lock = asyncio.Lock()
self._supervisor_task: asyncio.Task[None] | None = None
self._supervisor_running = False
self._supervisor_interval_s = 15.0
# Cap leftover Engine.IO teardown so an aborted websocket cannot wedge reconnect.
self._disconnect_timeout_s = max(0.05, min(5.0, float(connect_timeout_s)))
# True until the first successful connect, then after each disconnect notify.
self._disconnect_notified = True
# Reconnection is owned exclusively by our supervisor loop (with fresh-token
# resolution and a hard connect timeout); the built-in socket.io reconnect
# loop would bypass both and race with the supervisor, so it stays disabled.
self._sio: socketio.AsyncClient = self._make_sio()
self._register_sio_handlers()
# ---- Built-in handlers ----
async def _on_connect(self) -> None:
ns_sid = self.sid()
eng_sid = self.engine_sid()
log.info("WS connected, SID=%s", eng_sid)
log.info(
"WS connected, namespace_sid=%s, engine_sid=%s, namespaces=%s",
ns_sid,
eng_sid,
list(self._sio.namespaces),
)
for cb in list(self._on_connected):
try:
res = cb()
if asyncio.iscoroutine(res):
spawn(res, "on_connected_cb", log)
except Exception:
log.exception("Error in on_connected callback")
self._disconnect_notified = False
self._connected.set()
async def _on_disconnect(self) -> None:
log.info("WS disconnected")
self._connected.clear()
self._notify_disconnected()
async def _on_connect_error(self, data: Any | None = None) -> None:
log.warning("WS connect_error: %s", data)
was_connected = self._connected.is_set()
self._connected.clear()
if was_connected:
self._notify_disconnected()
def _notify_disconnected(self, *, force: bool = False) -> None:
"""Invoke disconnect callbacks (sync or async).
Args:
force: When True, notify even if a previous drop was already reported.
The supervisor reconnect loop and Socket.IO ``disconnect`` pass False
so a wedged client does not spam session-down callbacks.
"""
if self._disconnect_notified and not force:
return
self._disconnect_notified = True
for cb in list(self._on_disconnected):
try:
res = cb()
if asyncio.iscoroutine(res):
spawn(res, "on_disconnected_cb", log)
except Exception:
log.exception("Error in on_disconnected callback")
async def _on_reconnect(self) -> None:
log.info("WS reconnect OK")
async def _on_reconnect_attempt(self, number: int) -> None:
log.info("WS reconnect attempt #%s", number)
async def _on_reconnect_error(self, data: Any | None = None) -> None:
log.warning("WS reconnect_error: %s", data)
async def _on_error(self, data: Any) -> None:
log.error("WS ERROR: %s", data)
async def _on_message(self, data: Any) -> None:
log.debug("WS message → %s", data)
# --- Key domain events ---
async def _on_snapshot(self, payload: Any) -> None:
log.debug("WS EVENT snapshot → %s", payload)
self._dispatch("snapshot", payload)
async def _on_app_modules_parameters_change(self, payload: Any) -> None:
log.debug("WS EVENT app:modules:parameters:change → %s", payload)
self._dispatch("app:modules:parameters:change", payload)
# Fallback alt names occasionally seen in traces
async def _on_modules_parameters_change(self, payload: Any) -> None:
log.debug("WS EVENT modules:parameters:change → %s", payload)
self._dispatch("modules:parameters:change", payload)
async def _on_parameters_change(self, payload: Any) -> None:
log.debug("WS EVENT parameters:change → %s", payload)
self._dispatch("parameters:change", payload)
# Optional task-related events; kept for completeness/diagnostics.
async def _on_app_modules_task_created(self, p: Any) -> None:
log.debug("WS EVENT app:module:task:created → %s", p)
self._dispatch("app:module:task:created", p)
async def _on_app_modules_task_status_changed(self, p: Any) -> None:
log.debug("WS EVENT app:module:task:status:changed → %s", p)
self._dispatch("app:module:task:status:changed", p)
async def _on_app_modules_task_completed(self, p: Any) -> None:
log.debug("WS EVENT app:module:task:completed → %s", p)
self._dispatch("app:module:task:completed", p)
async def _on_app_module_connection_status_changed(self, p: Any) -> None:
log.debug("WS EVENT %s → %s", MODULE_CONNECTION_STATUS_CHANGED, p)
self._dispatch(MODULE_CONNECTION_STATUS_CHANGED, p)
# Numeric fallbacks observed in some builds
async def _on_ev60(self, p: Any) -> None:
log.debug("WS EVENT 60 → %s", p)
self._dispatch("app:module:task:status:changed", p)
async def _on_ev61(self, p: Any) -> None:
log.debug("WS EVENT 61 → %s", p)
self._dispatch("app:module:task:created", p)
async def _on_ev63(self, p: Any) -> None:
log.debug("WS EVENT 63 → %s", p)
self._dispatch("app:module:task:completed", p)
async def _on_module_memory_updated(self, p: Any) -> None:
log.debug("WS EVENT %s → %s", MODULE_MEMORY_UPDATED, p)
self._dispatch(MODULE_MEMORY_UPDATED, p)
# ---------------- Public API ----------------
[docs]
async def connect(self) -> None:
"""Open a Socket.IO connection with appropriate headers and wait for join."""
await self._ensure_connected(initial=True)
# Short grace period to ensure the namespace is fully established
await self._connected.wait()
await asyncio.sleep(0.1)
self._start_supervisor()
[docs]
async def force_reconnect(self) -> None:
"""Tear down a still-"connected" Socket.IO session and reconnect.
Used when ParamUpdates go silent while the client still reports up
(zombie Engine.IO that skipped disconnect callbacks). The SPA recovers
via built-in Socket.IO reconnect then ``connect`` →
``ModulesService.connect`` + REST ``/modules/parameters``; our supervisor
only acts when ``connected`` looks down, so this forces that same path
(``on_connected`` → gateway resubscribe + prime).
Waits for the namespace join (same as :meth:`connect`) so callers can
safely ``sid()`` / ``modules.connect`` immediately afterwards.
"""
log.warning("Forcing WS hard reconnect (zombie session recovery)")
self._notify_disconnected(force=True)
# Clear the namespace-joined bit so ``_ensure_connected`` does not treat a
# wedged ``sio.connected=True`` session as healthy.
self._connected.clear()
try:
await self._ensure_connected(initial=False)
try:
await asyncio.wait_for(self._connected.wait(), timeout=self._connect_timeout_s)
except TimeoutError:
log.warning("WS force reconnect: namespace join timed out")
return
await asyncio.sleep(0.1)
except Exception:
# Defensive: callers (gateway poll) must keep running after a failed force.
log.exception("WS force reconnect failed unexpectedly")
[docs]
async def hard_reset(self) -> None:
"""Abandon the Socket.IO client and open a brand-new transport session.
Stronger than :meth:`force_reconnect`: stops the supervisor, replaces the
``AsyncClient`` (handlers re-bound), then runs a full :meth:`connect`.
Used when disconnect/reconnect on the same client fails to restore live
``ParamUpdate`` traffic.
"""
log.warning("Hard-resetting WS transport (zombie session recovery)")
self._notify_disconnected(force=True)
self._supervisor_running = False
task = self._supervisor_task
self._supervisor_task = None
if task is not None and not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
self._connected.clear()
try:
await asyncio.wait_for(self._sio.disconnect(), timeout=self._disconnect_timeout_s)
except Exception:
log.debug("WS hard_reset disconnect ignored", exc_info=True)
self._replace_client()
await self.connect()
def _start_supervisor(self) -> None:
if self._supervisor_task is not None and not self._supervisor_task.done():
return
self._supervisor_running = True
task = asyncio.create_task(
self._connection_supervisor(),
name="ws_connection_supervisor",
)
self._supervisor_task = task
def _supervisor_done(done_task: asyncio.Task[None]) -> None:
with suppress(asyncio.CancelledError):
done_task.result()
if self._supervisor_task is done_task:
self._supervisor_task = None
task.add_done_callback(_supervisor_done)
def _make_sio(self) -> socketio.AsyncClient:
"""Create a Socket.IO client with library reconnect disabled."""
return socketio.AsyncClient(
reconnection=False,
logger=sio_log, # pyright: ignore[reportArgumentType] # route socket.io logs to a sub-logger
engineio_logger=eio_log, # route engine.io logs to a sub-logger
)
def _register_sio_handlers(self) -> None:
"""Bind namespace handlers on the current Socket.IO client."""
ns = self._namespace
self._sio.on("connect", self._on_connect, namespace=ns)
self._sio.on("disconnect", self._on_disconnect, namespace=ns)
self._sio.on("connect_error", self._on_connect_error, namespace=ns)
self._sio.on("reconnect", self._on_reconnect, namespace=ns)
self._sio.on("reconnect_attempt", self._on_reconnect_attempt, namespace=ns)
self._sio.on("reconnect_error", self._on_reconnect_error, namespace=ns)
self._sio.on("error", self._on_error, namespace=ns)
self._sio.on("message", self._on_message, namespace=ns)
self._sio.on("snapshot", self._on_snapshot, namespace=ns)
self._sio.on(
"app:modules:parameters:change",
self._on_app_modules_parameters_change,
namespace=ns,
)
self._sio.on(
"modules:parameters:change", # fallback alt name
self._on_modules_parameters_change,
namespace=ns,
)
self._sio.on(
"parameters:change",
self._on_parameters_change,
namespace=ns,
)
self._sio.on(
"app:module:task:created",
self._on_app_modules_task_created,
namespace=ns,
)
self._sio.on(
"app:module:task:status:changed",
self._on_app_modules_task_status_changed,
namespace=ns,
)
self._sio.on(
"app:module:task:completed",
self._on_app_modules_task_completed,
namespace=ns,
)
self._sio.on(
MODULE_CONNECTION_STATUS_CHANGED,
self._on_app_module_connection_status_changed,
namespace=ns,
)
self._sio.on(
MODULE_MEMORY_UPDATED,
self._on_module_memory_updated,
namespace=ns,
)
self._sio.on("60", self._on_ev60, namespace=ns)
self._sio.on("61", self._on_ev61, namespace=ns)
self._sio.on("63", self._on_ev63, namespace=ns)
def _abandon_client(self, old: socketio.AsyncClient) -> None:
"""Best-effort abort of a wedged Engine.IO client without blocking reconnect."""
eio = getattr(old, "eio", None)
disconnect = getattr(eio, "disconnect", None) if eio is not None else None
if not callable(disconnect):
return
try:
result = disconnect(abort=True)
except TypeError:
try:
result = disconnect()
except Exception:
log.exception("Failed to abort leftover Engine.IO client")
return
except Exception:
log.exception("Failed to abort leftover Engine.IO client")
return
if asyncio.iscoroutine(result):
spawn(result, "eio_abort_disconnect", log)
def _replace_client(self) -> None:
"""Swap in a fresh Socket.IO client after a hung/aborted transport."""
old = self._sio
self._sio = self._make_sio()
self._register_sio_handlers()
self._abandon_client(old)
async def _reset_transport(self) -> None:
"""Drop a half-open Engine.IO session so the next ``connect()`` is clean.
Engine.IO abort (``WSMsgType.CLOSED``) can deadlock ``disconnect()`` waiting
for read/write loops. Bound that wait, then replace the client if it is still
wedged so the supervisor can reconnect.
"""
self._connected.clear()
try:
await asyncio.wait_for(self._sio.disconnect(), timeout=self._disconnect_timeout_s)
except TimeoutError:
log.warning("WS transport disconnect timed out; replacing Socket.IO client")
self._replace_client()
except Exception:
log.warning("WS transport disconnect failed; replacing Socket.IO client", exc_info=True)
self._replace_client()
async def _resolve_token(self) -> str:
"""Return a token for the next connect attempt, refreshing via the provider if set."""
if self._token_provider is None:
return self._token
try:
# Bound the provider too: re-authentication can stall during the same
# outage, and it must not wedge the supervisor before connect() even runs.
token = await asyncio.wait_for(self._token_provider(), timeout=self._connect_timeout_s)
except Exception:
# Auth backend unreachable or stalled — retry with the last known token.
log.warning("WS token refresh failed; falling back to the previous token", exc_info=True)
return self._token
self._token = token
return token
async def _ensure_connected(self, *, initial: bool) -> None:
if self._sio.connected and self._connected.is_set():
return
async with self._connect_lock:
if self._sio.connected and self._connected.is_set():
return
# Engine.IO abort / connect timeout can leave ``connected=True`` without a
# namespace join. A second ``connect()`` then hangs or no-ops — drop first.
await self._reset_transport()
token = await self._resolve_token()
headers = {
"Authorization": f"Bearer {token}",
"Origin": self._origin,
"Referer": self._referer,
"Accept-Language": "pl-PL,pl;q=0.9",
"x-appversion": "1.1.78",
}
try:
# Hard timeout: a hung TCP/DNS handshake must not wedge the supervisor forever.
await asyncio.wait_for(
self._sio.connect(
self._io_base,
headers=headers,
transports=["pooling", "websocket"],
socketio_path=self._socket_path,
namespaces=[self._namespace],
),
timeout=self._connect_timeout_s,
)
except Exception as err:
await self._reset_transport()
if initial:
raise
if _is_expected_ws_reconnect_failure(err):
log.warning(
"WS supervisor reconnect failed (expected upstream/transient): %s",
format_expected_failure_reason(err),
)
else:
log.warning("WS supervisor reconnect failed", exc_info=True)
async def _connection_supervisor(self) -> None:
try:
while self._supervisor_running:
await asyncio.sleep(self._supervisor_interval_s)
if self._sio.connected and self._connected.is_set():
continue
log.warning("WS disconnected state detected; forcing reconnect")
# Engine.IO abort often skips the Socket.IO disconnect callback.
# Notify now so CloudSession goes down and REST-prime can run even if
# leftover disconnect() is still wedged.
self._notify_disconnected()
try:
await self._ensure_connected(initial=False)
except Exception:
# Defensive: the supervisor must never die on an unexpected error.
log.exception("WS supervisor iteration failed unexpectedly")
except asyncio.CancelledError:
raise
[docs]
def sid(self) -> str | None:
"""Return the namespace SID (``/ws``), if available."""
try:
return str(self._sio.get_sid(self._namespace))
except Exception:
return None
[docs]
def engine_sid(self) -> str | None:
"""Return the underlying Engine.IO SID (transport-level)."""
return getattr(self._sio, "sid", None)
[docs]
async def disconnect(self) -> None:
"""Close the Socket.IO connection if open."""
self._supervisor_running = False
task = self._supervisor_task
self._supervisor_task = None
if task is not None and not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
if self._sio.connected:
await self._sio.disconnect()
[docs]
def on_event(self, handler: EventDispatcher) -> None:
"""Register a single event dispatcher (gateway attaches here)."""
self._on_event = handler
[docs]
def add_on_connected(self, cb: ConnectedCb) -> None:
"""Register a callback to be called when the connection is established."""
self._on_connected.append(cb)
[docs]
def add_on_disconnected(self, cb: ConnectedCb) -> None:
"""Register a callback to be called when the Socket.IO session drops."""
self._on_disconnected.append(cb)
@property
def group_id(self) -> int | None:
"""Return the optional ``group_id`` included in subscription payloads."""
return self._group_id
@group_id.setter
def group_id(self, group_id: int | None) -> None:
"""Set an optional ``group_id`` to be included in subscription payloads."""
self._group_id = group_id
[docs]
async def subscribe(self, modules: list[str]) -> None:
"""Emit listen events for the provided devices (devids/modules)."""
self._modules = sorted(set(modules))
if not self._modules:
return
# Base variants: "modules" and alt. "devids"
base: _SubPayload = {"modules": self._modules}
base_alt: _SubPayload = {"devids": self._modules}
# Optionally include "group_id"
if self.group_id is not None:
base["group_id"] = self.group_id
base_alt["group_id"] = self.group_id
payloads: list[tuple[str, _SubPayload]] = [
("app:modules:parameters:listen", base),
("app:modules:parameters:listen", base_alt),
("app:modules:activity:quantity:listen", base),
("app:modules:activity:quantity:listen", base_alt),
("app:modules:alarms:quantity:listen", base),
("app:modules:alarms:quantity:listen", base_alt),
]
for ev, pl in payloads:
log.debug("EMIT %s %s", ev, pl)
try:
await self._sio.emit(ev, pl, namespace=self._namespace)
except Exception:
log.exception("Emit failed for %s", ev)
[docs]
async def resubscribe(self) -> None:
"""Re-emit subscription events after a reconnect."""
if self._modules:
await self.subscribe(self._modules)
# ---------------- Internal ----------------
def _dispatch(self, name: str, payload: Any) -> None:
"""Forward an event to the registered dispatcher (if any)."""
if self._on_event:
try:
self._on_event(name, payload)
except Exception:
log.exception("Error in on_event(%s) callback", name)