Abort individual subshell on exception

Author: ianthomas23

ipykernel/kernelbase.py

@@ -393,10 +393,15 @@ async def dispatch_shell(self, msg, /, stream=None, subshell_id: str | None = No
                 stream = self.shell_stream
 
         # Only abort execute requests
-        if self._aborting and msg_type == "execute_request":
-            self._send_abort_reply(self.shell_stream, msg, idents)
-            self._publish_status_and_flush("idle", "shell", stream)
-            return
+        if msg_type == "execute_request":
+            if subshell_id is None:
+                aborting = self._aborting  # type:ignore[unreachable]
+            else:
+                aborting = self.shell_channel_thread.manager.get_subshell_aborting(subshell_id)
+            if aborting:
+                self._send_abort_reply(self.shell_stream, msg, idents)
+                self._publish_status_and_flush("idle", "shell", stream)
+                return
 
         # Print some info about this message and leave a '--->' marker, so it's
         # easier to trace visually the message chain when debugging.  Each

ipykernel/subshell_manager.py

@@ -113,6 +113,10 @@ def get_shell_channel_stream(self, subshell_id: str | None) -> ZMQStream:
         with self._lock_cache:
             return self._cache[subshell_id].shell_channel_stream
 
+    def get_subshell_aborting(self, subshell_id: str) -> bool:
+        """Get the aborting flag of the specified subshell."""
+        return self._cache[subshell_id].aborting
+
     def list_subshell(self) -> list[str]:
         """Return list of current subshell ids.
 

tests/test_subshells.py

@@ -209,7 +209,7 @@ def test_shutdown_with_subshell():
         assert not km.is_alive()
 
 
-@pytest.mark.parametrize("are_subshells", [(False, True)])
+@pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)])
 def test_execute_stop_on_error(are_subshells):
     # Based on test_message_spec.py::test_execute_stop_on_error, testing that exception
     # in one subshell aborts execution queue in that subshell but not others.
@@ -221,7 +221,6 @@ def test_execute_stop_on_error(are_subshells):
 
         msg_ids = []
 
-        # msg = execute_request(kc, "print('ok')", subshell_ids[0])
         msg = execute_request(
             kc, "import asyncio; await asyncio.sleep(1); raise ValueError()", subshell_ids[0]
         )