Short-circuit to prevent internal deadlocks

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -405,12 +405,18 @@ async def _enqueue_message(
         self._process_messages.set()
 
     async def close(
-        self, reason: Exception | None = None, current_state: SessionState | None = None
+        self,
+        reason: Exception | None = None,
+        current_state: SessionState | None = None,
+        _wait_for_closed: bool = True,
     ) -> None:
         """Close the session and all associated streams."""
         if (current_state or self._state) in TerminalStates:
             start = datetime.now()
-            while (current_state or self._state) != SessionState.CLOSED:
+            while (
+                _wait_for_closed
+                and (current_state or self._state) != SessionState.CLOSED
+            ):
                 elapsed = (datetime.now() - start).total_seconds()
                 if elapsed >= SESSION_CLOSE_TIMEOUT_SEC:
                     logger.warning(
@@ -632,7 +638,7 @@ async def block_until_connected() -> None:
                 get_state=lambda: self._state,
                 get_ws=lambda: self._ws,
                 transition_no_connection=transition_no_connection,
-                close_session=self.close,
+                close_session=lambda err: self.close(err, _wait_for_closed=False),
                 assert_incoming_seq_bookkeeping=assert_incoming_seq_bookkeeping,
                 get_stream=lambda stream_id: self._streams.get(stream_id),
                 enqueue_message=self._enqueue_message,