Allow async cell to process Comm messages

ipykernel/comm/comm.py

@@ -4,6 +4,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 import uuid
+from threading import Lock
 from typing import Optional
 from warnings import warn
 
@@ -15,12 +16,32 @@
 from ipykernel.kernelbase import Kernel
 
 
+def _request_id(data):
+    if isinstance(data, dict):
+        content = data.get("content")
+        if isinstance(content, dict) and isinstance(content.get("id"), str):
+            return content["id"]
+    return None
+
+
 # this is the class that will be created if we do comm.create_comm
 class BaseComm(comm.base_comm.BaseComm):
     """The base class for comms."""
 
     kernel: Optional["Kernel"] = None
 
+    def __init__(self, *args, **kwargs):
+        self._reply_subshell_lock = Lock()
+        self._reply_subshell_ids = {}
+        super().__init__(*args, **kwargs)
+
+    def _reply_subshell_for(self, data, default):
+        request_id = _request_id(data)
+        with self._reply_subshell_lock:
+            if request_id is not None and request_id in self._reply_subshell_ids:
+                return self._reply_subshell_ids.pop(request_id)
+            return getattr(self, "_reply_subshell_id", default)
+
     def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
         """Helper for sending a comm message on IOPub"""
         if not Kernel.initialized():
@@ -34,12 +55,22 @@ def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
             self.kernel = Kernel.instance()
 
         assert self.kernel.session is not None
+        parent = self.kernel.get_parent()
+        if parent.get("header"):
+            # A comm can be used from a different subshell than the one that
+            # created it. Route the frontend reply to the loop that sent it.
+            subshell_id = parent["header"].get("subshell_id")
+            request_id = _request_id(data)
+            with self._reply_subshell_lock:
+                self._reply_subshell_id = subshell_id
+                if request_id is not None:
+                    self._reply_subshell_ids[request_id] = subshell_id
         self.kernel.session.send(
             self.kernel.iopub_socket,
             msg_type,
             content,
             metadata=json_clean(metadata),
-            parent=self.kernel.get_parent(),
+            parent=parent,
             ident=self.topic,
             buffers=buffers,
         )

ipykernel/kernelbase.py

