feat(notebook): implemented interupt execution and restart kernel

Author: d-biehl

package.json

@@ -1337,6 +1337,16 @@
         "shortTitle": "Robot Framework Notebook",
         "category": "RobotCode",
         "command": "robotcode.createNewNotebook"
+      },
+      {
+        "command": "robotcode.notebookEditor.restartKernel",
+        "title": "Restart Kernel",
+        "category": "RobotCode",
+        "shortTitle": "Restart",
+        "icon": {
+          "dark": "./resources/dark/restart-kernel.svg",
+          "light": "./resources/light/restart-kernel.svg"
+        }
       }
     ],
     "menus": {
@@ -1405,6 +1415,13 @@
           "group": "notebook",
           "when": "!virtualWorkspace"
         }
+      ],
+      "notebook/toolbar": [
+        {
+          "command": "robotcode.notebookEditor.restartKernel",
+          "group": "navigation/execute@1",
+          "when": "notebookKernel =~ /robotframework-repl/ && isWorkspaceTrusted"
+        }
       ]
     },
     "breakpoints": [
@@ -1926,4 +1943,4 @@
   "workspaces": [
     "docs"
   ]
-}
+}

packages/repl/src/robotcode/repl/base_interpreter.py

@@ -1,4 +1,5 @@
 import abc
+import signal
 from datetime import datetime
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union, cast
@@ -9,7 +10,7 @@
 from robot.output import Message as OutputMessage
 from robot.running import Keyword, TestCase, TestSuite
 from robot.running.context import EXECUTION_CONTEXTS
-from robot.running.signalhandler import _StopSignalMonitor
+from robot.running.signalhandler import STOP_SIGNAL_MONITOR, _StopSignalMonitor
 
 from robotcode.core.utils.path import normalized_path
 from robotcode.robot.utils import get_robot_version
@@ -19,18 +20,33 @@
     from robot import result, running
 
 
+class ExecutionInterrupted(ExecutionStatus):
+    pass
+
+
 def _register_signal_handler(self: Any, exsignum: Any) -> None:
     pass
 
 
+def _stop_signal_monitor_call(self: Any, signum: Any, frame: Any) -> None:
+    if self._running_keyword:
+        self._stop_execution_gracefully()
+
+
+def _stop_signal_monitor_stop_execution_gracefully(self: Any) -> None:
+    raise ExecutionInterrupted("Execution interrupted")
+
+
 _patched = False
 
 
 def _patch() -> None:
     global _patched
     if not _patched:
         # Monkey patching the _register_signal_handler method to disable robot's signal handling
-        _StopSignalMonitor._register_signal_handler = _register_signal_handler
+        # _StopSignalMonitor._register_signal_handler = _register_signal_handler
+        _StopSignalMonitor.__call__ = _stop_signal_monitor_call
+        _StopSignalMonitor._stop_execution_gracefully = _stop_signal_monitor_stop_execution_gracefully
 
     _patched = True
 
@@ -167,7 +183,10 @@ def run_keyword(self, kw: Keyword) -> Any:
         except ExecutionStatus:
             raise
         except BaseException as e:
-            self.log_message(str(e), "ERROR", timestamp=datetime.now())  # noqa: DTZ005
+            self.log_message(f"{type(e)}: {e}", "ERROR", timestamp=datetime.now())  # noqa: DTZ005
+
+    def interrupt(self) -> None:
+        signal.raise_signal(signal.SIGINT)
 
     def run(self) -> Any:
         self._logger.enabled = True
@@ -181,6 +200,8 @@ def run(self) -> Any:
                     break
                 except (SystemExit, KeyboardInterrupt):
                     break
+                except ExecutionInterrupted as e:
+                    self.log_message(str(e), "ERROR", timestamp=datetime.now())  # noqa: DTZ005
                 except ExecutionStatus:
                     pass
                 except BaseException as e:
@@ -190,9 +211,10 @@ def run(self) -> Any:
 
     def run_input(self) -> None:
         for kw in self.get_input():
-            if kw is None:
-                break
-            self.set_last_result(self.run_keyword(kw))
+            with STOP_SIGNAL_MONITOR:
+                if kw is None:
+                    break
+                self.set_last_result(self.run_keyword(kw))
 
     def set_last_result(self, result: Any) -> None:
         self.last_result = result

packages/repl_server/src/robotcode/repl_server/interpreter.py

