Differentiating between closing and closed

Author: blast-hardcheese

src/replit_river/v2/client_transport.py

@@ -54,8 +54,10 @@ async def get_or_create_session(self) -> Session:
         call ensure_connected on whatever session is active.
         """
         existing_session = self._session
-        if not existing_session or existing_session.is_closed():
+        if not existing_session or existing_session.is_terminal():
             logger.info("Creating new session")
+            if existing_session:
+                await existing_session.close()
             new_session = Session(
                 client_id=self._client_id,
                 server_id=self._server_id,

src/replit_river/v2/session.py

@@ -324,7 +324,7 @@ def unbind_connecting_task() -> None:
         if self._terminating_task:
             await self._terminating_task
 
-    def is_closed(self) -> bool:
+    def is_terminal(self) -> bool:
         """
         If the session is in a terminal state.
         Do not send messages, do not expect any more messages to be emitted,
@@ -402,12 +402,15 @@ async def close(
         self, reason: Exception | None = None, current_state: SessionState | None = None
     ) -> None:
         """Close the session and all associated streams."""
-        logger.info(
-            f"{self.session_id} closing session to {self._server_id}, ws: {self._ws}"
-        )
         if (current_state or self._state) in TerminalStates:
+            while (current_state or self._state) != SessionState.CLOSED:
+                logger.debug("Session already closing, waiting...")
+                await asyncio.sleep(0.2)
             # already closing
             return
+        logger.info(
+            f"{self.session_id} closing session to {self._server_id}, ws: {self._ws}"
+        )
         self._state = SessionState.CLOSING
 
         # We're closing, so we need to wake up...