feat: execute cells through persistent Jupyter kernel

src/ipython_shell/transport.py

@@ -1,5 +1,9 @@
+import asyncio
+from collections.abc import AsyncIterator
 from dataclasses import dataclass
 
+from jupyter_client import AsyncKernelManager
+
 
 @dataclass
 class KernelMessage:
@@ -18,3 +22,91 @@ class InputRequest:
     call_id: str
     prompt: str
     password: bool
+
+
+class JupyterTransport:
+    """Own one persistent ipykernel subprocess and its message channels."""
+
+    def __init__(self) -> None:
+        self.manager = AsyncKernelManager()
+        self.client = None
+        self._active_call: str | None = None
+        self._call_lock = asyncio.Lock()
+
+    async def start(self) -> None:
+        """Start the kernel process and wait until it accepts requests."""
+        await self.manager.start_kernel()
+        self.client = self.manager.client()
+        self.client.start_channels()
+        await self.client.wait_for_ready()
+
+    async def execute(self, code: str) -> str:
+        """Submit one cell and return its Jupyter message ID."""
+        if self.client is None:
+            raise RuntimeError("JupyterTransport has not been started")
+
+        async with self._call_lock:
+            if self._active_call is not None:
+                raise RuntimeError("JupyterTransport already has an active call")
+            self._active_call = self.client.execute(code, allow_stdin=True)
+            return self._active_call
+
+    async def messages_for(self, call_id: str) -> AsyncIterator[KernelMessage]:
+        """Yield decoded IOPub and shell messages for one execution."""
+        if self.client is None:
+            raise RuntimeError("JupyterTransport has not been started")
+        if self._active_call != call_id:
+            raise RuntimeError(f"Unknown or inactive call: {call_id}")
+
+        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())
+                done, pending = await asyncio.wait(
+                    {iopub_task, shell_task},
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                for task in pending:
+                    task.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)
+
+                for task in done:
+                    message = task.result()
+                    parent_id = message.get("parent_header", {}).get("msg_id")
+                    if parent_id != call_id:
+                        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 become ready 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
+                    return
+        finally:
+            if self._active_call == call_id:
+                self._active_call = None
+
+    async def shutdown(self) -> None:
+        """Stop channels and terminate the kernel process."""
+        if self.client is not None:
+            self.client.stop_channels()
+            self.client = None
+        await self.manager.shutdown_kernel(now=True)

tests/test_transport.py

@@ -1,6 +1,6 @@
 import unittest
 
-from ipython_shell.transport import KernelMessage
+from ipython_shell.transport import JupyterTransport, KernelMessage
 
 
 class TransportTests(unittest.TestCase):
@@ -17,5 +17,24 @@ class TransportTests(unittest.TestCase):
         self.assertEqual(message.buffers, [b"binary"])
 
 
+class AsyncTransportTests(unittest.IsolatedAsyncioTestCase):
+    async def test_transport_executes_and_returns_execute_reply(self):
+        transport = JupyterTransport()
+        await transport.start()
+        try:
+            call_id = await transport.execute("2 + 2")
+            messages = [
+                message async for message in transport.messages_for(call_id)
+            ]
+        finally:
+            await transport.shutdown()
+
+        self.assertTrue(
+            any(message.msg_type == "execute_result" for message in messages)
+        )
+        self.assertEqual(messages[-1].msg_type, "execute_reply")
+        self.assertEqual(messages[-1].content["status"], "ok")
+
+
 if __name__ == "__main__":
     unittest.main()