PR feedback

Author: blast-hardcheese

src/replit_river/client.py

@@ -226,9 +226,11 @@ def _trace_procedure(
     span_handle = _SpanHandle(span)
     try:
         yield span_handle
+        span_handle.set_status(StatusCode.OK)
     except GeneratorExit:
         # This error indicates the caller is done with the async generator
         # but messages are still left. This is okay, we do not consider it an error.
+        span_handle.set_status(StatusCode.OK)
         raise
     except RiverException as e:
         span.record_exception(e, escaped=True)
@@ -239,7 +241,6 @@ def _trace_procedure(
         span_handle.set_status(StatusCode.ERROR, f"{type(e).__name__}: {e}")
         raise e
     finally:
-        span_handle.set_status(StatusCode.OK)
         span.end()
 
 

src/replit_river/codegen/client.py

@@ -850,11 +850,11 @@ async def {name}(
     elif procedure.type == "subscription":
         match protocol_version:
             case "v1.1":
-                assert input_meta, "rpc expects input to be required"
+                assert input_meta, "subscription expects input to be required"
                 _, tpe, render_method = input_meta
                 binding = "input"
             case "v2.0":
-                assert init_meta, "rpc expects init to be required"
+                assert init_meta, "subscription expects init to be required"
                 _, tpe, render_method = init_meta
                 binding = "init"
             case other:
@@ -932,7 +932,7 @@ async def {name}(
                 ]
             )
         elif protocol_version == "v1.1":
-            assert input_meta, "Protocol v1 requires input to be defined"
+            assert input_meta, "upload requires input to be defined"
             _, input_type, render_input_method = input_meta
             current_chunks.extend(
                 [
@@ -1009,7 +1009,7 @@ async def {name}(
                 ]
             )
         elif protocol_version == "v1.1":
-            assert input_meta, "Protocol v1 requires input to be defined"
+            assert input_meta, "stream requires input to be defined"
             _, input_type, render_input_method = input_meta
             current_chunks.extend(
                 [

src/replit_river/common_session.py

@@ -102,13 +102,13 @@ async def buffered_message_sender(
             except WebsocketClosedException as e:
                 logger.debug(
                     "_buffered_message_sender: Connection closed while sending "
-                    "message %r, waiting for retry from buffer",
+                    "message %r, waiting for reconnect and retry from buffer",
                     type(e),
                     exc_info=e,
                 )
             except FailedSendingMessageException:
                 logger.error(
-                    "Failed sending message, waiting for retry from buffer",
+                    "Failed sending message, waiting for reconnect and retry from buffer",
                     exc_info=True,
                 )
             except Exception:

src/replit_river/message_buffer.py

@@ -68,6 +68,7 @@ async def close(self) -> None:
         Closes the message buffer and rejects any pending put operations.
         """
         self._closed = True
+        # Wake up block_until_message_available to permit graceful cleanup
         self._has_messages.set()
         async with self._space_available_cond:
             self._space_available_cond.notify_all()

src/replit_river/v2/client.py

@@ -181,9 +181,11 @@ def _trace_procedure(
     span_handle = _SpanHandle(span)
     try:
         yield span_handle
+        span_handle.set_status(StatusCode.OK)
     except GeneratorExit:
         # This error indicates the caller is done with the async generator
         # but messages are still left. This is okay, we do not consider it an error.
+        span_handle.set_status(StatusCode.OK)
         raise
     except RiverException as e:
         span.record_exception(e, escaped=True)
@@ -194,7 +196,6 @@ def _trace_procedure(
         span_handle.set_status(StatusCode.ERROR, f"{type(e).__name__}: {e}")
         raise e
     finally:
-        span_handle.set_status(StatusCode.OK)
         span.end()