Do not fire pre_execute/post_execute for comms dispatched concurrently during a cell

ipykernel/comm/comm.py

@@ -35,6 +35,25 @@ def __init__(self, *args, **kwargs):
         self._reply_subshell_ids = {}
         super().__init__(*args, **kwargs)
 
+    def handle_msg(self, msg: "comm.base_comm.MessageType") -> None:
+        """Handle a comm_msg message.
+
+        When a comm is dispatched concurrently while an async cell holds the
+        shell lock, the pre_execute/post_execute events are not fired for its
+        callback: the callback runs under the comm_msg's parent header, so any
+        output a post_execute hook publishes would be parented to the comm_msg
+        and dropped by frontends instead of reaching the cell, and the hook's
+        queue would be drained before the cell's own post_execute fires
+        (see ipython/ipykernel#1564). Comms handled while no cell is running
+        keep the previous behaviour.
+        """
+        kernel = Kernel.instance() if Kernel.initialized() else None
+        if kernel is not None and kernel._in_concurrent_comm_dispatch.get():
+            if self._msg_callback:
+                self._msg_callback(msg)
+            return
+        super().handle_msg(msg)
+
     def _reply_subshell_for(self, data, default):
         request_id = _request_id(data)
         with self._reply_subshell_lock:

ipykernel/kernelbase.py

@@ -233,6 +233,11 @@ def _default_ident(self):
     _shell_parent: ContextVar[dict[str, Any]]
     _shell_parent_ident: ContextVar[bytes]
     _shell_context: Context
+    # Set while a comm is dispatched concurrently during a running cell
+    # (see shell_main). Comm callbacks then run under the comm_msg's parent
+    # header, so firing pre_execute/post_execute for them would parent any
+    # hook output to the comm_msg instead of the cell.
+    _in_concurrent_comm_dispatch: ContextVar[bool]
     # Kept for backward-compatibility, accesses _control_parent_ident and _shell_parent_ident,
     # see https://github.com/jupyterlab/jupyterlab/issues/17785
     _parent_ident: Mapping[str, bytes]
@@ -352,6 +357,9 @@ def __init__(self, **kwargs):
         self._shell_parent.set({})
         self._shell_parent_ident = ContextVar("shell_parent_ident")
         self._shell_parent_ident.set(b"")
+        self._in_concurrent_comm_dispatch: ContextVar[bool] = ContextVar(
+            "in_concurrent_comm_dispatch", default=False
+        )
         self._shell_context = copy_context()
 
         # For backward compatibility so that _parent_ident["shell"] and _parent_ident["control"]
@@ -665,13 +673,21 @@ async def shell_main(self, subshell_id: str | None, msg):
                 # 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)
+                # The comm callback runs under the comm_msg's parent header, not
+                # the cell's execute_request. Mark the dispatch so BaseComm
+                # does not fire pre_execute/post_execute for it: any output a
+                # post_execute hook publishes would be parented to the comm_msg
+                # and dropped by frontends, and the hook's queue would be
+                # drained before the cell's own post_execute fires (#1564).
+                token = self._in_concurrent_comm_dispatch.set(True)
                 try:
                     comm_task = asyncio.create_task(
                         self.dispatch_shell(msg, subshell_id=subshell_id, concurrent=True),
                         context=copy_context(),
                     )
                     await comm_task
                 finally:
+                    self._in_concurrent_comm_dispatch.reset(token)
                     self.set_parent(shell_ident, shell_parent, channel="shell")
                 return
         async with asyncio_lock:

tests/test_comm.py

@@ -104,3 +104,37 @@ def test_comm_in_manager(ipkernel: IPythonKernel) -> None:
         comm = Comm()
 
     assert comm.comm_id in ipkernel.comm_manager.comms
