Replace proposed VSC API with our implementation (#15573)

package.json

@@ -2068,7 +2068,6 @@
         "notebookDeprecated",
         "notebookMessaging",
         "notebookMime",
-        "notebookCellExecutionState",
         "contribNotebookStaticPreloads",
         "portsAttributes",
         "quickPickSortByLabel",

src/interactive-window/editor-integration/codeGenerator.ts

@@ -3,14 +3,12 @@
 
 import {
     Disposable,
-    NotebookCellExecutionStateChangeEvent,
     NotebookDocument,
     Position,
     Range,
     TextDocumentChangeEvent,
     TextDocumentContentChangeEvent,
     Uri,
-    notebooks,
     workspace
 } from 'vscode';
 
@@ -21,6 +19,11 @@ import { uncommentMagicCommands } from './cellFactory';
 import { CellMatcher } from './cellMatcher';
 import { IGeneratedCode, IInteractiveWindowCodeGenerator, IGeneratedCodeStore, InteractiveCellMetadata } from './types';
 import { computeHash } from '../../platform/common/crypto';
+import {
+    NotebookCellExecutionState,
+    notebookCellExecutions,
+    type NotebookCellExecutionStateChangeEvent
+} from '../../platform/notebooks/cellExecutionStateService';
 
 // This class provides generated code for debugging jupyter cells. Call getGeneratedCode just before starting debugging to compute all of the
 // generated codes for cells & update the source maps in the python debugger.
@@ -39,7 +42,7 @@ export class CodeGenerator implements IInteractiveWindowCodeGenerator {
         disposables.push(this);
         // Watch document changes so we can update our generated code
         workspace.onDidChangeTextDocument(this.onChangedDocument, this, this.disposables);
-        notebooks.onDidChangeNotebookCellExecutionState(this.onDidCellStateChange, this, this.disposables);
+        notebookCellExecutions.onDidChangeNotebookCellExecutionState(this.onDidCellStateChange, this, this.disposables);
     }
 
     public dispose() {
@@ -107,7 +110,11 @@ export class CodeGenerator implements IInteractiveWindowCodeGenerator {
     }
 
     private onDidCellStateChange(e: NotebookCellExecutionStateChangeEvent) {
+        if (e.state !== NotebookCellExecutionState.Idle) {
+            return;
+        }
         if (
+            e.state !== NotebookCellExecutionState.Idle ||
             e.cell.notebook !== this.notebook ||
             !e.cell.executionSummary?.executionOrder ||
             this.cellIndexesCounted[e.cell.index]

src/interactive-window/editor-integration/codeGenerator.unit.test.ts

@@ -2,15 +2,14 @@
 // Licensed under the MIT License.
 
 import { assert } from 'chai';
-import { EventEmitter, NotebookCellExecutionStateChangeEvent, NotebookDocument, Position, Range, Uri } from 'vscode';
+import { NotebookDocument, Position, Range, Uri } from 'vscode';
 
 import { IConfigurationService, IDisposable, IWatchableJupyterSettings } from '../../platform/common/types';
 import { CodeGenerator } from './codeGenerator';
 import { MockDocumentManager } from '../../test/datascience/mockDocumentManager';
 import { anything, instance, mock, when } from 'ts-mockito';
 import { IGeneratedCodeStore, InteractiveCellMetadata } from './types';
 import { GeneratedCodeStorage } from './generatedCodeStorage';
-import { mockedVSCodeNamespaces } from '../../test/vscode-mock';
 import { dispose } from '../../platform/common/utils/lifecycle';
 
 // eslint-disable-next-line
@@ -21,7 +20,6 @@ suite.skip('Code Generator Unit Tests', () => {
     let pythonSettings: IWatchableJupyterSettings;
     let storage: IGeneratedCodeStore;
     let notebook: NotebookDocument;
-    let onDidChangeNotebookCellExecutionState: EventEmitter<NotebookCellExecutionStateChangeEvent>;
     let disposables: IDisposable[] = [];
     setup(() => {
         configurationService = mock<IConfigurationService>();
@@ -30,11 +28,6 @@ suite.skip('Code Generator Unit Tests', () => {
         when(configurationService.getSettings(anything())).thenReturn(instance(pythonSettings));
         documentManager = new MockDocumentManager();
         notebook = mock<NotebookDocument>();
-        onDidChangeNotebookCellExecutionState = new EventEmitter<NotebookCellExecutionStateChangeEvent>();
-        disposables.push(onDidChangeNotebookCellExecutionState);
-        when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn(
-            onDidChangeNotebookCellExecutionState.event
-        );
         when(notebook.uri).thenReturn();
         codeGenerator = new CodeGenerator(instance(configurationService), storage, instance(notebook), []);
         disposables.push(codeGenerator);

src/interactive-window/editor-integration/codeLensFactory.ts

@@ -2,17 +2,7 @@
 // Licensed under the MIT License.
 
 import { inject, injectable } from 'inversify';
-import {
-    CodeLens,
-    Command,
-    Event,
-    EventEmitter,
-    NotebookCellExecutionState,
-    NotebookCellExecutionStateChangeEvent,
-    Range,
-    TextDocument,
-    workspace
-} from 'vscode';
+import { CodeLens, Command, Event, EventEmitter, Range, TextDocument, workspace } from 'vscode';
 
 import { traceWarning, traceInfoIfCI, traceVerbose } from '../../platform/logging';
 
@@ -24,7 +14,11 @@ import { CodeLensCommands, Commands, InteractiveWindowView } from '../../platfor
 import { generateCellRangesFromDocument } from './cellFactory';
 import { CodeLensPerfMeasures, ICodeLensFactory, IGeneratedCode, IGeneratedCodeStorageFactory } from './types';
 import { StopWatch } from '../../platform/common/utils/stopWatch';
-import { notebooks } from 'vscode';
+import {
+    NotebookCellExecutionState,
+    notebookCellExecutions,
+    type NotebookCellExecutionStateChangeEvent
+} from '../../platform/notebooks/cellExecutionStateService';
 
 type CodeLensCacheData = {
     cachedDocumentVersion: number | undefined;
@@ -76,7 +70,11 @@ export class CodeLensFactory implements ICodeLensFactory {
         workspace.onDidCloseTextDocument(this.onClosedDocument, this, disposables);
         workspace.onDidGrantWorkspaceTrust(() => this.codeLensCache.clear(), this, disposables);
         this.configService.getSettings(undefined).onDidChange(this.onChangedSettings, this, disposables);
-        notebooks.onDidChangeNotebookCellExecutionState(this.onDidChangeNotebookCellExecutionState, this, disposables);
+        notebookCellExecutions.onDidChangeNotebookCellExecutionState(
+            this.onDidChangeNotebookCellExecutionState,
+            this,
+            disposables
+        );
         kernelProvider.onDidDisposeKernel(
             (kernel) => {
                 this.notebookData.delete(kernel.notebook.uri.toString());
@@ -387,7 +385,6 @@ export class CodeLensFactory implements ICodeLensFactory {
                         range.start.character
                     ]);
                 }
-                break;
             case Commands.RunCellAndAllBelowPalette:
             case Commands.RunCellAndAllBelow:
                 return this.generateCodeLens(range, Commands.RunCellAndAllBelow, runAllBelowTitle, [

src/interactive-window/editor-integration/codewatcher.unit.test.ts

@@ -5,17 +5,7 @@
 // Disable whitespace / multiline as we use that to pass in our fake file strings
 import { expect } from 'chai';
 import * as TypeMoq from 'typemoq';
-import {
-    CancellationTokenSource,
-    CodeLens,
-    Disposable,
-    EventEmitter,
-    NotebookCellExecutionStateChangeEvent,
-    Range,
-    Selection,
-    TextEditor,
-    Uri
-} from 'vscode';
+import { CancellationTokenSource, CodeLens, Disposable, EventEmitter, Range, Selection, TextEditor, Uri } from 'vscode';
 
 import { anything, instance, mock, when } from 'ts-mockito';
 import { IDebugService } from '../../platform/common/application/types';
@@ -121,16 +111,11 @@ suite('Code Watcher Unit Tests', () => {
         when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true);
         const trustedEvent = new EventEmitter<void>();
         when(mockedVSCodeNamespaces.workspace.onDidGrantWorkspaceTrust).thenReturn(trustedEvent.event);
-        const execStateChangeEvent = new EventEmitter<NotebookCellExecutionStateChangeEvent>();
-        when(mockedVSCodeNamespaces.notebooks.onDidChangeNotebookCellExecutionState).thenReturn(
-            execStateChangeEvent.event
-        );
         const storageFactory = mock<IGeneratedCodeStorageFactory>();
         const kernelProvider = mock<IKernelProvider>();
         const kernelDisposedEvent = new EventEmitter<IKernel>();
         when(kernelProvider.onDidDisposeKernel).thenReturn(kernelDisposedEvent.event);
         disposables.push(trustedEvent);
-        disposables.push(execStateChangeEvent);
         disposables.push(kernelDisposedEvent);
         const codeLensFactory = new CodeLensFactory(
             configService.object,

src/interactive-window/editor-integration/hoverProvider.ts

@@ -16,6 +16,10 @@ import { IKernel, IKernelProvider } from '../../kernels/types';
 import { IJupyterVariables } from '../../kernels/variables/types';
 import { IInteractiveWindowProvider } from '../types';
 import { getInteractiveCellMetadata } from '../helpers';
+import {
+    notebookCellExecutions,
+    type NotebookCellExecutionStateChangeEvent
+} from '../../platform/notebooks/cellExecutionStateService';
 
 /**
  * Provides hover support in python files based on the state of a jupyter kernel. Files that are
@@ -35,10 +39,11 @@ export class HoverProvider implements IExtensionSyncActivationService, vscode.Ho
         @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider
     ) {}
     public activate() {
-        this.onDidChangeNotebookCellExecutionStateHandler = vscode.notebooks.onDidChangeNotebookCellExecutionState(
-            (e) => this.delayer.trigger(() => this.onDidChangeNotebookCellExecutionState(e)),
-            this
-        );
+        this.onDidChangeNotebookCellExecutionStateHandler =
+            notebookCellExecutions.onDidChangeNotebookCellExecutionState(
+                (e) => this.delayer.trigger(() => this.onDidChangeNotebookCellExecutionState(e)),
+                this
+            );
         this.kernelProvider.onDidRestartKernel(() => this.runFiles.clear(), this, this.disposables);
     }
     public dispose() {
@@ -48,9 +53,7 @@ export class HoverProvider implements IExtensionSyncActivationService, vscode.Ho
             this.hoverProviderRegistration.dispose();
         }
     }
-    private async onDidChangeNotebookCellExecutionState(
-        e: vscode.NotebookCellExecutionStateChangeEvent
-    ): Promise<void> {
+    private async onDidChangeNotebookCellExecutionState(e: NotebookCellExecutionStateChangeEvent): Promise<void> {
         try {
             if (e.cell.notebook.notebookType !== InteractiveWindowView) {
                 return;

src/kernels/execution/cellExecution.ts

@@ -2,15 +2,7 @@
 // Licensed under the MIT License.
 
 import type * as KernelMessage from '@jupyterlab/services/lib/kernel/messages';
-import {
-    NotebookCell,
-    NotebookCellExecution,
-    workspace,
-    NotebookCellOutput,
-    NotebookCellExecutionState,
-    Event,
-    EventEmitter
-} from 'vscode';
+import { NotebookCell, NotebookCellExecution, workspace, NotebookCellOutput } from 'vscode';
 
 import type { Kernel } from '@jupyterlab/services';
 import { CellExecutionCreator } from './cellExecutionCreator';
@@ -39,6 +31,7 @@ import { ICellExecution } from './types';
 import { KernelError } from '../errors/kernelError';
 import { getCachedSysPrefix } from '../../platform/interpreter/helpers';
 import { getCellMetadata } from '../../platform/common/utils';
+import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService';
 
 /**
  * Factory for CellExecution objects.
@@ -74,9 +67,6 @@ export class CellExecution implements ICellExecution, IDisposable {
     public get result(): Promise<void> {
         return this._result.promise;
     }
-    public get preExecute(): Event<NotebookCell> {
-        return this._preExecuteEmitter.event;
-    }
     private readonly _result = createDeferred<void>();
 
     private started?: boolean;
@@ -94,10 +84,13 @@ 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 cellExecutionHandler?: CellExecutionMessageHandler;
     private session?: IKernelSession;
     private cancelRequested?: boolean;
+    private _executionOrder?: number;
+    public get executionOrder() {
+        return this._executionOrder;
+    }
     private constructor(
         public readonly cell: NotebookCell,
         private readonly codeOverride: string | undefined,
@@ -129,7 +122,6 @@ export class CellExecution implements ICellExecution, IDisposable {
             this,
             this.disposables
         );
-        NotebookCellStateTracker.setCellState(cell, NotebookCellExecutionState.Idle);
         if (this.canExecuteCell()) {
             this.execution = CellExecutionCreator.getOrCreate(
                 cell,
@@ -378,6 +370,7 @@ export class CellExecution implements ICellExecution, IDisposable {
         if (activeNotebookCellExecution.get(this.cell.notebook) === this.execution) {
             activeNotebookCellExecution.set(this.cell.notebook, undefined);
         }
+        this._executionOrder = this.execution?.executionOrder;
         NotebookCellStateTracker.setCellState(this.cell, NotebookCellExecutionState.Idle);
         this.execution = undefined;
     }
@@ -423,7 +416,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);
+            notebookCellExecutions.changeCellState(this.cell, NotebookCellExecutionState.Executing);
             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/cellExecutionCreator.ts

@@ -21,7 +21,10 @@ import { getNotebookTelemetryTracker } from '../telemetry/notebookTelemetry';
  * - Do something when 'end' is called
  */
 export class NotebookCellExecutionWrapper implements NotebookCellExecution {
-    public started: boolean = false;
+    public _started: boolean = false;
+    public get started() {
+        return this._started;
+    }
     private _startTime?: number;
     public errorInfo: CellExecutionError;
     /**
@@ -56,7 +59,7 @@ export class NotebookCellExecutionWrapper implements NotebookCellExecution {
     start(startTime?: number): void {
         // Allow this to be called more than once (so we can switch out a kernel during running a cell)
         if (!this.started) {
-            this.started = true;
+            this._started = true;
             this._impl.start(startTime);
             this._startTime = startTime;
             // We clear the output as soon as we start,

src/kernels/execution/cellExecutionMessageHandler.ts

@@ -449,8 +449,8 @@ export class CellExecutionMessageHandler implements IDisposable {
         this.previousResultsToRestore = { ...(this.cell.executionSummary || {}) };
         this.temporaryExecution = CellExecutionCreator.getOrCreate(this.cell, this.controller);
         this.temporaryExecution?.start();
-        if (this.previousResultsToRestore?.executionOrder && this.execution) {
-            this.execution.executionOrder = this.previousResultsToRestore.executionOrder;
+        if (this.previousResultsToRestore?.executionOrder && this.temporaryExecution) {
+            this.temporaryExecution.executionOrder = this.previousResultsToRestore.executionOrder;
         }
         return this.temporaryExecution;
     }

src/kernels/execution/cellExecutionQueue.ts

@@ -12,6 +12,7 @@ import { ICellExecution, ICodeExecution } from './types';
 import { CodeExecution } from './codeExecution';
 import { once } from '../../platform/common/utils/events';
 import { getCellMetadata } from '../../platform/common/utils';
+import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService';
 
 /**
  * A queue responsible for execution of cells.
@@ -25,8 +26,6 @@ export class CellExecutionQueue implements Disposable {
     }
     private cancelledOrCompletedWithErrors = false;
     private startedRunningCells = false;
-    private readonly _onPreExecute = new EventEmitter<NotebookCell>();
-    private readonly _onPostExecute = new EventEmitter<NotebookCell>();
     private disposables: Disposable[] = [];
     private lastCellExecution?: ICellExecution | ICodeExecution;
     /**
@@ -57,14 +56,6 @@ export class CellExecutionQueue implements Disposable {
         this.lastCellExecution?.dispose();
     }
 
-    public get onPreExecute() {
-        return this._onPreExecute.event;
-    }
-
-    public get onPostExecute() {
-        return this._onPostExecute.event;
-    }
-
     /**
      * Queue the cell for execution & start processing it immediately.
      */
@@ -101,7 +92,6 @@ 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);
             this.queueOfItemsToExecute.push(cellExecution);
 
             traceCellMessage(cell, 'User queued cell for execution');
@@ -245,7 +235,11 @@ export class CellExecutionQueue implements Disposable {
                 }
 
                 if (itemToExecute.type === 'cell') {
-                    this._onPostExecute.fire(itemToExecute.cell);
+                    notebookCellExecutions.changeCellState(
+                        itemToExecute.cell,
+                        NotebookCellExecutionState.Idle,
+                        itemToExecute.executionOrder
+                    );
                 }
             }