chore/schema handshake metadata pt2

pyproject.toml

@@ -26,6 +26,7 @@ nanoid = "^2.0.0"
 pydantic = {git = "https://github.com/pydantic/pydantic.git", rev = "f5d6acfe19fca38fad802458dab2b4c859182d7b"}
 websockets = "^12.0"
 pydantic-core = "^2.20.1"
+msgpack-types = "^0.3.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.0"

replit_river/client.py

@@ -25,7 +25,7 @@ def __init__(
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata: Optional[HandshakeType] = None,
+        handshake_metadata: HandshakeType,
     ) -> None:
         self._client_id = client_id
         self._server_id = server_id

replit_river/client_transport.py

@@ -51,7 +51,7 @@ def __init__(
         client_id: str,
         server_id: str,
         transport_options: TransportOptions,
-        handshake_metadata: Optional[HandshakeType] = None,
+        handshake_metadata: HandshakeType,
     ):
         super().__init__(
             transport_id=client_id,
@@ -226,7 +226,7 @@ async def websocket_closed_callback() -> None:
         try:
             await send_transport_message(
                 TransportMessage(
-                    from_=transport_id,
+                    from_=transport_id,  # type: ignore
                     to=to_id,
                     streamId=stream_id,
                     controlFlags=0,
@@ -276,7 +276,7 @@ async def _establish_handshake(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: Optional[HandshakeType],
+        handshake_metadata: HandshakeType,
         websocket: WebSocketCommonProtocol,
         old_session: Optional[ClientSession],
     ) -> Tuple[

replit_river/codegen/client.py

@@ -60,7 +60,7 @@ class RiverService(BaseModel):
 
 class RiverSchema(BaseModel):
     services: Dict[str, RiverService]
-    handshakeSchema: RiverConcreteType
+    handshakeSchema: Optional[RiverConcreteType]
 
 
 RiverSchemaFile = RootModel[RiverSchema]
@@ -266,10 +266,13 @@ def generate_river_client_module(
         "",
     ]
 
-    (handshake_type, handshake_chunks) = encode_type(
-        schema_root.handshakeSchema, "HandshakeSchema"
-    )
-    chunks.extend(handshake_chunks)
+    if schema_root.handshakeSchema is not None:
+        (handshake_type, handshake_chunks) = encode_type(
+            schema_root.handshakeSchema, "HandshakeSchema"
+        )
+        chunks.extend(handshake_chunks)
+    else:
+        handshake_type = "Literal[None]"
 
     for schema_name, schema in schema_root.services.items():
         current_chunks: List[str] = [

replit_river/messages.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Callable, Coroutine
 
-import msgpack  # type: ignore
+import msgpack
 import websockets
 from pydantic import ValidationError
 from pydantic_core import ValidationError as PydanticCoreValidationError
@@ -43,12 +43,11 @@ async def send_transport_message(
 ) -> None:
     logger.debug("sending a message %r to ws %s", msg, ws)
     try:
-        await ws.send(
-            prefix_bytes
-            + msgpack.packb(
-                msg.model_dump(by_alias=True, exclude_none=True), datetime=True
-            )
+        packed = msgpack.packb(
+            msg.model_dump(by_alias=True, exclude_none=True), datetime=True
         )
+        assert isinstance(packed, bytes)
+        await ws.send(prefix_bytes + packed)
     except websockets.exceptions.ConnectionClosed as e:
         await websocket_closed_callback()
         raise WebsocketClosedException("Websocket closed during send message") from e

replit_river/server_transport.py

@@ -87,7 +87,7 @@ async def _send_handshake_response(
         response_message = TransportMessage(
             streamId=request_message.streamId,
             id=nanoid.generate(),
-            from_=request_message.to,
+            from_=request_message.to,  # type: ignore
             to=request_message.from_,
             seq=0,
             ack=0,

replit_river/session.py

@@ -373,7 +373,7 @@ async def send_message(
         msg = TransportMessage(
             streamId=stream_id,
             id=nanoid.generate(),
-            from_=self._transport_id,
+            from_=self._transport_id,  # type: ignore
             to=self._to_id,
             seq=await self._seq_manager.get_seq_and_increment(),
             ack=await self._seq_manager.get_ack(),

tests/conftest.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator
-from typing import Any, AsyncGenerator, NoReturn
+from typing import Any, AsyncGenerator, Literal
 
 import nanoid  # type: ignore
 import pytest
@@ -36,7 +36,7 @@ def transport_message(
 ) -> TransportMessage:
     return TransportMessage(
         id=str(nanoid.generate()),
-        from_=from_,
+        from_=from_,  # type: ignore
         to=to,
         streamId=streamId,
         seq=seq,
@@ -139,11 +139,12 @@ async def client(
 ) -> AsyncGenerator[Client, None]:
     try:
         async with serve(server.serve, "localhost", 8765):
-            client: Client[NoReturn] = Client(
+            client: Client[Literal[None]] = Client(
                 "ws://localhost:8765",
                 client_id="test_client",
                 server_id="test_server",
                 transport_options=transport_options,
+                handshake_metadata=None,
             )
             try:
                 yield client