From b25b1e13ddba642c56aae0dee942f7efcc09acc8 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Mon, 21 Sep 2026 17:02:42 +0200 Subject: [PATCH 1/3] Add tests --- tests/test_kernel.py | 29 +++++++++++++++++++++++ tests/test_subshells.py | 51 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index fb9a7ad1f..27355baab 100644 --- a/tests/test_kernel.py +++ b/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 anyio, comm +reply = anyio.Future() +widget = comm.create_comm(target_name='comm-reply-test') +widget.on_msg(lambda msg: setattr(reply, 'return_value', msg['content']['data']['value'])) +with anyio.fail_after(5): + await reply.wait() +result = reply.return_value +assert result == 42 +""" + ) + while True: + msg = kc.get_iopub_msg(timeout=10) + 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"), [ diff --git a/tests/test_subshells.py b/tests/test_subshells.py index a8e1a1abf..402b1c664 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -133,6 +133,57 @@ 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 anyio, threading +reply = anyio.Future() +request_thread = threading.get_ident() +def on_reply(message): + global callback_thread + callback_thread = threading.get_ident() + reply.return_value = message['content']['data']['content']['value'] +widget.on_msg(on_reply) +widget.send({'method': 'custom', 'content': {'id': 'request-1', 'operation': 'get'}}) +with anyio.fail_after(2): + await reply.wait() +assert reply.return_value == 42 +assert callback_thread == request_thread +""", + None, + ) + while True: + outgoing = kc.get_iopub_msg(timeout=10) + 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): From 977cb71d26b19cc7e93048a2f3b6aaebe67bbac4 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Mon, 21 Sep 2026 16:22:47 +0200 Subject: [PATCH 2/3] Implement feature --- ipykernel/comm/comm.py | 33 ++++++++++++++++++++- ipykernel/kernelbase.py | 64 ++++++++++++++++++++++++++++++++--------- 2 files changed, 83 insertions(+), 14 deletions(-) diff --git a/ipykernel/comm/comm.py b/ipykernel/comm/comm.py index 9be5e23d0..54df0bbac 100644 --- a/ipykernel/comm/comm.py +++ b/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, ) diff --git a/ipykernel/kernelbase.py b/ipykernel/kernelbase.py index 3f709d0aa..be79337cc 100644 --- a/ipykernel/kernelbase.py +++ b/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) From 93fc609c5795559459c197955af25c0de7994bbb Mon Sep 17 00:00:00 2001 From: David Brochart Date: Mon, 21 Sep 2026 17:32:36 +0200 Subject: [PATCH 3/3] Use asyncio instead of anyio in tests --- tests/test_kernel.py | 12 ++++++------ tests/test_subshells.py | 15 +++++++++------ 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 27355baab..f5417db24 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -61,18 +61,18 @@ def test_simple_print(): def test_async_cell_waiting_for_comm_reply(): with new_kernel() as kc: msg_id = kc.execute( - """import anyio, comm -reply = anyio.Future() + """import asyncio, comm +reply = asyncio.get_running_loop().create_future() widget = comm.create_comm(target_name='comm-reply-test') -widget.on_msg(lambda msg: setattr(reply, 'return_value', msg['content']['data']['value'])) -with anyio.fail_after(5): - await reply.wait() -result = reply.return_value +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 diff --git a/tests/test_subshells.py b/tests/test_subshells.py index 402b1c664..6dcfaec35 100644 --- a/tests/test_subshells.py +++ b/tests/test_subshells.py @@ -144,24 +144,27 @@ def test_comm_reply_follows_requesting_subshell(): msg = execute_request( kc, - """import anyio, threading -reply = anyio.Future() + """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.return_value = message['content']['data']['content']['value'] + reply.set_result(message['content']['data']['content']['value']) widget.on_msg(on_reply) widget.send({'method': 'custom', 'content': {'id': 'request-1', 'operation': 'get'}}) -with anyio.fail_after(2): - await reply.wait() -assert reply.return_value == 42 +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"]