refactor: use a background task to process incoming responses

This makes sure that we keep getting events from the simulator in the background, even when there's on outstanding request.

Author: urish

src/wokwi_client/event_queue.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: 2025-present CodeMagic LTD
+#
+# SPDX-License-Identifier: MIT
+
+"""
+Queue-based event subscription helper for Transport.
+
+Usage:
+    with EventQueue(transport, "serial-monitor:data") as queue:
+        # Use the queue to get events, calling get() or get_nowait()
+        event = await queue.get()
+        # do something with the event
+        ...
+"""
+
+import asyncio
+
+from .protocol_types import EventMessage
+from .transport import Transport
+
+
+class EventQueue:
+    """A queue for events from a specific event type."""
+
+    def __init__(self, transport: Transport, event_type: str) -> None:
+        self._queue: asyncio.Queue[EventMessage] = asyncio.Queue()
+        self._transport = transport
+        self._event_type = event_type
+
+        def listener(event: EventMessage) -> None:
+            self._queue.put_nowait(event)
+
+        self._listener = listener
+        self._transport.add_event_listener(self._event_type, self._listener)
+
+    def close(self) -> None:
+        """Close the queue. This is useful when you want to stop listening for events."""
+        self._transport.remove_event_listener(self._event_type, self._listener)
+
+    async def get(self) -> EventMessage:
+        """Get an event from the queue. Blocks until an event is available."""
+        return await self._queue.get()
+
+    def get_nowait(self) -> EventMessage:
+        """Get an event from the queue. Raises QueueEmpty if no event is available."""
+        return self._queue.get_nowait()
+
+    def flush(self) -> None:
+        """Flush the queue. This is useful when you want to wait for all events to be processed."""
+        while not self._queue.empty():
+            self._queue.get_nowait()
+
+    def __enter__(self) -> "EventQueue":
+        return self
+
+    def __exit__(self, exc_type: type, exc_value: Exception, traceback: object) -> None:
+        self.close()

src/wokwi_client/serial.py

@@ -3,17 +3,14 @@
 # SPDX-License-Identifier: MIT
 
 from collections.abc import AsyncGenerator
-from typing import cast
 
-from .protocol_types import EventMessage
+from .event_queue import EventQueue
 from .transport import Transport
 
 
 async def monitor_lines(transport: Transport) -> AsyncGenerator[bytes, None]:
     await transport.request("serial-monitor:listen", {})
-    while True:
-        msg = await transport.recv()
-        if msg["type"] == "event":
-            event_msg = cast(EventMessage, msg)
-            if event_msg["event"] == "serial-monitor:data":
-                yield bytes(event_msg["payload"]["bytes"])
+    with EventQueue(transport, "serial-monitor:data") as queue:
+        while True:
+            event_msg = await queue.get()
+            yield bytes(event_msg["payload"]["bytes"])

src/wokwi_client/transport.py

@@ -2,22 +2,24 @@
 #
 # SPDX-License-Identifier: MIT
 
+import asyncio
 import json
 import os
 import warnings
-from typing import Any, Optional, cast
+from typing import Any, Callable, Optional, cast
 
 import websockets
 
 from .__version__ import get_version
 from .constants import (
     DEFAULT_WS_URL,
+    MSG_TYPE_EVENT,
     MSG_TYPE_HELLO,
     MSG_TYPE_RESPONSE,
     PROTOCOL_VERSION,
 )
 from .exceptions import ProtocolError, ServerError, WokwiError
-from .protocol_types import HelloMessage, IncomingMessage, ResponseMessage
+from .protocol_types import EventMessage, HelloMessage, IncomingMessage, ResponseMessage
 
 TRANSPORT_DEFAULT_WS_URL = os.getenv("WOKWI_CLI_SERVER", DEFAULT_WS_URL)
 
@@ -28,6 +30,10 @@ def __init__(self, token: str, url: str = TRANSPORT_DEFAULT_WS_URL):
         self._url = url
         self._next_id = 1
         self._ws: Optional[websockets.WebSocketClientProtocol] = None
