Adding start_serve_response to client_session

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -8,6 +8,7 @@
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span
+import websockets
 
 from replit_river.error_schema import (
     ERROR_CODE_CANCEL,
@@ -33,6 +34,12 @@
 
 
 class ClientSession(Session):
+    async def replace_with_new_websocket(
+        self, new_ws: websockets.WebSocketCommonProtocol
+    ) -> None:
+        await super().replace_with_new_websocket(new_ws)
+        await self.start_serve_responses()
+
     async def send_rpc(
         self,
         service_name: str,

src/replit_river/common_session.py

@@ -11,6 +11,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 class SendMessage(Protocol):
     async def __call__(
         self,