Properly close streams on exception (#167)

Why
===

We've had persistent timeout errors in AI-Infra, and I suspect that it's
related to not handling bumps in the connection correctly.

What changed
============

- WebSocket drops or send failures left _streams populated, so any
in-flight RPC hung until the session fully shut down. That meant clients
didn’t see an abort signal and could block indefinitely even though the
transport was already defunct.
- Added _abort_all_streams() in src/replit_river/session.py#L289 and
call it from both client_session.serve() and server_session.serve() on
ConnectionClosed, FailedSendingMessageException, or any other unexpected
exception (src/replit_river/client_session.py#L95,
src/replit_river/server_session.py#L82). This immediately closes every
active channel and clears _streams, ensuring callers are notified right
away when the socket dies so they can retry or surface an error.

Test plan
=========

CI/CD, ran against an internal branch with no issues 3x without flake.

Author: JimmyAustin

src/replit_river/client_session.py

@@ -99,17 +99,22 @@ async def serve(self) -> None:
             try:
                 await self._handle_messages_from_ws()
             except ConnectionClosed:
+                if self._should_abort_streams_after_transport_failure():
+                    await self.close()
                 if self._retry_connection_callback:
                     self._task_manager.create_task(self._retry_connection_callback())
 
                 await self._begin_close_session_countdown()
                 logger.debug("ConnectionClosed while serving", exc_info=True)
             except FailedSendingMessageException:
                 # Expected error if the connection is closed.
+                if self._should_abort_streams_after_transport_failure():
+                    await self.close()
                 logger.debug(
                     "FailedSendingMessageException while serving", exc_info=True
                 )
             except Exception:
+                await self.close()
                 logger.exception("caught exception at message iterator")
         except ExceptionGroup as eg:
             _, unhandled = eg.split(lambda e: isinstance(e, ConnectionClosed))

src/replit_river/session.py

@@ -286,6 +286,17 @@ async def close_websocket(
         if should_retry and self._retry_connection_callback:
             self._task_manager.create_task(self._retry_connection_callback())
 
+    def _should_abort_streams_after_transport_failure(self) -> bool:
+        return not self._transport_options.transparent_reconnect
+
+    def _abort_all_streams(self) -> None:
+        """Close all active stream channels, notifying any waiting consumers."""
+        if not self._streams:
+            return
+        for stream in self._streams.values():
+            stream.close()
+        self._streams.clear()
+
     async def close(self) -> None:
         """Close the session and all associated streams."""
         logger.info(
@@ -310,9 +321,7 @@ async def close(self) -> None:
 
             # TODO: unexpected_close should close stream differently here to
             # throw exception correctly.
-            for stream in self._streams.values():
-                stream.close()
-            self._streams.clear()
+            self._abort_all_streams()
 
             self._state = SessionState.CLOSED