-

Author: davidbrochart

ipykernel/kernelbase.py

@@ -430,7 +430,6 @@ async def shell_main(self, subshell_id: str | None):
 
     async def process_shell(self, socket=None):
         # socket=None is valid if kernel subshells are not supported.
-        _socket = t.cast(zmq_anyio.Socket, self.shell_socket if socket is None else socket)
         try:
             while True:
                 await self.process_shell_message(socket=socket)
@@ -567,7 +566,7 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None:
                 manager = self.shell_channel_thread.manager
                 self.shell_channel_thread.start_soon(self.shell_channel_thread_main)
                 self.shell_channel_thread.start_soon(
-                    partial(manager.listen_from_control, self.shell_main)
+                    partial(manager.listen_from_control, self.shell_main, self.shell_channel_thread)
                 )
                 self.shell_channel_thread.start_soon(manager.listen_from_subshells)
                 self.shell_channel_thread.start()

ipykernel/subshell_manager.py

@@ -118,6 +118,11 @@ async def get_control_other_socket(self, thread: BaseThread) -> zmq_anyio.Socket
             await thread.task_group.start(self._control_other_socket.start)
         return self._control_other_socket
 
+    async def get_control_shell_channel_socket(self, thread: BaseThread) -> zmq_anyio.Socket:
+        if not self._control_shell_channel_socket.started.is_set():
+            await thread.task_group.start(self._control_shell_channel_socket.start)
+        return self._control_shell_channel_socket
+
     def get_other_socket(self, subshell_id: str | None) -> zmq_anyio.Socket:
         """Return the other inproc pair socket for a subshell.
 
@@ -148,18 +153,17 @@ def list_subshell(self) -> list[str]:
         with self._lock_cache:
             return list(self._cache)
 
-    async def listen_from_control(self, subshell_task: t.Any) -> None:
+    async def listen_from_control(self, subshell_task: t.Any, thread: BaseThread) -> None:
         """Listen for messages on the control inproc socket, handle those messages and
         return replies on the same socket.  Runs in the shell channel thread.
         """
         assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
 
-        socket = self._control_shell_channel_socket
-        async with socket:
-            while True:
-                request = await socket.arecv_json().wait()
-                reply = await self._process_control_request(request, subshell_task)
-                await socket.asend_json(reply).wait()
+        socket = await self.get_control_shell_channel_socket(thread)
+        while True:
+            request = await socket.arecv_json().wait()
+            reply = await self._process_control_request(request, subshell_task)
+            await socket.asend_json(reply).wait()
 
     async def listen_from_subshells(self) -> None:
         """Listen for reply messages on inproc sockets of all subshells and resend
@@ -265,8 +269,8 @@ async def _listen_for_subshell_reply(
 
         shell_channel_socket = self._get_shell_channel_socket(subshell_id)
 
-        task_group.start_soon(shell_channel_socket.start)
-        await shell_channel_socket.started.wait()
+        if not shell_channel_socket.started.is_set():
+            await task_group.start(shell_channel_socket.start)
         try:
             while True:
                 msg = await shell_channel_socket.arecv_multipart(copy=False).wait()