feat/typed metadata (#72)

Why
===

From a discussion with Jacky, egged on by the revert of #62, we should
know the schema of metadata that we're threading through.

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

`river.Client` is now bound by `A`, the schema of the structure passed
in to `handshake_metadata`. If this agrees with the generated client's
expectations, we're good to go.

Test plan
=========

Manual testing

Author: blast-hardcheese

replit_river/client.py

@@ -1,6 +1,6 @@
 import logging
 from collections.abc import AsyncIterable, AsyncIterator
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Generic, Optional, TypeVar, Union
 
 from replit_river.client_transport import ClientTransport
 from replit_river.transport_options import TransportOptions
@@ -15,18 +15,21 @@
 logger = logging.getLogger(__name__)
 
 
-class Client:
+HandshakeType = TypeVar("HandshakeType")
+
+
+class Client(Generic[HandshakeType]):
     def __init__(
         self,
         websocket_uri: str,
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata: Optional[Any] = None,
+        handshake_metadata: Optional[HandshakeType] = None,
     ) -> None:
         self._client_id = client_id
         self._server_id = server_id
-        self._transport = ClientTransport(
+        self._transport = ClientTransport[HandshakeType](
             websocket_uri=websocket_uri,
             client_id=client_id,
             server_id=server_id,

replit_river/client_transport.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Any, Optional, Tuple
+from typing import Generic, Optional, Tuple, TypeVar
 
 import websockets
 from pydantic import ValidationError
@@ -41,14 +41,17 @@
 logger = logging.getLogger(__name__)
 
 
-class ClientTransport(Transport):
+HandshakeType = TypeVar("HandshakeType")
+
+
+class ClientTransport(Transport, Generic[HandshakeType]):
     def __init__(
         self,
         websocket_uri: str,
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata: Optional[Any] = None,
+        handshake_metadata: Optional[HandshakeType] = None,
     ):
         super().__init__(
             transport_id=client_id,
@@ -91,7 +94,7 @@ async def _establish_new_connection(
         old_session: Optional[ClientSession] = None,
     ) -> Tuple[
         WebSocketCommonProtocol,
-        ControlMessageHandshakeRequest,
+        ControlMessageHandshakeRequest[HandshakeType],
         ControlMessageHandshakeResponse,
     ]:
         """Build a new websocket connection with retry logic."""
@@ -204,11 +207,11 @@ async def _send_handshake_request(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: Optional[Any],
+        handshake_metadata: Optional[HandshakeType],
         websocket: WebSocketCommonProtocol,
         expected_session_state: ExpectedSessionState,
-    ) -> ControlMessageHandshakeRequest:
-        handshake_request = ControlMessageHandshakeRequest(
+    ) -> ControlMessageHandshakeRequest[HandshakeType]:
+        handshake_request = ControlMessageHandshakeRequest[HandshakeType](
             type="HANDSHAKE_REQ",
             protocolVersion=PROTOCOL_VERSION,
             sessionId=session_id,
@@ -273,10 +276,12 @@ async def _establish_handshake(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: Optional[Any],
+        handshake_metadata: Optional[HandshakeType],
         websocket: WebSocketCommonProtocol,
         old_session: Optional[ClientSession],
-    ) -> Tuple[ControlMessageHandshakeRequest, ControlMessageHandshakeResponse]:
+    ) -> Tuple[
+        ControlMessageHandshakeRequest[HandshakeType], ControlMessageHandshakeResponse
+    ]:
         try:
             handshake_request = await self._send_handshake_request(
                 transport_id=transport_id,

replit_river/codegen/client.py

@@ -60,6 +60,7 @@ class RiverService(BaseModel):
 
 class RiverSchema(BaseModel):
     services: Dict[str, RiverService]
+    handshakeSchema: RiverConcreteType
 
 
 RiverSchemaFile = RootModel[RiverSchema]
@@ -264,12 +265,21 @@ def generate_river_client_module(
         "import replit_river as river",
         "",
     ]
+
+    (handshake_type, handshake_chunks) = encode_type(
+        schema_root.handshakeSchema, "HandshakeSchema"
+    )
+    chunks.extend(handshake_chunks)
+
     for schema_name, schema in schema_root.services.items():
         current_chunks: List[str] = [
-            f"class {schema_name.title()}Service:",
-            "  def __init__(self, client: river.Client):",
-            "    self.client = client",
-            "",
+            dedent(
+                f"""\
+                  class {schema_name.title()}Service:
+                    def __init__(self, client: river.Client[{handshake_type}]):
+                      self.client = client
+                """
+            ),
         ]
         for name, procedure in schema.procedures.items():
             init_type: Optional[str] = None
@@ -309,13 +319,13 @@ def generate_river_client_module(
                                     .validate_python(
                                         x # type: ignore[arg-type]
                                     )
-                """.strip()
+                """.rstrip()
             parse_error_method = f"""\
                                 lambda x: TypeAdapter({error_type})
                                     .validate_python(
                                         x # type: ignore[arg-type]
                                     )
-                """.strip()
+                """.rstrip()
 
             if output_type == "None":
                 parse_output_method = "lambda x: None"
@@ -506,8 +516,12 @@ async def {name}(
 
     chunks.extend(
         [
-            f"class {client_name}:",
-            "  def __init__(self, client: river.Client):",
+            dedent(
+                f"""\
+                class {client_name}:
+                  def __init__(self, client: river.Client[{handshake_type}]):
+                """.rstrip()
+            )
         ]
     )
     for schema_name, schema in schema_root.services.items():

replit_river/rpc.py

@@ -7,6 +7,7 @@
     Callable,
     Coroutine,
     Dict,
+    Generic,
     Iterable,
     Literal,
     Mapping,
@@ -61,12 +62,15 @@ class ExpectedSessionState(BaseModel):
     nextSentSeq: Optional[int] = None
 
 
-class ControlMessageHandshakeRequest(BaseModel):
+HandshakeType = TypeVar("HandshakeType")
+
+
+class ControlMessageHandshakeRequest(BaseModel, Generic[HandshakeType]):
     type: Literal["HANDSHAKE_REQ"] = "HANDSHAKE_REQ"
     protocolVersion: str
     sessionId: str
     expectedSessionState: ExpectedSessionState
-    metadata: Optional[Any] = None
+    metadata: Optional[HandshakeType] = None
 
 
 class HandShakeStatus(BaseModel):

replit_river/server_transport.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Tuple
+from typing import Any, Tuple
 
 import nanoid  # type: ignore  # type: ignore
 from pydantic import ValidationError
@@ -113,11 +113,11 @@ async def websocket_closed_callback() -> None:
     async def _establish_handshake(
         self, request_message: TransportMessage, websocket: WebSocketCommonProtocol
     ) -> Tuple[
-        ControlMessageHandshakeRequest,
+        ControlMessageHandshakeRequest[Any],
         ControlMessageHandshakeResponse,
     ]:
         try:
-            handshake_request = ControlMessageHandshakeRequest(
+            handshake_request = ControlMessageHandshakeRequest[Any](
                 **request_message.payload
             )
             logger.debug('Got handshake request "%r"', handshake_request)

tests/conftest.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator
-from typing import Any, AsyncGenerator
+from typing import Any, AsyncGenerator, NoReturn
 
 import nanoid  # type: ignore
 import pytest
@@ -139,7 +139,7 @@ async def client(
 ) -> AsyncGenerator[Client, None]:
     try:
         async with serve(server.serve, "localhost", 8765):
-            client = Client(
+            client: Client[NoReturn] = Client(
                 "ws://localhost:8765",
                 client_id="test_client",
                 server_id="test_server",