@@ -108,11 +108,14 @@ def __init__(
         self.files = files
         self.has_input = Event()
         self.executed = Event()
+        self.no_execution = Event()
+        self.no_execution.set()
         self._code: List[str] = []
         self._success: Optional[bool] = None
         self._result_data: Optional[ResultData] = None
         self._result_data_stack: List[ResultData] = []
         self.collect_messages: bool = False
+        self._interrupted = False
         self._has_shutdown = False
         self._cell_errors: List[str] = []
 
@@ -122,11 +125,17 @@ def shutdown(self) -> None:
         self.has_input.set()
 
     def execute(self, source: str) -> ExecutionResult:
+        self.no_execution.wait()
+
+        self.no_execution.clear()
+
         self._result_data_stack = []
 
         self._success = None
         try:
             self._cell_errors = []
+            self._interrupted = False
+
             self._result_data = RootResultData()
 
             self.executed.clear()
@@ -159,6 +168,8 @@ def execute(self, source: str) -> ExecutionResult:
             )
         except BaseException as e:
             return ExecutionResult(False, [ExecutionOutput("application/vnd.code.notebook.stderr", str(e))])
+        finally:
+            self.no_execution.set()
 
     def get_input(self) -> Iterator[Optional[Keyword]]:
         while self._code:

packages/repl_server/src/robotcode/repl_server/protocol.py

@@ -17,9 +17,13 @@ def initialize(self, message: str) -> str:
         return "yeah initialized " + message
 
     @rpc_method(name="executeCell", threaded=True)
-    def execute_cell(self, source: str) -> Optional[ExecutionResult]:
+    def execute_cell(self, source: str, language_id: str) -> Optional[ExecutionResult]:
         return self.interpreter.execute(source)
 
+    @rpc_method(name="interrupt", threaded=True)
+    def interrupt(self) -> None:
+        self.interpreter.interrupt()
+
     @rpc_method(name="shutdown", threaded=True)
     def shutdown(self) -> None:
         try:

resources/dark/restart-kernel.svg

@@ -234,7 +234,14 @@ export class LanguageToolsManager {
           }
           if (folder === undefined) return;
 
