feat: add websockets client/server

Author: rachwalk

src/rai_core/rai/communication/http/api.py

@@ -2,8 +2,10 @@
 import threading
 from enum import IntFlag
 from typing import Callable, Optional, Any
+import json
 import logging
 
+import aiohttp
 from aiohttp import ClientSession, ClientTimeout, web
 
 
@@ -26,6 +28,9 @@ def __init__(
         self.host = host
         self.port = port
         self.mode = mode
+
+        self.routes: dict[str, list[str]] = {}
+
         self.app = web.Application()
         self.runner = web.AppRunner(self.app)
         self.loop = asyncio.new_event_loop()
@@ -34,6 +39,9 @@ def __init__(
         self._started_event = threading.Event()
         self.unresolved_futures = []
 
+        self.websockets: dict[str, set[web.WebSocketResponse]] = {}
+        self.ws_clients: dict[str, aiohttp.ClientWebSocketResponse] = {}
+
     def _start_loop(self):
         asyncio.set_event_loop(self.loop)
         self.loop.run_until_complete(self._start_server())
@@ -86,21 +94,117 @@ def register():
             self.app.router.add_route(method.upper(), path, handler)
 
         self.loop.call_soon_threadsafe(register)
+        if self.routes.get(path) is not None:
+            self.routes[path] = [method]
+        else:
+            self.routes[path].append(method)
 
     def add_websocket(self, path: str, handler_lambda):
-        if not (self.mode & HTTPConnectorMode.server):
-            return
+        """
+        In server mode:
+            `path` is the HTTP path (e.g. "/ws").
+            `handler_lambda(ws, request)` is called for each connection.
+
+        In client mode:
+            `path` is the full WebSocket URL (e.g. "ws://example.com/ws").
+            `handler_lambda(ws, msg)` is called for each incoming message.
+        """
+        # SERVER SIDE
+        if self.mode & HTTPConnectorMode.server:
+            if path not in self.websockets:
+                self.websockets[path] = set()
 
-        async def ws_handler(request):
-            ws = web.WebSocketResponse()
-            await ws.prepare(request)
-            await handler_lambda(ws, request)
-            return ws
+            async def ws_handler(request):
+                ws = web.WebSocketResponse()
+                await ws.prepare(request)
 
-        def register():
-            self.app.router.add_get(path, ws_handler)
+                # register this connection
+                self.websockets[path].add(ws)
 
-        self.loop.call_soon_threadsafe(register)
+                try:
+                    # user handler can read/write freely, e.g.:
+                    # async for msg in ws: ...
+                    await handler_lambda(ws, request)
+                finally:
+                    # ensure it is removed on close
+                    self.websockets[path].discard(ws)
+                    await ws.close()
+
+                return ws
+
+            def register_server():
+                self.app.router.add_get(path, ws_handler)
+
+            self.loop.call_soon_threadsafe(register_server)
+
+        # CLIENT SIDE
+        if self.mode & HTTPConnectorMode.client:
+            async def connect_client_ws():
+                assert self.client_session is not None, "ClientSession not initialized"
+                ws = await self.client_session.ws_connect(path)
+                self.ws_clients[path] = ws
+
+                try:
+                    async for msg in ws:
+                        # let user handler inspect/read messages and optionally write
+                        await handler_lambda(ws, msg)
+                finally:
+                    # clean up on close
+                    if self.ws_clients.get(path) is ws:
+                        del self.ws_clients[path]
+                    await ws.close()
+
+            def start_client():
+                asyncio.create_task(connect_client_ws())
+
+            self.loop.call_soon_threadsafe(start_client)
+
+    def publish_websocket(
+            self,
+            path: str,
+            payload: Optional[str | dict],
+    ):
+        """
+        Send `payload` over all WebSocket connections associated with `path`.
+
+        - For server mode: broadcasts to all connected clients on that route.
+        - For client mode: sends to the single client WebSocket created for that URL.
+        """
+        if payload is None:
+            msg = ""
+        elif isinstance(payload, dict):
+            msg = json.dumps(payload)
+        else:
+            msg = str(payload)
+
+        async def _publish():
+            # collect all websockets (server + client) associated with this key
+            server_conns = list(self.websockets.get(path, []))
+            client_ws = self.ws_clients.get(path)
+            all_conns = server_conns + ([client_ws] if client_ws is not None else [])
+
+            dead_server = []
+            dead_client = False
+
+            for ws in all_conns:
+                try:
+                    await ws.send_str(msg)
+                except Exception:
+                    # mark broken ones to be removed
+                    if ws in server_conns:
+                        dead_server.append(ws)
+                    elif ws is client_ws:
+                        dead_client = True
+
+            # cleanup broken server connections
+            for ws in dead_server:
+                self.websockets.get(path, set()).discard(ws)
+
+            # cleanup broken client connection
+            if dead_client and self.ws_clients.get(path) is client_ws:
+                del self.ws_clients[path]
+
+        asyncio.run_coroutine_threadsafe(_publish(), self.loop)
 
     def send_request(
         self,

src/rai_core/rai/communication/http/connectors/base.py

@@ -1,4 +1,6 @@
-from typing import Any, Callable, Optional, TypeVar
+import time
+import logging
+from typing import Any, Callable, Optional, TypeVar, Dict
 
 from rai.communication.base_connector import BaseConnector, BaseMessage
 from rai.communication.http.messages import HTTPMessage
@@ -19,19 +21,45 @@ def __init__(
         self._api = HTTPAPI(mode, host, port)
         self._api.run()
         self._services = []
+        self.last_msg: Dict[str, T] = {}
 
     def send_message(self, message: T, target: str, **kwargs: Optional[Any]) -> None:
-        self._api.send_request(
-            message.method,
-            target,
-            None,
-            payload=message.payload,
-            headers=message.headers,
-        )
+        if message.protocol == "http":
+             self._api.send_request(
+                 message.method,
+                 target,
+                 None,
+                 payload=message.payload,
+                 headers=message.headers,
+             )
+        else:
+        #    self._api.
 
     def receive_message(
         self, source: str, timeout_sec: float, **kwargs: Optional[Any]
     ) -> T:
+        msg = None
+        if self._api.routes.get(source, None) is not None:
+            # a GET method has already been added...
+        else:
+            def local_callback(payload: Any) -> None:
+                msg = payload
+            self._api.add_route(
+                "GET",
+                source,
+                self.general_callback
+            )
+
+        start_time = time.time()
+        # wait for the message to be received
+        while time.time() - start_time < timeout_sec:
+            if source in self.last_msg:
+                return self.last_msg[source]
+            time.sleep(0.1)
+        else:
+            raise TimeoutError(
+                f"Message from {source} not received in {timeout_sec} seconds"
+            )
         raise NotImplementedError("This method should be implemented by the subclass.")
 
     def _safe_callback_wrapper(self, callback: Callable[[T], None], message: T) -> None:
@@ -74,6 +102,10 @@ def create_service(
         **kwargs: Optional[Any],
     ) -> str:
         id_str = f"{method.upper()}_{service_name}"
+        if on_done is not None:
+            logging.warning(
+                f"not None on_done argument passed to create_service of {self.__class__}; will have no effect!"
+            )
         if id_str in self._services:
             raise HTTPAPIError(
                 f"Service {service_name} already has a {method.upper()} handler"

src/rai_core/rai/communication/http/messages.py

@@ -18,6 +18,7 @@
 
 
 class HTTPMessage(BaseMessage):
+    protocol: Literal["http", "websocket"]
     headers: Optional[dict] = None
     method: Literal[
         "GET", "HEAD", "OPTIONS", "TRACE", "PUT", "DELETE", "POST", "PATCH", "CONNECT"