Dunno

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -173,7 +173,7 @@ async def _handle_messages_from_ws(self) -> None:
                                 # The client is no longer interested in this stream,
                                 # just drop the message.
                                 pass
-                            except RuntimeError as e:
+                            except Exception as e:
                                 raise InvalidMessageException(e) from e
                     else:
                         raise InvalidMessageException(
@@ -244,7 +244,7 @@ async def send_rpc(
                     service_name,
                     procedure_name,
                 ) from e
-            except RuntimeError as e:
+            except Exception as e:
                 raise RiverException(ERROR_CODE_STREAM_CLOSED, str(e)) from e
             if not response.get("ok", False):
                 try:
@@ -330,7 +330,7 @@ async def send_upload(
                     service_name,
                     procedure_name,
                 ) from e
-            except RuntimeError as e:
+            except Exception as e:
                 raise RiverException(ERROR_CODE_STREAM_CLOSED, str(e)) from e
             if not response.get("ok", False):
                 try:
@@ -387,15 +387,13 @@ async def send_subscription(
                         )
                     continue
                 yield response_deserializer(item["payload"])
-        except (RuntimeError, ChannelClosed) as e:
+        except Exception as e:
             raise RiverServiceException(
                 ERROR_CODE_STREAM_CLOSED,
                 "Stream closed before response",
                 service_name,
                 procedure_name,
             ) from e
-        except Exception as e:
-            raise e
         finally:
             output.close()
 
@@ -490,17 +488,14 @@ async def _encode_stream() -> None:
                         )
                     continue
                 yield response_deserializer(item["payload"])
-        except (RuntimeError, ChannelClosed) as e:
+        except Exception as e:
+            logger.exception("There was a problem")
             raise RiverServiceException(
                 ERROR_CODE_STREAM_CLOSED,
                 "Stream closed before response",
                 service_name,
                 procedure_name,
             ) from e
-        except Exception as e:
-            raise e
-        finally:
-            output.close()
 
     async def send_close_stream(
         self,

src/replit_river/v2/session.py

@@ -878,7 +878,7 @@ async def _encode_stream() -> None:
             try:
                 async for result in output:
                     # Raise as early as we possibly can in case of an emission error
-                    if err := emitter_task.exception():
+                    if err := emitter_task.done() and emitter_task.exception():
                         raise err
                     if result.get("type") == "CLOSE":
                         break