archastro.phx_channel

 30class Socket:
 31    """
 32    Manages a WebSocket connection to a Phoenix server.
 33
 34    Usage::
 35
 36        socket = Socket("ws://localhost:4000/socket/websocket", params={"token": "..."})
 37        await socket.connect()
 38
 39        channel = socket.channel("room:lobby", {"user_id": "123"})
 40        resp = await channel.join()
 41
 42        await channel.push("new_msg", {"body": "hello"})
 43        channel.on("new_msg", lambda payload: print(payload))
 44
 45        await socket.disconnect()
 46    """
 47
 48    def __init__(
 49        self,
 50        url: str,
 51        *,
 52        params: dict[str, str] | None = None,
 53        heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL_S,
 54        timeout: float = DEFAULT_TIMEOUT_S,
 55        reconnect_backoff_ms: list[int] | None = None,
 56        auto_reconnect: bool = True,
 57    ):
 58        self._base_url = url
 59        self._params = params or {}
 60        self._heartbeat_interval = heartbeat_interval
 61        self._timeout = timeout
 62        self._backoff = reconnect_backoff_ms or DEFAULT_RECONNECT_BACKOFF_MS
 63        self._auto_reconnect = auto_reconnect
 64
 65        self._ws: ClientConnection | None = None
 66        self._ref = 0
 67        self._pending_heartbeat_ref: str | None = None
 68        self._channels: dict[str, Channel] = {}
 69        self._connected = False
 70        self._closing = False
 71
 72        self._heartbeat_task: asyncio.Task[None] | None = None
 73        self._receive_task: asyncio.Task[None] | None = None
 74        self._reconnect_attempt = 0
 75
 76        self._on_open_callbacks: list[Callable[[], Any]] = []
 77        self._on_close_callbacks: list[Callable[[int, str], Any]] = []
 78        self._on_error_callbacks: list[Callable[[Exception], Any]] = []
 79
 80    @property
 81    def is_connected(self) -> bool:
 82        return self._connected and self._ws is not None
 83
 84    def _make_ref(self) -> str:
 85        self._ref += 1
 86        return str(self._ref)
 87
 88    def _build_url(self) -> str:
 89        parsed = urlparse(self._base_url)
 90        params = {**self._params, "vsn": "2.0.0"}
 91        # Merge any existing query params
 92        existing_qs = parsed.query
 93        if existing_qs:
 94            qs = existing_qs + "&" + urlencode(params)
 95        else:
 96            qs = urlencode(params)
 97        return urlunparse(parsed._replace(query=qs))
 98
 99    # ─── Connection lifecycle ─────────────────────────────────
