feat: support Jupyter kernel transport input

src/ipython_shell/transport.py

@@ -32,6 +32,7 @@ class JupyterTransport:
         self.client = None
         self._active_call: str | None = None
         self._call_lock = asyncio.Lock()
+        self._waiting_for_input: str | None = None
 
     async def start(self) -> None:
         """Start the kernel process and wait until it accepts requests."""
@@ -51,8 +52,10 @@ class JupyterTransport:
             self._active_call = self.client.execute(code, allow_stdin=True)
             return self._active_call
 
-    async def messages_for(self, call_id: str) -> AsyncIterator[KernelMessage]:
+    async def messages_for(
-        """Yield decoded IOPub and shell messages for one execution."""
+        self, call_id: str
+    ) -> AsyncIterator[KernelMessage | InputRequest]:
+        """Yield decoded output, input, and completion messages for one call."""
         if self.client is None:
             raise RuntimeError("JupyterTransport has not been started")
         if self._active_call != call_id:
@@ -64,8 +67,9 @@ class JupyterTransport:
             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},
+                    {iopub_task, shell_task, stdin_task},
                     return_when=asyncio.FIRST_COMPLETED,
                 )
                 for task in pending:
@@ -78,6 +82,17 @@ class JupyterTransport:
                     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,
@@ -101,9 +116,21 @@ class JupyterTransport:
                     yield reply
                     return
         finally:
+            if self._waiting_for_input == call_id:
+                self._waiting_for_input = None
             if self._active_call == call_id:
                 self._active_call = None
 
+    async def reply_to_input(self, value: str) -> None:
+        """Reply to the active kernel input request without using process stdin."""
+        if self.client is None:
+            raise RuntimeError("JupyterTransport has not been started")
+        if self._waiting_for_input is None:
+            raise RuntimeError("JupyterTransport is not waiting for input")
+
+        self.client.input(value)
+        self._waiting_for_input = None
+
     async def shutdown(self) -> None:
         """Stop channels and terminate the kernel process."""
         if self.client is not None:

tests/test_transport.py

@@ -1,6 +1,7 @@
+import asyncio
 import unittest
 
-from ipython_shell.transport import JupyterTransport, KernelMessage
+from ipython_shell.transport import InputRequest, JupyterTransport, KernelMessage
 
 
 class TransportTests(unittest.TestCase):
@@ -35,6 +36,60 @@ class AsyncTransportTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(messages[-1].msg_type, "execute_reply")
         self.assertEqual(messages[-1].content["status"], "ok")
 
+    async def test_transport_preserves_mime_and_errors(self):
+        transport = JupyterTransport()
+        await transport.start()
+        try:
+            display_id = await transport.execute(
+                "from IPython.display import display\n"
+                "display({'text/plain': 'hello'}, raw=True)"
+            )
+            display_messages = [
+                message async for message in transport.messages_for(display_id)
+            ]
+
+            error_id = await transport.execute("raise ValueError('boom')")
+            error_messages = [
+                message async for message in transport.messages_for(error_id)
+            ]
+        finally:
+            await transport.shutdown()
+
+        display = next(
+            message
+            for message in display_messages
+            if message.msg_type == "display_data"
+        )
+        self.assertEqual(display.content["data"]["text/plain"], "hello")
+
+        error = next(
+            message for message in error_messages if message.msg_type == "error"
+        )
+        self.assertEqual(error.content["ename"], "ValueError")
+
+    async def test_transport_routes_input_reply_without_parent_stdin(self):
+        transport = JupyterTransport()
+        await transport.start()
+        try:
+            call_id = await transport.execute("answer = input('name? '); answer")
+            stream = transport.messages_for(call_id)
+
+            while True:
+                event = await asyncio.wait_for(anext(stream), timeout=5)
+                if isinstance(event, InputRequest):
+                    input_request = event
+                    break
+
+            self.assertEqual(input_request.prompt, "name? ")
+            await transport.reply_to_input("Ada")
+            messages = [message async for message in stream]
+        finally:
+            await transport.shutdown()
+
+        self.assertTrue(
+            any(message.msg_type == "execute_result" for message in messages)
+        )
+
 
 if __name__ == "__main__":
     unittest.main()