clear session if the session is closed to avoid using closed sessions (#81)

Why
===

when we get a `SESSION_STATE_MISMATCH`, we close the old session and
delete it, but the recreate fn still has a handle to it

_Describe what prompted you to make this change, link relevant
resources: Linear issues, Slack discussions, etc._

What changed
============

if the retry loop detects the old session to be closed, stop using it

_Describe what changed to a level of detail that someone with no context
with your PR could be able to review it_

Test plan
=========

_Describe what you did to test this change to a level of detail that
allows your reviewer to test it_

Author: jackyzha0

replit_river/client_transport.py

@@ -102,13 +102,19 @@ async def _establish_new_connection(
         max_retry = self._transport_options.connection_retry_options.max_retry
         client_id = self._client_id
         logger.info("Attempting to establish new ws connection")
+
         for i in range(max_retry):
             if i > 0:
                 logger.info(f"Retrying build handshake number {i} times")
             if not rate_limit.has_budget(client_id):
                 logger.debug("No retry budget for %s.", client_id)
                 break
             rate_limit.consume_budget(client_id)
+
+            # if the session is closed, we shouldn't use it
+            if old_session and not await old_session.is_session_open():
+                old_session = None
+
             try:
                 ws = await websockets.connect(self._websocket_uri)
                 session_id = (
@@ -333,6 +339,8 @@ async def _establish_handshake(
                 await self._delete_session(old_session)
 
             raise RiverException(
-                ERROR_HANDSHAKE, f"Handshake failed: {handshake_response.status.reason}"
+                ERROR_HANDSHAKE,
+                f"Handshake failed with code ${handshake_response.status.code}: "
+                + f"{handshake_response.status.reason}",
             )
         return handshake_request, handshake_response

replit_river/session.py

@@ -124,7 +124,7 @@ async def _begin_close_session_countdown(self) -> None:
         if self._close_session_after_time_secs is not None:
             # already in grace period, no need to set again
             return
-        logger.debug(
+        logger.info(
             "websocket closed from %s to %s begin grace period",
             self._transport_id,
             self._to_id,
@@ -133,6 +133,7 @@ async def _begin_close_session_countdown(self) -> None:
 
     async def serve(self) -> None:
         """Serve messages from the websocket."""
+        self._reset_session_close_countdown()
         try:
             async with asyncio.TaskGroup() as tg:
                 try:
@@ -226,7 +227,6 @@ async def replace_with_new_websocket(
             old_wrapper = self._ws_wrapper
             old_ws_id = old_wrapper.ws.id
             if new_ws.id != old_ws_id:
-                self._reset_session_close_countdown()
                 await old_wrapper.close()
             self._ws_wrapper = WebsocketWrapper(new_ws)
         await self._send_buffered_messages(new_ws)