Adding some more safeguards

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -1090,11 +1090,17 @@ async def _do_ensure_connected[HandshakeMetadata](
         get_current_time() + transport_options.handshake_timeout_ms / 1000
     )
 
-    while rate_limiter.has_budget(client_id):
+    while (
+        rate_limiter.has_budget(client_id)
+        and get_current_time() < handshake_deadline_ms
+    ):
         if (state := get_state()) in TerminalStates or state in ActiveStates:
             logger.info(f"_do_ensure_connected stopping due to state={state}")
             break
 
+        if (task := asyncio.current_task()) and task.cancelled():
+            break
+
         if attempt_count > 0:
             logger.info(f"Retrying build handshake number {attempt_count} times")
         attempt_count += 1

tests/v2/test_v2_session_lifecycle.py

@@ -55,6 +55,7 @@ def trigger_close(
         nonlocal ws_close
 
         closing_event = asyncio.Event()
+
         async def _do_close() -> None:
             signal_closing()
             await task_manager.cancel_all_tasks()
@@ -64,6 +65,7 @@ async def _do_close() -> None:
                 await ws.close()
             become_closed()
             closing_event.set()
+
         ws_close = asyncio.create_task(_do_close())
 
         return closing_event