Inline add_msg_to_stream

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -11,7 +11,6 @@
 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,
@@ -154,7 +153,14 @@ async def _handle_messages_from_ws(self) -> None:
                             # close message is not sent to the stream
                             pass
                         else:
-                            await add_msg_to_stream(msg, stream)
+                            try:
+                                await stream.put(msg.payload)
+                            except ChannelClosed:
+                                # The client is no longer interested in this stream,
+                                # just drop the message.
+                                pass
+                            except RuntimeError as e:
+                                raise InvalidMessageException(e) from e
                     else:
                         raise InvalidMessageException(
                             "Client should not receive stream open bit"

src/replit_river/codegen/typing.py

@@ -213,16 +213,20 @@ def extract_inner_type(value: TypeExpression) -> TypeName:
             raise ValueError(f"Unexpected literal type: {repr(value)}")
         case UnionTypeExpr(_):
             raise ValueError(
-                f"Attempting to extract from a union, currently not possible: {repr(value)}"
+                "Attempting to extract from a union, "
+                f"currently not possible: {repr(value)}"
             )
         case OpenUnionTypeExpr(_):
             raise ValueError(
-                f"Attempting to extract from a union, currently not possible: {repr(value)}"
+                "Attempting to extract from a union, "
+                f"currently not possible: {repr(value)}"
             )
         case TypeName(name):
             return TypeName(name)
         case NoneTypeExpr():
-            raise ValueError(f"Attempting to extract from a literal 'None': {repr(value)}")
+            raise ValueError(
+                f"Attempting to extract from a literal 'None': {repr(value)}",
+            )
         case other:
             assert_never(other)
 

src/replit_river/common_session.py

@@ -3,7 +3,6 @@
 import logging
 from typing import Any, Awaitable, Callable, Coroutine, Protocol
 
-from aiochannel import Channel, ChannelClosed
 from opentelemetry.trace import Span
 from websockets import WebSocketCommonProtocol
 from websockets.asyncio.client import ClientConnection
@@ -14,7 +13,6 @@
     send_transport_message,
 )
 from replit_river.rpc import ACK_BIT, TransportMessage
-from replit_river.seq_manager import InvalidMessageException
 
 logger = logging.getLogger(__name__)
 
@@ -181,17 +179,3 @@ async def buffered_message_sender(
                 logger.exception("Error attempting to send buffered messages")
                 message_enqueued.release()
                 break
-
-
-async def add_msg_to_stream(
-    msg: TransportMessage,
-    stream: Channel[Any],
-) -> None:
-    try:
-        await stream.put(msg.payload)
-    except ChannelClosed:
-        # The client is no longer interested in this stream,
-        # just drop the message.
-        pass
-    except RuntimeError as e:
-        raise InvalidMessageException(e) from e

src/replit_river/server_session.py

@@ -7,7 +7,6 @@
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from websockets.exceptions import ConnectionClosed
 
-from replit_river.common_session import add_msg_to_stream
 from replit_river.messages import (
     FailedSendingMessageException,
     parse_transport_msg,
@@ -151,7 +150,14 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                             # close message is not sent to the stream
                             pass
                         else:
-                            await add_msg_to_stream(msg, stream)
+                            try:
+                                await stream.put(msg.payload)
+                            except ChannelClosed:
+                                # The client is no longer interested in this stream,
+                                # just drop the message.
+                                pass
+                            except RuntimeError as e:
+                                raise InvalidMessageException(e) from e
                     else:
                         _stream = await self._open_stream_and_call_handler(msg, tg)
                         if not stream: