From 2f2f629fbccf37b4ed1f4f110953a7b04bf76dc1 Mon Sep 17 00:00:00 2001 From: Daniel Biehl Date: Mon, 16 Dec 2024 00:02:51 +0100 Subject: [PATCH] feat(notebook): implemented interupt execution and restart kernel --- package.json | 19 ++++- .../src/robotcode/repl/base_interpreter.py | 34 +++++++-- .../src/robotcode/repl_server/interpreter.py | 11 +++ .../src/robotcode/repl_server/protocol.py | 6 +- resources/dark/restart-kernel.svg | 3 + resources/light/restart-kernel.svg | 3 + .../extension/languageToolsManager.ts | 9 ++- vscode-client/extension/notebook.ts | 70 +++++++++++++------ vscode-client/extension/pythonmanger.ts | 12 +++- .../extension/testcontrollermanager.ts | 10 +-- 10 files changed, 139 insertions(+), 38 deletions(-) create mode 100644 resources/dark/restart-kernel.svg create mode 100644 resources/light/restart-kernel.svg diff --git a/package.json b/package.json index 59a9ea0b4..7fdcfd47d 100644 --- a/package.json +++ b/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" ] -} +} \ No newline at end of file diff --git a/packages/repl/src/robotcode/repl/base_interpreter.py b/packages/repl/src/robotcode/repl/base_interpreter.py index 8c73b392e..bbcc730ea 100644 --- a/packages/repl/src/robotcode/repl/base_interpreter.py +++ b/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,10 +20,23 @@ 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 @@ -30,7 +44,9 @@ 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 diff --git a/packages/repl_server/src/robotcode/repl_server/interpreter.py b/packages/repl_server/src/robotcode/repl_server/interpreter.py index 9e66a7af7..23f3d7e8a 100644 --- a/packages/repl_server/src/robotcode/repl_server/interpreter.py +++ b/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: diff --git a/packages/repl_server/src/robotcode/repl_server/protocol.py b/packages/repl_server/src/robotcode/repl_server/protocol.py index f3e817f60..17b0e7714 100644 --- a/packages/repl_server/src/robotcode/repl_server/protocol.py +++ b/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: diff --git a/resources/dark/restart-kernel.svg b/resources/dark/restart-kernel.svg new file mode 100644 index 000000000..cb8138e8c --- /dev/null +++ b/resources/dark/restart-kernel.svg @@ -0,0 +1,3 @@ + + + diff --git a/resources/light/restart-kernel.svg b/resources/light/restart-kernel.svg new file mode 100644 index 000000000..63e4973f3 --- /dev/null +++ b/resources/light/restart-kernel.svg @@ -0,0 +1,3 @@ + + + diff --git a/vscode-client/extension/languageToolsManager.ts b/vscode-client/extension/languageToolsManager.ts index 274c1a9c4..a6865808d 100644 --- a/vscode-client/extension/languageToolsManager.ts +++ b/vscode-client/extension/languageToolsManager.ts @@ -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("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})`}`, diff --git a/vscode-client/extension/notebook.ts b/vscode-client/extension/notebook.ts index 2be2cac56..eded44e59 100644 --- a/vscode-client/extension/notebook.ts +++ b/vscode-client/extension/notebook.ts @@ -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(); 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 { @@ -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("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( "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 { + 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(); + readonly clients = new Map(); _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); + }, + ); + } + }), ); } diff --git a/vscode-client/extension/pythonmanger.ts b/vscode-client/extension/pythonmanger.ts index 347627afd..ce580bdd9 100644 --- a/vscode-client/extension/pythonmanger.ts +++ b/vscode-client/extension/pythonmanger.ts @@ -167,13 +167,21 @@ export class PythonManager { public async executeRobotCode( folder: vscode.WorkspaceFolder, args: string[], + profiles?: string[], format?: string, noColor?: boolean, noPager?: boolean, stdioData?: string, token?: vscode.CancellationToken, ): Promise { - const { pythonCommand, final_args } = await this.buildRobotCodeCommand(folder, args, format, noColor, noPager); + const { pythonCommand, final_args } = await this.buildRobotCodeCommand( + folder, + args, + profiles, + format, + noColor, + noPager, + ); this.outputChannel.appendLine(`executeRobotCode: ${pythonCommand} ${final_args.join(" ")}`); @@ -237,6 +245,7 @@ export class PythonManager { public async buildRobotCodeCommand( folder: vscode.WorkspaceFolder, args: string[], + profiles?: string[], format?: string, noColor?: boolean, noPager?: boolean, @@ -256,6 +265,7 @@ export class PythonManager { ...(format ? ["--format", format] : []), ...(noColor ? ["--no-color"] : []), ...(noPager ? ["--no-pager"] : []), + ...(profiles !== undefined ? profiles.flatMap((v) => ["--profile", v]) : []), ...args, ]; return { pythonCommand, final_args }; diff --git a/vscode-client/extension/testcontrollermanager.ts b/vscode-client/extension/testcontrollermanager.ts index be5c607fd..dcbc73f2d 100644 --- a/vscode-client/extension/testcontrollermanager.ts +++ b/vscode-client/extension/testcontrollermanager.ts @@ -427,12 +427,8 @@ export class TestControllerManager { return (await this.languageClientsManager.pythonManager.executeRobotCode( folder, - [ - ...(profiles === undefined ? [] : profiles.flatMap((v) => ["--profile", v])), - ...(paths?.length ? paths.flatMap((v) => ["--default-path", v]) : ["--default-path", "."]), - "profiles", - "list", - ], + [...(paths?.length ? paths.flatMap((v) => ["--default-path", v]) : ["--default-path", "."]), "profiles", "list"], + profiles, "json", true, true, @@ -677,7 +673,6 @@ export class TestControllerManager { const result = (await this.languageClientsManager.pythonManager.executeRobotCode( folder, [ - ...(profiles === undefined ? [] : profiles.flatMap((v) => ["--profile", v])), ...(paths?.length ? paths.flatMap((v) => ["--default-path", v]) : ["--default-path", "."]), ...discoverArgs, ...mode_args, @@ -686,6 +681,7 @@ export class TestControllerManager { ...robotArgs, ...extraArgs, ], + profiles, "json", true, true,