diff --git a/src/ipython_shell/jupyter/transport.py b/src/ipython_shell/jupyter/transport.py index 320b332..fa3f809 100644 --- a/src/ipython_shell/jupyter/transport.py +++ b/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,59 +112,53 @@ 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() - parent_id = message.get("parent_header", {}).get("msg_id") - if parent_id != call_id: - continue + message = queued_message + parent_id = message.get("parent_header", {}).get("msg_id") + if parent_id != call_id: + continue - if message["msg_type"] == "input_request": - self._waiting_for_input = call_id - yield InputRequest( - call_id=call_id, - prompt=str(message.get("content", {}).get("prompt", "")), - password=bool( - message.get("content", {}).get("password", False) - ), - ) - continue - - decoded = KernelMessage( - msg_type=message["msg_type"], - parent_id=parent_id, - content=dict(message.get("content", {})), - buffers=[ - bytes(buffer) for buffer in message.get("buffers", []) - ], + if message["msg_type"] == "input_request": + self._waiting_for_input = call_id + yield InputRequest( + call_id=call_id, + prompt=str(message.get("content", {}).get("prompt", "")), + password=bool( + message.get("content", {}).get("password", False) + ), ) - 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. - reply = decoded - else: - yield decoded - if ( - decoded.msg_type == "status" - and decoded.content.get("execution_state") == "idle" - ): - idle = True + continue + + decoded = KernelMessage( + msg_type=message["msg_type"], + parent_id=parent_id, + content=dict(message.get("content", {})), + buffers=[ + bytes(buffer) for buffer in message.get("buffers", []) + ], + ) + 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 arrive close together. + reply = decoded + else: + yield decoded + if ( + decoded.msg_type == "status" + and decoded.content.get("execution_state") == "idle" + ): + idle = True if reply is not None and idle: yield reply @@ -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 diff --git a/tests/test_transport.py b/tests/test_transport.py index 5939bc2..586fc88 100644 --- a/tests/test_transport.py +++ b/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()