Combine uri and metadata factory (#87)

Why
===

It's unnecessary to have two functions. With the way our metadata and
websocket uri are setup, it's just added complexity

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

- Make `handshake_metadata_factory` and `websocket_uri_factory` a single
`uri_and_metadata_factory` that grabs both.
- Returns a new TypeDict `UriAndMetadata`
- Rename `HandshakeType` to `HandshakeMetadataType`
- `HandshakeMetadataType` `TypeVar` is reused (not sure if i'm doing
some python no-no here). Also moved it to the options module to avoid
circular imports

Author: masad-frost

replit_river/client.py

@@ -1,9 +1,13 @@
 import logging
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
-from typing import Any, Generic, Optional, TypeVar, Union
+from typing import Any, Generic, Optional, Union
 
 from replit_river.client_transport import ClientTransport
-from replit_river.transport_options import TransportOptions
+from replit_river.transport_options import (
+    HandshakeMetadataType,
+    TransportOptions,
+    UriAndMetadata,
+)
 
 from .rpc import (
     ErrorType,
@@ -15,26 +19,23 @@
 logger = logging.getLogger(__name__)
 
 
-HandshakeType = TypeVar("HandshakeType")
-
-
-class Client(Generic[HandshakeType]):
+class Client(Generic[HandshakeMetadataType]):
     def __init__(
         self,
-        websocket_uri_factory: Callable[[], Awaitable[str]],
+        uri_and_metadata_factory: Callable[
+            [], Awaitable[UriAndMetadata[HandshakeMetadataType]]
+        ],
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata_factory: Callable[[], Awaitable[HandshakeType]],
     ) -> None:
         self._client_id = client_id
         self._server_id = server_id
-        self._transport = ClientTransport[HandshakeType](
-            websocket_uri_factory=websocket_uri_factory,
+        self._transport = ClientTransport[HandshakeMetadataType](
+            uri_and_metadata_factory=uri_and_metadata_factory,
             client_id=client_id,
             server_id=server_id,
             transport_options=transport_options,
-            handshake_metadata_factory=handshake_metadata_factory,
         )
 
     async def close(self) -> None:

replit_river/client_transport.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import Awaitable, Callable
-from typing import Generic, Optional, Tuple, TypeVar
+from typing import Generic, Optional, Tuple
 
 import websockets
 from pydantic import ValidationError
@@ -37,35 +37,34 @@
     InvalidMessageException,
 )
 from replit_river.transport import Transport
-from replit_river.transport_options import TransportOptions
+from replit_river.transport_options import (
+    HandshakeMetadataType,
+    TransportOptions,
+    UriAndMetadata,
+)
 
 logger = logging.getLogger(__name__)
 
 
-HandshakeType = TypeVar("HandshakeType")
-
-
-class ClientTransport(Transport, Generic[HandshakeType]):
+class ClientTransport(Transport, Generic[HandshakeMetadataType]):
     def __init__(
         self,
-        websocket_uri_factory: Callable[[], Awaitable[str]],
+        uri_and_metadata_factory: Callable[[], Awaitable[UriAndMetadata]],
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata_factory: Callable[[], Awaitable[HandshakeType]],
     ):
         super().__init__(
             transport_id=client_id,
             transport_options=transport_options,
             is_server=False,
         )
-        self._websocket_uri_factory = websocket_uri_factory
+        self._uri_and_metadata_factory = uri_and_metadata_factory
         self._client_id = client_id
         self._server_id = server_id
         self._rate_limiter = LeakyBucketRateLimit(
             transport_options.connection_retry_options
         )
-        self._handshake_metadata_factory = handshake_metadata_factory
         # We want to make sure there's only one session creation at a time
         self._create_session_lock = asyncio.Lock()
 
@@ -121,7 +120,7 @@ async def _establish_new_connection(
         old_session: Optional[ClientSession] = None,
     ) -> Tuple[
         WebSocketCommonProtocol,
-        ControlMessageHandshakeRequest[HandshakeType],
+        ControlMessageHandshakeRequest[HandshakeMetadataType],
         ControlMessageHandshakeResponse,
     ]:
         """Build a new websocket connection with retry logic."""
@@ -147,9 +146,8 @@ async def _establish_new_connection(
                 old_session = None
 
             try:
-                websocket_uri = await self._websocket_uri_factory()
-                handshake_metadata = await self._handshake_metadata_factory()
-                ws = await websockets.connect(websocket_uri)
+                uri_and_metadata = await self._uri_and_metadata_factory()
+                ws = await websockets.connect(uri_and_metadata["uri"])
                 session_id = (
                     self.generate_nanoid()
                     if not old_session
@@ -164,7 +162,7 @@ async def _establish_new_connection(
                         self._transport_id,
                         self._server_id,
                         session_id,
-                        handshake_metadata,
+                        uri_and_metadata["metadata"],
                         ws,
                         old_session,
                     )
@@ -223,11 +221,11 @@ async def _send_handshake_request(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: Optional[HandshakeType],
+        handshake_metadata: Optional[HandshakeMetadataType],
         websocket: WebSocketCommonProtocol,
         expected_session_state: ExpectedSessionState,
-    ) -> ControlMessageHandshakeRequest[HandshakeType]:
-        handshake_request = ControlMessageHandshakeRequest[HandshakeType](
+    ) -> ControlMessageHandshakeRequest[HandshakeMetadataType]:
+        handshake_request = ControlMessageHandshakeRequest[HandshakeMetadataType](
             type="HANDSHAKE_REQ",
             protocolVersion=PROTOCOL_VERSION,
             sessionId=session_id,
@@ -292,11 +290,12 @@ async def _establish_handshake(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: HandshakeType,
+        handshake_metadata: HandshakeMetadataType,
         websocket: WebSocketCommonProtocol,
         old_session: Optional[ClientSession],
     ) -> Tuple[
-        ControlMessageHandshakeRequest[HandshakeType], ControlMessageHandshakeResponse
+        ControlMessageHandshakeRequest[HandshakeMetadataType],
+        ControlMessageHandshakeResponse,
     ]:
         try:
             handshake_request = await self._send_handshake_request(

replit_river/rpc.py

@@ -30,7 +30,10 @@
     stringify_exception,
 )
 from replit_river.task_manager import BackgroundTaskManager
-from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE
+from replit_river.transport_options import (
+    MAX_MESSAGE_BUFFER_SIZE,
+    HandshakeMetadataType,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -62,15 +65,12 @@ class ExpectedSessionState(BaseModel):
     nextSentSeq: Optional[int] = None
 
 
-HandshakeType = TypeVar("HandshakeType")
-
-
-class ControlMessageHandshakeRequest(BaseModel, Generic[HandshakeType]):
+class ControlMessageHandshakeRequest(BaseModel, Generic[HandshakeMetadataType]):
     type: Literal["HANDSHAKE_REQ"] = "HANDSHAKE_REQ"
     protocolVersion: str
     sessionId: str
     expectedSessionState: ExpectedSessionState
-    metadata: Optional[HandshakeType] = None
+    metadata: Optional[HandshakeMetadataType] = None
 
 
 class HandShakeStatus(BaseModel):

replit_river/transport_options.py

@@ -1,4 +1,5 @@
 import os
+from typing import Generic, TypedDict, TypeVar
 
 from pydantic import BaseModel
 
@@ -50,3 +51,11 @@ def create_from_env(cls) -> "TransportOptions":
             heartbeat_ms=heartbeat_ms,
             heartbeats_until_dead=heartbeats_to_dead,
         )
+
+
+HandshakeMetadataType = TypeVar("HandshakeMetadataType")
+
+
+class UriAndMetadata(TypedDict, Generic[HandshakeMetadataType]):
+    uri: str
+    metadata: HandshakeMetadataType

tests/conftest.py

@@ -8,6 +8,7 @@
 from websockets.server import serve
 
 from replit_river.client import Client
+from replit_river.client_transport import UriAndMetadata
 from replit_river.error_schema import RiverError
 from replit_river.rpc import (
     GrpcContext,
@@ -138,20 +139,19 @@ async def client(
     no_logging_error: NoErrors,
 ) -> AsyncGenerator[Client, None]:
 
-    async def websocket_uri_factory() -> str:
-        return "ws://localhost:8765"
-
-    async def handshake_metadata_factory() -> None:
-        return None
+    async def websocket_uri_factory() -> UriAndMetadata[None]:
+        return {
+            "uri": "ws://localhost:8765",
+            "metadata": None,
+        }
 
     try:
         async with serve(server.serve, "localhost", 8765):
-            client: Client[Literal[None]] = Client(
-                websocket_uri_factory,
+            client: Client[Literal[None]] = Client[None](
+                uri_and_metadata_factory=websocket_uri_factory,
                 client_id="test_client",
                 server_id="test_server",
                 transport_options=transport_options,
-                handshake_metadata_factory=handshake_metadata_factory,
             )
             try:
                 yield client