Permit sensible message_buffer state transitions

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -147,7 +147,7 @@ async def _handle_messages_from_ws(self) -> None:
                         case other:
                             assert_never(other)
 
-                    self._buffer.remove_old_messages(
+                    await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
                     self._reset_session_close_countdown()

src/replit_river/codegen/run.py

@@ -52,12 +52,6 @@ def main() -> None:
         default="v1.1",
         choices=["v1.1", "v2.0"],
     )
-    client.add_argument(
-        "--method-filter",
-        help="Only generate a subset of the specified methods",
-        action="store",
-        type=pathlib.Path,
-    )
     client.add_argument("schema", help="schema file")
     args = parser.parse_args()
 
@@ -82,11 +76,11 @@ def file_opener(path: Path) -> TextIO:
                 method_filter = set(x.strip() for x in handle.readlines())
 
         schema_to_river_client_codegen(
-            lambda: open(schema_path),
-            target_path,
-            args.client_name,
-            args.typed_dict_inputs,
-            file_opener,
+            read_schema=lambda: open(schema_path),
+            target_path=target_path,
+            client_name=args.client_name,
+            typed_dict_inputs=args.typed_dict_inputs,
+            file_opener=file_opener,
             method_filter=method_filter,
             protocol_version=args.protocol_version,
         )

src/replit_river/message_buffer.py

@@ -17,6 +17,7 @@ class MessageBuffer:
     def __init__(self, max_num_messages: int = MAX_MESSAGE_BUFFER_SIZE):
         self.max_size = max_num_messages
         self.buffer: list[TransportMessage] = []
+        self._has_messages = asyncio.Event()
         self._space_available_cond = asyncio.Condition()
         self._closed = False
 
@@ -35,6 +36,7 @@ def put(self, message: TransportMessage) -> None:
         if self._closed:
             raise MessageBufferClosedError("message buffer is closed")
         self.buffer.append(message)
+        self._has_messages.set()
 
     def get_next_sent_seq(self) -> int | None:
         if self.buffer:
@@ -47,16 +49,25 @@ def peek(self) -> TransportMessage | None:
             return None
         return self.buffer[0]
 
-    def remove_old_messages(self, min_seq: int) -> None:
+    async def remove_old_messages(self, min_seq: int) -> None:
         """Remove messages in the buffer with a seq number less than min_seq."""
         self.buffer = [msg for msg in self.buffer if msg.seq >= min_seq]
+        if self.buffer:
+            self._has_messages.set()
+        else:
+            self._has_messages.clear()
         async with self._space_available_cond:
             self._space_available_cond.notify_all()
 
-    def close(self) -> None:
+    async def block_until_message_available(self) -> None:
+        """Allow consumers to avoid spinning unnecessarily"""
+        await self._has_messages.wait()
+
+    async def close(self) -> None:
         """
         Closes the message buffer and rejects any pending put operations.
         """
         self._closed = True
+        self._has_messages.set()
         async with self._space_available_cond:
             self._space_available_cond.notify_all()

src/replit_river/server_session.py

@@ -136,7 +136,7 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                             pass
                         case other:
                             assert_never(other)
-                    self._buffer.remove_old_messages(
+                    await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
                     self._reset_session_close_countdown()

src/replit_river/session.py

@@ -303,7 +303,7 @@ async def close(self) -> None:
 
             await self.close_websocket(self._ws_wrapper, should_retry=False)
 
-            self._buffer.close()
+            await self._buffer.close()
 
             # Clear the session in transports
             await self._close_session_callback(self)

tests/test_message_buffer.py

@@ -45,7 +45,7 @@ async def put_messages() -> None:
         # Wait for the put call to return.
         await sync_events.get()
         assert len(buffer.buffer) == 1
-        buffer.remove_old_messages(i)
+        await buffer.remove_old_messages(i)
 
     await background_puts