+        self._event_listeners: dict[str, list[Callable[[EventMessage], Any]]] = {}
+        self._response_futures: dict[str, asyncio.Future[ResponseMessage]] = {}
+        self._recv_task: Optional[asyncio.Task[None]] = None
+        self._closed = False
 
     async def connect(self) -> dict[str, Any]:
         self._ws = await websockets.connect(
@@ -41,28 +47,85 @@ async def connect(self) -> dict[str, Any]:
         if hello["type"] != MSG_TYPE_HELLO or hello.get("protocolVersion") != PROTOCOL_VERSION:
             raise ProtocolError(f"Unsupported protocol handshake: {hello}")
         hello_msg = cast(HelloMessage, hello)
+        self._closed = False
+        # Start background message processor
+        self._recv_task = asyncio.create_task(self._background_recv())
         return {"version": hello_msg["appVersion"]}
 
     async def close(self) -> None:
+        self._closed = True
+        if self._recv_task:
+            self._recv_task.cancel()
+            try:
+                await self._recv_task
+            except asyncio.CancelledError:
+                pass
         if self._ws:
             await self._ws.close()
 
+    def add_event_listener(self, event_type: str, listener: Callable[[EventMessage], Any]) -> None:
+        """Register a listener for a specific event type."""
+        if event_type not in self._event_listeners:
+            self._event_listeners[event_type] = []
+        self._event_listeners[event_type].append(listener)
+
+    def remove_event_listener(
+        self, event_type: str, listener: Callable[[EventMessage], Any]
+    ) -> None:
+        """Remove a previously registered listener for a specific event type."""
+        if event_type in self._event_listeners:
+            self._event_listeners[event_type] = [
+                registered_listener
+                for registered_listener in self._event_listeners[event_type]
+                if registered_listener != listener
+            ]
+            if not self._event_listeners[event_type]:
+                del self._event_listeners[event_type]
+
+    async def _dispatch_event(self, event_msg: EventMessage) -> None:
+        listeners = self._event_listeners.get(event_msg["event"], [])
+        for listener in listeners:
+            result = listener(event_msg)
+            if hasattr(result, "__await__"):
+                await result
+
     async def request(self, command: str, params: dict[str, Any]) -> ResponseMessage:
         msg_id = str(self._next_id)
         self._next_id += 1
         if self._ws is None:
             raise WokwiError("Not connected")
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[ResponseMessage] = loop.create_future()
+        self._response_futures[msg_id] = future
         await self._ws.send(
             json.dumps({"type": "command", "command": command, "params": params, "id": msg_id})
         )
-        while True:
-            msg: IncomingMessage = await self._recv()
-            if msg["type"] == MSG_TYPE_RESPONSE and msg.get("id") == msg_id:
-                resp_msg = cast(ResponseMessage, msg)
-                if resp_msg.get("error"):
-                    result = resp_msg["result"]
-                    raise ServerError(result["message"])
-                return resp_msg
+        try:
+            resp_msg_resp = await future
+            if resp_msg_resp.get("error"):
+                result = resp_msg_resp["result"]
+                raise ServerError(result["message"])
+            return resp_msg_resp
+        finally:
+            del self._response_futures[msg_id]
+
+    async def _background_recv(self) -> None:
+        try:
+            while not self._closed and self._ws is not None:
+                msg: IncomingMessage = await self._recv()
+                if msg["type"] == MSG_TYPE_EVENT:
+                    resp_msg_event = cast(EventMessage, msg)
+                    await self._dispatch_event(resp_msg_event)
+                elif msg["type"] == MSG_TYPE_RESPONSE:
+                    resp_msg_resp = cast(ResponseMessage, msg)
+                    future = self._response_futures.get(resp_msg_resp["id"])
+                    if future is None or future.done():
+                        continue
+                    future.set_result(resp_msg_resp)
+        except (websockets.ConnectionClosed, asyncio.CancelledError):
+            pass
+        except Exception as e:
+            warnings.warn(f"Background recv error: {e}", RuntimeWarning)
 
     async def _recv(self) -> IncomingMessage:
         if self._ws is None:
@@ -87,6 +150,3 @@ async def _recv(self) -> IncomingMessage:
             )
             raise WokwiError(f"Server error {result['code']}: {result['message']}")
         return cast(IncomingMessage, message)
-
-    async def recv(self) -> IncomingMessage:
-        return await self._recv()