Remove is_server

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -5,10 +5,10 @@
 from typing import Any, AsyncGenerator, Callable, Coroutine
 
 import nanoid  # type: ignore
+import websockets
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span
-import websockets
 from websockets.exceptions import ConnectionClosed
 
 from replit_river.error_schema import (
@@ -72,6 +72,15 @@ def __init__(
             retry_connection_callback=retry_connection_callback,
         )
 
+        async def do_close_websocket() -> None:
+            await self.close_websocket(
+                self._ws_wrapper,
+                should_retry=True,
+            )
+            await self._begin_close_session_countdown()
+
+        self._setup_heartbeats_task(do_close_websocket)
+
     async def start_serve_responses(self) -> None:
         self._task_manager.create_task(self.serve())
 

src/replit_river/server_session.py

@@ -50,7 +50,7 @@ def __init__(
         websocket: websockets.WebSocketCommonProtocol,
         transport_options: TransportOptions,
         handlers: dict[tuple[str, str], tuple[str, GenericRpcHandlerBuilder]],
-        close_session_callback: Callable[["Session"], Coroutine[Any, Any, Any]],
+        close_session_callback: Callable[[Session], Coroutine[Any, Any, Any]],
         retry_connection_callback: (
             Callable[
                 [],
@@ -68,9 +68,17 @@ def __init__(
             close_session_callback=close_session_callback,
             retry_connection_callback=retry_connection_callback,
         )
-        self._is_server = True
         self._handlers = handlers
 
+        async def do_close_websocket() -> None:
+            await self.close_websocket(
+                self._ws_wrapper,
+                should_retry=False,
+            )
+            await self._begin_close_session_countdown()
+
+        self._setup_heartbeats_task(do_close_websocket)
+
     async def start_serve_responses(self) -> None:
         self._task_manager.create_task(self.serve())
 

src/replit_river/session.py

@@ -1,7 +1,7 @@
 import asyncio
 import enum
 import logging
-from typing import Any, Callable, Coroutine, Protocol
+from typing import Any, Awaitable, Callable, Coroutine, Protocol
 
 import nanoid  # type: ignore
 import websockets
@@ -62,7 +62,6 @@ class SessionState(enum.Enum):
 
 class Session:
     """Common functionality shared between client_session and server_session"""
-    _is_server: bool
 
     def __init__(
         self,
@@ -83,7 +82,6 @@ def __init__(
         self._transport_id = transport_id
         self._to_id = to_id
         self.session_id = session_id
-        self._is_server = False
         self._transport_options = transport_options
 
         # session state, only modified during closing
@@ -108,16 +106,10 @@ def __init__(
         self._buffer = MessageBuffer(self._transport_options.buffer_size)
         self._task_manager = BackgroundTaskManager()
 
-        self._setup_heartbeats_task()
-
-    def _setup_heartbeats_task(self) -> None:
-        async def do_close_websocket() -> None:
-            await self.close_websocket(
-                self._ws_wrapper,
-                should_retry=not self._is_server,
-            )
-            await self._begin_close_session_countdown()
-
+    def _setup_heartbeats_task(
+        self,
+        do_close_websocket: Callable[[], Awaitable[None]],
+    ) -> None:
         def increment_and_get_heartbeat_misses() -> int:
             self._heartbeat_misses += 1
             return self._heartbeat_misses