100
101    async def connect(self) -> None:
102        """Connect to the Phoenix server."""
103        self._closing = False
104        self._reconnect_attempt = 0
105        await self._do_connect()
106
107    async def _do_connect(self) -> None:
108        url = self._build_url()
109        try:
110            self._ws = await websockets.connect(url)
111            self._connected = True
112            self._reconnect_attempt = 0
113            logger.info("Connected to %s", self._base_url)
114
115            for cb in self._on_open_callbacks:
116                cb()
117
118            # Start heartbeat and receive loops
119            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
120            self._receive_task = asyncio.create_task(self._receive_loop())
121
122            # Rejoin any channels that were previously joined
123            for channel in self._channels.values():
124                if channel._state in ("joined", "errored"):
125                    asyncio.create_task(channel._rejoin())
126
127        except Exception as exc:
128            logger.error("Connection failed: %s", exc)
129            self._connected = False
130            for cb in self._on_error_callbacks:
131                cb(exc)
132            if self._auto_reconnect and not self._closing:
133                await self._schedule_reconnect()
134
135    async def disconnect(self) -> None:
136        """Gracefully disconnect from the server."""
137        self._closing = True
138        self._connected = False
139
140        if self._heartbeat_task:
141            self._heartbeat_task.cancel()
142            self._heartbeat_task = None
143        if self._receive_task:
144            self._receive_task.cancel()
145            self._receive_task = None
146
147        if self._ws:
148            await self._ws.close()
149            self._ws = None
150
151        logger.info("Disconnected")
152
153    async def _schedule_reconnect(self) -> None:
154        if self._closing:
155            return
156        idx = min(self._reconnect_attempt, len(self._backoff) - 1)
157        delay_ms = self._backoff[idx]
158        self._reconnect_attempt += 1
159        logger.info("Reconnecting in %dms (attempt %d)", delay_ms, self._reconnect_attempt)
160        await asyncio.sleep(delay_ms / 1000)
161        if not self._closing:
162            await self._do_connect()
163
164    # ─── Channel management ───────────────────────────────────
165
166    def channel(self, topic: str, params: dict[str, Any] | None = None) -> Channel:
167        """Create a channel for the given topic."""
168        if topic in self._channels:
169            return self._channels[topic]
170        ch = Channel(self, topic, params or {})
171        self._channels[topic] = ch
172        return ch
173
174    def _remove_channel(self, topic: str) -> None:
175        self._channels.pop(topic, None)
176
177    # ─── Send ─────────────────────────────────────────────────
178
179    async def _send(
180        self,
181        join_ref: str | None,
182        ref: str | None,
183        topic: str,
184        event: str,
185        payload: Any,
186    ) -> None:
187        if not self._ws or not self._connected:
188            raise ConnectionError("Socket is not connected")
189        msg = json.dumps([join_ref, ref, topic, event, payload])
190        await self._ws.send(msg)
191
192    # ─── Receive loop ─────────────────────────────────────────
193
194    async def _receive_loop(self) -> None:
195        try:
196            assert self._ws is not None
197            async for raw in self._ws:
198                if isinstance(raw, bytes):
199                    raw = raw.decode("utf-8")
200                try:
201                    msg = json.loads(raw)
202                except json.JSONDecodeError:
203                    logger.warning("Non-JSON message: %s", raw[:100])
204                    continue
205
206                if not isinstance(msg, list) or len(msg) != 5:
207                    logger.warning("Malformed message: %s", raw[:100])
208                    continue
209
210                join_ref, ref, topic, event, payload = msg
211                self._dispatch(join_ref, ref, topic, event, payload)
212
213        except websockets.ConnectionClosed as exc:
214            logger.info("Connection closed: code=%s reason=%s", exc.code, exc.reason)
215            self._connected = False
216            for cb in self._on_close_callbacks:
217                cb(exc.code, exc.reason)
218            if self._auto_reconnect and not self._closing:
219                await self._schedule_reconnect()
220        except asyncio.CancelledError:
221            pass
222        except Exception as exc:
223            logger.error("Receive error: %s", exc)
224            self._connected = False
225            for cb in self._on_error_callbacks:
226                cb(exc)
227            if self._auto_reconnect and not self._closing:
228                await self._schedule_reconnect()
229
230    def _dispatch(
231        self,
232        join_ref: str | None,
233        ref: str | None,
234        topic: str,
235        event: str,
236        payload: Any,
237    ) -> None:
238        # Heartbeat reply
239        if ref and ref == self._pending_heartbeat_ref:
240            self._pending_heartbeat_ref = None
241            return
242
243        # Route to channel
244        channel = self._channels.get(topic)
245        if channel:
246            channel._on_message(join_ref, ref, event, payload)
247
248    # ─── Heartbeat ────────────────────────────────────────────
249
250    async def _heartbeat_loop(self) -> None:
251        try:
252            while self._connected:
253                await asyncio.sleep(self._heartbeat_interval)
254                if not self._connected:
255                    break
256                if self._pending_heartbeat_ref is not None:
257                    # Previous heartbeat not acknowledged — connection is dead
258                    logger.warning("Heartbeat timeout — closing connection")
259                    self._pending_heartbeat_ref = None
260                    if self._ws:
261                        await self._ws.close(1000, "heartbeat timeout")
262                    break
263                ref = self._make_ref()
264                self._pending_heartbeat_ref = ref
265                try:
266                    await self._send(None, ref, "phoenix", "heartbeat", {})
267                except Exception:
268                    break
269        except asyncio.CancelledError:
270            pass
271
272    # ─── Event callbacks ──────────────────────────────────────
273
274    def on_open(self, callback: Callable[[], Any]) -> None:
275        self._on_open_callbacks.append(callback)
276
277    def on_close(self, callback: Callable[[int, str], Any]) -> None:
278        self._on_close_callbacks.append(callback)
279
280    def on_error(self, callback: Callable[[Exception], Any]) -> None:
281        self._on_error_callbacks.append(callback)

 21class Channel:
 22    """
 23    Represents a single Phoenix Channel subscription on a topic.
 24
 25    Created via ``socket.channel("topic", params)``.
 26    Call ``await channel.join()`` to subscribe.
 27    """
 28
 29    def __init__(self, socket: Socket, topic: str, params: dict[str, Any]):
 30        self._socket = socket
 31        self._topic = topic
 32        self._params = params
 33
 34        self._state: str = "closed"  # closed | joining | joined | leaving | errored
 35        self._join_ref: str | None = None
 36        self._join_push_ref: str | None = None
 37
 38        self._event_handlers: dict[str, list[Callable[..., Any]]] = {}
 39        self._pending_replies: dict[str, asyncio.Future[dict[str, Any]]] = {}
 40        self._push_buffer: list[tuple[str, Any, asyncio.Future[dict[str, Any]]]] = []
 41        # Inbound pushes that arrived before any handler was registered.
 42        # Server scenarios can `autoPush` in response to a join frame, and
 43        # that push can land on the asyncio queue before the caller has
 44        # had a chance to register a handler post-join. We replay these
 45        # to the first matching handler registered within the window.
 46        # Bounded so a forgotten handler doesn't leak unbounded memory.
 47        self._pending_pushes: dict[str, list[Any]] = {}
 48        self._pending_pushes_cap: int = 32
 49
 50    @property
 51    def topic(self) -> str:
 52        return self._topic
 53
 54    @property
 55    def state(self) -> str:
 56        return self._state
 57
 58    @property
 59    def is_joined(self) -> bool:
 60        return self._state == "joined"
 61
 62    # ─── Join / Leave ─────────────────────────────────────────
 63
 64    async def join(
 65        self,
 66        payload: dict[str, Any] | None = None,
 67        *,
 68        timeout: float = DEFAULT_TIMEOUT_S,
 69    ) -> dict[str, Any]:
 70        """
 71        Join the channel. Returns the join response payload.
 72
 73        ``payload`` overrides the params passed to ``socket.channel(topic, params)``
 74        when provided — generated SDK channel classes pass the join payload
 75        directly here, while hand-written callers typically rely on the params
 76        captured at channel creation time.
 77
 78        Raises ``TimeoutError`` if the server doesn't reply within ``timeout``.
 79        Raises ``ChannelError`` if the server rejects the join.
 80        Raises ``ChannelError`` if the channel is already joined — callers
 81        should ``leave()`` before re-joining. Silently returning an empty
 82        response would let a double-call (e.g. two ``LiveDocChannel.join_*``
 83        invocations reusing the same underlying channel) masquerade as a
 84        successful join while dropping the real server response.
 85        """
 86        if self._state == "joined":
 87            raise ChannelError(f"Channel {self._topic} is already joined; call leave() first")
 88
 89        self._state = "joining"
 90        ref = self._socket._make_ref()
 91        self._join_ref = ref
 92        self._join_push_ref = ref
 93
 94        future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
 95        self._pending_replies[ref] = future
 96
 97        params = self._params if payload is None else payload
 98        await self._socket._send(ref, ref, self._topic, "phx_join", params)
 99
