Make ensure_connected callable based on our own internal state

Author: blast-hardcheese

src/replit_river/v2/client_transport.py

@@ -57,16 +57,15 @@ async def get_or_create_session(self) -> Session:
                 transport_options=self._transport_options,
                 close_session_callback=self._delete_session,
                 retry_connection_callback=self._retry_connection,
+                uri_and_metadata_factory=self._uri_and_metadata_factory,
+                rate_limiter=self._rate_limiter,
+                client_id=self._client_id,
             )
 
             self._session = new_session
             existing_session = new_session
 
-        await existing_session.ensure_connected(
-            client_id=self._client_id,
-            rate_limiter=self._rate_limiter,
-            uri_and_metadata_factory=self._uri_and_metadata_factory,
-        )
+        await existing_session.ensure_connected()
         return existing_session
 
     async def _retry_connection(self) -> Session:

src/replit_river/v2/session.py

@@ -119,7 +119,7 @@ class _IgnoreMessage:
     pass
 
 
-class Session:
+class Session[HandshakeMetadata]:
     _transport_id: str
     _to_id: str
     session_id: str
@@ -132,6 +132,12 @@ class Session:
     _connecting_task: asyncio.Task[None] | None
     _wait_for_connected: asyncio.Event
 
+    _client_id: str
+    _rate_limiter: LeakyBucketRateLimit
+    _uri_and_metadata_factory: Callable[
+        [], Awaitable[UriAndMetadata[HandshakeMetadata]]
+    ]
+
     # ws state
     _ws: ClientConnection | None
     _heartbeat_misses: int
@@ -161,6 +167,11 @@ def __init__(
         session_id: str,
         transport_options: TransportOptions,
         close_session_callback: CloseSessionCallback,
+        client_id: str,
+        rate_limiter: LeakyBucketRateLimit,
+        uri_and_metadata_factory: Callable[
+            [], Awaitable[UriAndMetadata[HandshakeMetadata]]
+        ],
         retry_connection_callback: RetryConnectionCallback | None = None,
     ) -> None:
         self._transport_id = transport_id
@@ -175,6 +186,10 @@ def __init__(
         self._connecting_task = None
         self._wait_for_connected = asyncio.Event()
 
+        self._client_id = client_id
+        self._rate_limiter = rate_limiter
+        self._uri_and_metadata_factory = uri_and_metadata_factory
+
         # ws state
         self._ws = None
         self._heartbeat_misses = 0
@@ -204,14 +219,7 @@ def __init__(
         self._start_close_session_checker()
         self._start_buffered_message_sender()
 
-    async def ensure_connected[HandshakeMetadata](
-        self,
-        client_id: str,
-        rate_limiter: LeakyBucketRateLimit,
-        uri_and_metadata_factory: Callable[
-            [], Awaitable[UriAndMetadata[HandshakeMetadata]]
-        ],
-    ) -> None:
+    async def ensure_connected(self) -> None:
         """
         Either return immediately or establish a websocket connection and return
         once we can accept messages.
@@ -279,12 +287,12 @@ def finalize_attempt() -> None:
             self._connecting_task = asyncio.create_task(
                 _do_ensure_connected(
                     transport_id=self._transport_id,
-                    client_id=client_id,
+                    client_id=self._client_id,
                     to_id=self._to_id,
                     session_id=self.session_id,
                     max_retry=self._transport_options.connection_retry_options.max_retry,
-                    rate_limiter=rate_limiter,
-                    uri_and_metadata_factory=uri_and_metadata_factory,
+                    rate_limiter=self._rate_limiter,
+                    uri_and_metadata_factory=self._uri_and_metadata_factory,
                     get_next_sent_seq=get_next_sent_seq,
                     get_current_ack=lambda: self.ack,
                     get_current_time=self._get_current_time,