Improvements in notebook cell execution display

subhash-arabhi · Achal1607 · commit cf445a5c4108 · 2025-10-10T19:56:41.000+05:30
- Corrected execution counter
- Introduced mimeTypeHandler and executionSummary and refactored code
- Added new commands to localisation english file
- Added tests for mimeType and executionSummary of notebook code
- New and unexecuted cells will not display success or failure
diff --git a/vscode/package.json b/vscode/package.json
@@ -483,7 +483,7 @@
 			},
 			{
 				"command": "jdk.jshell.project",
-				"title": "Open JShell in context of a Project",
+				"title": "%jdk.jshell.project%",
 				"category": "Java"
 			},
 			{
@@ -591,7 +591,7 @@
 			},
 			{
 				"command": "jdk.notebook.new",
-				"title": "Create New Notebook...",
+				"title": "%jdk.notebook.new%",
 				"category": "Java",
 				"icon": "$(new-file)"
 			}
diff --git a/vscode/package.nls.json b/vscode/package.nls.json
@@ -6,11 +6,13 @@
     "jdk.workspace.new": "New File from Template...",
     "jdk.workspace.newproject": "New Project...",
     "jdk.java.goto.super.implementation": "Go to Super Implementation",
+    "jdk.jshell.project": "Open JShell in context of a Project",
     "jdk.open.type": "Open Type...",
     "jdk.foundProjects.deleteEntry": "Delete",
     "jdk.Edit.org.openide.actions.DeleteAction": "Delete",
     "workbench.action.debug.run": "Run Without Debugging",
     "workbench.action.debug.start": "Start Debugging",
+    "jdk.notebook.new": "Create New Notebook...",
     "jdk.project.run": "Run Project Without Debugging",
     "jdk.project.debug": "Debug Project",
     "jdk.project.test": "Test Project",
diff --git a/vscode/src/notebooks/codeCellExecution.ts b/vscode/src/notebooks/codeCellExecution.ts
@@ -1,9 +1,10 @@
 import { commands, NotebookCell, NotebookCellExecution, NotebookCellOutput, NotebookController } from "vscode";
 import { LOGGER } from "../logger";
 import { NotebookCellExecutionResult } from "../lsp/protocol";
-import { createErrorOutputItem, createOutputItem } from "./utils";
+import { createErrorOutputItem } from "./utils";
 import { nbCommands } from "../commands/commands";
 import { mimeTypes } from "./constants";
