From cebad6f57442bde06713700c5880c2bf60ad3426 Mon Sep 17 00:00:00 2001
From: zhenthebuilder <zhen@repl.it>
Date: Mon, 22 Apr 2024 19:10:49 -0700
Subject: [PATCH] Fix

---
 pyproject.toml          |  2 +-
 replit_river/session.py | 14 ++++++--------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 495b0a6..fa62c98 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name="replit-river"
-version="0.1.7-beta.1"
+version="0.1.7-beta.2"
 description="Replit river toolkit for Python"
 authors = ["Replit <eng@replit.com>"]
 license = "LICENSE"
diff --git a/replit_river/session.py b/replit_river/session.py
index c0a4fed..7d9caea 100644
--- a/replit_river/session.py
+++ b/replit_river/session.py
@@ -285,7 +285,6 @@ async def _heartbeat(
     async def _send_buffered_messages(
         self, websocket: websockets.WebSocketCommonProtocol
     ) -> None:
-        logging.debug(f"Sending buffered messages to {self._to_id}")
         buffered_messages = list(self._buffer.buffer)
         for msg in buffered_messages:
             try:
@@ -338,6 +337,12 @@ async def send_message(
             serviceName=service_name,
             procedureName=procedure_name,
         )
+        try:
+            await self._buffer.put(msg)
+        except Exception:
+            # We should close the session when there are too many messages in buffer
+            await self.close(True)
+            return
         try:
             await self._send_transport_message(
                 msg,
@@ -353,13 +358,6 @@ async def send_message(
             logging.error(
                 f"Failed sending message : {e}, waiting for retry from buffer"
             )
-        finally:
-            try:
-                await self._buffer.put(msg)
-            except Exception:
-                # We should close the session when there are too many messages in buffer
-                await self.close(True)
-                return
 
     async def _send_responses_from_output_stream(
         self,