Prototype for re-attach

Author: DonJayamanne

src/kernels/execution/cellExecution.ts

@@ -44,9 +44,21 @@ export class CellExecutionFactory {
         private readonly requestListener: CellExecutionMessageHandlerService
     ) {}
 
-    public create(cell: NotebookCell, code: string | undefined, metadata: Readonly<KernelConnectionMetadata>) {
+    public create(
+        cell: NotebookCell,
+        code: string | undefined,
+        metadata: Readonly<KernelConnectionMetadata>,
+        resumeExecutionMsgId?: string
+    ) {
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
-        return CellExecution.fromCell(cell, code, metadata, this.controller, this.requestListener);
+        return CellExecution.fromCell(
+            cell,
+            code,
+            metadata,
+            this.controller,
+            this.requestListener,
+            resumeExecutionMsgId
+        );
     }
 }
 
@@ -88,7 +100,8 @@ export class CellExecution implements IDisposable {
         private readonly codeOverride: string | undefined,
         private readonly kernelConnection: Readonly<KernelConnectionMetadata>,
         private readonly controller: IKernelController,
-        private readonly requestListener: CellExecutionMessageHandlerService
+        private readonly requestListener: CellExecutionMessageHandlerService,
+        private readonly resumeExecutionMsgId?: string
     ) {
         workspace.onDidCloseTextDocument(
             (e) => {
@@ -134,11 +147,55 @@ export class CellExecution implements IDisposable {
         code: string | undefined,
         metadata: Readonly<KernelConnectionMetadata>,
         controller: IKernelController,
-        requestListener: CellExecutionMessageHandlerService
+        requestListener: CellExecutionMessageHandlerService,
+        resumeExecutionMsgId?: string
     ) {
-        return new CellExecution(cell, code, metadata, controller, requestListener);
+        return new CellExecution(cell, code, metadata, controller, requestListener, resumeExecutionMsgId);
+    }
+    public async resume(session: IKernelConnectionSession) {
+        if (this.cancelHandled) {
+            traceCellMessage(this.cell, 'Not resuming as it was cancelled');
+            return;
+        }
+        traceCellMessage(this.cell, 'Start resuming execution');
+        traceInfoIfCI(`Cell Exec (resuming) contents ${this.cell.document.getText().substring(0, 50)}...`);
+        if (!this.canExecuteCell()) {
+            // End state is bool | undefined not optional. Undefined == not success or failure
+            this.execution?.end(undefined);
+            this.execution = undefined;
+            this._result.resolve();
+            return;
+        }
+        if (this.started) {
+            traceCellMessage(this.cell, 'Cell has already been started yet CellExecution.Start invoked again');
+            traceError(`Cell has already been started yet CellExecution.Start invoked again ${this.cell.index}`);
+            // TODO: Send telemetry this should never happen, if it does we have problems.
+            return this.result;
+        }
+        this.started = true;
+
+        this.startTime = new Date().getTime();
+        activeNotebookCellExecution.set(this.cell.notebook, this.execution);
+        this.execution?.start(this.startTime);
+        NotebookCellStateTracker.setCellState(this.cell, NotebookCellExecutionState.Executing);
+        this.stopWatch.reset();
+
+        this.cellExecutionHandler = this.requestListener.registerListenerForResumingExecution(this.cell, {
+            kernel: session.kernel!,
+            cellExecution: this.execution!,
+            msg_id: this.resumeExecutionMsgId!,
+            onErrorHandlingExecuteRequestIOPubMessage: (error) => {
+                traceError(`Cell (index = ${this.cell.index}) execution completed with errors (2).`, error);
+                // If not a restart error, then tell the subscriber
+                this.completedWithErrors(error);
+            }
+        });
+        this.cellExecutionHandler.completed.finally(() => this.completedSuccessfully());
     }
     public async start(session: IKernelConnectionSession) {
+        if (this.resumeExecutionMsgId) {
+            return this.resume(session);
+        }
         if (this.cancelHandled) {
             traceCellMessage(this.cell, 'Not starting as it was cancelled');
             return;
@@ -348,7 +405,7 @@ export class CellExecution implements IDisposable {
             traceError(`Cell execution failed without request, for cell Index ${this.cell.index}`, ex);
             return this.completedWithErrors(ex);
         }
-        this.cellExecutionHandler = this.requestListener.registerListener(this.cell, {
+        this.cellExecutionHandler = this.requestListener.registerListenerForExecution(this.cell, {
             kernel: session.kernel!,
             cellExecution: this.execution!,
             request: this.request,

src/kernels/execution/cellExecutionMessageHandler.ts

@@ -41,6 +41,7 @@ import { IKernelController, ITracebackFormatter } from '../../kernels/types';
 import { handleTensorBoardDisplayDataOutput } from './executionHelpers';
 import { Identifiers, WIDGET_MIMETYPE } from '../../platform/common/constants';
 import { CellOutputDisplayIdTracker } from './cellDisplayIdTracker';
+import { createDeferred } from '../../platform/common/utils/async';
 
 // Helper interface for the set_next_input execute reply payload
 interface ISetNextInputPayload {
@@ -166,18 +167,26 @@ export class CellExecutionMessageHandler implements IDisposable {
      * or for any subsequent requests as a result of outputs sending custom messages.
      */
     private readonly ownedRequestMsgIds = new Set<string>();
+    private readonly _completed = createDeferred<void>();
+    public get completed() {
+        return this._completed.promise;
+    }
     constructor(
         public readonly cell: NotebookCell,
         private readonly applicationService: IApplicationShell,
         private readonly controller: IKernelController,
         private readonly context: IExtensionContext,
         private readonly formatters: ITracebackFormatter[],
         private readonly kernel: Kernel.IKernelConnection,
-        request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg>,
-        cellExecution: NotebookCellExecution
+        private readonly request:
+            | Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg>
+            | undefined,
+        cellExecution: NotebookCellExecution,
+        executionMessageId: string
     ) {
-        this.executeRequestMessageId = request.msg.header.msg_id;
-        this.ownedRequestMsgIds.add(request.msg.header.msg_id);
+        this._completed.promise.catch(noop);
+        this.executeRequestMessageId = executionMessageId;
+        this.ownedRequestMsgIds.add(executionMessageId);
         workspace.onDidChangeNotebookDocument(
             (e) => {
                 if (!isJupyterNotebook(e.notebook)) {
@@ -203,27 +212,29 @@ export class CellExecutionMessageHandler implements IDisposable {
         this.kernel.anyMessage.connect(this.onKernelAnyMessage, this);
         this.kernel.iopubMessage.connect(this.onKernelIOPubMessage, this);
 
-        request.onIOPub = () => {
-            // Cell has been deleted or the like.
-            if (this.cell.document.isClosed && !this.completedExecution) {
-                request.dispose();
-            }
-        };
-        request.onReply = (msg) => {
-            // Cell has been deleted or the like.
-            if (this.cell.document.isClosed) {
-                request.dispose();
-                return;
-            }
-            this.handleReply(msg);
-        };
-        request.onStdin = this.handleInputRequest.bind(this);
-        request.done
-            .finally(() => {
-                this.completedExecution = true;
-                this.endCellExecution();
-            })
-            .catch(noop);
+        if (request) {
+            request.onIOPub = () => {
+                // Cell has been deleted or the like.
+                if (this.cell.document.isClosed && !this.completedExecution) {
+                    request.dispose();
+                }
+            };
+            request.onReply = (msg) => {
+                // Cell has been deleted or the like.
+                if (this.cell.document.isClosed) {
+                    request.dispose();
+                    return;
+                }
+                this.handleReply(msg);
+            };
+            request.onStdin = this.handleInputRequest.bind(this);
+            request.done
+                .finally(() => {
+                    this.completedExecution = true;
+                    this.endCellExecution();
+                })
+                .catch(noop);
+        }
     }
     /**
      * This method is called when all execution has been completed (successfully or failed).
@@ -254,12 +265,24 @@ export class CellExecutionMessageHandler implements IDisposable {
         this.prompts.clear();
         this.clearLastUsedStreamOutput();
         this.execution = undefined;
+        this._completed.resolve();
     }
     private onKernelAnyMessage(_: unknown, { direction, msg }: Kernel.IAnyMessageArgs) {
         if (this.cell.document.isClosed) {
             return this.endCellExecution();
         }
-
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        console.log((msg as any).msg_type);
+        if (
+            !this.request &&
+            'msg_type' in msg &&
+            (msg.msg_type === 'kernel_info_reply' ||
+                msg.msg_type === 'execute_input' ||
+                msg.msg_type === 'execute_reply')
+        ) {
+            this.completedExecution = true;
+            return this.endCellExecution();
+        }
         // We're only interested in messages after execution has completed.
         // See https://github.com/microsoft/vscode-jupyter/issues/9503 for more information.
         if (direction !== 'send' || !this.completedExecution) {

src/kernels/execution/cellExecutionMessageHandlerService.ts

@@ -2,12 +2,13 @@
 // Licensed under the MIT License.
 
 import type { Kernel, KernelMessage } from '@jupyterlab/services';
-import { NotebookCell, NotebookCellExecution, NotebookDocument, workspace } from 'vscode';
+import { Memento, NotebookCell, NotebookCellExecution, NotebookDocument, workspace } from 'vscode';
 import { IKernelController, ITracebackFormatter } from '../../kernels/types';
 import { IApplicationShell } from '../../platform/common/application/types';
 import { disposeAllDisposables } from '../../platform/common/helpers';
 import { IDisposable, IExtensionContext } from '../../platform/common/types';
 import { CellExecutionMessageHandler } from './cellExecutionMessageHandler';
+import { noop } from '../../platform/common/utils/misc';
 
 /**
  * Allows registering a CellExecutionMessageHandler for a given execution.
@@ -20,7 +21,8 @@ export class CellExecutionMessageHandlerService {
         private readonly appShell: IApplicationShell,
         private readonly controller: IKernelController,
         private readonly context: IExtensionContext,
-        private readonly formatters: ITracebackFormatter[]
+        private readonly formatters: ITracebackFormatter[],
+        private readonly workspaceStorage: Memento
     ) {
         workspace.onDidChangeNotebookDocument(
             (e) => {
@@ -42,7 +44,7 @@ export class CellExecutionMessageHandlerService {
             this.notebook.getCells().forEach((cell) => this.messageHandlers.get(cell)?.dispose());
         }
     }
-    public registerListener(
+    public registerListenerForExecution(
         cell: NotebookCell,
         options: {
             kernel: Kernel.IKernelConnection;
@@ -54,6 +56,12 @@ export class CellExecutionMessageHandlerService {
         this.notebook = cell.notebook;
         // Always dispose any previous handlers & create new ones.
         this.messageHandlers.get(cell)?.dispose();
+        this.workspaceStorage
+            .update(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, {
+                index: cell.index,
+                msg_id: options.request?.msg.header.msg_id
+            })
+            .then(noop, noop);
         const handler = new CellExecutionMessageHandler(
             cell,
             this.appShell,
@@ -62,7 +70,41 @@ export class CellExecutionMessageHandlerService {
             this.formatters,
             options.kernel,
             options.request,
-            options.cellExecution
+            options.cellExecution,
+            options.request.msg.header.msg_id
+        );
+        // This object must be kept in memory has it monitors the kernel messages.
+        this.messageHandlers.set(cell, handler);
+        return handler;
+    }
+    public registerListenerForResumingExecution(
+        cell: NotebookCell,
+        options: {
+            kernel: Kernel.IKernelConnection;
+            msg_id: string;
+            cellExecution: NotebookCellExecution;
+            onErrorHandlingExecuteRequestIOPubMessage: (error: Error) => void;
+        }
+    ): CellExecutionMessageHandler {
+        this.notebook = cell.notebook;
+        // Always dispose any previous handlers & create new ones.
+        this.messageHandlers.get(cell)?.dispose();
+        this.workspaceStorage
+            .update(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, {
+                index: cell.index,
+                msg_id: options.msg_id
+            })
+            .then(noop, noop);
+        const handler = new CellExecutionMessageHandler(
+            cell,
+            this.appShell,
+            this.controller,
+            this.context,
+            this.formatters,
+            options.kernel,
+            undefined,
+            options.cellExecution,
+            options.msg_id
         );
         // This object must be kept in memory has it monitors the kernel messages.
         this.messageHandlers.set(cell, handler);

src/kernels/execution/cellExecutionQueue.ts

@@ -77,6 +77,25 @@ export class CellExecutionQueue implements Disposable {
         // Start executing the cells.
         this.startExecutingCells();
     }
+
+    /**
+     * Queue the cell for execution & start processing it immediately.
+     */
+    public resumeCell(cell: NotebookCell, msg_id: string): void {
+        const existingCellExecution = this.queueOfCellsToExecute.find((item) => item.cell === cell);
+        if (existingCellExecution) {
+            traceCellMessage(cell, 'Use existing cell execution');
+            return;
+        }
+        const cellExecution = this.executionFactory.create(cell, '', this.metadata, msg_id);
+        this.disposables.push(cellExecution);
+        this.queueOfCellsToExecute.push(cellExecution);
+
+        traceCellMessage(cell, 'User queued cell for execution');
+
+        // Start executing the cells.
+        this.startExecutingCells();
+    }
     /**
      * Cancel all cells that have been queued & wait for them to complete.
      * @param {boolean} [forced=false]

src/kernels/kernel.ts

@@ -684,7 +684,6 @@ abstract class BaseKernel implements IBaseKernel {
         // So that we don't have problems with ipywidgets, always register the default ipywidgets comm target.
         // Restart sessions and retries might make this hard to do correctly otherwise.
         session.registerCommTarget(Identifiers.DefaultCommTarget, noop);
-
         if (this.kernelConnectionMetadata.kind === 'connectToLiveRemoteKernel') {
             // As users can have IPyWidgets at any point in time, we need to determine the version of ipywidgets
             // This must happen early on as the state of the kernel needs to be synced with the Kernel in the webview (renderer)