From 9f27836a65bf5adfed0a38b8713ed61f2c5928ab Mon Sep 17 00:00:00 2001
From: klioen <klioen@users.noreply.github.com>
Date: Mon, 7 Sep 2026 12:50:53 +0000
Subject: [PATCH] fix(server): don't silently swallow WebSocket message
 processing errors

A single malformed or unexpected message (e.g. a missing field in a
Jupyter kernel message) currently bubbles out of _process_message, which
terminates the whole receive loop and then marks every ongoing execution
as WebSocketError + UnexpectedEndOfExecution. The original exception is
only logged as a one-line message without stack trace, so the root cause
is effectively swallowed while all in-flight executions get a misleading
'connection lost' error.

Isolate per-message processing so that:
- a per-message failure logs the full stack trace and a preview of the
  raw message instead of being silently dropped,
- the affected execution is notified with a MessageProcessingError +
  EndOfExecution so it doesn't hang waiting for results,
- the receive loop keeps running for subsequent messages,
- only connection-level failures (ConnectionClosedError/WebSocketException)
  still terminate the loop and cancel all ongoing executions.
---
 template/server/messaging.py | 38 +++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index c51f8b21..9d88222b 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -426,7 +426,43 @@ async def _receive_message(self):
 
         try:
             async for message in self._ws:
-                await self._process_message(json.loads(message))
+                try:
+                    data = json.loads(message)
+                    await self._process_message(data)
+                except (ConnectionClosedError, WebSocketException):
+                    # Connection-level failures are handled below and must
+                    # terminate the receive loop so ongoing executions are
+                    # cancelled instead of hanging forever.
+                    raise
+                except Exception as e:
+                    # A single malformed or unexpected message must not kill
+                    # the receive loop nor be silently swallowed. Log the full
+                    # stack trace together with the raw message so it can be
+                    # diagnosed, and notify the affected execution so it does
+                    # not hang waiting for results.
+                    logger.exception(
+                        "Error while processing WebSocket message: %s (message: %s)",
+                        e,
+                        message[:500],
+                    )
+                    parent_msg_id = None
+                    try:
+                        parent_msg_id = json.loads(message).get(
+                            "parent_header", {}
+                        ).get("msg_id")
+                    except Exception:
+                        pass
+                    if parent_msg_id:
+                        execution = self._executions.get(parent_msg_id)
+                        if execution:
+                            await execution.queue.put(
+                                Error(
+                                    name="MessageProcessingError",
+                                    value=f"Failed to process kernel message: {e}",
+                                    traceback="",
+                                )
+                            )
+                            await execution.queue.put(EndOfExecution())
         except Exception as e:
             logger.error(f"WebSocket received error while receiving messages: {str(e)}")
         finally: