Propagating stream_id into logs

Author: blast-hardcheese

src/replit_river/error_schema.py

@@ -86,8 +86,10 @@ class SessionClosedRiverServiceException(RiverException):
     def __init__(
         self,
         message: str,
+        streamId: str,
     ) -> None:
         super().__init__(SYNTHETIC_ERROR_CODE_SESSION_CLOSED, message)
+        self.streamId = streamId
 
 
 def exception_from_message(code: str) -> type[RiverServiceException]:

src/replit_river/task_manager.py

@@ -2,7 +2,11 @@
 import logging
 from typing import Coroutine, Set
 
-from replit_river.error_schema import ERROR_CODE_STREAM_CLOSED, RiverException
+from replit_river.error_schema import (
+    ERROR_CODE_STREAM_CLOSED,
+    RiverException,
+    SessionClosedRiverServiceException,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,6 +41,13 @@ async def cancel_task(
             # If we cancel the task manager we will get called here as well,
             # if we want to handle the cancellation differently we can do it here.
             logger.debug("Task was cancelled %r", task_to_remove)
+        except SessionClosedRiverServiceException as e:
+            logger.warning(
+                "Session was closed",
+                extra={
+                    "stream_id": e.streamId,
+                },
+            )
         except RiverException as e:
             if e.code == ERROR_CODE_STREAM_CLOSED:
                 # Task is cancelled
@@ -76,6 +87,14 @@ def _task_done_callback(
             ):
                 # Task is cancelled
                 pass
+            elif isinstance(exception, SessionClosedRiverServiceException):
+                # Session is closed, don't bother logging
+                logger.info(
+                    "Session closed",
+                    extra={
+                        "stream_id": exception.streamId,
+                    },
+                )
             else:
                 logger.error(
                     "Exception on cancelling task",

src/replit_river/v2/session.py

@@ -353,6 +353,7 @@ async def _enqueue_message(
             # session is closing / closed, raise
             raise SessionClosedRiverServiceException(
                 "river session is closed, dropping message",
+                stream_id,
             )
 
         # Begin critical section: Avoid any await between here and _send_buffer.append
@@ -448,14 +449,15 @@ async def do_close() -> None:
 
             await self._task_manager.cancel_all_tasks()
 
-            for stream_meta in self._streams.values():
+            for stream_id, stream_meta in self._streams.items():
                 stream_meta["output"].close()
                 # Wake up backpressured writers
                 try:
                     stream_meta["error_channel"].put_nowait(
                         reason
                         or SessionClosedRiverServiceException(
                             "river session is closed",
+                            stream_id,
                         )
                     )
                 except ChannelFull:
@@ -1023,12 +1025,14 @@ async def _encode_stream() -> None:
                 # ... block the outer function until the emitter is finished emitting,
                 #     possibly raising a terminal exception.
                 await emitter_task
-            except asyncio.CancelledError:
+            except asyncio.CancelledError as e:
                 await self._send_cancel_stream(
                     stream_id=stream_id,
                     message="Stream cancelled",
                     span=span,
                 )
+                if emitter_task.done() and (err := emitter_task.exception()):
+                    raise e from err
                 raise
             except Exception as e:
                 await self._send_cancel_stream(
@@ -1316,6 +1320,7 @@ async def _recv_from_ws(
                     # the outer loop.
                     await transition_no_connection()
                     break
+                msg: TransportMessage | str | None = None
                 try:
                     msg = parse_transport_msg(message)
                     logger.debug(
@@ -1395,19 +1400,27 @@ async def _recv_from_ws(
                         stream_meta["output"].close()
                 except OutOfOrderMessageException:
                     logger.exception("Out of order message, closing connection")
+                    stream_id = "unknown"
+                    if isinstance(msg, TransportMessage):
+                        stream_id = msg.streamId
                     close_session(
                         SessionClosedRiverServiceException(
-                            "Out of order message, closing connection"
+                            "Out of order message, closing connection",
+                            stream_id,
                         )
                     )
                     continue
                 except InvalidMessageException:
                     logger.exception(
                         "Got invalid transport message, closing session",
                     )
+                    stream_id = "unknown"
+                    if isinstance(msg, TransportMessage):
+                        stream_id = msg.streamId
                     close_session(
                         SessionClosedRiverServiceException(
-                            "Out of order message, closing connection"
+                            "Out of order message, closing connection",
+                            stream_id,
                         )
                     )
                     continue