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

Author: krassowski

ipykernel/kernelbase.py

@@ -594,14 +594,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 KeyError:
+                    self._send_unknown_subshell_reply(idents, msg3, subshell_id)
+                    return
                 assert socket is not None
                 socket.send_multipart(msg, copy=False)
             except Exception:
@@ -1378,6 +1382,46 @@ def _send_abort_reply(self, stream, msg, idents):
             ident=idents,
         )
 
+    def _send_unknown_subshell_reply(self, idents, msg, subshell_id) -> 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: unknown subshell_id %r",
+            msg_type,
+            msg["header"]["msg_id"],
+            subshell_id,
+        )
+        self._publish_status("busy", "shell", parent=msg)
+        content = {
+            "status": "error",
+            "ename": "KeyError",
+            "evalue": f"Unknown subshell_id {subshell_id!r}",
+            "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."""

tests/test_subshells.py

@@ -391,3 +391,45 @@ 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)
+        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 subshell_id in reply["content"]["evalue"]
+
+        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"])