-          const { pythonCommand, final_args } = await this.pythonManager.buildRobotCodeCommand(folder, ["repl"]);
+          const config = vscode.workspace.getConfiguration(CONFIG_SECTION, folder);
+          const profiles = config.get<string[]>("profiles", []);
+
+          const { pythonCommand, final_args } = await this.pythonManager.buildRobotCodeCommand(
+            folder,
+            ["repl"],
+            profiles,
+          );
           vscode.window
             .createTerminal({
               name: `Robot REPL${vscode.workspace.workspaceFolders?.length === 1 ? "" : ` (${folder.name})`}`,

resources/light/restart-kernel.svg

@@ -6,6 +6,7 @@ import { PythonManager } from "./pythonmanger";
 import * as cp from "child_process";
 import * as rpc from "vscode-jsonrpc/node";
 import { withTimeout } from "./utils";
+import { CONFIG_SECTION } from "./config";
 
 interface RawNotebook {
   cells: RawNotebookCell[];
@@ -161,7 +162,7 @@ export class ReplServerClient {
 
   connection: rpc.MessageConnection | undefined;
   childProcess: cp.ChildProcessWithoutNullStreams | undefined;
-  private _cancelationTokenSource: vscode.CancellationTokenSource | undefined;
+  private _cancelationTokenSources = new Map<vscode.CancellationTokenSource, vscode.NotebookCell>();
 
   dispose(): void {
     this.exitClient().finally(() => {});
@@ -198,10 +199,11 @@ export class ReplServerClient {
   }
 
   cancelCurrentExecution(): void {
-    if (this._cancelationTokenSource) {
-      this._cancelationTokenSource.cancel();
-      this._cancelationTokenSource.dispose();
+    for (const [token] of this._cancelationTokenSources) {
+      token.cancel();
+      token.dispose();
     }
+    this._cancelationTokenSources.clear();
   }
 
   async ensureInitialized(): Promise<void> {
@@ -223,10 +225,14 @@ export class ReplServerClient {
 
     const transport = await rpc.createClientPipeTransport(pipeName, "utf-8");
 
+    const config = vscode.workspace.getConfiguration(CONFIG_SECTION, folder);
+    const profiles = config.get<string[]>("profiles", []);
+
     const { pythonCommand, final_args } = await this.pythonManager.buildRobotCodeCommand(
       folder,
       //["-v", "--debugpy", "--debugpy-wait-for-client", "repl-server", "--pipe", pipeName],
       ["repl-server", "--pipe", pipeName, "--source", this.document.uri.fsPath],
+      profiles,
       undefined,
       true,
       true,
@@ -279,16 +285,16 @@ export class ReplServerClient {
     this.connection = connection;
   }
 
-  async executeCell(source: string): Promise<{ success?: boolean; output: vscode.NotebookCellOutput }> {
-    this._cancelationTokenSource = new vscode.CancellationTokenSource();
-
+  async executeCell(cell: vscode.NotebookCell): Promise<{ success?: boolean; output: vscode.NotebookCellOutput }> {
+    const _cancelationTokenSource = new vscode.CancellationTokenSource();
+    this._cancelationTokenSources.set(_cancelationTokenSource, cell);
     try {
       await this.ensureInitialized();
 
       const result = await this.connection?.sendRequest<ExecutionResult>(
         "executeCell",
-        { source },
-        this._cancelationTokenSource.token,
+        { source: cell.document.getText(), language_id: cell.document.languageId },
+        _cancelationTokenSource.token,
       );
 
       return {
@@ -306,10 +312,14 @@ export class ReplServerClient {
         ),
       };
     } finally {
-      this._cancelationTokenSource.dispose();
-      this._cancelationTokenSource = undefined;
+      this._cancelationTokenSources.delete(_cancelationTokenSource);
+      _cancelationTokenSource.dispose();
     }
   }
+
+  async interrupt(): Promise<void> {
+    await this.connection?.sendRequest("interrupt");
+  }
 }
 
 export class REPLNotebookController {
@@ -320,7 +330,7 @@ export class REPLNotebookController {
   readonly description = "A Robot Framework REPL notebook controller";
   readonly supportsExecutionOrder = true;
   readonly controller: vscode.NotebookController;
-  readonly _clients = new Map<vscode.NotebookDocument, ReplServerClient>();
+  readonly clients = new Map<vscode.NotebookDocument, ReplServerClient>();
 
   _outputChannel: vscode.OutputChannel | undefined;
 
@@ -343,18 +353,23 @@ export class REPLNotebookController {
     this.controller.supportsExecutionOrder = true;
     this.controller.description = "Robot Framework REPL";
     this.controller.interruptHandler = async (notebook: vscode.NotebookDocument) => {
-      this._clients.get(notebook)?.dispose();
-      this._clients.delete(notebook);
+      this.clients.get(notebook)?.interrupt();
     };
     this._disposables = vscode.Disposable.from(
       this.controller,
       vscode.workspace.onDidCloseNotebookDocument((document) => {
-        this._clients.get(document)?.dispose();
-        this._clients.delete(document);
+        this.disposeDocument(document);
       }),
     );
   }
 
+  disposeDocument(notebook: vscode.NotebookDocument): void {
+    const client = this.clients.get(notebook);
+    client?.interrupt();
+    client?.dispose();
+    this.clients.delete(notebook);
+  }
+
   outputChannel(): vscode.OutputChannel {
     if (!this._outputChannel) {
       this._outputChannel = vscode.window.createOutputChannel("RobotCode REPL");
@@ -363,18 +378,18 @@ export class REPLNotebookController {
   }
 
   dispose(): void {
-    for (const client of this._clients.values()) {
+    for (const client of this.clients.values()) {
       client.dispose();
     }
     this._disposables.dispose();
   }
 
   private getClient(document: vscode.NotebookDocument): ReplServerClient {
-    let client = this._clients.get(document);
+    let client = this.clients.get(document);
     if (!client) {
       client = new ReplServerClient(document, this.extensionContext, this.pythonManager, this.outputChannel());
       this.finalizeRegistry.register(document, client);
-      this._clients.set(document, client);
+      this.clients.set(document, client);
     }
     return client;
   }
@@ -398,8 +413,7 @@ export class REPLNotebookController {
 
     execution.start(Date.now());
     try {
-      const source = cell.document.getText();
-      const result = await client.executeCell(source);
+      const result = await client.executeCell(cell);
       if (result !== undefined) {
         success = result.success;
 
@@ -445,6 +459,20 @@ export class NotebookManager {
         );
         await vscode.commands.executeCommand("vscode.openWith", newNotebook.uri, "robotframework-repl");
       }),
+      vscode.commands.registerCommand("robotcode.notebookEditor.restartKernel", () => {
+        const notebook = vscode.window.activeNotebookEditor?.notebook;
+        if (notebook) {
+          vscode.window.withProgress(
+            {
+              location: vscode.ProgressLocation.Notification,
+              title: "Restarting kernel...",
+            },
+            async (_progress, _token) => {
+              this._notebookController.disposeDocument(notebook);
+            },
+          );
+        }
+      }),
     );
   }