VS Code API replacement for cell execution events

src/kernels/execution/cellExecution.ts

@@ -74,8 +74,11 @@ export class CellExecution implements ICellExecution, IDisposable {
     public get result(): Promise<void> {
         return this._result.promise;
     }
-    public get preExecute(): Event<NotebookCell> {
-        return this._preExecuteEmitter.event;
+    public get onWillExecute(): Event<void> {
+        return this._onWillExecute.event;
+    }
+    public get onDidExecute(): Event<void> {
+        return this._onDidExecute.event;
     }
     private readonly _result = createDeferred<void>();
 
@@ -94,7 +97,8 @@ export class CellExecution implements ICellExecution, IDisposable {
     private disposed?: boolean;
     private request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg> | undefined;
     private readonly disposables: IDisposable[] = [];
-    private _preExecuteEmitter = new EventEmitter<NotebookCell>();
+    private _onWillExecute = new EventEmitter<void>();
+    private _onDidExecute = new EventEmitter<void>();
     private cellExecutionHandler?: CellExecutionMessageHandler;
     private session?: IKernelSession;
     private cancelRequested?: boolean;
@@ -159,6 +163,7 @@ export class CellExecution implements ICellExecution, IDisposable {
     }
     public async start(session: IKernelSession) {
         this.session = session;
+        void this._result.promise.finally(() => this._onDidExecute.fire());
         if (this.resumeExecution?.msg_id) {
             return this.resume(session, this.resumeExecution);
         }
@@ -423,7 +428,7 @@ export class CellExecution implements ICellExecution, IDisposable {
         const kernelConnection = session.kernel;
         try {
             // At this point we're about to ACTUALLY execute some code. Fire an event to indicate that
-            this._preExecuteEmitter.fire(this.cell);
+            this._onWillExecute.fire();
             traceVerbose(`Cell Index:${this.cell.index} sent to kernel`);
             // For Jupyter requests, silent === don't output, while store_history === don't update execution count
             // https://jupyter-client.readthedocs.io/en/stable/api/client.html#jupyter_client.KernelClient.execute

src/kernels/execution/cellExecutionQueue.ts

@@ -101,7 +101,8 @@ export class CellExecutionQueue implements Disposable {
             const cellExecution = this.executionFactory.create(cell, codeOverride, this.metadata);
             executionItem = cellExecution;
             this.disposables.push(cellExecution);
-            cellExecution.preExecute((c) => this._onPreExecute.fire(c), this, this.disposables);
+            cellExecution.onWillExecute(() => this._onPreExecute.fire(cellExecution.cell), this, this.disposables);
+            cellExecution.onDidExecute(() => this._onPostExecute.fire(cellExecution.cell), this, this.disposables);
             this.queueOfItemsToExecute.push(cellExecution);
 
             traceCellMessage(cell, 'User queued cell for execution');
@@ -243,10 +244,6 @@ export class CellExecutionQueue implements Disposable {
                 if (index >= 0) {
                     this.queueOfItemsToExecute.splice(index, 1);
                 }
-
-                if (itemToExecute.type === 'cell') {
-                    this._onPostExecute.fire(itemToExecute.cell);
-                }
             }
 
             let cellErrorsAllowed = false;

src/notebooks/controllers/controllerRegistration.ts

@@ -25,6 +25,7 @@ import { PythonEnvironmentFilter } from '../../platform/interpreter/filter/filte
 import {
     IConnectionDisplayDataProvider,
     IControllerRegistration,
+    INotebookCellExecutionStateService,
     InteractiveControllerIdSuffix,
     IVSCodeNotebookController,
     IVSCodeNotebookControllerUpdateEvent
@@ -360,7 +361,10 @@ export class ControllerRegistration implements IControllerRegistration, IExtensi
                         this.extensionChecker,
                         this.serviceContainer,
                         this.serviceContainer.get<IConnectionDisplayDataProvider>(IConnectionDisplayDataProvider),
-                        this.serviceContainer.get<IJupyterVariables>(IJupyterVariables, Identifiers.KERNEL_VARIABLES)
+                        this.serviceContainer.get<IJupyterVariables>(IJupyterVariables, Identifiers.KERNEL_VARIABLES),
+                        this.serviceContainer.get<INotebookCellExecutionStateService>(
+                            INotebookCellExecutionStateService
+                        )
                     );
                     // Hook up to if this NotebookController is selected or de-selected
                     const controllerDisposables: IDisposable[] = [];

src/notebooks/controllers/notebookCellExecutionStateService.ts

@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { EventEmitter, type Event, type NotebookCell } from 'vscode';
+import { DisposableBase, DisposableStore } from '../../platform/common/utils/lifecycle';
+import { inject, injectable } from 'inversify';
+import { IKernelProvider, type IKernel, type INotebookKernelExecution } from '../../kernels/types';
+import {
+    NotebookCellExecutionState,
+    type INotebookCellExecutionStateService,
+    type NotebookCellExecutionStateChangeEvent
+} from './types';
+import { IDisposableRegistry } from '../../platform/common/types';
+
+@injectable()
+export class NotebookCellExecutionStateService extends DisposableBase implements INotebookCellExecutionStateService {
+    /**
+     * An {@link Event} which fires when the execution state of a cell has changed.
+     */
+    // todo@API this is an event that is fired for a property that cells don't have and that makes me wonder
+    // how a correct consumer works, e.g the consumer could have been late and missed an event?
+    private readonly _onDidChangeNotebookCellExecutionState = this._register(
+        new EventEmitter<NotebookCellExecutionStateChangeEvent>()
+    );
+
+    public get onDidChangeNotebookCellExecutionState() {
+        return this._onDidChangeNotebookCellExecutionState.event;
+    }
+    private readonly trackedExecutions = new WeakSet<INotebookKernelExecution>();
+    private readonly kernelExecutionMaps = new WeakMap<IKernel, INotebookKernelExecution>();
+    private readonly kernelDisposables = new WeakMap<IKernel, DisposableStore>();
+    constructor(
+        @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry
+    ) {
+        super();
+        disposables.push(this);
+        this._register(kernelProvider.onDidCreateKernel(this.monitorKernelExecutionEvents, this));
+        this._register(kernelProvider.onDidStartKernel(this.monitorKernelExecutionEvents, this));
+        this._register(kernelProvider.onDidDisposeKernel((k) => this.kernelDisposables.get(k)?.dispose(), this));
+    }
+    setPendingState(cell: NotebookCell): void {
+        this.triggerStateChange(cell, NotebookCellExecutionState.Pending);
+    }
+
+    private monitorKernelExecutionEvents(kernel: IKernel) {
+        const execution = this.kernelProvider.getKernelExecution(kernel);
+        if (this.trackedExecutions.has(execution)) {
+            return;
+        }
+        const disposableStore = this.kernelDisposables.get(kernel) || new DisposableStore();
+        this.kernelDisposables.set(kernel, disposableStore);
+        this._register(disposableStore);
+
+        this.kernelExecutionMaps.set(kernel, execution);
+        disposableStore.add(
+            execution.onPreExecute((cell) => this.triggerStateChange(cell, NotebookCellExecutionState.Executing), this)
+        );
+        disposableStore.add(
+            execution.onPostExecute((cell) => this.triggerStateChange(cell, NotebookCellExecutionState.Idle), this)
+        );
+    }
+
+    private triggerStateChange(cell: NotebookCell, state: NotebookCellExecutionState) {
+        this._onDidChangeNotebookCellExecutionState.fire({ cell, state: state });
+    }
+}

src/notebooks/controllers/serviceRegistry.node.ts

@@ -10,11 +10,13 @@ import { KernelSourceCommandHandler } from './kernelSource/kernelSourceCommandHa
 import { LocalNotebookKernelSourceSelector } from './kernelSource/localNotebookKernelSourceSelector.node';
 import { LocalPythonEnvNotebookKernelSourceSelector } from './kernelSource/localPythonEnvKernelSourceSelector.node';
 import { RemoteNotebookKernelSourceSelector } from './kernelSource/remoteNotebookKernelSourceSelector';
+import { NotebookCellExecutionStateService } from './notebookCellExecutionStateService';
 import {
     IConnectionDisplayDataProvider,
     IControllerRegistration,
     ILocalNotebookKernelSourceSelector,
     ILocalPythonNotebookKernelSourceSelector,
+    INotebookCellExecutionStateService,
     IRemoteNotebookKernelSourceSelector
 } from './types';
 
@@ -37,6 +39,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
         ILocalPythonNotebookKernelSourceSelector,
         LocalPythonEnvNotebookKernelSourceSelector
     );
+    serviceManager.addSingleton<INotebookCellExecutionStateService>(
+        INotebookCellExecutionStateService,
+        NotebookCellExecutionStateService
+    );
     serviceManager.addBinding(ILocalPythonNotebookKernelSourceSelector, IExtensionSyncActivationService);
     serviceManager.addSingleton<IExtensionSyncActivationService>(
         IExtensionSyncActivationService,

src/notebooks/controllers/types.ts

@@ -141,3 +141,42 @@ export const IConnectionDisplayDataProvider = Symbol('IConnectionDisplayData');
 export interface IConnectionDisplayDataProvider {
     getDisplayData(connection: KernelConnectionMetadata): IConnectionDisplayData;
 }
+
+/**
+ * The execution state of a notebook cell.
+ */
+export enum NotebookCellExecutionState {
+    /**
+     * The cell is idle.
+     */
+    Idle = 1,
+    /**
+     * Execution for the cell is pending.
+     */
+    Pending = 2,
+    /**
+     * The cell is currently executing.
+     */
+    Executing = 3
+}
+
+/**
+ * An event describing a cell execution state change.
+ */
+export interface NotebookCellExecutionStateChangeEvent {
+    /**
+     * The {@link NotebookCell cell} for which the execution state has changed.
+     */
+    readonly cell: vscode.NotebookCell;
+
+    /**
+     * The new execution state of the cell.
+     */
+    readonly state: NotebookCellExecutionState;
+}
+
+export const INotebookCellExecutionStateService = Symbol('INotebookCellExecutionStateService');
+export interface INotebookCellExecutionStateService {
+    onDidChangeNotebookCellExecutionState: vscode.Event<NotebookCellExecutionStateChangeEvent>;
+    setPendingState(cell: vscode.NotebookCell): void;
+}

src/notebooks/controllers/vscodeNotebookController.ts

@@ -59,7 +59,12 @@ import { DisplayOptions } from '../../kernels/displayOptions';
 import { getNotebookMetadata, isJupyterNotebook, updateNotebookMetadata } from '../../platform/common/utils';
 import { ConsoleForegroundColors } from '../../platform/logging/types';
 import { KernelConnector } from './kernelConnector';
-import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types';
+import {
+    IConnectionDisplayData,
+    IConnectionDisplayDataProvider,
+    IVSCodeNotebookController,
+    type INotebookCellExecutionStateService
+} from './types';
 import { isCancellationError } from '../../platform/common/cancellation';
 import { CellExecutionCreator } from '../../kernels/execution/cellExecutionCreator';
 import {
@@ -159,7 +164,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
         extensionChecker: IPythonExtensionChecker,
         serviceContainer: IServiceContainer,
         displayDataProvider: IConnectionDisplayDataProvider,
-        jupyterVariables: IJupyterVariables
+        jupyterVariables: IJupyterVariables,
+        notebookCellExecutionStateService: INotebookCellExecutionStateService
     ): IVSCodeNotebookController {
         return new VSCodeNotebookController(
             kernelConnection,
@@ -173,7 +179,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
             extensionChecker,
             serviceContainer,
             displayDataProvider,
-            jupyterVariables
+            jupyterVariables,
+            notebookCellExecutionStateService
         );
     }
     constructor(
@@ -188,7 +195,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
         private readonly extensionChecker: IPythonExtensionChecker,
         private serviceContainer: IServiceContainer,
         private readonly displayDataProvider: IConnectionDisplayDataProvider,
-        jupyterVariables: IJupyterVariables
+        jupyterVariables: IJupyterVariables,
+        private readonly notebookCellExecutionStateService: INotebookCellExecutionStateService
     ) {
         trackControllerCreation(kernelConnection.id, kernelConnection.interpreter?.id);
         disposableRegistry.push(this);
@@ -392,6 +400,7 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont
             return;
         }
         traceInfo(`Handle Execution of Cells ${cells.map((c) => c.index)} for ${getDisplayPath(notebook.uri)}`);
+        cells.forEach((cell) => this.notebookCellExecutionStateService.setPendingState(cell));
         await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(notebook.uri, this.connection);
         telemetryTracker?.stop();
         // Notebook is trusted. Continue to execute cells

src/notebooks/controllers/vscodeNotebookController.unit.test.ts

@@ -43,6 +43,7 @@ import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock
 import { IJupyterVariables } from '../../kernels/variables/types';
 import { Environment, PythonExtension } from '@vscode/python-extension';
 import { crateMockedPythonApi, whenResolveEnvironment } from '../../kernels/helpers.unit.test';
+import { NotebookCellExecutionStateService } from './notebookCellExecutionStateService';
 
 suite(`Notebook Controller`, function () {
     let controller: NotebookController;
@@ -144,6 +145,10 @@ suite(`Notebook Controller`, function () {
     });
     teardown(() => (disposables = dispose(disposables)));
     function createController(viewType: 'jupyter-notebook' | 'interactive') {
+        const notebookCellExecutionStateService = new NotebookCellExecutionStateService(
+            instance(kernelProvider),
+            disposables
+        );
         new VSCodeNotebookController(
             instance(kernelConnection),
             '1',
@@ -156,7 +161,8 @@ suite(`Notebook Controller`, function () {
             instance(extensionChecker),
             instance(serviceContainer),
             displayDataProvider,
-            jupyterVariables
+            jupyterVariables,
+            notebookCellExecutionStateService
         );
         notebook = new TestNotebookDocument(undefined, viewType);
     }

src/test/datascience/notebook/kernelCrashes.vscode.test.ts

@@ -36,6 +36,7 @@ import { IPlatformService } from '../../../platform/common/platform/types';
 import { IInterpreterService } from '../../../platform/interpreter/contracts';
 import { ConnectionDisplayDataProvider } from '../../../notebooks/controllers/connectionDisplayData.node';
 import { IJupyterVariables } from '../../../kernels/variables/types';
+import { INotebookCellExecutionStateService } from '../../../notebooks/controllers/types';
 
 const codeToKillKernel = dedent`
 import IPython
@@ -125,7 +126,8 @@ suite('VSCode Notebook Kernel Error Handling - @kernelCore', function () {
                 extensionChecker,
                 api.serviceContainer,
                 displayDataProvider,
-                jupyterVariables
+                jupyterVariables,
+                api.serviceContainer.get<INotebookCellExecutionStateService>(INotebookCellExecutionStateService)
             );
             disposables.push(interpreterController);