Refactored StdioServer and StdioClient to pass and use captured stdout buffer, ensuring isolated binary protocol handling on stdout.

Author: m-messer

lf_toolkit/io/stdio_server.py

@@ -1,5 +1,6 @@
 import sys
 
+from typing import BinaryIO
 from typing import Optional
 
 import anyio
@@ -15,9 +16,10 @@
 
 class StdioClient(StreamIO):
 
-    def __init__(self):
+    def __init__(self, stdout_buffer: BinaryIO):
+        self._stdout_buffer = stdout_buffer
         self.stream = StapledByteStream(
-            FileWriteStream(sys.stdout.buffer),
+            FileWriteStream(stdout_buffer),
             FileReadStream(sys.stdin.buffer),
         )
 
@@ -26,7 +28,7 @@ async def read(self, size: int) -> bytes:
 
     async def write(self, data: bytes):
         await self.stream.send(data)
-        await anyio.to_thread.run_sync(sys.stdout.buffer.flush)
+        await anyio.to_thread.run_sync(self._stdout_buffer.flush)
 
     async def close(self):
         await self.stream.aclose()
@@ -35,15 +37,20 @@ async def close(self):
 class StdioServer(StreamServer):
 
     _client: StdioClient
+    _stdout_buffer: BinaryIO
 
     def __init__(self, handler: Optional[Handler] = None):
         super().__init__(handler)
-
+        # Capture the real stdout buffer before redirecting sys.stdout.
+        # Any print() in user code after this point goes to stderr,
+        # keeping the binary Content-Length-framed protocol on fd 1 clean.
+        self._stdout_buffer = sys.stdout.buffer
+        sys.stdout = sys.stderr
 
     def wrap_io(self, client: StreamIO) -> StreamIO:
         return PrefixStreamIO(client)
 
     async def run(self):
         print("StdioServer started", file=sys.stderr, flush=True)
-        self._client = StdioClient()
-        await self._handle_client(self._client)
+        self._client = StdioClient(self._stdout_buffer)
+        await self._handle_client(self._client)