Merge branch 'main' into th-stricter-error-types

Author: airportyh

scripts/lint/src/lint/__init__.py

@@ -11,7 +11,7 @@ def raise_err(code: int) -> None:
 
 def main() -> None:
     fix = ["--fix"] if "--fix" in sys.argv else []
-    raise_err(os.system(" ".join(["ruff", "check", "src"] + fix)))
-    raise_err(os.system("ruff format src"))
+    raise_err(os.system(" ".join(["ruff", "check", "src", "scripts", "tests"] + fix)))
+    raise_err(os.system("ruff format src scripts tests"))
     raise_err(os.system("mypy src"))
     raise_err(os.system("pyright src"))

scripts/parity/check_parity.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Literal, TypedDict, TypeVar, Union
+from typing import Any, Callable, Literal, TypedDict, TypeVar
 
 import pyd
 import tyd
@@ -85,35 +85,37 @@ def testAgenttoollanguageserverOpendocumentInput() -> None:
     )
 
 
-kind_type = Union[
-    Literal[1],
-    Literal[2],
-    Literal[3],
-    Literal[4],
-    Literal[5],
-    Literal[6],
-    Literal[7],
-    Literal[8],
-    Literal[9],
-    Literal[10],
-    Literal[11],
-    Literal[12],
-    Literal[13],
-    Literal[14],
-    Literal[15],
-    Literal[16],
-    Literal[17],
-    Literal[18],
-    Literal[19],
-    Literal[20],
-    Literal[21],
-    Literal[22],
-    Literal[23],
-    Literal[24],
-    Literal[25],
-    Literal[26],
-    None,
-]
+kind_type = (
+    Literal[
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        26,
+    ]
+    | None
+)
 
 
 def testAgenttoollanguageserverGetcodesymbolInput() -> None:

scripts/parity/gen.py

@@ -1,6 +1,6 @@
 import random
 import string
-from typing import Callable, Optional, TypeVar
+from typing import Callable, TypeVar
 
 A = TypeVar("A")
 
@@ -37,7 +37,7 @@ def gen_choice(choices: list[A]) -> Callable[[], A]:
     return lambda: random.choice(choices)
 
 
-def gen_opt(gen_x: Callable[[], A]) -> Callable[[], Optional[A]]:
+def gen_opt(gen_x: Callable[[], A]) -> Callable[[], A | None]:
     return lambda: gen_x() if gen_bool() else None
 
 

src/replit_river/__init__.py