+
+
+def test_handle_msg_skips_events_when_dispatched_concurrently(
+    kernel: Kernel, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """handle_msg must not fire pre/post_execute for a concurrently dispatched comm.
+
+    Regression test for ipython/ipykernel#1564: when a comm_msg is dispatched
+    while an async cell holds the shell lock, its callback runs under the
+    comm_msg's parent header, so output a post_execute hook publishes would be
+    parented to the comm_msg and dropped by frontends.
+    """
+    Kernel._instance = kernel  # type: ignore[attr-defined]
+    try:
+        events: list[str] = []
+        shell = unittest.mock.Mock()
+        shell.events.trigger.side_effect = lambda name: events.append(name)
+        monkeypatch.setattr("IPython.get_ipython", lambda: shell)
+        with pytest.deprecated_call():
+            comm = Comm(kernel=kernel, target_name="test")
+        comm.on_msg(lambda msg: events.append("callback"))
+
+        token = kernel._in_concurrent_comm_dispatch.set(True)
+        try:
+            comm.handle_msg({})
+        finally:
+            kernel._in_concurrent_comm_dispatch.reset(token)
+        assert events == ["callback"]
+
+        events.clear()
+        comm.handle_msg({})
+        assert events == ["pre_execute", "callback", "post_execute"]
+    finally:
+        Kernel.clear_instance()

tests/test_subshells.py

@@ -493,3 +493,70 @@ def test_comm_close_on_deleted_subshell():
         kc.shell_channel.send(msg)
 
         wait_for_idle(kc, msg["header"]["msg_id"])
+
+
+@pytest.mark.parametrize("use_subshell", [True, False])
+def test_comm_msg_during_async_cell_does_not_fire_post_execute(use_subshell, request):
+    """A comm_msg arriving during an async cell must not fire post_execute.
+
+    Regression test for #1564: the hook runs under the comm_msg's parent
+    header, so output it publishes is parented to the comm_msg and dropped by
+    frontends, and the hook's queue is drained before the cell's own
+    post_execute fires. Both the subshell and the plain (no-subshell) routes
+    go through the concurrent dispatch in shell_main.
+    """
+    if request.config.getvalue("--cov"):
+        pytest.skip("Skip time-sensitive subshell tests if measuring coverage")
+
+    with new_kernel() as kc:
+        setup = (
+            "import time\n"
+            "from IPython.display import display\n"
+            "import comm\n"
+            "_pending = []\n"
+            "def _post_execute():\n"
+            "    while _pending:\n"
+            "        display(_pending.pop())\n"
+            "get_ipython().events.register('post_execute', _post_execute)\n"
+            "_comm = comm.create_comm(target_name='repro')\n"
+            "_comm.on_msg(lambda msg: None)\n"
+            "print('COMM_ID', _comm.comm_id)\n"
+        )
+        comm_id = execute_request_subshell_id(kc, setup, None).split()[1]
+
+        subshell_id = create_subshell_helper(kc)["subshell_id"] if use_subshell else None
+        try:
+            cell = kc.session.msg(
+                "execute_request",
+                {"code": "_pending.append('PAYLOAD')\nimport asyncio\nawait asyncio.sleep(3)\n"},
+            )
+            kc.shell_channel.send(cell)
+            time.sleep(1.0)
+            comm_msg = kc.session.msg(
+                "comm_msg",
+                {"comm_id": comm_id, "data": {"method": "custom", "content": {}}},
+            )
+            if subshell_id:
+                comm_msg["header"]["subshell_id"] = subshell_id
+            kc.shell_channel.send(comm_msg)
+
+            cell_id = cell["header"]["msg_id"]
+            display_parent = None
+            while True:
+                msg = kc.get_iopub_msg(timeout=TIMEOUT)
+                if msg["msg_type"] == "display_data":
+                    display_parent = msg["parent_header"].get("msg_id")
+                if (
+                    msg["msg_type"] == "status"
+                    and msg["content"]["execution_state"] == "idle"
+                    and msg["parent_header"].get("msg_id") == cell_id
+                ):
+                    break
+        finally:
+            if subshell_id:
+                delete_subshell_helper(kc, subshell_id)
+
+        assert display_parent is not None, "no display_data observed"
+        assert display_parent == cell_id, (
+            "display_data parented to the comm_msg, not the cell's execute_request"
+        )