resolve merge

Author: jackyzha0

src/replit_river/server_transport.py

@@ -110,10 +110,6 @@ async def handshake_to_get_session(
     async def close(self) -> None:
         await self._close_all_sessions()
 
-    async def _get_existing_session(self, to_id: str) -> ServerSession | None:
-        async with self._session_lock:
-            return self._sessions.get(to_id)
-
     async def _get_or_create_session(
         self,
         transport_id: str,
@@ -251,57 +247,68 @@ async def _establish_handshake(
             )
             raise InvalidMessageException("handshake request to wrong server")
 
-        old_session = await self._get_existing_session(request_message.from_)
-        client_next_expected_seq = (
-            handshake_request.expectedSessionState.nextExpectedSeq
-        )
-        client_next_sent_seq = handshake_request.expectedSessionState.nextSentSeq or 0
-        if old_session and old_session.session_id == handshake_request.sessionId:
-            # check invariants
-            # ordering must be correct
-            our_next_seq = await old_session.get_next_sent_seq()
-            our_ack = await old_session.get_next_expected_seq()
+        async with self._session_lock:
+            old_session = self._sessions.get(request_message.from_)
+            client_next_expected_seq = (
+                handshake_request.expectedSessionState.nextExpectedSeq
+            )
+            client_next_sent_seq = (
+                handshake_request.expectedSessionState.nextSentSeq or 0
+            )
+            if old_session and old_session.session_id == handshake_request.sessionId:
+                # check invariants
+                # ordering must be correct
+                our_next_seq = await old_session.get_next_sent_seq()
+                our_ack = await old_session.get_next_expected_seq()
 
-            if client_next_sent_seq > our_ack:
-                message = (
-                    "client is in the future: "
-                    f"server wanted {our_ack} but client has {client_next_sent_seq}"
-                )
-                await self._send_handshake_response(
-                    request_message,
-                    HandShakeStatus(ok=False, reason=message),
-                    websocket,
-                )
-                raise SessionStateMismatchException(message)
+                if client_next_sent_seq > our_ack:
+                    message = (
+                        "client is in the future: "
+                        f"server wanted {our_ack} but client has {client_next_sent_seq}"
+                    )
+                    await self._send_handshake_response(
+                        request_message,
+                        HandShakeStatus(ok=False, reason=message),
+                        websocket,
+                    )
+                    raise SessionStateMismatchException(message)
 
-            if our_next_seq > client_next_expected_seq:
-                message = (
-                    "server is in the future: "
-                    f"client wanted {client_next_expected_seq} "
-                    f"but server has {our_next_seq}"
-                )
+                if our_next_seq > client_next_expected_seq:
+                    message = (
+                        "server is in the future: "
+                        f"client wanted {client_next_expected_seq} "
+                        f"but server has {our_next_seq}"
+                    )
+                    await self._send_handshake_response(
+                        request_message,
+                        HandShakeStatus(ok=False, reason=message),
+                        websocket,
+                    )
+                    raise SessionStateMismatchException(message)
+            elif old_session:
+                # we have an old session but the session id is different
+                # just delete the old session
+                await old_session.close()
+                old_session = None
+
+            if not old_session and (
+                client_next_sent_seq > 0 or client_next_expected_seq > 0
+            ):
+                message = "client is trying to resume a session but we don't have it"
                 await self._send_handshake_response(
                     request_message,
                     HandShakeStatus(ok=False, reason=message),
                     websocket,
                 )
                 raise SessionStateMismatchException(message)
-        elif old_session:
-            # we have an old session but the session id is different
-            # just delete the old session
-            await old_session.close()
-            old_session = None
 
-        if not old_session and (
-            client_next_sent_seq > 0 or client_next_expected_seq > 0
-        ):
-            message = "client is trying to resume a session but we don't have it"
-            await self._send_handshake_response(
+            # from this point on, we're committed to connecting
+            session_id = handshake_request.sessionId
+            handshake_response = await self._send_handshake_response(
                 request_message,
-                HandShakeStatus(ok=False, reason=message),
+                HandShakeStatus(ok=True, sessionId=session_id),
                 websocket,
             )
-            raise SessionStateMismatchException(message)
 
             return handshake_request, handshake_response