Skipping heartbeats when we know we're closed

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -47,6 +47,7 @@ async def setup_heartbeat(
     heartbeat_ms: float,
     heartbeats_until_dead: int,
     get_state: Callable[[], SessionState],
+    get_connected: Callable[[], bool],
     get_closing_grace_period: Callable[[], float | None],
     close_websocket: Callable[[], Awaitable[None]],
     send_message: SendMessage,
@@ -56,6 +57,11 @@ async def setup_heartbeat(
     while True:
         await asyncio.sleep(heartbeat_ms / 1000)
         state = get_state()
+        if not get_connected():
+            logger.debug(
+                "Websocket is not connected, not sending heartbeat"
+            )
+            continue
         if state != SessionState.ACTIVE:
             logger.debug(
                 "Session is closed, no need to send heartbeat, state : "

src/replit_river/v2/session.py

@@ -115,6 +115,7 @@ def increment_and_get_heartbeat_misses() -> int:
                 self._transport_options.heartbeat_ms,
                 self._transport_options.heartbeats_until_dead,
                 lambda: self._state,
+                lambda: self._ws_connected,
                 lambda: self._close_session_after_time_secs,
                 close_websocket=do_close_websocket,
                 send_message=self.send_message,