100        try:
101            result = await asyncio.wait_for(future, timeout=timeout)
102        except TimeoutError:
103            self._pending_replies.pop(ref, None)
104            self._state = "errored"
105            raise TimeoutError(f"Join timed out for {self._topic}") from None
106
107        status = result.get("status")
108        if status == "ok":
109            self._state = "joined"
110            logger.info("Joined %s", self._topic)
111            # Flush buffered pushes
112            await self._flush_push_buffer()
113            return result.get("response", {})
114        else:
115            self._state = "errored"
116            raise ChannelError(f"Join rejected for {self._topic}: {result.get('response', {})}")
117
118    async def leave(self, timeout: float = DEFAULT_TIMEOUT_S) -> None:
119        """Leave the channel."""
120        if self._state == "closed":
121            return
122
123        self._state = "leaving"
124        ref = self._socket._make_ref()
125        future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
126        self._pending_replies[ref] = future
127
128        await self._socket._send(self._join_ref, ref, self._topic, "phx_leave", {})
129
130        try:
131            await asyncio.wait_for(future, timeout=timeout)
132        except TimeoutError:
133            pass  # Leave is best-effort
134        finally:
135            self._state = "closed"
136            self._join_ref = None
137            self._socket._remove_channel(self._topic)
138            logger.info("Left %s", self._topic)
139
140    async def _rejoin(self) -> None:
141        """Rejoin after reconnection."""
142        if self._state in ("closed", "leaving"):
143            return
144        self._state = "closed"
145        self._join_ref = None
146        try:
147            await self.join()
148        except Exception as exc:
149            logger.warning("Rejoin failed for %s: %s", self._topic, exc)
150            self._state = "errored"
151
152    # ─── Push / Reply ─────────────────────────────────────────
153
154    async def push(
155        self, event: str, payload: Any = None, timeout: float = DEFAULT_TIMEOUT_S
156    ) -> dict[str, Any]:
157        """
158        Push an event to the channel and wait for a reply.
159
160        Returns the reply ``{"status": ..., "response": ...}``.
161        Raises ``TimeoutError`` if no reply within ``timeout``.
162        """
163        if payload is None:
164            payload = {}
165
166        if self._state != "joined":
167            # Buffer the push for when we rejoin
168            future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
169            self._push_buffer.append((event, payload, future))
170            return await asyncio.wait_for(future, timeout=timeout)
171
172        return await self._do_push(event, payload, timeout)
173
174    async def _do_push(self, event: str, payload: Any, timeout: float) -> dict[str, Any]:
175        ref = self._socket._make_ref()
176        future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
177        self._pending_replies[ref] = future
178
179        await self._socket._send(self._join_ref, ref, self._topic, event, payload)
180
181        try:
182            result = await asyncio.wait_for(future, timeout=timeout)
183        except TimeoutError:
184            self._pending_replies.pop(ref, None)
185            raise TimeoutError(f"Push '{event}' timed out on {self._topic}") from None
186
187        return result
188
189    async def _flush_push_buffer(self) -> None:
190        buffer = self._push_buffer[:]
191        self._push_buffer.clear()
192        for event, payload, future in buffer:
193            try:
194                result = await self._do_push(event, payload, DEFAULT_TIMEOUT_S)
195                if not future.done():
196                    future.set_result(result)
197            except Exception as exc:
198                if not future.done():
199                    future.set_exception(exc)
200
201    # ─── Event handlers ──────────────────────────────────────
202
203    def on(self, event: str, callback: Callable[..., Any]) -> Callable[[], None]:
204        """
205        Register a callback for a channel event.
206
207        Returns an unsubscribe function. If pushes for this event already
208        arrived before any handler was registered (e.g. server-side
209        autoPush fired on join), they're replayed to this callback in
210        arrival order so the caller can't miss them due to scheduling.
211        """
212        handlers = self._event_handlers.setdefault(event, [])
213        handlers.append(callback)
214
215        pending = self._pending_pushes.pop(event, None)
216        if pending:
217            for payload in pending:
218                try:
219                    callback(payload)
220                except Exception:
221                    logger.exception("Error in handler for %s:%s (replayed)", self._topic, event)
222
223        def unsubscribe() -> None:
224            handlers.remove(callback)
225
226        return unsubscribe
227
228    # ─── Internal message dispatch ────────────────────────────
229
230    def _on_message(
231        self,
232        join_ref: str | None,
233        ref: str | None,
234        event: str,
235        payload: Any,
236    ) -> None:
237        # Ignore messages from stale join
238        if join_ref is not None and join_ref != self._join_ref:
239            return
240
241        if event == "phx_reply":
242            self._handle_reply(ref, payload)
243        elif event == "phx_close":
244            self._handle_close()
245        elif event == "phx_error":
246            self._handle_error(payload)
247        else:
248            # User event — dispatch to handlers, or buffer if none yet.
249            handlers = self._event_handlers.get(event)
250            if handlers:
251                for handler in handlers:
252                    try:
253                        handler(payload)
254                    except Exception:
255                        logger.exception("Error in handler for %s:%s", self._topic, event)
256            else:
257                buf = self._pending_pushes.setdefault(event, [])
258                if len(buf) < self._pending_pushes_cap:
259                    buf.append(payload)
260                else:
261                    # Drop the oldest to bound memory for orphan events.
262                    buf.pop(0)
263                    buf.append(payload)
264
265    def _handle_reply(self, ref: str | None, payload: Any) -> None:
266        if ref and ref in self._pending_replies:
267            future = self._pending_replies.pop(ref)
268            if not future.done():
269                future.set_result(payload)
270
271    def _handle_close(self) -> None:
272        logger.info("Channel closed: %s", self._topic)
273        self._state = "closed"
274        self._trigger_event("phx_close", {})
275
276    def _handle_error(self, payload: Any) -> None:
277        logger.warning("Channel error: %s — %s", self._topic, payload)
278        self._state = "errored"
279        self._trigger_event("phx_error", payload)
280
281    def _trigger_event(self, event: str, payload: Any) -> None:
282        handlers = self._event_handlers.get(event, [])
283        for handler in handlers:
284            try:
285                handler(payload)
286            except Exception:
287                logger.exception("Error in handler for %s:%s", self._topic, event)

 31class HarnessServiceClient:
 32    """Thin client over the harness service's HTTP + WebSocket surfaces.
 33
 34    One instance per test is typical::
 35
 36        client = HarnessServiceClient(ws_url, control_url)
 37        try:
 38            await client.reset()
 39            await client.register_scenario({
 40                "topic": "doc:doc_42",
 41                "onJoin": [{"type": "autoReply"}],
 42            })
 43            socket = await client.open_socket()
 44            channel = await LiveDocChannel.join_document(
 45                socket, "doc_42", user_id="user_1"
 46            )
 47            assert channel.join_response is not None
 48        finally:
 49            await client.close()
 50    """
 51
 52    def __init__(self, ws_url: str, control_url: str, *, request_timeout: float = 5.0):
 53        self.ws_url = ws_url
 54        self.control_url = control_url.rstrip("/")
 55        self._http = httpx.AsyncClient(timeout=request_timeout)
 56        self._sockets: list[Socket] = []
 57
 58    async def close(self) -> None:
 59        """Disconnect every socket opened through this client + close HTTP."""
 60        for s in self._sockets:
 61            try:
 62                await s.disconnect()
 63            except Exception:
 64                pass
 65        self._sockets.clear()
 66        await self._http.aclose()
 67
 68    # ─── HTTP control ────────────────────────────────────────────
 69
 70    async def reset(self) -> None:
 71        """Clear every scenario, observation, and handler error on the server."""
 72        r = await self._http.post(f"{self.control_url}/reset")
 73        r.raise_for_status()
 74
 75    async def register_scenario(self, scenario: dict[str, Any]) -> None:
 76        """Register a scenario for an exact topic.
 77
 78        See the TS ``ScenarioRequest`` / ``ScenarioAction`` types for the
 79        JSON shape — the server validates and returns 400 on invalid bodies.
 80        """
 81        r = await self._http.post(f"{self.control_url}/scenarios", json=scenario)
 82        if r.status_code != 201:
 83            raise HarnessServiceError(f"register_scenario failed: {r.status_code} {r.text}")
 84
 85    async def register_stream_scenario(self, scenario: dict[str, Any]) -> None:
 86        """Register a scenario for an SSE streaming route (``"METHOD /path"``).
 87
 88        See the TS ``StreamScenarioRequest`` / ``StreamAction`` types for the
 89        JSON shape — the server validates and returns 400 on invalid bodies.
 90        """
 91        r = await self._http.post(f"{self.control_url}/stream-scenarios", json=scenario)
 92        if r.status_code != 201:
 93            raise HarnessServiceError(f"register_stream_scenario failed: {r.status_code} {r.text}")
 94
 95    async def observations(
 96        self, topic: str | None = None, event: str | None = None
 97    ) -> list[dict[str, Any]]:
 98        """Fetch inbound frames the server validated, optionally filtered."""
 99        params: dict[str, str] = {}
100        if topic is not None:
101            params["topic"] = topic
102        if event is not None:
103            params["event"] = event
104        r = await self._http.get(f"{self.control_url}/observations", params=params)
105        r.raise_for_status()
106        return r.json()
107
108    async def handler_errors(self) -> list[dict[str, Any]]:
109        """Fetch scenario handler errors recorded by the server."""
110        r = await self._http.get(f"{self.control_url}/handler-errors")
111        r.raise_for_status()
112        return r.json()
113
114    # ─── Socket lifecycle ────────────────────────────────────────
115
116    async def open_socket(self, *, auto_reconnect: bool = False) -> Socket:
117        """Open a fresh Phoenix socket to the service's WebSocket endpoint.
118
119        Every call produces a new connection — the generated SDK receives
120        the same ``Socket`` it would talk to in production. ``auto_reconnect``
121        defaults to ``False`` so a disconnected test surfaces immediately
122        rather than silently retrying in the background.
123        """
124        socket = Socket(self.ws_url, auto_reconnect=auto_reconnect)
125        await socket.connect()
126        self._sockets.append(socket)
127        return socket