fix: keep Jupyter channel readers alive

src/ipython_shell/jupyter/transport.py

@@ -1,6 +1,7 @@
 import asyncio
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass
+from typing import Any
 
 from jupyter_client import AsyncKernelManager
 
@@ -33,14 +34,65 @@ class JupyterTransport:
         self._active_call: str | None = None
         self._call_lock = asyncio.Lock()
         self._waiting_for_input: str | None = None
+        self._message_queue: asyncio.Queue[
+            tuple[str, dict[str, Any] | BaseException]
+        ] | None = None
+        self._reader_tasks: set[asyncio.Task[None]] = set()
 
     async def start(self) -> None:
-        """Start the kernel process and wait until it accepts requests."""
+        """Start the kernel and its persistent channel readers."""
+        if self.client is not None:
+            return
+
         await self.manager.start_kernel()
         self.client = self.manager.client()
         self.client.start_channels()
         await self.client.wait_for_ready()
 
+        self._message_queue = asyncio.Queue()
+        self._reader_tasks = {
+            asyncio.create_task(
+                self._read_channel("iopub", self.client.get_iopub_msg),
+                name="jupyter-iopub-reader",
+            ),
+            asyncio.create_task(
+                self._read_channel("shell", self.client.get_shell_msg),
+                name="jupyter-shell-reader",
+            ),
+            asyncio.create_task(
+                self._read_channel("stdin", self.client.get_stdin_msg),
+                name="jupyter-stdin-reader",
+            ),
+        }
+
+    async def _read_channel(
+        self,
+        channel: str,
+        get_message: Callable[[], Awaitable[dict[str, Any]]],
+    ) -> None:
+        """Read one ZMQ channel continuously into the transport queue."""
+        if self._message_queue is None:
+            raise RuntimeError("JupyterTransport has not been started")
+
+        try:
+            while True:
+                message = await get_message()
+                await self._message_queue.put((channel, message))
+        except asyncio.CancelledError:
+            raise
+        except BaseException as error:
+            # Surface reader failures instead of leaving messages_for() stuck.
+            await self._message_queue.put((channel, error))
+
+    async def _stop_readers(self) -> None:
+        """Cancel channel readers once, during transport shutdown."""
+        readers = self._reader_tasks
+        self._reader_tasks = set()
+        for reader in readers:
+            reader.cancel()
+        await asyncio.gather(*readers, return_exceptions=True)
+        self._message_queue = None
+
     async def execute(self, code: str) -> str:
         """Submit one cell and return its Jupyter message ID."""
         if self.client is None:
@@ -60,24 +112,18 @@ class JupyterTransport:
             raise RuntimeError("JupyterTransport has not been started")
         if self._active_call != call_id:
             raise RuntimeError(f"Unknown or inactive call: {call_id}")
+        if self._message_queue is None:
+            raise RuntimeError("JupyterTransport message readers are not running")
 
         reply: KernelMessage | None = None
         idle = False
         try:
             while True:
-                iopub_task = asyncio.create_task(self.client.get_iopub_msg())
-                shell_task = asyncio.create_task(self.client.get_shell_msg())
-                stdin_task = asyncio.create_task(self.client.get_stdin_msg())
-                done, pending = await asyncio.wait(
-                    {iopub_task, shell_task, stdin_task},
-                    return_when=asyncio.FIRST_COMPLETED,
-                )
-                for task in pending:
-                    task.cancel()
-                await asyncio.gather(*pending, return_exceptions=True)
+                _, queued_message = await self._message_queue.get()
+                if isinstance(queued_message, BaseException):
+                    raise queued_message
 
-                for task in done:
-                    message = task.result()
+                message = queued_message
                 parent_id = message.get("parent_header", {}).get("msg_id")
                 if parent_id != call_id:
                     continue
@@ -104,7 +150,7 @@ class JupyterTransport:
                 if decoded.msg_type == "execute_reply":
                     # The shell reply and IOPub messages use different
                     # channels. Buffer the reply so it remains terminal
-                        # even when both channels become ready together.
+                    # even when both channels arrive close together.
                     reply = decoded
                 else:
                     yield decoded
@@ -124,7 +170,7 @@ class JupyterTransport:
                 self._active_call = None
 
     async def reply_to_input(self, value: str) -> None:
-        """Reply to the active kernel input request without using process stdin."""
+        """Reply to the active kernel input request without process stdin."""
         if self.client is None:
             raise RuntimeError("JupyterTransport has not been started")
         if self._waiting_for_input is None:
@@ -134,7 +180,8 @@ class JupyterTransport:
         self._waiting_for_input = None
 
     async def shutdown(self) -> None:
-        """Stop channels and terminate the kernel process."""
+        """Stop readers, channels, and the kernel process."""
+        await self._stop_readers()
         if self.client is not None:
             self.client.stop_channels()
             self.client = None

tests/test_transport.py

@@ -23,6 +23,26 @@ class TransportTests(unittest.TestCase):
 
 
 class AsyncTransportTests(unittest.IsolatedAsyncioTestCase):
+    async def test_transport_keeps_channel_readers_alive_during_a_call(self):
+        transport = JupyterTransport()
+        await transport.start()
+        try:
+            readers = set(transport._reader_tasks)
+            call_id = await transport.execute("2 + 2")
+            messages = [
+                message async for message in transport.messages_for(call_id)
+            ]
+
+            self.assertEqual(len(readers), 3)
+            self.assertEqual(transport._reader_tasks, readers)
+            self.assertTrue(all(not reader.done() for reader in readers))
+        finally:
+            await transport.shutdown()
+
+        self.assertTrue(
+            any(message.msg_type == "execute_result" for message in messages)
+        )
+
     async def test_transport_executes_and_returns_execute_reply(self):
         transport = JupyterTransport()
         await transport.start()