continuum-task-server-sdk 1.1.0__py3-none-any.whl

continuum_task_server/websocket.py

@@ -0,0 +1,852 @@
+"""Synchronous STOMP-over-WebSocket client for the Queue API worker protocol."""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import time
+import uuid
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Protocol
+from urllib.parse import urljoin, urlparse
+from uuid import UUID
+
+from . import _stomp
+from ._http import drop_none
+from .client import _encode_input
+from .exceptions import (
+    AmbiguousCommandError,
+    CommandTimeoutError,
+    ContinuumError,
+    ProtocolError,
+    TransportError,
+    error_for_status,
+)
+from .models import (
+    EnqueueAndSubscribeWaitResult,
+    EventAckResult,
+    QueueEvent,
+    QueueEventType,
+    QueueItem,
+    TaskStatus,
+    WaitMode,
+    WaitSubscriptionResult,
+    WorkSubscriptionResult,
+)
+
+logger = logging.getLogger("continuum_task_server")
+
+Connector = Callable[[str, dict[str, str], float], "SocketLike"]
+EventCallback = Callable[[QueueEvent], None]
+
+
+class SocketLike(Protocol):
+    def send(self, data: str) -> None: ...
+
+    def recv(self) -> str | bytes: ...
+
+    def close(self) -> None: ...
+
+
+@dataclass
+class WebSocketOptions:
+    """Tunables for :class:`ContinuumWebSocketClient`."""
+
+    endpoint: str = "/ws/queue"
+    connect_timeout: float = 10.0
+    request_timeout: float = 30.0
+    heartbeat_ms: int = 10_000
+    reconnect_backoff: float = 1.0
+    max_reconnect_backoff: float = 30.0
+    auto_reconnect: bool = True
+
+
+def derive_websocket_url(base_url: str, endpoint: str = "/ws/queue") -> str:
+    """Map an HTTP(S) Queue API base URL to the native STOMP WebSocket URL.
+
+    ``http://`` stays ``ws://`` and ``https://`` stays ``wss://``. The path
+    ``endpoint`` is appended to the origin (not to extra HTTP path prefixes).
+    """
+    if not base_url:
+        raise ValueError("base_url must not be empty")
+    parsed = urlparse(base_url)
+    scheme = parsed.scheme.lower()
+    if scheme in ("https", "wss"):
+        ws_scheme = "wss"
+    elif scheme in ("http", "ws"):
+        ws_scheme = "ws"
+    else:
+        raise ValueError(f"unsupported URL scheme: {parsed.scheme!r}")
+    path = endpoint if endpoint.startswith("/") else "/" + endpoint
+    netloc = parsed.netloc
+    if not netloc:
+        raise ValueError(f"base_url has no host: {base_url!r}")
+    return urljoin(f"{ws_scheme}://{netloc}", path)
+
+
+class _Pending:
+    def __init__(self) -> None:
+        self.event = threading.Event()
+        self.payload: dict[str, Any] | None = None
+        self.error: Exception | None = None
+
+
+class ContinuumWebSocketClient:
+    """Synchronous Queue API STOMP session.
+
+    Authenticate the WebSocket upgrade with ``Api-Key``. Commands are correlated
+    by ``requestId``. ``work.available`` hints are not reservations; ``claim``
+    remains the arbiter. Durable ``queue.completed`` deliveries must be
+    acknowledged with :meth:`ack_event`.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        *,
+        options: WebSocketOptions | None = None,
+        connector: Connector | None = None,
+    ) -> None:
+        if not base_url:
+            raise ValueError("base_url must not be empty")
+        if not api_key:
+            raise ValueError("api_key must not be empty")
+        self._base_url = base_url.rstrip("/")
+        self._api_key = api_key
+        self._options = options or WebSocketOptions()
+        self._connector = connector or _default_connector
+        self._socket: SocketLike | None = None
+        self._send_lock = threading.Lock()
+        self._state_lock = threading.Lock()
+        self._pending: dict[str, _Pending] = {}
+        self._events: list[QueueEvent] = []
+        self._events_cv = threading.Condition()
+        self._event_callbacks: list[EventCallback] = []
+        self._work_subscriptions: set[str] = set()
+        self._wait_subscribers: set[str] = set()
+        self._seen_completed: OrderedDict[str, None] = OrderedDict()
+        self._connected = threading.Event()
+        self._closed = threading.Event()
+        self._receiver: threading.Thread | None = None
+        self._heartbeat_thread: threading.Thread | None = None
+        self._heartbeat_stop = threading.Event()
+        self._decoder = _stomp.StompDecoder()
+        self._reconnect_attempt = 0
+        self._client_heartbeat_ms = self._options.heartbeat_ms
+        self._generation = 0
+
+    @property
+    def connected(self) -> bool:
+        return self._connected.is_set() and not self._closed.is_set()
+
+    def connect(self) -> None:
+        """Open the WebSocket and complete the STOMP handshake.
+
+        Idempotent while already connected. Blocks until ``CONNECTED`` or
+        ``connect_timeout``.
+        """
+        if self._closed.is_set():
+            self._closed.clear()
+        if self._connected.is_set() and self._socket is not None:
+            return
+        self._start_receiver()
+        if not self._connected.wait(timeout=self._options.connect_timeout):
+            raise TransportError("timed out connecting to Queue API WebSocket")
+
+    def close(self) -> None:
+        self._closed.set()
+        self._heartbeat_stop.set()
+        self._connected.clear()
+        self._fail_pending(TransportError("WebSocket client closed"))
+        sock = self._socket
+        self._socket = None
+        if sock is not None:
+            try:
+                sock.close()
+            except Exception:
+                logger.debug("error closing websocket", exc_info=True)
+        with self._events_cv:
+            self._events_cv.notify_all()
+        if self._receiver is not None and self._receiver is not threading.current_thread():
+            self._receiver.join(timeout=5.0)
+        if (
+            self._heartbeat_thread is not None
+            and self._heartbeat_thread is not threading.current_thread()
+        ):
+            self._heartbeat_thread.join(timeout=2.0)
+        self._receiver = None
+        self._heartbeat_thread = None
+
+    def __enter__(self) -> ContinuumWebSocketClient:
+        self.connect()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def on_event(self, callback: EventCallback) -> Callable[[], None]:
+        """Register a callback invoked on the receiver thread. Returns unsubscribe."""
+        self._event_callbacks.append(callback)
+
+        def _off() -> None:
+            try:
+                self._event_callbacks.remove(callback)
+            except ValueError:
+                pass
+
+        return _off
+
+    def next_event(self, timeout: float | None = None) -> QueueEvent | None:
+        """Block until the next event, or return ``None`` on timeout / close."""
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with self._events_cv:
+            while True:
+                if self._events:
+                    return self._events.pop(0)
+                if self._closed.is_set():
+                    return None
+                remaining = None
+                if deadline is not None:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        return None
+                self._events_cv.wait(timeout=remaining)
+
+    def enqueue(
+        self,
+        *,
+        task_name: str | None = None,
+        task_item_name: str | None = None,
+        input_data: Any = None,
+        priority: int | None = None,
+        parent_id: UUID | str | None = None,
+        request_id: str | None = None,
+        idempotency_token: str | None = None,
+    ) -> QueueItem:
+        if task_name is None and task_item_name is None:
+            raise ValueError("at least one of task_name or task_item_name is required")
+        body = drop_none(
+            {
+                "parent_id": str(parent_id) if parent_id else None,
+                "taskName": task_name,
+                "taskItemName": task_item_name,
+                "priority": priority,
+                "inputData": _encode_input(input_data),
+                "idempotencyToken": idempotency_token,
+            }
+        )
+        payload = self._command("/app/enqueue", body, request_id=request_id, retry=False)
+        item = payload.get("item")
+        if item is None:
+            raise ContinuumError("enqueue succeeded without an item")
+        return QueueItem.model_validate(item)
+
+    def enqueue_and_subscribe_wait(
+        self,
+        *,
+        subscriber_id: str,
+        task_name: str | None = None,
+        task_item_name: str | None = None,
+        input_data: Any = None,
+        priority: int | None = None,
+        parent_id: UUID | str | None = None,
+        mode: WaitMode | str | None = None,
+        correlation: str | None = None,
+        request_id: str | None = None,
+        idempotency_token: str | None = None,
+    ) -> EnqueueAndSubscribeWaitResult:
+        if task_name is None and task_item_name is None:
+            raise ValueError("at least one of task_name or task_item_name is required")
+        if not subscriber_id or not str(subscriber_id).strip():
+            raise ValueError("subscriber_id is required")
+        mode_value = mode.value if isinstance(mode, WaitMode) else mode
+        body = drop_none(
+            {
+                "parent_id": str(parent_id) if parent_id else None,
+                "taskName": task_name,
+                "taskItemName": task_item_name,
+                "priority": priority,
+                "inputData": _encode_input(input_data),
+                "subscriberId": subscriber_id,
+                "mode": mode_value,
+                "correlation": correlation,
+                "idempotencyToken": idempotency_token,
+            }
+        )
+        payload = self._command(
+            "/app/enqueueAndSubscribeWait", body, request_id=request_id, retry=False
+        )
+        result = EnqueueAndSubscribeWaitResult.model_validate(payload)
+        if result.item is None:
+            raise ContinuumError("enqueueAndSubscribeWait succeeded without an item")
+        with self._state_lock:
+            self._wait_subscribers.add(subscriber_id)
+        return result
+
+    def claim(self, task_name: str, *, request_id: str | None = None) -> QueueItem | None:
+        payload = self._command(
+            "/app/claim",
+            {"taskName": task_name},
+            request_id=request_id,
+            retry=False,
+        )
+        item = payload.get("item")
+        if item is None:
+            return None
+        return QueueItem.model_validate(item)
+
+    def heartbeat(self, queue_item_id: UUID | str, *, request_id: str | None = None) -> QueueItem:
+        payload = self._command(
+            "/app/heartbeat",
+            {"queueItemId": str(queue_item_id)},
+            request_id=request_id,
+            retry=True,
+        )
+        return QueueItem.model_validate(payload["item"])
+
+    def update_status(
+        self,
+        queue_item_id: UUID | str,
+        status: TaskStatus,
+        *,
+        output_data: Any = None,
+        request_id: str | None = None,
+    ) -> QueueItem:
+        body: dict[str, Any] = {
+            "queueItemId": str(queue_item_id),
+            "status": status.value if isinstance(status, TaskStatus) else str(status),
+        }
+        encoded = _encode_input(output_data)
+        if encoded is not None:
+            body["outputData"] = encoded
+        payload = self._command("/app/status", body, request_id=request_id, retry=True)
+        return QueueItem.model_validate(payload["item"])
+
+    def subscribe_work(
+        self, task_name: str, *, request_id: str | None = None
+    ) -> WorkSubscriptionResult:
+        payload = self._command(
+            "/app/subscribeWork",
+            {"taskName": task_name},
+            request_id=request_id,
+            retry=True,
+        )
+        with self._state_lock:
+            self._work_subscriptions.add(task_name)
+        return WorkSubscriptionResult.model_validate(payload)
+
+    def unsubscribe_work(
+        self, task_name: str, *, request_id: str | None = None
+    ) -> WorkSubscriptionResult:
+        payload = self._command(
+            "/app/unsubscribeWork",
+            {"taskName": task_name},
+            request_id=request_id,
+            retry=True,
+        )
+        with self._state_lock:
+            self._work_subscriptions.discard(task_name)
+        return WorkSubscriptionResult.model_validate(payload)
+
+    def subscribe_wait(
+        self,
+        subscriber_id: str,
+        *,
+        mode: WaitMode | str | None = None,
+        correlation: str | None = None,
+        queue_item_ids: list[UUID | str] | None = None,
+        request_id: str | None = None,
+    ) -> WaitSubscriptionResult:
+        ids = [str(qid) for qid in queue_item_ids] if queue_item_ids else None
+        mode_value = mode.value if isinstance(mode, WaitMode) else mode
+        payload = self._command(
+            "/app/subscribeWait",
+            drop_none(
+                {
+                    "subscriberId": subscriber_id,
+                    "mode": mode_value,
+                    "correlation": correlation,
+                    "queueItemIds": ids,
+                }
+            ),
+            request_id=request_id,
+            retry=True,
+        )
+        with self._state_lock:
+            self._wait_subscribers.add(subscriber_id)
+        return WaitSubscriptionResult.model_validate(payload)
+
+    def bind_wait(
+        self, subscriber_id: str, *, request_id: str | None = None
+    ) -> WaitSubscriptionResult:
+        """Reconnect helper: bind this socket and replay unacked deliveries."""
+        return self.subscribe_wait(subscriber_id, request_id=request_id)
+
+    def ack_event(
+        self,
+        subscriber_id: str,
+        event_id: UUID | str,
+        *,
+        request_id: str | None = None,
+    ) -> EventAckResult:
+        payload = self._command(
+            "/app/ackEvent",
+            {"subscriberId": subscriber_id, "eventId": str(event_id)},
+            request_id=request_id,
+            retry=True,
+        )
+        return EventAckResult.model_validate(payload)
+
+    def _command(
+        self,
+        destination: str,
+        body: dict[str, Any],
+        *,
+        request_id: str | None,
+        retry: bool,
+    ) -> dict[str, Any]:
+        if self._closed.is_set():
+            raise TransportError("WebSocket client is closed")
+        deadline = time.monotonic() + self._options.request_timeout
+        last_error: Exception | None = None
+        sent = False
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                if last_error is not None:
+                    raise last_error
+                raise CommandTimeoutError(f"timed out waiting to send {destination}")
+            try:
+                self._wait_connected(min(remaining, self._options.connect_timeout))
+                rid = request_id or str(uuid.uuid4())
+                payload = dict(body)
+                payload["requestId"] = rid
+                pending = _Pending()
+                with self._state_lock:
+                    self._pending[rid] = pending
+                self._send_json(destination, payload)
+                sent = True
+                wait_for = max(0.01, deadline - time.monotonic())
+                if not pending.event.wait(timeout=wait_for):
+                    with self._state_lock:
+                        self._pending.pop(rid, None)
+                    raise CommandTimeoutError(f"timed out waiting for reply to {destination}")
+                if pending.error is not None:
+                    raise pending.error
+                assert pending.payload is not None
+                return pending.payload
+            except AmbiguousCommandError:
+                raise
+            except CommandTimeoutError:
+                raise
+            except TransportError as exc:
+                last_error = exc
+                if sent and not retry:
+                    raise AmbiguousCommandError(
+                        f"{destination} was sent but the reply was lost; outcome is unknown"
+                    ) from exc
+                if not self._options.auto_reconnect or self._closed.is_set():
+                    raise
+                request_id = None
+                sent = False
+                time.sleep(min(0.05, max(0.0, deadline - time.monotonic())))
+
+    def _wait_connected(self, timeout: float) -> None:
+        if self._connected.wait(timeout=timeout) and self._socket is not None:
+            return
+        if self._closed.is_set():
+            raise TransportError("WebSocket client is closed")
+        raise TransportError("WebSocket is not connected")
+
+    def _send_json(self, destination: str, payload: dict[str, Any]) -> None:
+        frame = _stomp.StompFrame(
+            command="SEND",
+            headers={
+                "destination": destination,
+                "content-type": "application/json",
+            },
+            body=json.dumps(payload, default=_json_default),
+        )
+        self._send_raw(_stomp.encode_frame(frame))
+
+    def _send_raw(self, data: str) -> None:
+        sock = self._socket
+        if sock is None:
+            raise TransportError("WebSocket is not connected")
+        with self._send_lock:
+            try:
+                sock.send(data)
+            except Exception as exc:
+                raise TransportError(f"WebSocket send failed: {exc}") from exc
+
+    def _start_receiver(self) -> None:
+        if self._receiver is not None and self._receiver.is_alive():
+            return
+        self._receiver = threading.Thread(
+            target=self._receiver_loop, name="continuum-stomp-recv", daemon=True
+        )
+        self._receiver.start()
+
+    def _receiver_loop(self) -> None:
+        backoff = self._options.reconnect_backoff
+        while not self._closed.is_set():
+            try:
+                self._open_and_handshake()
+                self._reconnect_attempt = 0
+                backoff = self._options.reconnect_backoff
+                self._restore_subscriptions()
+                self._connected.set()
+                self._read_until_disconnect()
+            except Exception:
+                logger.info("WebSocket session ended", exc_info=True)
+            self._teardown_socket()
+            self._connected.clear()
+            self._fail_pending(TransportError("WebSocket disconnected"))
+            if self._closed.is_set() or not self._options.auto_reconnect:
+                with self._events_cv:
+                    self._events_cv.notify_all()
+                return
+            self._reconnect_attempt += 1
+            logger.info(
+                "Reconnecting to Queue API WebSocket in %.1fs (attempt %d)",
+                backoff,
+                self._reconnect_attempt,
+            )
+            if self._closed.wait(timeout=backoff):
+                return
+            backoff = min(backoff * 2, self._options.max_reconnect_backoff)
+
+    def _open_and_handshake(self) -> None:
+        url = derive_websocket_url(self._base_url, self._options.endpoint)
+        socket = self._connector(
+            url,
+            {"Api-Key": self._api_key},
+            self._options.connect_timeout,
+        )
+        self._socket = socket
+        self._decoder = _stomp.StompDecoder()
+        self._generation += 1
+        cx = max(0, int(self._options.heartbeat_ms))
+        connect = _stomp.StompFrame(
+            command="CONNECT",
+            headers={
+                "accept-version": "1.2,1.1,1.0",
+                "heart-beat": f"{cx},{cx}",
+                "host": urlparse(url).hostname or "localhost",
+            },
+        )
+        self._send_raw(_stomp.encode_frame(connect))
+        connected = self._recv_until_command("CONNECTED", timeout=self._options.connect_timeout)
+        self._client_heartbeat_ms = _negotiate_heartbeat(
+            cx, connected.headers.get("heart-beat", "0,0")
+        )
+        self._subscribe_destinations(connected.headers.get("session"))
+        self._start_heartbeat()
+        logger.info("WebSocket connected to %s", url)
+
+    def _subscribe_destinations(self, session_id: str | None) -> None:
+        """Subscribe to user destinations and Spring's resolved session queues."""
+        self._send_raw(
+            _stomp.encode_frame(
+                _stomp.StompFrame(
+                    command="SUBSCRIBE",
+                    headers={"id": "replies", "destination": "/user/queue/replies"},
+                )
+            )
+        )
+        self._send_raw(
+            _stomp.encode_frame(
+                _stomp.StompFrame(
+                    command="SUBSCRIBE",
+                    headers={"id": "events", "destination": "/user/queue/events"},
+                )
+            )
+        )
+        if not session_id:
+            return
+        self._send_raw(
+            _stomp.encode_frame(
+                _stomp.StompFrame(
+                    command="SUBSCRIBE",
+                    headers={
+                        "id": "replies-session",
+                        "destination": f"/queue/replies-user{session_id}",
+                    },
+                )
+            )
+        )
+        self._send_raw(
+            _stomp.encode_frame(
+                _stomp.StompFrame(
+                    command="SUBSCRIBE",
+                    headers={
+                        "id": "events-session",
+                        "destination": f"/queue/events-user{session_id}",
+                    },
+                )
+            )
+        )
+
+    def _restore_subscriptions(self) -> None:
+        # Runs on the receiver thread; must recv replies inline (not via
+        # ``_command``, which waits for this same thread to dispatch).
+        with self._state_lock:
+            work = list(self._work_subscriptions)
+            waits = list(self._wait_subscribers)
+        for task_name in work:
+            try:
+                self._command_inline("/app/subscribeWork", {"taskName": task_name})
+            except ContinuumError:
+                logger.error("Failed to restore work subscription %s", task_name, exc_info=True)
+        for subscriber_id in waits:
+            try:
+                self._command_inline("/app/subscribeWait", {"subscriberId": subscriber_id})
+            except ContinuumError:
+                logger.error("Failed to restore wait subscriber %s", subscriber_id, exc_info=True)
+
+    def _command_inline(self, destination: str, body: dict[str, Any]) -> dict[str, Any]:
+        rid = str(uuid.uuid4())
+        payload = dict(body)
+        payload["requestId"] = rid
+        self._send_json(destination, payload)
+        deadline = time.monotonic() + self._options.request_timeout
+        sock = self._socket
+        if sock is None:
+            raise TransportError("WebSocket is not connected")
+        while time.monotonic() < deadline:
+            try:
+                data = sock.recv()
+            except Exception as exc:
+                if _is_timeout(exc):
+                    continue
+                raise TransportError(f"WebSocket recv failed: {exc}") from exc
+            if data is None or data == "":
+                raise TransportError("WebSocket closed during subscription restore")
+            for frame in self._decoder.feed(data):
+                if frame is None:
+                    continue
+                if frame.command == "ERROR":
+                    raise _error_frame(frame)
+                if frame.command != "MESSAGE":
+                    continue
+                destination_header = frame.headers.get("destination", "")
+                try:
+                    parsed = json.loads(frame.body) if frame.body else {}
+                except json.JSONDecodeError as exc:
+                    raise ProtocolError(f"invalid JSON: {exc}") from exc
+                if "events" in destination_header:
+                    self._handle_event(parsed)
+                    continue
+                if str(parsed.get("requestId")) != rid:
+                    self._handle_reply(parsed)
+                    continue
+                status = int(parsed.get("status", 0))
+                if status != 0:
+                    message = parsed.get("message") or f"command failed with status {status}"
+                    raise error_for_status(status, message)
+                return parsed
+        raise CommandTimeoutError(f"timed out restoring {destination}")
+
+    def _read_until_disconnect(self) -> None:
+        sock = self._socket
+        if sock is None:
+            return
+        while not self._closed.is_set():
+            try:
+                data = sock.recv()
+            except Exception as exc:
+                if _is_timeout(exc):
+                    continue
+                raise TransportError(f"WebSocket recv failed: {exc}") from exc
+            if data is None or data == "":
+                raise TransportError("WebSocket closed by peer")
+            self._dispatch_bytes(data)
+
+    def _recv_until_command(self, command: str, *, timeout: float) -> _stomp.StompFrame:
+        sock = self._socket
+        if sock is None:
+            raise TransportError("WebSocket is not connected")
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            try:
+                data = sock.recv()
+            except Exception as exc:
+                if _is_timeout(exc):
+                    continue
+                raise TransportError(f"WebSocket recv failed: {exc}") from exc
+            if data is None or data == "":
+                raise TransportError("WebSocket closed during handshake")
+            for frame in self._decoder.feed(data):
+                if frame is None:
+                    continue
+                if frame.command == "ERROR":
+                    raise _error_frame(frame)
+                if frame.command == command:
+                    return frame
+                self._handle_frame(frame)
+        raise TransportError(f"timed out waiting for STOMP {command}")
+
+    def _dispatch_bytes(self, data: str | bytes) -> None:
+        for frame in self._decoder.feed(data):
+            if frame is None:
+                continue
+            self._handle_frame(frame)
+
+    def _handle_frame(self, frame: _stomp.StompFrame) -> None:
+        if frame.command == "ERROR":
+            exc = _error_frame(frame)
+            self._fail_pending(exc)
+            raise exc
+        if frame.command != "MESSAGE":
+            return
+        destination = frame.headers.get("destination", "")
+        try:
+            payload = json.loads(frame.body) if frame.body else {}
+        except json.JSONDecodeError as exc:
+            raise ProtocolError(f"invalid JSON on {destination}: {exc}") from exc
+        if "replies" in destination:
+            self._handle_reply(payload)
+        elif "events" in destination:
+            self._handle_event(payload)
+
+    def _handle_reply(self, payload: dict[str, Any]) -> None:
+        request_id = payload.get("requestId")
+        pending: _Pending | None = None
+        if request_id is not None:
+            with self._state_lock:
+                pending = self._pending.pop(str(request_id), None)
+        status = int(payload.get("status", 0))
+        if status != 0:
+            message = payload.get("message") or f"command failed with status {status}"
+            error = error_for_status(status, message)
+            if pending is not None:
+                pending.error = error
+                pending.event.set()
+            return
+        if pending is not None:
+            pending.payload = payload
+            pending.event.set()
+
+    def _handle_event(self, payload: dict[str, Any]) -> None:
+        event_body = payload.get("event")
+        if not event_body:
+            return
+        event = QueueEvent.model_validate(event_body)
+        if event.event_type == QueueEventType.COMPLETED:
+            key = str(event.event_id)
+            if key in self._seen_completed:
+                return
+            self._seen_completed[key] = None
+            while len(self._seen_completed) > 1024:
+                self._seen_completed.popitem(last=False)
+        for callback in list(self._event_callbacks):
+            try:
+                callback(event)
+            except Exception:
+                logger.exception("event callback failed")
+        with self._events_cv:
+            self._events.append(event)
+            self._events_cv.notify_all()
+
+    def _fail_pending(self, error: Exception) -> None:
+        with self._state_lock:
+            pending = list(self._pending.values())
+            self._pending.clear()
+        for item in pending:
+            if item.error is None:
+                item.error = error
+            item.event.set()
+
+    def _teardown_socket(self) -> None:
+        self._heartbeat_stop.set()
+        sock = self._socket
+        self._socket = None
+        if sock is not None:
+            try:
+                sock.close()
+            except Exception:
+                logger.debug("error closing websocket after disconnect", exc_info=True)
+        thread = self._heartbeat_thread
+        self._heartbeat_thread = None
+        if thread is not None and thread is not threading.current_thread():
+            thread.join(timeout=1.0)
+
+    def _start_heartbeat(self) -> None:
+        self._heartbeat_stop = threading.Event()
+        if self._client_heartbeat_ms <= 0:
+            return
+        interval = self._client_heartbeat_ms / 1000.0
+
+        def _loop() -> None:
+            while not self._heartbeat_stop.wait(timeout=interval):
+                if self._closed.is_set() or self._socket is None:
+                    return
+                try:
+                    self._send_raw(_stomp.encode_heartbeat())
+                except Exception:
+                    return
+
+        self._heartbeat_thread = threading.Thread(
+            target=_loop, name="continuum-stomp-heartbeat", daemon=True
+        )
+        self._heartbeat_thread.start()
+
+
+def _negotiate_heartbeat(client_cx_ms: int, server_header: str) -> int:
+    """Return the client send interval in milliseconds, or 0 to disable."""
+    parts = [p.strip() for p in server_header.split(",")]
+    server_cy = 0
+    if len(parts) >= 2:
+        try:
+            server_cy = int(parts[1])
+        except ValueError:
+            server_cy = 0
+    if client_cx_ms <= 0 or server_cy <= 0:
+        return 0
+    return max(client_cx_ms, server_cy)
+
+
+def _is_timeout(exc: BaseException) -> bool:
+    name = type(exc).__name__.lower()
+    return "timeout" in name or isinstance(exc, TimeoutError)
+
+
+def _error_frame(frame: _stomp.StompFrame) -> ContinuumError:
+    message = frame.headers.get("message") or frame.body or "STOMP ERROR"
+    lower = message.lower()
+    if "unauthorized" in lower or "401" in lower:
+        return error_for_status(401, message)
+    return ProtocolError(message, body=frame.body)
+
+
+def _json_default(obj: Any) -> Any:
+    import datetime as _dt
+    import enum as _enum
+
+    if isinstance(obj, UUID):
+        return str(obj)
+    if isinstance(obj, _dt.datetime):
+        return obj.isoformat()
+    if isinstance(obj, _enum.Enum):
+        return obj.value
+    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+
+def _default_connector(url: str, headers: dict[str, str], timeout: float) -> SocketLike:
+    try:
+        import websocket
+    except ImportError as exc:  # pragma: no cover
+        raise TransportError(
+            "websocket-client is required for WebSocket mode; install continuum-task-server-sdk"
+        ) from exc
+    header = [f"{key}: {value}" for key, value in headers.items()]
+    try:
+        sock = websocket.create_connection(url, header=header, timeout=timeout)
+        sock.settimeout(1.0)
+        return sock
+    except Exception as exc:
+        raise TransportError(f"WebSocket handshake failed for {url}: {exc}") from exc