Missing wait_for_connected.clear()

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -462,6 +462,8 @@ def get_ws() -> ClientConnection | None:
             return None
 
         async def block_until_connected() -> None:
+            if self._state in TerminalStates:
+                return
             logger.debug("block_until_connected")
             await self._wait_for_connected.wait()
             logger.debug("block_until_connected released!")
@@ -486,6 +488,7 @@ async def transition_no_connection() -> None:
             if self._state in TerminalStates:
                 return
             self._state = SessionState.NO_CONNECTION
+            self._wait_for_connected.clear()
             if self._ws:
                 self._task_manager.create_task(self._ws.close())
                 self._ws = None