Just avoid calling add_msg_to_stream if it should not be called

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -144,7 +144,15 @@ async def _handle_messages_from_ws(self) -> None:
                             raise IgnoreMessageException(
                                 "no stream for message, ignoring"
                             )
-                        await add_msg_to_stream(msg, stream)
+
+                        if (
+                            msg.controlFlags & STREAM_CLOSED_BIT != 0
+                            and msg.payload.get("type", None) == "CLOSE"
+                        ):
+                            # close message is not sent to the stream
+                            pass
+                        else:
+                            await add_msg_to_stream(msg, stream)
                     else:
                         raise InvalidMessageException(
                             "Client should not receive stream open bit"

src/replit_river/common_session.py

@@ -119,12 +119,6 @@ async def add_msg_to_stream(
     msg: TransportMessage,
     stream: Channel[Any],
 ) -> None:
-    if (
-        msg.controlFlags & STREAM_CLOSED_BIT != 0
-        and msg.payload.get("type", None) == "CLOSE"
-    ):
-        # close message is not sent to the stream
-        return
     try:
         await stream.put(msg.payload)
     except ChannelClosed:

src/replit_river/server_session.py

@@ -143,7 +143,15 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                             raise IgnoreMessageException(
                                 "no stream for message, ignoring"
                             )
-                        await add_msg_to_stream(msg, stream)
+
+                        if (
+                            msg.controlFlags & STREAM_CLOSED_BIT != 0
+                            and msg.payload.get("type", None) == "CLOSE"
+                        ):
+                            # close message is not sent to the stream
+                            pass
+                        else:
+                            await add_msg_to_stream(msg, stream)
                     else:
                         _stream = await self._open_stream_and_call_handler(msg, tg)
                         if not stream:

src/replit_river/v2/client_session.py

@@ -140,7 +140,15 @@ async def _handle_messages_from_ws(self) -> None:
                             raise IgnoreMessageException(
                                 "no stream for message, ignoring"
                             )
-                        await add_msg_to_stream(msg, stream)
+
+                        if (
+                            msg.controlFlags & STREAM_CLOSED_BIT != 0
+                            and msg.payload.get("type", None) == "CLOSE"
+                        ):
+                            # close message is not sent to the stream
+                            pass
+                        else:
+                            await add_msg_to_stream(msg, stream)
                     else:
                         raise InvalidMessageException(
                             "Client should not receive stream open bit"