@@ -1,7 +1,7 @@
 from .client import Client
 from .error_schema import RiverError
 from .rpc import (
-    GenericRpcHandler,
+    GenericRpcHandlerBuilder,
     GrpcContext,
     rpc_method_handler,
     stream_method_handler,
@@ -15,7 +15,7 @@
     "Server",
     "GrpcContext",
     "RiverError",
-    "GenericRpcHandler",
+    "GenericRpcHandlerBuilder",
     "rpc_method_handler",
     "subscription_method_handler",
     "upload_method_handler",

src/replit_river/client.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Generator, Generic, Literal, Optional, Union
+from typing import Any, AsyncGenerator, Generator, Generic, Literal
 
 from opentelemetry import trace
 from opentelemetry.trace import Span, SpanKind, Status, StatusCode
@@ -119,9 +119,9 @@ async def send_upload(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
@@ -148,7 +148,7 @@ async def send_subscription(
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncGenerator[Union[ResponseType, RiverError], None]:
+    ) -> AsyncGenerator[ResponseType | RiverError, None]:
         with _trace_procedure(
             "subscription", service_name, procedure_name
         ) as span_handle:
@@ -170,13 +170,13 @@ async def send_stream(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
-    ) -> AsyncGenerator[Union[ResponseType, RiverError], None]:
+    ) -> AsyncGenerator[ResponseType | RiverError, None]:
         with _trace_procedure("stream", service_name, procedure_name) as span_handle:
             session = await self._transport.get_or_create_session()
             async for msg in session.send_stream(
@@ -204,8 +204,8 @@ class _SpanHandle:
 
     def set_status(
         self,
-        status: Union[Status, StatusCode],
-        description: Optional[str] = None,
+        status: Status | StatusCode,
+        description: str | None = None,
     ) -> None:
         if self.did_set_status:
             return

src/replit_river/client_session.py

@@ -2,13 +2,16 @@
 import logging
 from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Callable, Optional, Union
+from typing import Any, AsyncGenerator, Callable, Coroutine
 
 import nanoid  # type: ignore
+import websockets
 from aiochannel import Channel
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span
+from websockets.exceptions import ConnectionClosed
 
+from replit_river.common_session import add_msg_to_stream
 from replit_river.error_schema import (
     ERROR_CODE_CANCEL,
     ERROR_CODE_STREAM_CLOSED,
@@ -17,10 +20,20 @@
     StreamClosedRiverServiceException,
     exception_from_message,
 )
+from replit_river.messages import (
+    FailedSendingMessageException,
+    parse_transport_msg,
+)
+from replit_river.seq_manager import (
+    IgnoreMessageException,
+    InvalidMessageException,
+    OutOfOrderMessageException,
+)
 from replit_river.session import Session
-from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE
+from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE, TransportOptions
 
 from .rpc import (
+    ACK_BIT,
     STREAM_CLOSED_BIT,
     STREAM_OPEN_BIT,
     ErrorType,
@@ -33,6 +46,129 @@
 
 
 class ClientSession(Session):
+    def __init__(
+        self,
+        transport_id: str,
+        to_id: str,
+        session_id: str,
+        websocket: websockets.WebSocketCommonProtocol,
+        transport_options: TransportOptions,
+        close_session_callback: Callable[[Session], Coroutine[Any, Any, Any]],
+        retry_connection_callback: (
+            Callable[
+                [],
+                Coroutine[Any, Any, Any],
+            ]
+            | None
+        ) = None,
+    ) -> None:
+        super().__init__(
+            transport_id=transport_id,
+            to_id=to_id,
+            session_id=session_id,
+            websocket=websocket,
+            transport_options=transport_options,
+            close_session_callback=close_session_callback,
+            retry_connection_callback=retry_connection_callback,
+        )
+
+        async def do_close_websocket() -> None:
+            await self.close_websocket(
+                self._ws_wrapper,
+                should_retry=True,
+            )
+            await self._begin_close_session_countdown()
+
+        self._setup_heartbeats_task(do_close_websocket)
+
+    async def start_serve_responses(self) -> None:
+        self._task_manager.create_task(self.serve())
+
+    async def serve(self) -> None:
+        """Serve messages from the websocket."""
+        self._reset_session_close_countdown()
+        try:
+            try:
+                await self._handle_messages_from_ws()
+            except ConnectionClosed:
+                if self._retry_connection_callback:
+                    self._task_manager.create_task(self._retry_connection_callback())
+
+                await self._begin_close_session_countdown()
+                logger.debug("ConnectionClosed while serving", exc_info=True)
+            except FailedSendingMessageException:
+                # Expected error if the connection is closed.
+                logger.debug(
+                    "FailedSendingMessageException while serving", exc_info=True
+                )
+            except Exception:
+                logger.exception("caught exception at message iterator")
+        except ExceptionGroup as eg:
+            _, unhandled = eg.split(lambda e: isinstance(e, ConnectionClosed))
+            if unhandled:
+                raise ExceptionGroup(
+                    "Unhandled exceptions on River server", unhandled.exceptions
+                )
+
+    async def _handle_messages_from_ws(self) -> None:
+        logger.debug(
+            "%s start handling messages from ws %s",
+            "client",
+            self._ws_wrapper.id,
+        )
+        try:
+            ws_wrapper = self._ws_wrapper
+            async for message in ws_wrapper.ws:
+                try:
+                    if not await ws_wrapper.is_open():
+                        # We should not process messages if the websocket is closed.
+                        break
+                    msg = parse_transport_msg(message, self._transport_options)
+
+                    logger.debug(f"{self._transport_id} got a message %r", msg)
+
+                    # Update bookkeeping
+                    await self._seq_manager.check_seq_and_update(msg)
+                    await self._buffer.remove_old_messages(
+                        self._seq_manager.receiver_ack,
+                    )
+                    self._reset_session_close_countdown()
+
+                    if msg.controlFlags & ACK_BIT != 0:
+                        continue
+                    async with self._stream_lock:
+                        stream = self._streams.get(msg.streamId, None)
+                    if msg.controlFlags & STREAM_OPEN_BIT == 0:
+                        if not stream:
+                            logger.warning("no stream for %s", msg.streamId)
+                            raise IgnoreMessageException(
+                                "no stream for message, ignoring"
+                            )
+                        await add_msg_to_stream(msg, stream)
+                    else:
+                        raise InvalidMessageException(
+                            "Client should not receive stream open bit"
+                        )
+
+                    if msg.controlFlags & STREAM_CLOSED_BIT != 0:
+                        if stream:
+                            stream.close()
+                        async with self._stream_lock:
+                            del self._streams[msg.streamId]
+                except IgnoreMessageException:
+                    logger.debug("Ignoring transport message", exc_info=True)
+                    continue
+                except OutOfOrderMessageException:
+                    logger.exception("Out of order message, closing connection")
+                    await ws_wrapper.close()
+                    return
+                except InvalidMessageException:
+                    logger.exception("Got invalid transport message, closing session")
+                    await self.close()
+                    return
+        except ConnectionClosed as e:
+            raise e
+
     async def send_rpc(
         self,
         service_name: str,
@@ -102,9 +238,9 @@ async def send_upload(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
@@ -194,7 +330,7 @@ async def send_subscription(
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
         span: Span,
-    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
+    ) -> AsyncGenerator[ResponseType | ErrorType, None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.
@@ -241,14 +377,14 @@ async def send_stream(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
         span: Span,
-    ) -> AsyncGenerator[Union[ResponseType, ErrorType], None]:
+    ) -> AsyncGenerator[ResponseType | ErrorType, None]:
         """Sends a subscription request to the server.
 
         Expects the input and output be messages that will be msgpacked.