Corrected execution counter

Author: subhash-arabhi

vscode/src/notebooks/kernel.ts

@@ -24,7 +24,7 @@ import { isNbCommandRegistered } from '../commands/utils';
 import { nbCommands } from '../commands/commands';
 import { createErrorOutput, createOutputItem } from './utils';
 import { NotebookCellExecutionResult } from '../lsp/protocol';
-import { NotebookCell, NotebookController, NotebookDocument, Disposable, notebooks, commands, NotebookCellOutput } from 'vscode';
+import { NotebookCell, NotebookController, NotebookDocument, Disposable, notebooks, commands, NotebookCellOutput, workspace } from 'vscode';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { l10n } from '../localiser';
 import { ijnbConstants, ipynbConstants, supportLanguages } from './constants';
@@ -86,6 +86,7 @@ export class IJNBKernel implements Disposable {
     const cellId = cell.document.uri.toString();
     const sourceCode = cell.document.getText();
     const codeCellExecution = new CodeCellExecution(controller.id, notebookId, cell);
+    this.getIncementedExecutionCounter(notebookId);
     try {
       this.cellControllerIdMap.set(cellId, codeCellExecution);
       const client: LanguageClient = await globalState.getClientPromise().client;
@@ -141,15 +142,15 @@ export class IJNBKernel implements Disposable {
 
   private handleUnkownLanguageTypeExecution = async (notebookId: string, cell: NotebookCell, controller: NotebookController) => {
     const exec = controller.createNotebookCellExecution(cell);
-    exec.executionOrder = this.getExecutionCounterAndIncrement(notebookId);
+    exec.executionOrder = this.getIncementedExecutionCounter(notebookId);
     exec.start(Date.now());
     await exec.replaceOutput(createErrorOutput(new Error(`Doesn't support ${cell.document.languageId} execution`)));
-    exec.end(true, Date.now());
+    exec.end(false, Date.now());
   }
 
-  private getExecutionCounterAndIncrement = (notebookId: string) => {
-    const next = IJNBKernel.executionCounter.get(notebookId) ?? 1;
-    IJNBKernel.executionCounter.set(notebookId, next + 1);
+  private getIncementedExecutionCounter = (notebookId: string) => {
+    const next = (IJNBKernel.executionCounter.get(notebookId) ?? 0) + 1;
+    IJNBKernel.executionCounter.set(notebookId, next);
     return next;
   }
 
@@ -160,18 +161,24 @@ export class IJNBKernel implements Disposable {
   private handleMarkdownCellExecution = async (notebookId: string, cell: NotebookCell, controller: NotebookController) => {
     const exec = controller.createNotebookCellExecution(cell);
     const mimeType = 'text/markdown';
-    exec.executionOrder = this.getExecutionCounterAndIncrement(notebookId);
+    exec.executionOrder = this.getIncementedExecutionCounter(notebookId);
     try {
       exec.start(Date.now());
       await exec.replaceOutput([
         new NotebookCellOutput([createOutputItem(cell.document.getText(), mimeType)]),
       ]);
+      exec.end(true, Date.now());
     } catch (error) {
       await exec.replaceOutput(createErrorOutput(error as Error));
-    } finally {
-      exec.end(true, Date.now());
+      exec.end(false, Date.now());
     }
   }
+
+  cleanUpKernel = workspace.onDidCloseNotebookDocument(doc => {
+  if (doc.notebookType === ijnbConstants.NOTEBOOK_TYPE) {
+    IJNBKernel.executionCounter.delete(doc.uri.toString());
+  }
+});
 }
 
 export const notebookKernel = new IJNBKernel();

vscode/src/notebooks/register.ts

@@ -31,4 +31,5 @@ export const registerNotebookProviders = (context: ExtensionContext) => {
         ));
 
     context.subscriptions.push(notebookKernel);
+    context.subscriptions.push(notebookKernel.cleanUpKernel);
 }

vscode/src/notebooks/utils.ts

@@ -31,7 +31,7 @@ import {
   IMetadata
 } from './types';
 import { randomUUID } from 'crypto';
-import { isError, isString } from '../utils';
+import { isString } from '../utils';
 import { mimeTypes } from './constants';
 
 export function base64ToUint8Array(base64: string): Uint8Array {