Notebook kernel preloads

src/vs/workbench/api/browser/mainThreadNotebook.ts

@@ -8,7 +8,7 @@ import { MainContext, MainThreadNotebookShape, NotebookExtensionDescription, IEx
 import { Disposable } from 'vs/base/common/lifecycle';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { INotebookService, IMainNotebookController } from 'vs/workbench/contrib/notebook/common/notebookService';
-import { INotebookTextModel, INotebookMimeTypeSelector, NOTEBOOK_DISPLAY_ORDER, NotebookCellOutputsSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellEditType, CellKind, INotebookKernelInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookMimeTypeSelector, NOTEBOOK_DISPLAY_ORDER, NotebookCellOutputsSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellEditType, CellKind, INotebookKernelInfo, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
@@ -123,8 +123,8 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		this._notebookService.unregisterNotebookRenderer(handle);
 	}
 
-	async $registerNotebookProvider(extension: NotebookExtensionDescription, viewType: string, hasKernelSupport: boolean): Promise<void> {
-		let controller = new MainThreadNotebookController(this._proxy, this, viewType, hasKernelSupport);
+	async $registerNotebookProvider(extension: NotebookExtensionDescription, viewType: string, kernel: INotebookKernelInfoDto | undefined): Promise<void> {
+		let controller = new MainThreadNotebookController(this._proxy, this, viewType, kernel);
 		this._notebookProviders.set(viewType, controller);
 		this._notebookService.registerNotebookController(viewType, extension, controller);
 		return;
@@ -206,7 +206,7 @@ export class MainThreadNotebookController implements IMainNotebookController {
 		private readonly _proxy: ExtHostNotebookShape,
 		private _mainThreadNotebook: MainThreadNotebooks,
 		private _viewType: string,
-		readonly hasKernelSupport: boolean
+		readonly kernel: INotebookKernelInfoDto | undefined
 	) {
 	}
 

src/vs/workbench/api/common/extHost.protocol.ts

@@ -51,7 +51,7 @@ import { TunnelDto } from 'vs/workbench/api/common/extHostTunnelService';
 import { TunnelOptions } from 'vs/platform/remote/common/tunnel';
 import { Timeline, TimelineChangeEvent, TimelineOptions, TimelineProviderDescriptor, InternalTimelineOptions } from 'vs/workbench/contrib/timeline/common/timeline';
 import { revive } from 'vs/base/common/marshalling';
-import { INotebookMimeTypeSelector, IOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookMimeTypeSelector, IOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { CallHierarchyItem } from 'vs/workbench/contrib/callHierarchy/common/callHierarchy';
 import { Dto } from 'vs/base/common/types';
 import { ISerializableEnvironmentVariableCollection } from 'vs/workbench/contrib/terminal/common/environmentVariable';
@@ -690,7 +690,7 @@ export type NotebookCellOutputsSplice = [
 ];
 
 export interface MainThreadNotebookShape extends IDisposable {
-	$registerNotebookProvider(extension: NotebookExtensionDescription, viewType: string, hasKernelSupport: boolean): Promise<void>;
+	$registerNotebookProvider(extension: NotebookExtensionDescription, viewType: string, kernelInfoDto: INotebookKernelInfoDto | undefined): Promise<void>;
 	$unregisterNotebookProvider(viewType: string): Promise<void>;
 	$registerNotebookRenderer(extension: NotebookExtensionDescription, type: string, selectors: INotebookMimeTypeSelector, handle: number, preloads: UriComponents[]): Promise<void>;
 	$unregisterNotebookRenderer(handle: number): Promise<void>;

src/vs/workbench/api/common/extHostNotebook.ts

@@ -722,7 +722,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		}
 
 		this._notebookContentProviders.set(viewType, { extension, provider });
-		this._proxy.$registerNotebookProvider({ id: extension.identifier, location: extension.extensionLocation }, viewType, !!provider.kernel);
+		this._proxy.$registerNotebookProvider({ id: extension.identifier, location: extension.extensionLocation }, viewType, provider.kernel ? { id: viewType, label: provider.kernel.label, extensionLocation: extension.extensionLocation, preloads: provider.kernel.preloads } : undefined);
 		return new VSCodeDisposable(() => {
 			this._notebookContentProviders.delete(viewType);
 			this._proxy.$unregisterNotebookProvider(viewType);

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -39,7 +39,7 @@ import { CellDragAndDropController, CodeCellRenderer, MarkdownCellRenderer, Note
 import { CodeCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/codeCellViewModel';
 import { NotebookEventDispatcher, NotebookLayoutChangedEvent } from 'vs/workbench/contrib/notebook/browser/viewModel/eventDispatcher';
 import { CellViewModel, IModelDecorationsChangeAccessor, INotebookEditorViewState, NotebookViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
-import { CellKind, IOutput, INotebookKernelInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellKind, IOutput, INotebookKernelInfo, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 import { Webview } from 'vs/workbench/contrib/webview/browser/webview';
 import { getExtraColor } from 'vs/workbench/contrib/welcome/walkThrough/common/walkThroughUtils';
@@ -48,6 +48,7 @@ import { ILayoutService } from 'vs/platform/layout/browser/layoutService';
 import { generateUuid } from 'vs/base/common/uuid';
 import { Memento, MementoObject } from 'vs/workbench/common/memento';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
+import { URI } from 'vs/base/common/uri';
 
 const $ = DOM.$;
 
@@ -403,7 +404,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		const provider = this.notebookService.getContributedNotebookProviders(this.viewModel!.uri)[0];
 		const availableKernels = this.notebookService.getContributedNotebookKernels(textModel.viewType, textModel.uri);
 
-		if (provider.hasKernelSupport && availableKernels.length > 0) {
+		if (provider.kernel && availableKernels.length > 0) {
 			this.notebookHasMultipleKernels!.set(true);
 		} else if (availableKernels.length > 1) {
 			this.notebookHasMultipleKernels!.set(true);
@@ -411,13 +412,23 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 			this.notebookHasMultipleKernels!.set(false);
 		}
 
-		if (provider && provider.hasKernelSupport) {
+		if (provider && provider.kernel) {
 			// it has a builtin kernel, don't automatically choose a kernel
+			this.loadKernelPreloads(provider.providerExtensionLocation, provider.kernel);
 			return;
 		}
 
 		// the provider doesn't have a builtin kernel, choose a kernel
 		this.activeKernel = availableKernels[0];
+		if (this.activeKernel) {
+			this.loadKernelPreloads(this.activeKernel.extensionLocation, this.activeKernel);
+		}
+	}
+
+	private loadKernelPreloads(extensionLocation: URI, kernel: INotebookKernelInfoDto) {
+		if (kernel.preloads) {
+			this.webview?.updateKernelPreloads([extensionLocation], kernel.preloads.map(preload => URI.revive(preload)));
+		}
 	}
 
 	private updateForMetadata(): void {
@@ -1079,7 +1090,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 
 				if (this._activeKernel) {
 					await this.notebookService.executeNotebook2(this.notebookViewModel!.viewType, this.notebookViewModel!.uri, this._activeKernel.id, tokenSource.token);
-				} else if (provider.hasKernelSupport) {
+				} else if (provider.kernel) {
 					return await this.notebookService.executeNotebook(viewType, notebookUri, true, tokenSource.token);
 				} else {
 					return await this.notebookService.executeNotebook(viewType, notebookUri, false, tokenSource.token);
@@ -1126,7 +1137,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 
 				if (this._activeKernel) {
 					return await this.notebookService.executeNotebookCell2(viewType, notebookUri, cell.handle, this._activeKernel.id, tokenSource.token);
-				} else if (provider.hasKernelSupport) {
+				} else if (provider.kernel) {
 					return await this.notebookService.executeNotebookCell(viewType, notebookUri, cell.handle, true, tokenSource.token);
 				} else {
 					return await this.notebookService.executeNotebookCell(viewType, notebookUri, cell.handle, false, tokenSource.token);

src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts

@@ -133,6 +133,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 						displayName: notebookContribution.displayName,
 						selector: notebookContribution.selector || [],
 						providerDisplayName: extension.description.isBuiltin ? nls.localize('builtinProviderDisplayName', "Built-in") : extension.description.displayName || extension.description.identifier.value,
+						providerExtensionLocation: extension.description.extensionLocation
 					}));
 				}
 			}
@@ -176,7 +177,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 
 	registerNotebookController(viewType: string, extensionData: NotebookExtensionDescription, controller: IMainNotebookController) {
 		this._notebookProviders.set(viewType, { extensionData, controller });
-		this.notebookProviderInfoStore.get(viewType)!.hasKernelSupport = controller.hasKernelSupport;
+		this.notebookProviderInfoStore.get(viewType)!.kernel = controller.kernel;
 		this._onDidChangeViewTypes.fire();
 	}
 

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -132,8 +132,10 @@ export class BackLayerWebView extends Disposable {
 	hiddenInsetMapping: Set<IOutput> = new Set();
 	reversedInsetMapping: Map<string, IOutput> = new Map();
 	preloadsCache: Map<string, boolean> = new Map();
+	kernelPreloadsCache: Map<string, boolean> = new Map();
 	localResourceRootsCache: URI[] | undefined = undefined;
 	rendererRootsCache: URI[] = [];
+	kernelRootsCache: URI[] = [];
 	private readonly _onMessage = this._register(new Emitter<any>());
 	public readonly onMessage: Event<any> = this._onMessage.event;
 	private _initalized: Promise<void>;
@@ -730,6 +732,30 @@ ${loaderJs}
 		}, 50);
 	}
 
+	updateKernelPreloads(extensionLocations: URI[], preloads: URI[]) {
+		if (this._disposed) {
+			return;
+		}
+
+		let resources: string[] = [];
+		preloads = preloads.map(preload => {
+			if (this.environmentService.isExtensionDevelopment && (preload.scheme === 'http' || preload.scheme === 'https')) {
+				return preload;
+			}
+			return asWebviewUri(this.workbenchEnvironmentService, this.id, preload);
+		});
+
+		preloads.forEach(e => {
+			if (!this.kernelPreloadsCache.has(e.toString())) {
+				resources.push(e.toString());
+				this.kernelPreloadsCache.set(e.toString(), true);
+			}
+		});
+
+		this.kernelRootsCache = [...extensionLocations, ...this.kernelRootsCache];
+		this._updatePreloads(resources);
+	}
+
 	updateRendererPreloads(preloads: ReadonlySet<number>) {
 		if (this._disposed) {
 			return;
@@ -758,7 +784,11 @@ ${loaderJs}
 		});
 
 		this.rendererRootsCache = extensionLocations;
-		const mixedResourceRoots = [...(this.localResourceRootsCache || []), ...this.rendererRootsCache];
+		this._updatePreloads(resources);
+	}
+
+	private _updatePreloads(resources: string[]) {
+		const mixedResourceRoots = [...(this.localResourceRootsCache || []), ...this.rendererRootsCache, ...this.kernelRootsCache];
 
 		this.webview.contentOptions = {
 			allowMultipleAPIAcquire: true,

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -115,6 +115,13 @@ export interface INotebookKernelInfo {
 	executeNotebook(viewType: string, uri: URI, handle: number | undefined, token: CancellationToken): Promise<void>;
 }
 
+export interface INotebookKernelInfoDto {
+	id: string;
+	label: string,
+	extensionLocation: URI;
+	preloads?: UriComponents[];
+}
+
 export interface INotebookSelectors {
 	readonly filenamePattern?: string;
 }
@@ -494,3 +501,4 @@ export interface INotebookEditorModel extends IEditorModel {
 	isDirty(): boolean;
 	save(): Promise<boolean>;
 }
+

src/vs/workbench/contrib/notebook/common/notebookProvider.ts

@@ -6,6 +6,7 @@
 import * as glob from 'vs/base/common/glob';
 import { URI } from 'vs/base/common/uri';
 import { basename } from 'vs/base/common/resources';
+import { INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 
 export interface NotebookSelector {
 	readonly filenamePattern?: string;
@@ -18,18 +19,21 @@ export class NotebookProviderInfo {
 	readonly displayName: string;
 	readonly selector: readonly NotebookSelector[];
 	readonly providerDisplayName: string;
-	hasKernelSupport: boolean = false;
+	readonly providerExtensionLocation: URI;
+	kernel?: INotebookKernelInfoDto;
 
 	constructor(descriptor: {
 		readonly id: string;
 		readonly displayName: string;
 		readonly selector: readonly NotebookSelector[];
 		readonly providerDisplayName: string;
+		readonly providerExtensionLocation: URI;
 	}) {
 		this.id = descriptor.id;
 		this.displayName = descriptor.displayName;
 		this.selector = descriptor.selector;
 		this.providerDisplayName = descriptor.providerDisplayName;
+		this.providerExtensionLocation = descriptor.providerExtensionLocation;
 	}
 
 	matches(resource: URI): boolean {

src/vs/workbench/contrib/notebook/common/notebookService.ts

@@ -8,7 +8,7 @@ import { URI } from 'vs/base/common/uri';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
 import { NotebookExtensionDescription } from 'vs/workbench/api/common/extHost.protocol';
 import { Event } from 'vs/base/common/event';
-import { INotebookTextModel, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookTextModel, INotebookMimeTypeSelector, INotebookRendererInfo, NotebookDocumentMetadata, ICellDto2, INotebookKernelInfo, INotebookKernelInfoDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
@@ -17,7 +17,7 @@ import { INotebookEditorModelManager } from 'vs/workbench/contrib/notebook/commo
 export const INotebookService = createDecorator<INotebookService>('notebookService');
 
 export interface IMainNotebookController {
-	hasKernelSupport: boolean;
+	kernel: INotebookKernelInfoDto | undefined;
 	createNotebook(viewType: string, uri: URI, forBackup: boolean, forceReload: boolean): Promise<NotebookTextModel | undefined>;
 	executeNotebook(viewType: string, uri: URI, useAttachedKernel: boolean, token: CancellationToken): Promise<void>;
 	onDidReceiveMessage(uri: URI, message: any): void;