wip

Author: DonJayamanne

src/kernels/execution/cellExecutionMessageHandlerService.ts

@@ -75,6 +75,26 @@ export class CellExecutionMessageHandlerService {
         );
         // This object must be kept in memory has it monitors the kernel messages.
         this.messageHandlers.set(cell, handler);
+        handler.completed.finally(() => {
+            const info = this.workspaceStorage.get<
+                | {
+                      index: number;
+                      msg_id: string;
+                  }
+                | undefined
+            >(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, undefined);
+            if (
+                !info ||
+                info.index !== cell.index ||
+                cell.document.isClosed ||
+                info?.msg_id !== options.request?.msg.header.msg_id
+            ) {
+                return;
+            }
+            this.workspaceStorage
+                .update(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, undefined)
+                .then(noop, noop);
+        });
         return handler;
     }
     public registerListenerForResumingExecution(

src/notebooks/controllers/vscodeNotebookController.ts

@@ -247,6 +247,7 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
         );
     }
     private readonly pendingOutuptsOfNotebooks = new WeakSet<NotebookDocument>();
+    private readonly restoredConnections = new WeakSet<NotebookDocument>();
     public async restoreOutput(notebook: NotebookDocument) {
         console.error('Done');
         const kernel = await this.connectToKernel(notebook, new DisplayOptions(true));
@@ -270,6 +271,10 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
         console.error('Done', kernel.uri, kernelExecution.pendingCells.length);
     }
     public async restoreConnection(notebook: NotebookDocument) {
+        if (this.restoredConnections.has(notebook)) {
+            return;
+        }
+        this.restoredConnections.add(notebook);
         console.error('Done');
         const kernel = await this.connectToKernel(notebook, new DisplayOptions(true));
         const kernelExecution = this.kernelProvider.getKernelExecution(kernel);
@@ -309,13 +314,20 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
             lastExecutionInfo &&
             typeof lastExecutionInfo.index === 'number'
         ) {
+            let resumed = false;
+            kernel.session.kernel.statusChanged.connect((_, status) => {
+                console.log(status);
+            });
             kernel.session.kernel.anyMessage.connect((_, msg) => {
+                if (msg.direction === 'send' || resumed) {
+                    return;
+                }
                 if (
-                    msg.direction === 'recv' &&
                     msg.msg.parent_header &&
                     'msg_id' in msg.msg.parent_header &&
                     msg.msg.parent_header.msg_id === lastExecutionInfo.msg_id
                 ) {
+                    resumed = true;
                     kernelExecution
                         .resumeCellExecution(notebook.cellAt(lastExecutionInfo.index), lastExecutionInfo.msg_id)
                         .catch(noop);