diff --git a/jupyter_server_documents/kernels/kernel_client.py b/jupyter_server_documents/kernels/kernel_client.py index 9fd05c6..7c42e1d 100644 --- a/jupyter_server_documents/kernels/kernel_client.py +++ b/jupyter_server_documents/kernels/kernel_client.py @@ -80,7 +80,7 @@ async def stop_listening(self): """ # If the listening task isn't defined yet # do nothing. - if not self._listening_task: + if not hasattr(self, '_listening_task') or not self._listening_task: return # Attempt to cancel the task. @@ -93,6 +93,9 @@ async def stop_listening(self): # Log any exceptions that were raised. except Exception as err: self.log.error(err) + finally: + # Clear the task reference + self._listening_task = None _listening_task: t.Optional[t.Awaitable] = Any(allow_none=True) diff --git a/jupyter_server_documents/kernels/kernel_manager.py b/jupyter_server_documents/kernels/kernel_manager.py index 00083da..a59e443 100644 --- a/jupyter_server_documents/kernels/kernel_manager.py +++ b/jupyter_server_documents/kernels/kernel_manager.py @@ -127,8 +127,9 @@ async def connect(self): await asyncio.sleep(0.1) async def disconnect(self): - await self.main_client.stop_listening() - self.main_client.stop_channels() + if self.main_client: + await self.main_client.stop_listening() + self.main_client.stop_channels() async def broadcast_state(self): """Broadcast state to all listeners""" diff --git a/jupyter_server_documents/kernels/multi_kernel_manager.py b/jupyter_server_documents/kernels/multi_kernel_manager.py index 722ef9d..bc26806 100644 --- a/jupyter_server_documents/kernels/multi_kernel_manager.py +++ b/jupyter_server_documents/kernels/multi_kernel_manager.py @@ -7,4 +7,11 @@ def start_watching_activity(self, kernel_id): pass def stop_buffering(self, kernel_id): - pass \ No newline at end of file + pass + + # NOTE: Since we disable watching activity and buffering here, + # this method needs to be forked and remove code related to these things. + async def restart_kernel(self, kernel_id, now=False): + """Restart a kernel by kernel_id""" + self._check_kernel_id(kernel_id) + await self.pinned_superclass._async_restart_kernel(self, kernel_id, now=now) \ No newline at end of file