Add public API event for kernel post-initialization

src/api.d.ts

@@ -233,10 +233,17 @@ export interface Kernel {
     executeCode(code: string, token: CancellationToken): AsyncIterable<Output>;
 }
 export interface Kernels {
+    /**
+     * Event fired when a kernel becomes available for execution on a resource.
+     */
+    onDidInitialize: Event<{
+        uri: Uri;
+    }>;
+
     /**
      * Gets an the kernel associated with a given resource.
      * For instance if the resource is a notebook, then get the kernel associated with the given Notebook document.
-     * Only kernels which have already been started by the user and belonging to Notebooks that are currently opened will be returned.
+     * Only kernels which are ready for code execution and belonging to Notebooks that are currently opened will be returned.
      */
     getKernel(uri: Uri): Thenable<Kernel | undefined>;
 }

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;
     }
@@ -176,6 +179,7 @@ abstract class BaseKernel implements IBaseKernel {
     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>();
@@ -788,6 +792,9 @@ abstract class BaseKernel implements IBaseKernel {
                 logger.trace('End running kernel initialization, session is idle');
             }
             kernelIdle?.stop();
+
+            // Post initialization event is only emitted on successful initialization
+            this._onPostInitialized.fire();
         } finally {
             postInitialization?.stop();
         }

src/kernels/kernelProvider.base.ts

@@ -21,7 +21,7 @@
 /**
  * Provides kernels to the system. Generally backed by a URI or a notebook object.
  */
 export abstract class BaseCoreKernelProvider implements IKernelProvider {
     protected readonly executions = new WeakMap<IKernel, INotebookKernelExecution>();
 
     /**
@@ -36,6 +36,7 @@
     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 @@
         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 @@
     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(
@@ -175,7 +180,7 @@
     }
 }
 
 export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelProvider {
     /**
      * The life time of kernels not tied to a notebook will be managed by callers of the API.
      * Where as if a kernel is tied to a notebook, then the kernel dies along with notebooks.
@@ -187,6 +192,7 @@
     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 @@
         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 @@
     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;
     }

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);

src/kernels/kernelProvider.node.unit.test.ts

@@ -101,20 +101,24 @@
             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 @@
                 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 @@
             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 @@
             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');
         }
@@ -211,7 +222,7 @@
                 (_: unknown, defaultValue: unknown) => defaultValue
             );
 
             kernelProvider = new KernelProvider(
                 asyncDisposables,
                 disposables,
                 instance(sessionCreator),
@@ -223,7 +234,7 @@
                 instance(workspaceMemento),
                 instance(replTracker)
             );
             thirdPartyKernelProvider = new ThirdPartyKernelProvider(
                 asyncDisposables,
                 disposables,
                 instance(sessionCreator),

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');
         }

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;
@@ -506,6 +507,8 @@ export interface IBaseKernelProvider<T extends IBaseKernel> extends IAsyncDispos
     onDidRestartKernel: Event<T>;
     onDidDisposeKernel: Event<T>;
     onKernelStatusChanged: Event<{ status: KernelMessage.Status; kernel: T }>;
+    onDidPostInitializeKernel: Event<T>;
+    onDidDispose: Event<void>;
 }
 
 /**

src/standalone/api/kernels/index.ts

@@ -1,16 +1,18 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { Uri, workspace } from 'vscode';
+import { Uri, workspace, EventEmitter } from 'vscode';
 import { Kernel, Kernels } from '../../../api';
 import { ServiceContainer } from '../../../platform/ioc/container';
-import { IKernel, IKernelProvider, isRemoteConnection } from '../../../kernels/types';
+import { IKernel, IKernelProvider } from '../../../kernels/types';
 import { createKernelApiForExtension as createKernelApiForExtension } from './kernel';
 import { Telemetry, sendTelemetryEvent } from '../../../telemetry';
 import { JVSC_EXTENSION_ID } from '../../../platform/common/constants';
 import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../../../kernels/telemetry/helper';
+import { IDisposableRegistry } from '../../../platform/common/types';
 
 const kernelCache = new WeakMap<IKernel, Kernel>();
+let _onDidInitialize: EventEmitter<{ uri: Uri }> | undefined = undefined;
 
 export function getKernelsApi(extensionId: string): Kernels {
     return {
@@ -29,18 +31,6 @@ export function getKernelsApi(extensionId: string): Kernels {
                 });
                 return;
             }
-            const execution = kernelProvider.getKernelExecution(kernel);
-            if (!isRemoteConnection(kernel.kernelConnectionMetadata) && execution.executionCount === 0) {
-                // For local kernels, execution count must be greater than 0,
-                // As we pre-warms kernels (i.e. we start kernels even though the user may not have executed any code).
-                // The only way to determine whether users executed code is to look at the execution count
-                sendTelemetryEvent(Telemetry.NewJupyterKernelsApiUsage, undefined, {
-                    extensionId,
-                    pemUsed: 'getKernel',
-                    accessAllowed
-                });
-                return;
-            }
             if (extensionId !== JVSC_EXTENSION_ID) {
                 void initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
                     kernel.resourceUri,
@@ -50,6 +40,24 @@ export function getKernelsApi(extensionId: string): Kernels {
             let wrappedKernel = kernelCache.get(kernel) || createKernelApiForExtension(extensionId, kernel);
             kernelCache.set(kernel, wrappedKernel);
             return wrappedKernel;
+        },
+        get onDidInitialize() {
+            // We can cache the event emitter for subsequent calls.
+            if (!_onDidInitialize) {
+                const kernelProvider = ServiceContainer.instance.get<IKernelProvider>(IKernelProvider);
+                const disposableRegistry = ServiceContainer.instance.get<IDisposableRegistry>(IDisposableRegistry);
+                _onDidInitialize = new EventEmitter<{ uri: Uri }>();
+
+                disposableRegistry.push(
+                    kernelProvider.onDidPostInitializeKernel((e) => {
+                        _onDidInitialize?.fire({ uri: e.uri });
+                    }),
+                    _onDidInitialize,
+                    { dispose: () => (_onDidInitialize = undefined) }
+                );
+            }
+
+            return _onDidInitialize.event;
         }
     };
 }