Incorporated Jacky's feedback

Author: JimmyAustin

src/replit_river/client_session.py

@@ -100,7 +100,7 @@ async def serve(self) -> None:
                 await self._handle_messages_from_ws()
             except ConnectionClosed:
                 if self._should_abort_streams_after_transport_failure():
-                    self._abort_all_streams()
+                    await self.close()
                 if self._retry_connection_callback:
                     self._task_manager.create_task(self._retry_connection_callback())
 
@@ -109,12 +109,12 @@ async def serve(self) -> None:
             except FailedSendingMessageException:
                 # Expected error if the connection is closed.
                 if self._should_abort_streams_after_transport_failure():
-                    self._abort_all_streams()
+                    await self.close()
                 logger.debug(
                     "FailedSendingMessageException while serving", exc_info=True
                 )
             except Exception:
-                self._abort_all_streams()
+                await self.close()
                 logger.exception("caught exception at message iterator")
         except ExceptionGroup as eg:
             _, unhandled = eg.split(lambda e: isinstance(e, ConnectionClosed))

src/replit_river/server_session.py

@@ -87,8 +87,6 @@ async def serve(self) -> None:
                 try:
                     await self._handle_messages_from_ws(tg)
                 except ConnectionClosed:
-                    if self._should_abort_streams_after_transport_failure():
-                        self._abort_all_streams()
                     if self._retry_connection_callback:
                         self._task_manager.create_task(
                             self._retry_connection_callback()
@@ -98,13 +96,10 @@ async def serve(self) -> None:
                     logger.debug("ConnectionClosed while serving", exc_info=True)
                 except FailedSendingMessageException:
                     # Expected error if the connection is closed.
-                    if self._should_abort_streams_after_transport_failure():
-                        self._abort_all_streams()
                     logger.debug(
                         "FailedSendingMessageException while serving", exc_info=True
                     )
                 except Exception:
-                    self._abort_all_streams()
                     logger.exception("caught exception at message iterator")
         except ExceptionGroup as eg:
             _, unhandled = eg.split(lambda e: isinstance(e, ConnectionClosed))

src/replit_river/server_transport.py

@@ -153,35 +153,18 @@ async def _get_or_create_session(
                         close_session_callback=self._delete_session,
                     )
                 else:
-                    if not await old_session.is_session_open():
-                        logger.info(
-                            'Session "%s" is not active, creating replacement '
-                            "session %s instead of reusing",
-                            to_id,
-                            session_id,
-                        )
-                        new_session = ServerSession(
-                            transport_id,
-                            to_id,
-                            session_id,
-                            websocket,
-                            self._transport_options,
-                            self._handlers,
-                            close_session_callback=self._delete_session,
-                        )
-                    else:
-                        # If the instance id is the same, we reuse the session and assign
-                        # a new websocket to it.
-                        logger.debug(
-                            'Reuse old session with "%s" using new ws: %s',
-                            to_id,
-                            websocket.id,
-                        )
-                        try:
-                            await old_session.replace_with_new_websocket(websocket)
-                            new_session = old_session
-                        except FailedSendingMessageException as e:
-                            raise e
+                    # If the instance id is the same, we reuse the session and assign
+                    # a new websocket to it.
+                    logger.debug(
+                        'Reuse old session with "%s" using new ws: %s',
+                        to_id,
+                        websocket.id,
+                    )
+                    try:
+                        await old_session.replace_with_new_websocket(websocket)
+                        new_session = old_session
+                    except FailedSendingMessageException as e:
+                        raise e
 
             self._sessions[new_session._to_id] = new_session
 
@@ -328,6 +311,5 @@ async def _establish_handshake(
 
     async def _delete_session(self, session: Session) -> None:
         async with self._session_lock:
-            existing_session = self._sessions.get(session._to_id)
-            if existing_session is session:
+            if session._to_id in self._sessions:
                 del self._sessions[session._to_id]

src/replit_river/session.py

@@ -287,7 +287,7 @@ async def close_websocket(
             self._task_manager.create_task(self._retry_connection_callback())
 
     def _should_abort_streams_after_transport_failure(self) -> bool:
-        return self._retry_connection_callback is None
+        return not self._transport_options.transparent_reconnect
 
     def _abort_all_streams(self) -> None:
         """Close all active stream channels, notifying any waiting consumers."""