Add public API event for kernel post-initialization (#16214)

* event validated * update * update * remove unused event * PR * WIP * test WIP * test WIP * PR * PR * PR * PR * fix tests * proposed fix * fix tests * update test * clean
microsoft · Nov 19, 2024 · dbd78ca · dbd78ca
1 parent 5a3dac9
commit dbd78ca
Show file tree

Hide file tree

Showing 12 changed files with 293 additions and 24 deletions.
diff --git a/src/api.proposed.kernelStartHook.d.ts b/src/api.proposed.kernelStartHook.d.ts
@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import type { Event, Uri } from 'vscode';
+
+declare module './api' {
+    export interface Kernels {
+        /**
+         * Event fired when a kernel is started or restarted by a user on a resource.
+         */
+        onDidStart: Event<{
+            uri: Uri;
+            kernel: Kernel;
+        }>;
+    }
+}
diff --git a/src/extension.node.ts b/src/extension.node.ts
@@ -124,7 +124,10 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             createJupyterServerCollection: () => {
                 throw new Error('Not Implemented');
             },
-            kernels: { getKernel: () => Promise.resolve(undefined) }
+            kernels: {
+                getKernel: () => Promise.resolve(undefined),
+                onDidStart: () => ({ dispose: noop })
+            }
         };
     }
 }

diff --git a/src/extension.web.ts b/src/extension.web.ts
@@ -117,7 +117,10 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
             createJupyterServerCollection: () => {
                 throw new Error('Not Implemented');
             },
-            kernels: { getKernel: () => Promise.resolve(undefined) }
+            kernels: {
+                getKernel: () => Promise.resolve(undefined),
+                onDidStart: () => ({ dispose: noop })
+            }
         };
     }
 }

