From e3fb2d7c546acf817359fe7b572af46ad3a8515e Mon Sep 17 00:00:00 2001 From: Dikshant Neupane Date: Sat, 31 Oct 2026 11:43:36 +0000 Subject: [PATCH 1/2] Do not fire pre_execute/post_execute for comms dispatched concurrently during a cell During an async cell, comm messages are dispatched concurrently while the cell holds the shell lock. Firing pre_execute/post_execute for those callbacks parents any hook output to the comm_msg instead of the cell's execute_request, so frontends drop it and the hook's queue is drained before the cell's own post_execute fires (see #1564). A context variable marks the concurrent dispatch, and BaseComm.handle_msg skips the two events only there; comms handled while no cell is running keep the previous behaviour. --- ipykernel/comm/comm.py | 19 ++++++++++++ ipykernel/kernelbase.py | 16 ++++++++++ tests/test_comm.py | 34 +++++++++++++++++++++ tests/test_subshells.py | 67 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 136 insertions(+) diff --git a/ipykernel/comm/comm.py b/ipykernel/comm/comm.py index 54df0bbac..b61f06adf 100644 --- a/ipykernel/comm/comm.py +++ b/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: diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index e6705cfc8..08d176843 100644 --- a/ipykernel/kernelbase.py +++ b/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,6 +673,13 @@ 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), @@ -672,6 +687,7 @@ async def shell_main(self, subshell_id: str | None, msg): ) 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: diff --git a/tests/test_comm.py b/tests/test_comm.py index c9d4df77f..29606f2f0 100644 --- a/tests/test_comm.py +++ b/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() diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 6dcfaec35..5d6f8ab14 100644 --- a/tests/test_subshells.py +++ b/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" + ) From 51ff96ab1f9d6af2c77d60b0fdb0efb5d17ce062 Mon Sep 17 00:00:00 2001 From: Dikshant Neupane Date: Sat, 31 Oct 2026 11:49:54 +0000 Subject: [PATCH 2/2] Quote the MessageType annotation so old comm versions still import The annotation was evaluated at import time, so environments with a comm version that has no base_comm.MessageType failed to import the module. Quoting it makes the evaluation lazy; the runtime behaviour is unchanged. --- ipykernel/comm/comm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ipykernel/comm/comm.py b/ipykernel/comm/comm.py index b61f06adf..053bbc57c 100644 --- a/ipykernel/comm/comm.py +++ b/ipykernel/comm/comm.py @@ -35,7 +35,7 @@ def __init__(self, *args, **kwargs): self._reply_subshell_ids = {} super().__init__(*args, **kwargs) - def handle_msg(self, msg: comm.base_comm.MessageType) -> None: + 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