From 42a1f7b890bd22f5e84030dcfccd2be82094fc83 Mon Sep 17 00:00:00 2001
From: Marcus Messer <marcus@marcusm.co.uk>
Date: Thu, 24 Sep 2026 15:00:51 +0100
Subject: [PATCH] Keep the worker serving after a handler raises

A handler exception permanently killed the worker. jsonrpc_handler passed
the exception object as JSON-RPC `data`, which ujson cannot serialize, so
building the error response raised TypeError out of dispatch(); the serve
loop caught it and broke out, closing the client. The process stayed
resident but never read stdin again, so every later request on that worker
timed out in the shim's RPC send.

Pass only the exception message, and split the serve loop so a failed
dispatch no longer ends the session. Read and write failures still stop
serving, since a partial frame leaves no safe point to resume from.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 lf_toolkit/io/rpc_handler.py |  5 ++-
 lf_toolkit/io/stream_io.py   | 74 ++++++++++++++++++++++++------------
 2 files changed, 54 insertions(+), 25 deletions(-)

diff --git a/lf_toolkit/io/rpc_handler.py b/lf_toolkit/io/rpc_handler.py
index bc47a71..e9940a0 100644
--- a/lf_toolkit/io/rpc_handler.py
+++ b/lf_toolkit/io/rpc_handler.py
@@ -32,6 +32,9 @@ async def wrapped(req: dict):
             result = await handler.handle(name, {"params": req})
             return Success(result)
         except Exception as e:
-            return Error(0, str(e), e)
+            # Pass only the message: the exception object is not JSON
+            # serializable, so sending it as `data` makes serializing the
+            # error response raise, which tears down the serve loop.
+            return Error(0, str(e))
 
     return wrapped
diff --git a/lf_toolkit/io/stream_io.py b/lf_toolkit/io/stream_io.py
index 179d7e4..2f0cb86 100644
--- a/lf_toolkit/io/stream_io.py
+++ b/lf_toolkit/io/stream_io.py
@@ -1,3 +1,6 @@
+import sys
+import traceback
+
 from abc import ABC
 from abc import abstractmethod
 
@@ -82,29 +85,52 @@ def wrap_io(self, client: StreamIO) -> StreamIO:
     async def _handle_client(self, client: StreamIO):
         io = self.wrap_io(client)
 
-        while True:
-            try:
-                import sys
-                print("waiting for data...", file=sys.stderr, flush=True)
-                data = await io.read(4096)
-                print(f"got data: {data[:80]}", file=sys.stderr, flush=True)
-
-                if not data:
-                    break
-
-                print("dispatching...", file=sys.stderr, flush=True)
-                response = await self.dispatch(data.decode("utf-8"))
-                print(f"got response: {str(response)[:80]}", file=sys.stderr, flush=True)
-
-                await io.write(response.encode("utf-8"))
-                print("wrote response", file=sys.stderr, flush=True)
-            except anyio.EndOfStream:
-                break
-            except anyio.ClosedResourceError:
-                break
-            except Exception as e:
-                import traceback
-                traceback.print_exc(file=sys.stderr)
-                break
+        serving = True
+        while serving:
+            serving = await self._serve_once(io)
 
         await client.close()
+
+    async def _serve_once(self, io: StreamIO) -> bool:
+        """Read one request, dispatch it and write the response.
+
+        Returns True when the session can carry on, and False once the
+        stream has ended or is no longer safe to read from.
+        """
+        try:
+            print("waiting for data...", file=sys.stderr, flush=True)
+            data = await io.read(4096)
+            print(f"got data: {data[:80]}", file=sys.stderr, flush=True)
+        except (anyio.EndOfStream, anyio.ClosedResourceError):
+            return False
+        except Exception:
+            # A read failure may have consumed part of a frame, so there is
+            # no point in the stream we can safely resume from.
+            traceback.print_exc(file=sys.stderr)
+            return False
+
+        if not data:
+            return False
+
+        try:
+            print("dispatching...", file=sys.stderr, flush=True)
+            response = await self.dispatch(data.decode("utf-8"))
+            print(f"got response: {str(response)[:80]}", file=sys.stderr, flush=True)
+        except Exception:
+            # One bad request must not end the session: the frame was read in
+            # full, so the stream is still aligned and the next request can be
+            # served. The caller gets no reply for this one and will time out,
+            # which beats every later request on this worker timing out too.
+            traceback.print_exc(file=sys.stderr)
+            return True
+
+        try:
+            await io.write(response.encode("utf-8"))
+            print("wrote response", file=sys.stderr, flush=True)
+        except (anyio.EndOfStream, anyio.ClosedResourceError):
+            return False
+        except Exception:
+            traceback.print_exc(file=sys.stderr)
+            return False
+
+        return True