From dbd78ca681a0a6eb0b44670ae1ebadc9ddcb3bd8 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 18 Nov 2024 19:16:57 -0800 Subject: [PATCH] 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 --- src/api.proposed.kernelStartHook.d.ts | 16 ++ src/extension.node.ts | 5 +- src/extension.web.ts | 5 +- src/kernels/kernel.ts | 22 ++- src/kernels/kernelProvider.base.ts | 10 ++ src/kernels/kernelProvider.node.ts | 14 ++ src/kernels/kernelProvider.node.unit.test.ts | 11 ++ src/kernels/kernelProvider.web.ts | 2 + src/kernels/kernelProvider.web.unit.test.ts | 11 ++ src/kernels/types.ts | 8 + .../api/kernels/api.vscode.common.test.ts | 158 +++++++++++++++++- src/standalone/api/kernels/index.ts | 55 +++--- 12 files changed, 293 insertions(+), 24 deletions(-) create mode 100644 src/api.proposed.kernelStartHook.d.ts diff --git a/src/api.proposed.kernelStartHook.d.ts b/src/api.proposed.kernelStartHook.d.ts new file mode 100644 index 00000000000..5377ee50f08 --- /dev/null +++ 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 index 1857d5e46ce..4e519ab38f3 100644 --- a/src/extension.node.ts +++ b/src/extension.node.ts @@ -124,7 +124,10 @@ export async function activate(context: IExtensionContext): Promise { 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 index add9063bd7c..9ffca564c40 100644 --- a/src/extension.web.ts +++ b/src/extension.web.ts @@ -117,7 +117,10 @@ export async function activate(context: IExtensionContext): Promise { 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 index b6a552e9c49..a63893ad1fb 100644 --- a/src/kernels/kernel.ts +++ b/src/kernels/kernel.ts @@ -129,6 +129,9 @@ abstract class BaseKernel implements IBaseKernel { get onStarted(): Event { return this._onStarted.event; } + get onPostInitialized(): Event { + return this._onPostInitialized.event; + } get onDisposed(): Event { 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(); private readonly _onStatusChanged = new EventEmitter(); private readonly _onRestarted = new EventEmitter(); private readonly _onStarted = new EventEmitter(); + private readonly _onPostInitialized = new EventEmitter(); private readonly _onDisposed = new EventEmitter(); private _jupyterSessionPromise?: Promise; private readonly hookedSessionForEvents = new WeakSet(); @@ -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 index 4c60412f704..f3dd4609bc3 100644 --- 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(); protected readonly _onDidDisposeKernel = new EventEmitter(); protected readonly _onKernelStatusChanged = new EventEmitter<{ status: KernelMessage.Status; kernel: IKernel }>(); + protected readonly _onDidPostInitializeKernel = new EventEmitter(); public readonly onKernelStatusChanged = this._onKernelStatusChanged.event; public get kernels() { const kernels = new Set(); @@ -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 { @@ -73,6 +75,9 @@ export abstract class BaseCoreKernelProvider implements IKernelProvider { public get onDidCreateKernel(): Event { return this._onDidCreateKernel.event; } + public get onDidPostInitializeKernel(): Event { + 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(); protected readonly _onDidCreateKernel = new EventEmitter(); protected readonly _onDidDisposeKernel = new EventEmitter(); + protected readonly _onDidPostInitializeKernel = new EventEmitter(); 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 { @@ -227,6 +234,9 @@ export abstract class BaseThirdPartyKernelProvider implements IThirdPartyKernelP public get onDidCreateKernel(): Event { return this._onDidCreateKernel.event; } + public get onDidPostInitializeKernel(): Event { + 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 index 0dd08513ba8..158b4b59593 100644 --- 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 index 81138ccd7d6..f5ba56cb83d 100644 --- 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(); const onStatusChanged = new EventEmitter(); const onRestartedEvent = new EventEmitter(); + const onPostInitializedEvent = new EventEmitter(); const onDisposedEvent = new EventEmitter(); 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 index 299ec64b4ec..6a65b3f8c00 100644 --- 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 index 05d5789f7c8..6c8cc09cda2 100644 --- 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(); const onStatusChanged = new EventEmitter(); const onRestartedEvent = new EventEmitter(); + const onPostInitializedEvent = new EventEmitter(); const onDisposedEvent = new EventEmitter(); 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 index 9d8d40dc727..b9bc32cf8a4 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -355,6 +355,7 @@ export interface IBaseKernel extends IAsyncDisposable { readonly onDisposed: Event; readonly onStarted: Event; readonly onRestarted: Event; + readonly onPostInitialized: Event; readonly restarting: Promise; 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; interrupt(): Promise; restart(): Promise; @@ -506,6 +513,7 @@ export interface IBaseKernelProvider extends IAsyncDispos onDidRestartKernel: Event; onDidDisposeKernel: Event; onKernelStatusChanged: Event<{ status: KernelMessage.Status; kernel: T }>; + onDidPostInitializeKernel: Event; } /** diff --git a/src/standalone/api/kernels/api.vscode.common.test.ts b/src/standalone/api/kernels/api.vscode.common.test.ts index d29e011f4fc..9236e2e68ff 100644 --- a/src/standalone/api/kernels/api.vscode.common.test.ts +++ b/src/standalone/api/kernels/api.vscode.common.test.ts @@ -2,7 +2,15 @@ // Licensed under the MIT License. import { assert } from 'chai'; -import { CancellationTokenSource, NotebookCellOutputItem, NotebookDocument, commands, window, workspace } from 'vscode'; +import { + CancellationTokenSource, + NotebookCell, + NotebookCellOutputItem, + NotebookDocument, + commands, + window, + workspace +} from 'vscode'; import { logger } from '../../../platform/logging'; import { IDisposable } from '../../../platform/common/types'; import { @@ -32,6 +40,7 @@ import { KernelError } from '../../../kernels/errors/kernelError'; import { JVSC_EXTENSION_ID } from '../../../platform/common/constants'; import { escapeStringToEmbedInPythonCode } from '../../../kernels/chat/generator'; import { notebookCellExecutions } from '../../../platform/notebooks/cellExecutionStateService'; +import { noop } from '../../../test/core'; suiteMandatory('Kernel API Tests @typescript', function () { const disposables: IDisposable[] = []; @@ -106,7 +115,12 @@ suiteMandatory('Kernel API Tests @typescript', function () { // Even after starting a kernel the API should not return anything, // as no code has been executed against this kernel. - await realKernel.start(); + await realKernel.start({ + disableUI: true, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); assert.isUndefined(await kernels.getKernel(notebook.uri)); // Ensure user has executed some code against this kernel. @@ -185,6 +199,146 @@ suiteMandatory('Kernel API Tests @typescript', function () { 'Traceback does not contain original error' ); }); + testMandatory('Kernel start event is not triggered by silent executions', async function () { + let startEventCounter = 0; + // Register event listener to track invocations + disposables.push( + kernels.onDidStart(() => { + startEventCounter++; + }) + ); + + await realKernel.start({ + disableUI: true, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); + assert.equal(startEventCounter, 0, 'Kernel start event was triggered for a non-user kernel start'); + }); + testMandatory('Kernel start event is triggered when started with UI enabled', async function () { + let startEventCounter = 0; + // Register event listener to track invocations + disposables.push( + kernels.onDidStart(() => { + startEventCounter++; + }) + ); + + await realKernel.start({ + disableUI: false, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); + assert.equal(startEventCounter, 1, 'Kernel start event was not triggered for a user kernel start'); + + // If we call start again with UI enabled, we shouldn't fire additional events + await realKernel.start({ + disableUI: false, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); + assert.equal(startEventCounter, 1, 'Multiple start calls should not fire more events'); + }); + testMandatory('Kernel start event is triggered when kernel restarts', async function () { + let startEventCounter = 0; + // Register event listener to track invocations + disposables.push( + kernels.onDidStart(() => { + startEventCounter++; + }) + ); + + await realKernel.start({ + disableUI: true, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); + await realKernel.restart(); + assert.equal(startEventCounter, 1, 'Kernel start event should be fired exactly once after restarting'); + await realKernel.restart(); + assert.equal(startEventCounter, 2, 'Kernel start event should be fired more than once for restarts'); + }); + testMandatory( + 'Kernel start event is triggered when user executes code and the event execution runs first', + async function () { + // Register event listener to track invocations + const source = new CancellationTokenSource(); + let startEventCounter = 0; + disposables.push( + kernels.onDidStart(async ({ kernel }) => { + const codeToRun = + startEventCounter === 0 ? `let foo = ${startEventCounter}` : `foo = ${startEventCounter}`; + startEventCounter++; + + // This is needed for the async generator to get executed. + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _out of kernel.executeCode(codeToRun, source.token)) { + } + }) + ); + await insertCodeCell('console.log(foo)', { index: 0, language: 'typescript' }); + + await realKernel.start({ + disableUI: true, + onDidChangeDisableUI: () => ({ + dispose: noop + }) + }); + assert.equal(startEventCounter, 0, 'Kernel start event was triggered for a non-user kernel start'); + const cell = notebook.cellAt(0)!; + const executionOrderSet = createDeferred(); + const eventHandler = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.cell.executionSummary?.executionOrder) { + executionOrderSet.resolve(); + } + }); + disposables.push(eventHandler); + await Promise.all([runCell(cell), waitForExecutionCompletedSuccessfully(cell), executionOrderSet.promise]); + + // Validate the cell execution output is equal to the expected value of "foo = 0" + const expectedMime = NotebookCellOutputItem.stdout('').mime; + assert.equal( + await decodeFirstOutput(cell, expectedMime), + '0', + 'Invalid output, kernel start hook should execute code first' + ); + + const kernel = await kernels.getKernel(notebook.uri); + if (!kernel) { + throw new Error('Kernel not found'); + } + + // Start event counter should only be 1 after the initial user cell execution + assert.equal(startEventCounter, 1, 'Kernel start event was not triggered for a user kernel start'); + + // Running the same cell again should not fire additional events + await Promise.all([runCell(cell), waitForExecutionCompletedSuccessfully(cell), executionOrderSet.promise]); + assert.equal( + await decodeFirstOutput(cell, expectedMime), + '0', + 'Invalid output, kernel start hook should only execute once' + ); + assert.equal(startEventCounter, 1, 'Start event should not be triggered more than once'); + } + ); + + async function decodeFirstOutput(cell: NotebookCell, expectedMimetype: string) { + return ( + cell.outputs + .flatMap((output) => output.items) + .map((item) => { + if (item.mime === expectedMimetype) { + const output = new TextDecoder().decode(item.data).trim(); + return output; + } + }) + .find((item) => item !== undefined) ?? '' + ); + } async function waitForOutput( executionResult: AsyncIterable, diff --git a/src/standalone/api/kernels/index.ts b/src/standalone/api/kernels/index.ts index c0277c8e438..b84065cb414 100644 --- a/src/standalone/api/kernels/index.ts +++ b/src/standalone/api/kernels/index.ts @@ -1,16 +1,24 @@ // 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 { DATA_WRANGLER_EXTENSION_ID, JVSC_EXTENSION_ID } from '../../../platform/common/constants'; import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../../../kernels/telemetry/helper'; +import { IDisposableRegistry } from '../../../platform/common/types'; const kernelCache = new WeakMap(); +let _onDidStart: EventEmitter<{ uri: Uri; kernel: Kernel }> | undefined = undefined; + +function getWrappedKernel(kernel: IKernel, extensionId: string) { + let wrappedKernel = kernelCache.get(kernel) || createKernelApiForExtension(extensionId, kernel); + kernelCache.set(kernel, wrappedKernel); + return wrappedKernel; +} export function getKernelsApi(extensionId: string): Kernels { return { @@ -21,19 +29,8 @@ export function getKernelsApi(extensionId: string): Kernels { const notebook = workspace.notebookDocuments.find((item) => item.uri.toString() === uri.toString()); const kernel = kernelProvider.get(notebook || uri); // We are only interested in returning kernels that have been started by the user. - if (!kernel || !kernel.startedAtLeastOnce) { - sendTelemetryEvent(Telemetry.NewJupyterKernelsApiUsage, undefined, { - extensionId, - pemUsed: 'getKernel', - accessAllowed - }); - 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 + // Returning started kernels is not sufficient as we also pre-warm kernels (i.e. we start kernels even though the user may not have executed any code). + if (!kernel || !kernel.startedAtLeastOnce || !kernel.userStartedKernel) { sendTelemetryEvent(Telemetry.NewJupyterKernelsApiUsage, undefined, { extensionId, pemUsed: 'getKernel', @@ -47,9 +44,29 @@ export function getKernelsApi(extensionId: string): Kernels { kernel.kernelConnectionMetadata ); } - let wrappedKernel = kernelCache.get(kernel) || createKernelApiForExtension(extensionId, kernel); - kernelCache.set(kernel, wrappedKernel); - return wrappedKernel; + return getWrappedKernel(kernel, extensionId); + }, + get onDidStart() { + if (![JVSC_EXTENSION_ID, DATA_WRANGLER_EXTENSION_ID].includes(extensionId)) { + throw new Error(`Proposed API is not supported for extension ${extensionId}`); + } + + // We can cache the event emitter for subsequent calls. + if (!_onDidStart) { + const kernelProvider = ServiceContainer.instance.get(IKernelProvider); + const disposableRegistry = ServiceContainer.instance.get(IDisposableRegistry); + _onDidStart = new EventEmitter<{ uri: Uri; kernel: Kernel }>(); + + disposableRegistry.push( + kernelProvider.onDidPostInitializeKernel((kernel) => { + _onDidStart?.fire({ uri: kernel.uri, kernel: getWrappedKernel(kernel, extensionId) }); + }), + _onDidStart, + { dispose: () => (_onDidStart = undefined) } + ); + } + + return _onDidStart.event; } }; }