+import { MimeTypeHandler } from "./mimeTypeHandler";
 
 export class CodeCellExecution {
     private controller?: NotebookController;
@@ -32,7 +33,7 @@ export class CodeCellExecution {
         diagnostics: string[] | undefined,
         errorDiagnostics: string[] | undefined,
         metadata: any,
-        executionOrder: number) => {
+        executionOrder: number | undefined) => {
 
         if (!this.isExecutionStarted) {
             this.handleExecutionStart(executionOrder);
@@ -74,7 +75,7 @@ export class CodeCellExecution {
             await this.execution!.replaceOutputItems(createErrorOutputItem(updatedData), this.output);
         } else {
             await this.execution!.replaceOutputItems(
-                createOutputItem(updatedData, mimeType),
+                new MimeTypeHandler(mimeType).makeOutputItem(updatedData),
                 this.output
             );
         }
@@ -97,7 +98,7 @@ export class CodeCellExecution {
         }
     }
 
-    private handleExecutionStart = async (executionOrder: number) => {
+    private handleExecutionStart = async (executionOrder: number | undefined) => {
         if (this.controller) {
             this.execution = this.controller.createNotebookCellExecution(this.cell);
             this.isExecutionStarted = true;
diff --git a/vscode/src/notebooks/executionSummary.ts b/vscode/src/notebooks/executionSummary.ts
@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2024, Oracle and/or its affiliates.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export interface ExecutionSummaryData {
+  executionOrder?: number | null;
+  success?: boolean;
+}
+
+export class ExecutionSummary {
+  constructor(
+    public executionOrder: number | null = null,
+    public success: boolean = false
+  ) {}
+
+  static fromMetadata(
+    meta?: ExecutionSummaryData,
+    fallbackExecCount?: number | null
+  ): ExecutionSummary {
+    const order =
+      meta?.executionOrder != null
+        ? meta.executionOrder
+        : fallbackExecCount ?? null;
+    const success = meta?.success ?? false;
+    return new ExecutionSummary(order, success);
+  }
+
+  toJSON(): ExecutionSummaryData {
+    return {
+      executionOrder: this.executionOrder,
+      success: this.success,
+    };
+  }
+}
diff --git a/vscode/src/notebooks/kernel.ts b/vscode/src/notebooks/kernel.ts
@@ -22,15 +22,16 @@
 import { globalState } from '../globalState';
 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';
 import { LOGGER } from '../logger';
 import { CodeCellExecution } from './codeCellExecution';
-import { isError } from '../utils';
+import { isError, isString } from '../utils';
+import { MimeTypeHandler } from './mimeTypeHandler';
+import { createErrorOutput } from './utils';
 
 export class IJNBKernel implements Disposable {
   private readonly controllers: NotebookController[] = [];
@@ -86,6 +87,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.getIncrementedExecutionCounter(notebookId);
     try {
       this.cellControllerIdMap.set(cellId, codeCellExecution);
       const client: LanguageClient = await globalState.getClientPromise().client;
@@ -141,37 +143,43 @@ 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.getIncrementedExecutionCounter(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 getIncrementedExecutionCounter = (notebookId: string) => {
+    const next = (IJNBKernel.executionCounter.get(notebookId) ?? 0) + 1;
+    IJNBKernel.executionCounter.set(notebookId, next);
     return next;
   }
 
   private getExecutionCounter = (notebookId: string) => {
-    return IJNBKernel.executionCounter.get(notebookId) ?? 1;
+    return IJNBKernel.executionCounter.get(notebookId);
   }
 
   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.getIncrementedExecutionCounter(notebookId);
     try {
       exec.start(Date.now());
       await exec.replaceOutput([
-        new NotebookCellOutput([createOutputItem(cell.document.getText(), mimeType)]),
+        new NotebookCellOutput([new MimeTypeHandler(mimeType).makeOutputItem(cell.document.getText())]),
       ]);
+      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();
diff --git a/vscode/src/notebooks/mimeTypeHandler.ts b/vscode/src/notebooks/mimeTypeHandler.ts
@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2024, Oracle and/or its affiliates.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Buffer } from 'buffer';
+import * as vscode from 'vscode';
+import { IMimeBundle } from './types';
+import { mimeTypes } from './constants';
+
+export type DataOrBytes = string | Uint8Array;
+
+export class MimeTypeHandler {
+  constructor(public readonly value: string) {}
+  get isText(): boolean {
+    return this.value === mimeTypes.TEXT;
+  }
+  get isImage(): boolean {
+    return this.value.startsWith('image/');
+  }
+
+  static toBytes(data: DataOrBytes): Uint8Array {
+    if (typeof data === 'string') {
+      return Buffer.from(data, 'base64');
+    }
+    return data;
+  }
+
+  static toString(data: DataOrBytes): string {
+    if (typeof data === 'string') {
+      return data;
+    }
+    return new TextDecoder().decode(data);
+  }
+
+    makeOutputItem(data: DataOrBytes): vscode.NotebookCellOutputItem {
+    if (this.isImage) {
+      const bytes = MimeTypeHandler.toBytes(data);
+      return new vscode.NotebookCellOutputItem(bytes, this.value);
+    }
+    const text = MimeTypeHandler.toString(data);
+    return vscode.NotebookCellOutputItem.text(text, this.value);
+  }
+
+    static itemsFromBundle(bundle: IMimeBundle): vscode.NotebookCellOutputItem[] {
+    return Object.entries(bundle).flatMap(([mime, data]) => {
+      const mt = new MimeTypeHandler(mime);
+      if (mt.isText || mt.isImage) {
+        const payload = Array.isArray(data) ? data.join('') : data;
+        return [mt.makeOutputItem(payload)];
+      }
+      return [];
+    });
+  }
+}
diff --git a/vscode/src/notebooks/register.ts b/vscode/src/notebooks/register.ts
@@ -31,4 +31,5 @@ export const registerNotebookProviders = (context: ExtensionContext) => {
         ));
 
     context.subscriptions.push(notebookKernel);
+    context.subscriptions.push(notebookKernel.cleanUpKernel);
 }
diff --git a/vscode/src/notebooks/utils.ts b/vscode/src/notebooks/utils.ts
@@ -31,8 +31,11 @@ import {
   IMetadata
 } from './types';
 import { randomUUID } from 'crypto';
-import { isError, isString } from '../utils';
+import { isString } from '../utils';
 import { mimeTypes } from './constants';
+import { MimeTypeHandler } from './mimeTypeHandler';
+import { ExecutionSummary } from './executionSummary';
+
 
 export function base64ToUint8Array(base64: string): Uint8Array {
   if (typeof Buffer !== 'undefined' && typeof Buffer.from === 'function') {
@@ -56,15 +59,6 @@ export function uint8ArrayToBase64(data: Uint8Array): string {
   return btoa(binary);
 }
 
-export function createOutputItem(data: string | Uint8Array, mimeType: string): vscode.NotebookCellOutputItem {
-  if (mimeType.startsWith('image/')) {
-    const bytes = typeof data === 'string' ? base64ToUint8Array(data) : data;
-    return new vscode.NotebookCellOutputItem(bytes, mimeType);
-  }
-  const text = typeof data === 'string' ? data : new TextDecoder().decode(data);
-  return vscode.NotebookCellOutputItem.text(text, mimeType);
-}
-
 export const createErrorOutput = (err: string | Error) => {
   return new vscode.NotebookCellOutput([createErrorOutputItem(err)]);
 }
@@ -86,15 +80,16 @@ export function parseCell(cell: ICell): vscode.NotebookCellData {
   const cellData = new vscode.NotebookCellData(kind, value, language);
   cellData.metadata = { id: cell.id, ...cell.metadata };
   if (cell.cell_type === 'code') {
-
-    const metaExec = (cell.metadata as IMetadata).executionSummary;
-    const executionOrder = metaExec?.executionOrder ?? cell.execution_count ?? undefined;
-    const success = metaExec?.success ?? undefined;
-
-    cellData.executionSummary = {
-      executionOrder,
-      success,
-    };
+    const execSummary = ExecutionSummary.fromMetadata(
+      (cell.metadata as IMetadata).executionSummary,
+      cell.execution_count ?? null
+    );
+    if (execSummary.executionOrder) {
+      cellData.executionSummary = {
+        executionOrder: execSummary.executionOrder ?? undefined,
+        success: execSummary.success,
+      };
+    }
 
     if (Array.isArray(cell.outputs)) {
       const outputs = cell.outputs.flatMap((out: IOutput) => {
@@ -119,7 +114,7 @@ export function parseOutput(raw: IOutput): vscode.NotebookCellOutput[] {
     case 'stream':
       outputs.push(
         new vscode.NotebookCellOutput([
-          vscode.NotebookCellOutputItem.text(Array.isArray(raw.text) ? raw.text.join('') : raw.text),
+          new MimeTypeHandler(mimeTypes.TEXT).makeOutputItem(Array.isArray(raw.text) ? raw.text.join('') : raw.text),
         ])
       );
       break;
@@ -138,19 +133,8 @@ export function parseOutput(raw: IOutput): vscode.NotebookCellOutput[] {
 
     case 'display_data':
     case 'execute_result':
-      const items: vscode.NotebookCellOutputItem[] = [];
-      const bundle = raw.data || {};
-      for (const mime in bundle) {
-        const data = bundle[mime];
-        if (mime === mimeTypes.TEXT) {
-          const text = Array.isArray(data) ? data.join('') : String(data);
-          items.push(vscode.NotebookCellOutputItem.text(text));
-        } else if ((mime as string).startsWith('image/')) {
-          const b64 = Array.isArray(data) ? data.join('') : String(data);
-          const bytes = base64ToUint8Array(b64);
-          items.push(new vscode.NotebookCellOutputItem(bytes, mime));
-        }
-      }
+      const bundle = raw.data ?? {};
+      const items = MimeTypeHandler.itemsFromBundle(bundle);
       if (items.length) {
         outputs.push(new vscode.NotebookCellOutput(items, raw.metadata));
       }
@@ -166,9 +150,10 @@ export function serializeCell(cell: vscode.NotebookCellData): ICell {
   if (cell.kind === vscode.NotebookCellKind.Code) {
     const exec = cell.executionSummary ?? {};
     const executionCount = exec.executionOrder ?? null;
-    const success = exec.success ?? false;
+    const success = exec.success;
 
-    const metadata = { ...baseMeta, executionSummary: { executionOrder: executionCount, success } };
+    const execSummary = new ExecutionSummary(executionCount, success);
+    const metadata = executionCount ? { ...baseMeta, executionSummary: execSummary.toJSON() } : {};
 
     const outputs: IOutput[] = (cell.outputs || []).map((output): IOutput => {
       const data: IMimeBundle = {};
diff --git a/vscode/src/test/unit/mocks/vscode/mockVscode.ts b/vscode/src/test/unit/mocks/vscode/mockVscode.ts
@@ -19,13 +19,15 @@ import { URI } from './uri';
 import { mockWindowNamespace } from './namespaces/window';
 import { mockEnvNamespace } from './namespaces/env';
 import { mockedEnums } from './vscodeHostedTypes';
+import { NotebookCellOutputItem } from './notebookCellOutputItem';
 
 type VSCode = typeof vscode;
 const mockedVSCode: Partial<VSCode> = {};
 
 const mockedVscodeClassesAndTypes = () => {
     mockedVSCode.Uri = URI as any;
     mockedVSCode.ViewColumn = mockedEnums.viewColumn;
+    mockedVSCode.NotebookCellOutputItem = NotebookCellOutputItem as any;
 }
 
 const mockNamespaces = () => {
diff --git a/vscode/src/test/unit/mocks/vscode/notebookCellOutputItem.ts b/vscode/src/test/unit/mocks/vscode/notebookCellOutputItem.ts
diff --git a/vscode/src/test/unit/notebooks/executionSummary.unit.test.ts b/vscode/src/test/unit/notebooks/executionSummary.unit.test.ts
diff --git a/vscode/src/test/unit/notebooks/mimeTypeHandler.unit.test.ts b/vscode/src/test/unit/notebooks/mimeTypeHandler.unit.test.ts

Original file line number	Diff line number	Diff line change
`@@ -483,7 +483,7 @@`
`483`	`483`	`},`
`484`	`484`	`{`
`485`	`485`	`"command": "jdk.jshell.project",`
`486`		`- "title": "Open JShell in context of a Project",`
	`486`	`+ "title": "%jdk.jshell.project%",`
`487`	`487`	`"category": "Java"`
`488`	`488`	`},`
`489`	`489`	`{`
`@@ -591,7 +591,7 @@`
`591`	`591`	`},`
`592`	`592`	`{`
`593`	`593`	`"command": "jdk.notebook.new",`
`594`		`- "title": "Create New Notebook...",`
	`594`	`+ "title": "%jdk.notebook.new%",`
`595`	`595`	`"category": "Java",`
`596`	`596`	`"icon": "$(new-file)"`
`597`	`597`	`}`
Original file line number	Diff line number	Diff line change
`@@ -31,4 +31,5 @@ export const registerNotebookProviders = (context: ExtensionContext) => {`
`31`	`31`	`));`
`32`	`32`
`33`	`33`	`context.subscriptions.push(notebookKernel);`
	`34`	`+ context.subscriptions.push(notebookKernel.cleanUpKernel);`
`34`	`35`	`}`