Reply with busy -> error -> idle if subshell unknown

ipykernel/kernelbase.py

@@ -59,6 +59,7 @@
 
 from ._version import kernel_protocol_version
 from .iostream import OutStream
+from .subshell_manager import UnknownSubshellError
 from .utils import LazyDict, _async_in_context
 
 psutil: t.Any | None = None
@@ -594,14 +595,18 @@ async def shell_channel_thread_main(self, msg):
 
             # deserialize only the header to get subshell_id
             # Keep original message to send to subshell_id unmodified.
-            _, msg2 = self.session.feed_identities(msg, copy=False)
+            idents, msg2 = self.session.feed_identities(msg, copy=False)
             try:
                 msg3 = self.session.deserialize(msg2, content=False, copy=False)
                 subshell_id = msg3["header"].get("subshell_id")
 
                 # Find inproc pair socket to use to send message to correct subshell.
                 subshell_manager = self.shell_channel_thread.manager
-                socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
+                try:
+                    socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
+                except UnknownSubshellError as err:
+                    self._send_unknown_subshell_reply(idents, msg3, err)
+                    return
                 assert socket is not None
                 socket.send_multipart(msg, copy=False)
             except Exception:
@@ -1378,6 +1383,41 @@ def _send_abort_reply(self, stream, msg, idents):
             ident=idents,
         )
 
+    def _send_unknown_subshell_reply(self, idents, msg, err: UnknownSubshellError) -> None:
+        """Send an error reply to a request addressed to a subshell that is not there.
+
+        Runs in the shell channel thread, so it writes to the shell socket
+        directly instead of going through a subshell.
+
+        The busy and idle status messages matter as much as the reply here.
+        A client tracks the completion of a request by the idle status that
+        carries it as parent, and for message types that have no reply, such as
+        the comm messages, that status is all it has to go on.
+        """
+        if not self.session:
+            return
+        msg_type = msg["header"]["msg_type"]
+        self.log.warning("Cannot handle %s %s: %s", msg_type, msg["header"]["msg_id"], err)
+        self._publish_status("busy", "shell", parent=msg)
+        content = {
+            "status": "error",
+            "ename": type(err).__name__,
+            "evalue": str(err),
+            "traceback": [],
+        }
+        md = self.init_metadata(msg)
+        md = self.finish_metadata(msg, md, content)
+        md.update({"status": "error"})
+        self.session.send(
+            self.shell_stream,
+            msg_type.rsplit("_", 1)[0] + "_reply",
+            metadata=md,
+            content=content,
+            parent=msg,
+            ident=idents,
+        )
+        self._publish_status("idle", "shell", parent=msg)
+
     def _no_raw_input(self):
         """Raise StdinNotImplementedError if active frontend doesn't support
         stdin."""

ipykernel/subshell_manager.py

@@ -19,6 +19,20 @@
 from .utils import _async_in_context
 
 
+class UnknownSubshellError(KeyError):
+    """A subshell_id that does not name an existing subshell.
+
+    .. versionadded:: 7.4
+    """
+
+    def __init__(self, subshell_id: str) -> None:
+        super().__init__(subshell_id)
+        self.subshell_id = subshell_id
+
+    def __str__(self) -> str:
+        return f"Unknown subshell_id {self.subshell_id!r}"
+
+
 class SubshellManager:
     """A manager of subshells.
 
@@ -87,13 +101,24 @@ def close(self) -> None:
         self._main_to_shell_channel.close()
         self._shell_channel_to_main.close()
 
+    def _get_subshell(self, subshell_id: str) -> SubshellThread:
+        """Return the thread of the specified subshell.
+
+        The caller must hold ``_lock_cache``.  Raises ``UnknownSubshellError`` if
+        there is no such subshell.
+        """
+        try:
+            return self._cache[subshell_id]
+        except KeyError:
+            raise UnknownSubshellError(subshell_id) from None
+
     def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketPair:
         """Return the inproc socket pair used to send messages from the shell channel
         to a particular subshell or main shell."""
         if subshell_id is None:
             return self._shell_channel_to_main
         with self._lock_cache:
-            return self._cache[subshell_id].shell_channel_to_subshell
+            return self._get_subshell(subshell_id).shell_channel_to_subshell
 
     def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
         """Return the socket used by a particular subshell or main shell to send
@@ -102,7 +127,7 @@ def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.S
         if subshell_id is None:
             return self._main_to_shell_channel.from_socket
         with self._lock_cache:
