Translating a little more of PROTOCOL into code

Author: blast-hardcheese

src/replit_river/rpc.py

@@ -46,8 +46,10 @@
 GenericRpcHandlerBuilder = Callable[
     [str, Channel[Any], Channel[Any]], Coroutine[None, None, None]
 ]
-ACK_BIT = 0b00001
-STREAM_OPEN_BIT = 0b00010
+ACK_BIT_TYPE = Literal[0b00001]
+ACK_BIT: ACK_BIT_TYPE = 0b00001
+STREAM_OPEN_BIT_TYPE = Literal[0b00010]
+STREAM_OPEN_BIT: STREAM_OPEN_BIT_TYPE = 0b00010
 
 # these codes are retriable
 # if the server sends a response with one of these codes,

src/replit_river/v2/client_session.py

@@ -2,7 +2,7 @@
 import logging
 from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Callable, Coroutine
+from typing import Any, AsyncGenerator, Callable, Coroutine, Literal
 
 import nanoid  # type: ignore
 import websockets
@@ -37,8 +37,10 @@
 from replit_river.session import Session
 from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE, TransportOptions
 
-STREAM_CANCEL_BIT = 0b00100  # Synonymous with the cancel bit in v2
-STREAM_CLOSED_BIT = 0b01000  # Synonymous with the cancel bit in v2
+STREAM_CANCEL_BIT_TYPE = Literal[0b00100]
+STREAM_CANCEL_BIT: STREAM_CANCEL_BIT_TYPE = 0b00100
+STREAM_CLOSED_BIT_TYPE = Literal[0b01000]
+STREAM_CLOSED_BIT: STREAM_CLOSED_BIT_TYPE = 0b01000
 
 
 logger = logging.getLogger(__name__)

src/replit_river/v2/schema.py

@@ -1,6 +1,9 @@
 from typing import Any, Literal, NotRequired, TypeAlias, TypedDict
 
-from replit_river.rpc import ExpectedSessionState
+from grpc.aio import BaseError
+
+from replit_river.rpc import ACK_BIT, ACK_BIT_TYPE, ExpectedSessionState
+from replit_river.v2.client_session import STREAM_CANCEL_BIT, STREAM_CANCEL_BIT_TYPE
 
 
 class ControlClose(TypedDict):
@@ -33,7 +36,39 @@ class ControlHandshakeResponse(TypedDict):
     type: Literal["HANDSHAKE_RESP"]
     status: HandshakeOK | HandshakeError
 
+# This is sent when the server encounters an internal error
+# i.e. an invariant has been violated
+class BaseErrorStructure(TypedDict):
+  # This should be a defined literal to make sure errors are easily differentiated
+  # code: str  # Supplied by implementations
+  # This can be any string
+  message: str
+  # Any extra metadata
+  extra: NotRequired[Any]
+
+# When a client sends a malformed request. This can be
+# for a variety of reasons which would  be included
+# in the message.
+class InvalidRequestError(BaseErrorStructure):
+  code: Literal['INVALID_REQUEST']
+
+# This is sent when an exception happens in the handler of a stream.
+class UncaughtError(BaseErrorStructure):
+  code: Literal['UNCAUGHT_ERROR']
+
+# This is sent when one side wishes to cancel the stream
+# abruptly from user-space. Handling this is up to the procedure
+# implementation or the caller.
+class CancelError(BaseErrorStructure):
+  code: Literal['CANCEL']
+
+ProtocolError: TypeAlias = UncaughtError | InvalidRequestError | CancelError;
 
 Control: TypeAlias = (
     ControlClose | ControlAck | ControlHandshakeRequest | ControlHandshakeResponse
 )
+
+ValidPairings = (
+    tuple[ACK_BIT_TYPE, ControlAck] |
+    tuple[STREAM_CANCEL_BIT_TYPE, ProtocolError]
+)