@@ -419,7 +419,9 @@ def should_handle(self, stream, msg, idents):
         """
         return True
 
-    async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
+    async def dispatch_shell(
+        self, msg, /, subshell_id: str | None = None, *, concurrent: bool = False
+    ):
         """dispatch shell requests"""
         if len(msg) == 1 and msg[0].buffer == b"stop aborting":
             # Dummy "stop aborting" message to stop aborting execute requests on this subshell.
@@ -450,10 +452,12 @@ async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
 
         # Set the parent message for side effects.
         self.set_parent(idents, msg, channel="shell")
-        self._publish_status("busy", "shell")
+        if not concurrent:
+            self._publish_status("busy", "shell")
 
         msg_type = msg["header"]["msg_type"]
-        assert msg["header"].get("subshell_id") == subshell_id
+        if msg_type not in {"comm_msg", "comm_close"}:
+            assert msg["header"].get("subshell_id") == subshell_id
 
         if self._supports_kernel_subshells:
             stream = self.shell_channel_thread.manager.get_subshell_to_shell_channel_socket(
@@ -483,7 +487,8 @@ async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
         if inspect.isawaitable(should_handle):
             should_handle = await should_handle
         if not should_handle:
-            self._publish_status_and_flush("idle", "shell", stream)
+            if not concurrent:
+                self._publish_status_and_flush("idle", "shell", stream)
             self.log.debug("Not handling %s:%s", msg_type, msg["header"].get("msg_id"))
             return
 
@@ -492,10 +497,11 @@ async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
             self.log.warning("Unknown message type: %r", msg_type)
         else:
             self.log.debug("%s: %s", msg_type, msg)
-            try:
-                self.pre_handler_hook()
-            except Exception:
-                self.log.debug("Unable to signal in pre_handler_hook:", exc_info=True)
+            if not concurrent:
+                try:
+                    self.pre_handler_hook()
+                except Exception:
+                    self.log.debug("Unable to signal in pre_handler_hook:", exc_info=True)
             try:
                 result = handler(stream, idents, msg)
                 if inspect.isawaitable(result):
@@ -506,16 +512,18 @@ async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
                 # Ctrl-c shouldn't crash the kernel here.
                 self.log.error("KeyboardInterrupt caught in kernel.")
             finally:
-                try:
-                    self.post_handler_hook()
-                except Exception:
-                    self.log.debug("Unable to signal in post_handler_hook:", exc_info=True)
+                if not concurrent:
+                    try:
+                        self.post_handler_hook()
+                    except Exception:
+                        self.log.debug("Unable to signal in post_handler_hook:", exc_info=True)
 
         if sys.stdout is not None:
             sys.stdout.flush()
         if sys.stderr is not None:
             sys.stderr.flush()
-        self._publish_status_and_flush("idle", "shell", stream)
+        if not concurrent:
+            self._publish_status_and_flush("idle", "shell", stream)
 
     def pre_handler_hook(self):
         """Hook to execute before calling message handler"""
@@ -600,6 +608,16 @@ async def shell_channel_thread_main(self, msg):
                 msg3 = self.session.deserialize(msg2, content=False, copy=False)
                 subshell_id = msg3["header"].get("subshell_id")
 
+                if msg3["header"]["msg_type"] in {"comm_msg", "comm_close"} and hasattr(
+                    self, "comm_manager"
+                ):
+                    content = self.session.unpack(msg3["content"])
+                    comm = self.comm_manager.get_comm(content.get("comm_id"))
+                    if comm is not None:
+                        route = getattr(comm, "_reply_subshell_for", None)
+                        if route is not None:
+                            subshell_id = route(content.get("data"), subshell_id)
+
                 # Find inproc pair socket to use to send message to correct subshell.
                 subshell_manager = self.shell_channel_thread.manager
                 try:
@@ -635,6 +653,26 @@ async def shell_main(self, subshell_id: str | None, msg):
         # async cells at the same time which would be a nice feature to have but is an API
         # change.
         assert asyncio_lock is not None
+        if asyncio_lock.locked() and self.session is not None:
+            try:
+                _, frames = self.session.feed_identities(msg, copy=False)
+                header = self.session.deserialize(frames, content=False, copy=False)["header"]
+            except Exception:
+                header = {}
+            if header.get("msg_type") in {"comm_open", "comm_msg", "comm_close"}:
+                # A running async cell may be waiting for a widget reply on this
+                # channel. Dispatch comms without waiting for the cell's lock.
+                shell_parent = self.get_parent("shell")
+                shell_ident = self._get_shell_context_var(self._shell_parent_ident)
+                try:
+                    comm_task = asyncio.create_task(
+                        self.dispatch_shell(msg, subshell_id=subshell_id, concurrent=True),
+                        context=copy_context(),
+                    )
+                    await comm_task
+                finally:
+                    self.set_parent(shell_ident, shell_parent, channel="shell")
+                return
         async with asyncio_lock:
             await self.dispatch_shell(msg, subshell_id=subshell_id)
 

tests/test_kernel.py

@@ -58,6 +58,35 @@ def test_simple_print():
         _check_master(kc, expected=True)
 
 
+def test_async_cell_waiting_for_comm_reply():
+    with new_kernel() as kc:
+        msg_id = kc.execute(
+            """import asyncio, comm
+reply = asyncio.get_running_loop().create_future()
+widget = comm.create_comm(target_name='comm-reply-test')
+widget.on_msg(lambda msg: reply.set_result(msg['content']['data']['value']))
+result = await asyncio.wait_for(reply, 5)
+assert result == 42
+"""
+        )
+        while True:
+            msg = kc.get_iopub_msg(timeout=10)
+            if msg["msg_type"] == "error" and msg["parent_header"].get("msg_id") == msg_id:
+                raise AssertionError("\n".join(msg["content"]["traceback"]))
+            if msg["msg_type"] == "comm_open" and msg["parent_header"].get("msg_id") == msg_id:
+                comm_id = msg["content"]["comm_id"]
+                break
+
+        next_msg_id = kc.execute("assert result == 42")
+        kc.shell_channel.send(
+            kc.session.msg("comm_msg", {"comm_id": comm_id, "data": {"value": 42}})
+        )
+        reply_msg = get_reply(kc, msg_id, timeout=10)
+        assert reply_msg["content"]["status"] == "ok", reply_msg["content"]
+        next_reply = get_reply(kc, next_msg_id, timeout=10)
+        assert next_reply["content"]["status"] == "ok", next_reply["content"]
+
+
 @pytest.mark.parametrize(
     ("code", "expect_error_status"),
     [

tests/test_subshells.py

@@ -133,6 +133,60 @@ def test_thread_ids():
         delete_subshell_helper(kc, subshell_id)
 
 
+def test_comm_reply_follows_requesting_subshell():
+    with new_kernel() as kc:
+        subshell_id = create_subshell_helper(kc)["subshell_id"]
+        comm_id = execute_request_subshell_id(
+            kc,
+            "import comm; widget = comm.create_comm(target_name='comm-thread-test'); print(widget.comm_id)",
+            subshell_id,
+        )
+
+        msg = execute_request(
+            kc,
+            """import asyncio, threading
+reply = asyncio.get_running_loop().create_future()
+request_thread = threading.get_ident()
+def on_reply(message):
+    global callback_thread
+    callback_thread = threading.get_ident()
+    reply.set_result(message['content']['data']['content']['value'])
+widget.on_msg(on_reply)
+widget.send({'method': 'custom', 'content': {'id': 'request-1', 'operation': 'get'}})
+assert await asyncio.wait_for(reply, 2) == 42
+assert callback_thread == request_thread
+""",
+            None,
+        )
+        while True:
+            outgoing = kc.get_iopub_msg(timeout=10)
+            if (
+                outgoing["msg_type"] == "error"
+                and outgoing["parent_header"].get("msg_id") == msg["header"]["msg_id"]
+            ):
+                raise AssertionError("\n".join(outgoing["content"]["traceback"]))
+            if (
+                outgoing["msg_type"] == "comm_msg"
+                and outgoing["parent_header"].get("msg_id") == msg["header"]["msg_id"]
+            ):
+                break
+        assert outgoing["content"]["comm_id"] == comm_id
+
+        response = kc.session.msg(
+            "comm_msg",
+            {
+                "comm_id": comm_id,
+                "data": {"method": "custom", "content": {"id": "request-1", "value": 42}},
+            },
+        )
+        response["header"]["subshell_id"] = subshell_id
+        kc.shell_channel.send(response)
+        reply_msg = get_reply(kc, msg["header"]["msg_id"], timeout=5)
+        assert reply_msg["content"]["status"] == "ok", reply_msg["content"]
+        wait_for_idle(kc, msg["header"]["msg_id"])
+        delete_subshell_helper(kc, subshell_id)
+
+
 @pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)])
 @pytest.mark.parametrize("overlap", [True, False])
 def test_run_concurrently_sequence(are_subshells, overlap, request):