diff --git a/src/kernels/kernel.ts b/src/kernels/kernel.ts
@@ -129,6 +129,9 @@ abstract class BaseKernel implements IBaseKernel {
     get onStarted(): Event<void> {
         return this._onStarted.event;
     }
+    get onPostInitialized(): Event<void> {
+        return this._onPostInitialized.event;
+    }
     get onDisposed(): Event<void> {
         return this._onDisposed.event;
     }
@@ -138,6 +141,9 @@ abstract class BaseKernel implements IBaseKernel {
     get startedAtLeastOnce() {
         return this._startedAtLeastOnce;
     }
+    get userStartedKernel() {
+        return !this.startupUI.disableUI;
+    }
     private _info?: KernelMessage.IInfoReplyMsg['content'];
     private _startedAtLeastOnce?: boolean;
     get info(): KernelMessage.IInfoReplyMsg['content'] | undefined {
@@ -172,10 +178,12 @@ abstract class BaseKernel implements IBaseKernel {
     private _disposed?: boolean;
     private _disposing?: boolean;
     private _ignoreJupyterSessionDisposedErrors?: boolean;
+    private _postInitializedOnStart?: boolean;
     private readonly _onDidKernelSocketChange = new EventEmitter<void>();
     private readonly _onStatusChanged = new EventEmitter<KernelMessage.Status>();
     private readonly _onRestarted = new EventEmitter<void>();
     private readonly _onStarted = new EventEmitter<void>();
+    private readonly _onPostInitialized = new EventEmitter<void>();
     private readonly _onDisposed = new EventEmitter<void>();
     private _jupyterSessionPromise?: Promise<IKernelSession>;
     private readonly hookedSessionForEvents = new WeakSet<IKernelSession>();
@@ -246,7 +254,16 @@ abstract class BaseKernel implements IBaseKernel {
             this.startCancellation.dispose();
             this.startCancellation = new CancellationTokenSource();
         }
-        return this.startJupyterSession(options);
+        return this.startJupyterSession(options).then((result) => {
+            // If we started and the UI is no longer disabled (ie., a user executed a cell)
+            // then we can signal that the kernel was created and can be used by third-party extensions.
+            // We also only want to fire off a single event here.
+            if (!options?.disableUI && !this._postInitializedOnStart) {
+                this._onPostInitialized.fire();
+                this._postInitializedOnStart = true;
+            }
+            return result;
+        });
     }
     /**
      * Interrupts the execution of cells.
@@ -409,6 +426,9 @@ abstract class BaseKernel implements IBaseKernel {
 
             // Indicate a restart occurred if it succeeds
             this._onRestarted.fire();
+
+            // Also signal that the kernel post initialization completed.
+            this._onPostInitialized.fire();
         } catch (ex) {
             logger.error(`Failed to restart kernel ${getDisplayPath(this.uri)}`, ex);
             throw ex;

diff --git a/src/kernels/kernelProvider.base.ts b/src/kernels/kernelProvider.base.ts
@@ -36,6 +36,7 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider {
     protected readonly _onDidCreateKernel = new EventEmitter<IKernel>();
     protected readonly _onDidDisposeKernel = new EventEmitter<IKernel>();
     protected readonly _onKernelStatusChanged = new EventEmitter<{ status: KernelMessage.Status; kernel: IKernel }>();
+    protected readonly _onDidPostInitializeKernel = new EventEmitter<IKernel>();
     public readonly onKernelStatusChanged = this._onKernelStatusChanged.event;
     public get kernels() {
         const kernels = new Set<IKernel>();
@@ -58,6 +59,7 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider {
         disposables.push(this._onKernelStatusChanged);
         disposables.push(this._onDidStartKernel);
         disposables.push(this._onDidCreateKernel);
+        disposables.push(this._onDidPostInitializeKernel);
     }
 
     public get onDidDisposeKernel(): Event<IKernel> {
@@ -73,6 +75,9 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider {
     public get onDidCreateKernel(): Event<IKernel> {
         return this._onDidCreateKernel.event;
     }
+    public get onDidPostInitializeKernel(): Event<IKernel> {
+        return this._onDidPostInitializeKernel.event;
+    }
     public get(uriOrNotebook: Uri | NotebookDocument | string): IKernel | undefined {
         if (isUri(uriOrNotebook)) {
             const notebook = workspace.notebookDocuments.find(
@@ -187,6 +192,7 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP
     protected readonly _onDidStartKernel = new EventEmitter<IThirdPartyKernel>();
     protected readonly _onDidCreateKernel = new EventEmitter<IThirdPartyKernel>();
     protected readonly _onDidDisposeKernel = new EventEmitter<IThirdPartyKernel>();
+    protected readonly _onDidPostInitializeKernel = new EventEmitter<IThirdPartyKernel>();
     protected readonly _onKernelStatusChanged = new EventEmitter<{
         status: KernelMessage.Status;
         kernel: IThirdPartyKernel;
@@ -213,6 +219,7 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP
         disposables.push(this._onKernelStatusChanged);
         disposables.push(this._onDidStartKernel);
         disposables.push(this._onDidCreateKernel);
+        disposables.push(this._onDidPostInitializeKernel);
     }
 
     public get onDidDisposeKernel(): Event<IThirdPartyKernel> {
@@ -227,6 +234,9 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP
     public get onDidCreateKernel(): Event<IThirdPartyKernel> {
         return this._onDidCreateKernel.event;
     }
+    public get onDidPostInitializeKernel(): Event<IThirdPartyKernel> {
+        return this._onDidPostInitializeKernel.event;
+    }
     public get(uri: Uri | string): IThirdPartyKernel | undefined {
         return this.kernelsByUri.get(uri.toString())?.kernel || this.kernelsById.get(uri.toString())?.kernel;
     }

diff --git a/src/kernels/kernelProvider.node.ts b/src/kernels/kernelProvider.node.ts
@@ -97,6 +97,13 @@ export class KernelProvider extends BaseCoreKernelProvider {
             this,
             this.disposables
         );
+        kernel.onPostInitialized(
+            () => {
+                this._onDidPostInitializeKernel.fire(kernel);
+            },
+            this,
+            this.disposables
+        );
 
         this.executions.set(kernel, new NotebookKernelExecution(kernel, this.context, this.formatters, notebook));
         this.asyncDisposables.push(kernel);
@@ -152,6 +159,13 @@ export class ThirdPartyKernelProvider extends BaseThirdPartyKernelProvider {
             this,
             this.disposables
         );
+        kernel.onPostInitialized(
+            () => {
+                this._onDidPostInitializeKernel.fire(kernel);
+            },
+            this,
+            this.disposables
+        );
         this.asyncDisposables.push(kernel);
         this.storeKernel(uri, options, kernel);
         this.deleteMappingIfKernelIsDisposed(uri, kernel);

diff --git a/src/kernels/kernelProvider.node.unit.test.ts b/src/kernels/kernelProvider.node.unit.test.ts
@@ -101,20 +101,24 @@ suite('Jupyter Session', () => {
             const kernelStarted = createEventHandler(kernelProvider, 'onDidStartKernel', disposables);
             const kernelDisposed = createEventHandler(kernelProvider, 'onDidDisposeKernel', disposables);
             const kernelRestarted = createEventHandler(kernelProvider, 'onDidRestartKernel', disposables);
+            const kernelPostInitialized = createEventHandler(kernelProvider, 'onDidPostInitializeKernel', disposables);
             const kernelStatusChanged = createEventHandler(kernelProvider, 'onKernelStatusChanged', disposables);
             const notebook = new TestNotebookDocument(undefined, 'jupyter-notebook');
             const onStarted = new EventEmitter<void>();
             const onStatusChanged = new EventEmitter<void>();
             const onRestartedEvent = new EventEmitter<void>();
+            const onPostInitializedEvent = new EventEmitter<void>();
             const onDisposedEvent = new EventEmitter<void>();
             disposables.push(onStatusChanged);
             disposables.push(onRestartedEvent);
+            disposables.push(onPostInitializedEvent);
             disposables.push(onStarted);
             disposables.push(onDisposedEvent);
             if (kernelProvider instanceof KernelProvider) {
                 sinon.stub(Kernel.prototype, 'onStarted').get(() => onStarted.event);
                 sinon.stub(Kernel.prototype, 'onStatusChanged').get(() => onStatusChanged.event);
                 sinon.stub(Kernel.prototype, 'onRestarted').get(() => onRestartedEvent.event);
+                sinon.stub(Kernel.prototype, 'onPostInitialized').get(() => onPostInitializedEvent.event);
                 sinon.stub(Kernel.prototype, 'onDisposed').get(() => onDisposedEvent.event);
                 const kernel = kernelProvider.getOrCreate(notebook, {
                     controller,
@@ -126,6 +130,7 @@ suite('Jupyter Session', () => {
                 sinon.stub(ThirdPartyKernel.prototype, 'onStarted').get(() => onStarted.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onStatusChanged').get(() => onStatusChanged.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onRestarted').get(() => onRestartedEvent.event);
+                sinon.stub(ThirdPartyKernel.prototype, 'onPostInitialized').get(() => onPostInitializedEvent.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onDisposed').get(() => onDisposedEvent.event);
                 const kernel = kernelProvider.getOrCreate(notebook.uri, {
                     metadata: instance(metadata),
@@ -138,6 +143,10 @@ suite('Jupyter Session', () => {
             assert.isFalse(kernelStarted.fired, 'IKernelProvider.onDidStartKernel should not be fired');
             assert.isFalse(kernelStatusChanged.fired, 'IKernelProvider.onKernelStatusChanged should not be fired');
             assert.isFalse(kernelRestarted.fired, 'IKernelProvider.onDidRestartKernel should not have fired');
+            assert.isFalse(
+                kernelPostInitialized.fired,
+                'IKernelProvider.onDidPostInitializeKernel should not have fired'
+            );
             assert.isFalse(kernelDisposed.fired, 'IKernelProvider.onDidDisposeKernel should not have fired');
 
             onStarted.fire();
@@ -146,6 +155,8 @@ suite('Jupyter Session', () => {
             assert.isTrue(kernelStatusChanged.fired, 'IKernelProvider.onKernelStatusChanged not fired');
             onRestartedEvent.fire();
             assert.isTrue(kernelRestarted.fired, 'IKernelProvider.onKernelRestarted not fired');
+            onPostInitializedEvent.fire();
+            assert.isTrue(kernelPostInitialized.fired, 'IKernelProvider.onDidPostInitializeKernel not fired');
             onDisposedEvent.fire();
             assert.isTrue(kernelDisposed.fired, 'IKernelProvider.onDisposedEvent not fired');
         }

diff --git a/src/kernels/kernelProvider.web.ts b/src/kernels/kernelProvider.web.ts
@@ -80,6 +80,7 @@ export class KernelProvider extends BaseCoreKernelProvider {
             this.workspaceStorage
         ) as IKernel;
         kernel.onRestarted(() => this._onDidRestartKernel.fire(kernel), this, this.disposables);
+        kernel.onPostInitialized(() => this._onDidPostInitializeKernel.fire(kernel), this, this.disposables);
         kernel.onDisposed(() => this._onDidDisposeKernel.fire(kernel), this, this.disposables);
         kernel.onStarted(() => this._onDidStartKernel.fire(kernel), this, this.disposables);
         kernel.onStatusChanged(
@@ -132,6 +133,7 @@ export class ThirdPartyKernelProvider extends BaseThirdPartyKernelProvider {
             this.workspaceStorage
         );
         kernel.onRestarted(() => this._onDidRestartKernel.fire(kernel), this, this.disposables);
+        kernel.onPostInitialized(() => this._onDidPostInitializeKernel.fire(kernel), this, this.disposables);
         kernel.onDisposed(() => this._onDidDisposeKernel.fire(kernel), this, this.disposables);
         kernel.onStarted(() => this._onDidStartKernel.fire(kernel), this, this.disposables);
         kernel.onStatusChanged(

diff --git a/src/kernels/kernelProvider.web.unit.test.ts b/src/kernels/kernelProvider.web.unit.test.ts
@@ -82,20 +82,24 @@ suite('Jupyter Session', () => {
             const kernelStarted = createEventHandler(kernelProvider, 'onDidStartKernel', disposables);
             const kernelDisposed = createEventHandler(kernelProvider, 'onDidDisposeKernel', disposables);
             const kernelRestarted = createEventHandler(kernelProvider, 'onDidRestartKernel', disposables);
+            const kernelPostInitialized = createEventHandler(kernelProvider, 'onDidPostInitializeKernel', disposables);
             const kernelStatusChanged = createEventHandler(kernelProvider, 'onKernelStatusChanged', disposables);
             const notebook = new TestNotebookDocument(undefined, 'jupyter-notebook');
             const onStarted = new EventEmitter<void>();
             const onStatusChanged = new EventEmitter<void>();
             const onRestartedEvent = new EventEmitter<void>();
+            const onPostInitializedEvent = new EventEmitter<void>();
             const onDisposedEvent = new EventEmitter<void>();
             disposables.push(onStatusChanged);
             disposables.push(onRestartedEvent);
+            disposables.push(onPostInitializedEvent);
             disposables.push(onStarted);
             disposables.push(onDisposedEvent);
             if (kernelProvider instanceof KernelProvider) {
                 sinon.stub(Kernel.prototype, 'onStarted').get(() => onStarted.event);
                 sinon.stub(Kernel.prototype, 'onStatusChanged').get(() => onStatusChanged.event);
                 sinon.stub(Kernel.prototype, 'onRestarted').get(() => onRestartedEvent.event);
+                sinon.stub(Kernel.prototype, 'onPostInitialized').get(() => onPostInitializedEvent.event);
                 sinon.stub(Kernel.prototype, 'onDisposed').get(() => onDisposedEvent.event);
                 const kernel = kernelProvider.getOrCreate(notebook, {
                     controller,
@@ -107,6 +111,7 @@ suite('Jupyter Session', () => {
                 sinon.stub(ThirdPartyKernel.prototype, 'onStarted').get(() => onStarted.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onStatusChanged').get(() => onStatusChanged.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onRestarted').get(() => onRestartedEvent.event);
+                sinon.stub(ThirdPartyKernel.prototype, 'onPostInitialized').get(() => onPostInitializedEvent.event);
                 sinon.stub(ThirdPartyKernel.prototype, 'onDisposed').get(() => onDisposedEvent.event);
                 const kernel = kernelProvider.getOrCreate(notebook.uri, {
                     metadata: instance(metadata),
@@ -119,6 +124,10 @@ suite('Jupyter Session', () => {
             assert.isFalse(kernelStarted.fired, 'IKernelProvider.onDidStartKernel should not be fired');
             assert.isFalse(kernelStatusChanged.fired, 'IKernelProvider.onKernelStatusChanged should not be fired');
             assert.isFalse(kernelRestarted.fired, 'IKernelProvider.onDidRestartKernel should not have fired');
+            assert.isFalse(
+                kernelPostInitialized.fired,
+                'IKernelProvider.onDidPostInitializeKernel should not have fired'
+            );
             assert.isFalse(kernelDisposed.fired, 'IKernelProvider.onDidDisposeKernel should not have fired');
 
             onStarted.fire();
@@ -127,6 +136,8 @@ suite('Jupyter Session', () => {
             assert.isTrue(kernelStatusChanged.fired, 'IKernelProvider.onKernelStatusChanged not fired');
             onRestartedEvent.fire();
             assert.isTrue(kernelRestarted.fired, 'IKernelProvider.onKernelRestarted not fired');
+            onPostInitializedEvent.fire();
+            assert.isTrue(kernelPostInitialized.fired, 'IKernelProvider.onDidPostInitializeKernel not fired');
             onDisposedEvent.fire();
             assert.isTrue(kernelDisposed.fired, 'IKernelProvider.onDisposedEvent not fired');
         }

diff --git a/src/kernels/types.ts b/src/kernels/types.ts
@@ -355,6 +355,7 @@ export interface IBaseKernel extends IAsyncDisposable {
     readonly onDisposed: Event<void>;
     readonly onStarted: Event<void>;
     readonly onRestarted: Event<void>;
+    readonly onPostInitialized: Event<void>;
     readonly restarting: Promise<void>;
     readonly status: KernelMessage.Status;
     readonly disposed: boolean;
@@ -381,6 +382,12 @@ export interface IBaseKernel extends IAsyncDisposable {
      * This flag will tell us whether a real kernel was or is active.
      */
     readonly startedAtLeastOnce?: boolean;
+    /**
+     * A kernel might be started (e.g., for pre-warming or other internal reasons) but this does not
+     * necessarily correlate with whether it was started by a user.
+     * This flag will tell us whether the kernel has been started explicitly through user action from the UI.
+     */
+    readonly userStartedKernel: boolean;
     start(options?: IDisplayOptions): Promise<IKernelSession>;
     interrupt(): Promise<void>;
     restart(): Promise<void>;
@@ -506,6 +513,7 @@ export interface IBaseKernelProvider<T extends IBaseKernel> extends IAsyncDispos
     onDidRestartKernel: Event<T>;
     onDidDisposeKernel: Event<T>;
     onKernelStatusChanged: Event<{ status: KernelMessage.Status; kernel: T }>;
+    onDidPostInitializeKernel: Event<T>;
 }
 
 /**