Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add public API event for kernel post-initialization #16214

Open
wants to merge 17 commits into
base: main
Choose a base branch
from
16 changes: 16 additions & 0 deletions src/api.proposed.kernelStartHook.d.ts
Original file line number Diff line number Diff line change
@@ -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;
}>;
}
}
5 changes: 4 additions & 1 deletion src/extension.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
}
};
}
}
Expand Down
5 changes: 4 additions & 1 deletion src/extension.web.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
}
};
}
}
Expand Down
19 changes: 18 additions & 1 deletion src/kernels/kernel.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -172,10 +175,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>();
Expand Down Expand Up @@ -246,7 +251,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;
pwang347 marked this conversation as resolved.
Show resolved Hide resolved
}
return result;
});
}
/**
* Interrupts the execution of cells.
Expand Down Expand Up @@ -409,6 +423,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;
Expand Down
10 changes: 10 additions & 0 deletions src/kernels/kernelProvider.base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>();
Expand All @@ -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> {
Expand All @@ -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(
Expand Down Expand Up @@ -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;
Expand All @@ -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> {
Expand All @@ -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;
}
Expand Down
14 changes: 14 additions & 0 deletions src/kernels/kernelProvider.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
11 changes: 11 additions & 0 deletions src/kernels/kernelProvider.node.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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),
Expand All @@ -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();
Expand All @@ -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');
}
Expand Down
2 changes: 2 additions & 0 deletions src/kernels/kernelProvider.web.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
11 changes: 11 additions & 0 deletions src/kernels/kernelProvider.web.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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),
Expand All @@ -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();
Expand All @@ -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');
}
Expand Down
2 changes: 2 additions & 0 deletions src/kernels/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -506,6 +507,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>;
}

/**
Expand Down
Loading
Loading