-            return self._cache[subshell_id].subshell_to_shell_channel.from_socket
+            return self._get_subshell(subshell_id).subshell_to_shell_channel.from_socket
 
     def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
         """Return the socket used by the shell channel to send messages to a particular
@@ -113,12 +138,12 @@ def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.S
     def get_subshell_aborting(self, subshell_id: str) -> bool:
         """Get the boolean aborting flag of the specified subshell."""
         with self._lock_cache:
-            return self._cache[subshell_id].aborting
+            return self._get_subshell(subshell_id).aborting
 
     def get_subshell_asyncio_lock(self, subshell_id: str) -> asyncio.Lock:
         """Return the asyncio lock belonging to the specified subshell."""
         with self._lock_cache:
-            return self._cache[subshell_id].asyncio_lock
+            return self._get_subshell(subshell_id).asyncio_lock
 
     def list_subshell(self) -> list[str]:
         """Return list of current subshell ids.
@@ -141,7 +166,7 @@ def set_on_recv_callback(self, on_recv_callback):
     def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None:
         """Set the aborting flag of the specified subshell."""
         with self._lock_cache:
-            self._cache[subshell_id].aborting = aborting
+            self._get_subshell(subshell_id).aborting = aborting
 
     def subshell_id_from_thread_id(self, thread_id: int) -> str | None:
         """Return subshell_id of the specified thread_id.
@@ -185,14 +210,15 @@ def _create_subshell(self) -> str:
     def _delete_subshell(self, subshell_id: str) -> None:
         """Delete subshell identified by subshell_id.
 
-        Raises key error if subshell_id not in cache.
+        Raises ``UnknownSubshellError`` if subshell_id not in cache.
         """
         assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
 
         with self._lock_cache:
-            subshell_threwad = self._cache.pop(subshell_id)
+            subshell_thread = self._get_subshell(subshell_id)
+            del self._cache[subshell_id]
 
-        self._stop_subshell(subshell_threwad)
+        self._stop_subshell(subshell_thread)
 
     def _process_control_request(
         self,

tests/test_subshells.py

@@ -391,3 +391,51 @@ def test_silent_flag_in_subshells():
             # Ensure subshell is always deleted
             if subshell_id:
                 delete_subshell_helper(kc, subshell_id)
+
+
+def test_unknown_subshell_id():
+    # A request for a subshell that does not exist is answered with an error rather
+    # than dropped, so that a client waiting for it does not wait forever.
+    with new_kernel() as kc:
+        subshell_id = create_subshell_helper(kc)["subshell_id"]
+        delete_subshell_helper(kc, subshell_id)
+
+        # Deleting it again names the missing subshell in the same way.
+        content = delete_subshell_helper(kc, subshell_id)
+        assert content["status"] == "error"
+        assert content["evalue"] == f"Unknown subshell_id {subshell_id!r}"
+        flush_channels(kc)
+
+        msg = execute_request(kc, "a = 1", subshell_id)
+        msg_id = msg["header"]["msg_id"]
+
+        reply = get_reply(kc, msg_id, TIMEOUT)
+        assert reply["content"]["status"] == "error"
+        assert reply["content"]["ename"] == "UnknownSubshellError"
+        assert reply["content"]["evalue"] == f"Unknown subshell_id {subshell_id!r}"
+
+        states = []
+        while True:
+            iopub_msg = kc.get_iopub_msg(timeout=TIMEOUT)
+            if iopub_msg["parent_header"].get("msg_id") != msg_id:
+                continue
+            assert iopub_msg["msg_type"] == "status"
+            states.append(iopub_msg["content"]["execution_state"])
+            if states[-1] == "idle":
+                break
+        assert states == ["busy", "idle"]
+
+
+def test_comm_close_on_deleted_subshell():
+    # A comm message has no reply of its own, so the idle status is the only thing
+    # that tells the client the kernel is done with it.
+    with new_kernel() as kc:
+        subshell_id = create_subshell_helper(kc)["subshell_id"]
+        delete_subshell_helper(kc, subshell_id)
+        flush_channels(kc)
+
+        msg = kc.session.msg("comm_close", {"comm_id": "comm-1", "data": {}})
+        msg["header"]["subshell_id"] = subshell_id
+        kc.shell_channel.send(msg)
+
+        wait_for_idle(kc, msg["header"]["msg_id"])