Upgrade to the new websocket impl

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -6,6 +6,7 @@
 from aiochannel import Channel, ChannelClosed
 from opentelemetry.trace import Span
 from websockets import WebSocketCommonProtocol
+from websockets.asyncio.client import ClientConnection
 
 from replit_river.messages import (
     FailedSendingMessageException,
@@ -128,7 +129,7 @@ async def check_to_close_session(
 
 
 async def buffered_message_sender(
-    get_ws: Callable[[], WebSocketCommonProtocol | None],
+    get_ws: Callable[[], WebSocketCommonProtocol | ClientConnection | None],
     websocket_closed_callback: Callable[[], Coroutine[Any, Any, None]],
     get_next_pending: Callable[[], TransportMessage | None],
     commit: Callable[[TransportMessage], None],

src/replit_river/messages.py

@@ -8,6 +8,7 @@
 from websockets import (
     WebSocketCommonProtocol,
 )
+from websockets.asyncio.client import ClientConnection
 
 from replit_river.rpc import (
     TransportMessage,
@@ -31,7 +32,7 @@ class FailedSendingMessageException(Exception):
 
 async def send_transport_message(
     msg: TransportMessage,
-    ws: WebSocketCommonProtocol,
+    ws: WebSocketCommonProtocol | ClientConnection,  # legacy | asyncio
     websocket_closed_callback: Callable[[], Coroutine[Any, Any, None]],
 ) -> None:
     logger.debug("sending a message %r to ws %s", msg, ws)

src/replit_river/v2/client.py

@@ -58,27 +58,6 @@ def translate_unknown_error(
             return RiverUnknownError(code=ERROR_CODE_UNKNOWN, message="Unknown error")
 
 
-# Client[HandshakeSchema](
-#     uri_and_metadata_factory=uri_and_metadata_factory,
-#     client_id=self.client_id,
-#     server_id="SERVER",
-#     transport_options=TransportOptions(
-#         session_disconnect_grace_ms=settings.RIVER_SESSION_DISCONNECT_GRACE_MS,
-#         heartbeat_ms=settings.RIVER_HEARTBEAT_MS,
-#         heartbeats_until_dead=settings.RIVER_HEARTBEATS_UNTIL_DEAD,
-#         connection_retry_options=ConnectionRetryOptions(
-#             base_interval_ms=settings.RIVER_CONNECTION_BASE_INTERVAL_MS,
-#             max_jitter_ms=settings.RIVER_CONNECTION_MAX_JITTER_MS,
-#             max_backoff_ms=settings.RIVER_CONNECTION_MAX_BACKOFF_MS,
-#             attempt_budget_capacity=self.attempt_budget_capacity,
-#             budget_restore_interval_ms=
-#               settings.RIVER_CONNECTION_BUDGET_RESTORE_INTERVAL_MS,
-#             max_retry=self.max_retry_count,
-#         ),
-#     ),
-# )
-
-
 class Client(Generic[HandshakeMetadataType]):
     def __init__(
         self,

src/replit_river/v2/client_session.py

@@ -2,15 +2,17 @@
 import logging
 from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Callable, Literal
+from typing import Any, AsyncGenerator, Callable, Literal, cast
 
 import nanoid
 import websockets
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span
+from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed
 from websockets.frames import CloseCode
+from websockets.legacy.protocol import WebSocketCommonProtocol
 
 from replit_river.common_session import buffered_message_sender
 from replit_river.error_schema import (
@@ -58,7 +60,7 @@ def __init__(
         transport_id: str,
         to_id: str,
         session_id: str,
-        websocket: websockets.WebSocketCommonProtocol,
+        websocket: ClientConnection,
         transport_options: TransportOptions,
         close_session_callback: CloseSessionCallback,
         retry_connection_callback: RetryConnectionCallback | None = None,
@@ -101,7 +103,11 @@ def get_next_pending() -> TransportMessage | None:
 
         self._task_manager.create_task(
             buffered_message_sender(
-                get_ws=lambda: self._ws_unwrapped if self.is_websocket_open() else None,
+                get_ws=lambda: (
+                    cast(WebSocketCommonProtocol | ClientConnection, self._ws_unwrapped)
+                    if self.is_websocket_open()
+                    else None
+                ),
                 websocket_closed_callback=self._begin_close_session_countdown,
                 get_next_pending=get_next_pending,
                 commit=commit,

src/replit_river/v2/client_transport.py

@@ -6,9 +6,9 @@
 import nanoid
 import websockets
 from pydantic import ValidationError
-from websockets import (
-    WebSocketCommonProtocol,
-)
+import websockets.asyncio.client
+from websockets import WebSocketCommonProtocol
+from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed
 
 from replit_river.error_schema import (
@@ -129,7 +129,7 @@ async def _establish_new_connection(
         self,
         old_session: ClientSession | None = None,
     ) -> tuple[
-        WebSocketCommonProtocol,
+        ClientConnection,
         ControlMessageHandshakeRequest[HandshakeMetadataType],
         ControlMessageHandshakeResponse,
     ]:
@@ -159,7 +159,7 @@ async def _establish_new_connection(
 
             try:
                 uri_and_metadata = await self._uri_and_metadata_factory()
-                ws = await websockets.connect(uri_and_metadata["uri"])
+                ws = await websockets.asyncio.client.connect(uri_and_metadata["uri"])
                 session_id: str
                 if old_session:
                     session_id = old_session.session_id
@@ -228,7 +228,7 @@ async def _send_handshake_request(
         self,
         session_id: str,
         handshake_metadata: HandshakeMetadataType | None,
-        websocket: WebSocketCommonProtocol,
+        websocket: ClientConnection,
         expected_session_state: ExpectedSessionState,
     ) -> ControlMessageHandshakeRequest[HandshakeMetadataType]:
         handshake_request = ControlMessageHandshakeRequest[HandshakeMetadataType](
@@ -266,7 +266,7 @@ async def websocket_closed_callback() -> None:
             ) from e
 
     async def _get_handshake_response_msg(
-        self, websocket: WebSocketCommonProtocol
+        self, websocket: ClientConnection
     ) -> TransportMessage:
         while True:
             try:
@@ -295,7 +295,7 @@ async def _establish_handshake(
         self,
         session_id: str,
         handshake_metadata: HandshakeMetadataType,
-        websocket: WebSocketCommonProtocol,
+        websocket: ClientConnection,
         old_session: ClientSession | None,
     ) -> tuple[
         ControlMessageHandshakeRequest[HandshakeMetadataType],

src/replit_river/v2/session.py

@@ -8,6 +8,7 @@
 from aiochannel import Channel
 from opentelemetry.trace import Span, use_span
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from websockets.asyncio.client import ClientConnection
 from websockets.frames import CloseCode
 
 from replit_river.common_session import (
@@ -47,7 +48,7 @@ class Session:
 
     # ws state
     _ws_connected: bool
-    _ws_unwrapped: websockets.WebSocketCommonProtocol | None
+    _ws_unwrapped: ClientConnection | None
     _heartbeat_misses: int
     _retry_connection_callback: RetryConnectionCallback | None
 
@@ -66,7 +67,7 @@ def __init__(
         transport_id: str,
         to_id: str,
         session_id: str,
-        websocket: websockets.WebSocketCommonProtocol,
+        websocket: ClientConnection,
         transport_options: TransportOptions,
         close_session_callback: CloseSessionCallback,
         retry_connection_callback: RetryConnectionCallback | None = None,
@@ -162,7 +163,7 @@ async def _begin_close_session_countdown(self) -> None:
         self._ws_connected = False
 
     async def replace_with_new_websocket(
-        self, new_ws: websockets.WebSocketCommonProtocol
+        self, new_ws: ClientConnection
     ) -> None:
         if self._ws_unwrapped and new_ws.id != self._ws_unwrapped.id:
             self._task_manager.create_task(