extract exthost/mainthread proxy kernels.

src/vs/workbench/api/browser/extensionHost.contribution.ts

@@ -64,6 +64,7 @@ import './mainThreadWorkspace';
 import './mainThreadComments';
 import './mainThreadNotebook';
 import './mainThreadNotebookKernels';
+import './mainThreadNotebookProxyKernels';
 import './mainThreadNotebookDocumentsAndEditors';
 import './mainThreadNotebookRenderers';
 import './mainThreadInteractive';

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

@@ -15,9 +15,9 @@ import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/ext
 import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/notebookEditorService';
 import { INotebookCellExecution, INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
-import { IResolvedNotebookKernel, INotebookKernelChangeEvent, INotebookKernelService, INotebookProxyKernel, INotebookProxyKernelChangeEvent, ProxyKernelState, NotebookKernelType } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
+import { IResolvedNotebookKernel, INotebookKernelChangeEvent, INotebookKernelService, NotebookKernelType } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 import { SerializableObjectWithBuffers } from 'vs/workbench/services/extensions/common/proxyIdentifier';
-import { ExtHostContext, ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, ICellExecutionCompleteDto, INotebookKernelDto2, INotebookProxyKernelDto, MainContext, MainThreadNotebookKernelsShape } from '../common/extHost.protocol';
+import { ExtHostContext, ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, ICellExecutionCompleteDto, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape } from '../common/extHost.protocol';
 
 abstract class MainThreadKernel implements IResolvedNotebookKernel {
 	readonly type: NotebookKernelType.Resolved = NotebookKernelType.Resolved;
@@ -98,58 +98,6 @@ abstract class MainThreadKernel implements IResolvedNotebookKernel {
 	abstract cancelNotebookCellExecution(uri: URI, cellHandles: number[]): Promise<void>;
 }
 
-abstract class MainThreadProxyKernel implements INotebookProxyKernel {
-	readonly type: NotebookKernelType.Proxy = NotebookKernelType.Proxy;
-	protected readonly _onDidChange = new Emitter<INotebookProxyKernelChangeEvent>();
-	readonly onDidChange: Event<INotebookProxyKernelChangeEvent> = this._onDidChange.event;
-	readonly id: string;
-	readonly viewType: string;
-	readonly extension: ExtensionIdentifier;
-	label: string;
-	description?: string;
-	detail?: string;
-	kind?: string;
-	supportedLanguages: string[] = [];
-	connectionState: ProxyKernelState;
-
-	constructor(data: INotebookProxyKernelDto) {
-		this.id = data.id;
-		this.viewType = data.notebookType;
-		this.extension = data.extensionId;
-
-		this.label = data.label;
-		this.description = data.description;
-		this.detail = data.detail;
-		this.kind = data.kind;
-
-		this.connectionState = ProxyKernelState.Disconnected;
-	}
-
-	update(data: Partial<INotebookProxyKernel>) {
-		const event: INotebookProxyKernelChangeEvent = Object.create(null);
-		if (data.label !== undefined) {
-			this.label = data.label;
-			event.label = true;
-		}
-		if (data.description !== undefined) {
-			this.description = data.description;
-			event.description = true;
-		}
-		if (data.detail !== undefined) {
-			this.detail = data.detail;
-			event.detail = true;
-		}
-		if (data.kind !== undefined) {
-			this.kind = data.kind;
-			event.kind = true;
-		}
-
-		this._onDidChange.fire(event);
-	}
-
-	abstract resolveKernel(): Promise<string | null>;
-}
-
 @extHostNamedCustomer(MainContext.MainThreadNotebookKernels)
 export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape {
 
@@ -157,7 +105,6 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 	private readonly _disposables = new DisposableStore();
 
 	private readonly _kernels = new Map<number, [kernel: MainThreadKernel, registraion: IDisposable]>();
-	private readonly _proxyKernels = new Map<number, [kernel: MainThreadProxyKernel, registraion: IDisposable]>();
 	private readonly _proxy: ExtHostNotebookKernelsShape;
 
 	private readonly _executions = new Map<number, INotebookCellExecution>();
@@ -297,40 +244,6 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 		}
 	}
 
-	// -- Proxy kernel
-
-	async $addProxyKernel(handle: number, data: INotebookProxyKernelDto): Promise<void> {
-		const that = this;
-		const proxyKernel = new class extends MainThreadProxyKernel {
-			async resolveKernel(): Promise<string | null> {
-				this.connectionState = ProxyKernelState.Initializing;
-				this._onDidChange.fire({ connectionState: true });
-				const delegateKernel = await that._proxy.$resolveKernel(handle);
-				this.connectionState = ProxyKernelState.Connected;
-				this._onDidChange.fire({ connectionState: true });
-				return delegateKernel;
-			}
-		}(data);
-
-		const listener = this._notebookKernelService.onDidChangeSelectedNotebooks(e => {
-			if (e.oldKernel === proxyKernel.id) {
-				this._proxy.$acceptNotebookAssociation(handle, e.notebook, false);
-			} else if (e.newKernel === proxyKernel.id) {
-				this._proxy.$acceptNotebookAssociation(handle, e.notebook, true);
-			}
-		});
-
-		const registration = this._notebookKernelService.registerProxyKernel(proxyKernel);
-		this._proxyKernels.set(handle, [proxyKernel, combinedDisposable(listener, registration)]);
-	}
-
-	$updateProxyKernel(handle: number, data: Partial<INotebookProxyKernelDto>): void {
-		const tuple = this._proxyKernels.get(handle);
-		if (tuple) {
-			tuple[0].update(data);
-		}
-	}
-
 	// --- execution
 
 	$createExecution(handle: number, controllerId: string, rawUri: UriComponents, cellHandle: number): void {

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

@@ -0,0 +1,122 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Emitter, Event } from 'vs/base/common/event';
+import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
+import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
+import { INotebookKernelService, INotebookProxyKernel, INotebookProxyKernelChangeEvent, ProxyKernelState, NotebookKernelType } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
+import { ExtHostContext, ExtHostNotebookProxyKernelsShape, INotebookProxyKernelDto, MainContext, MainThreadNotebookProxyKernelsShape } from '../common/extHost.protocol';
+
+abstract class MainThreadProxyKernel implements INotebookProxyKernel {
+	readonly type: NotebookKernelType.Proxy = NotebookKernelType.Proxy;
+	protected readonly _onDidChange = new Emitter<INotebookProxyKernelChangeEvent>();
+	readonly onDidChange: Event<INotebookProxyKernelChangeEvent> = this._onDidChange.event;
+	readonly id: string;
+	readonly viewType: string;
+	readonly extension: ExtensionIdentifier;
+	readonly preloadProvides: string[] = [];
+	label: string;
+	description?: string;
+	detail?: string;
+	kind?: string;
+	supportedLanguages: string[] = [];
+	connectionState: ProxyKernelState;
+
+	constructor(data: INotebookProxyKernelDto) {
+		this.id = data.id;
+		this.viewType = data.notebookType;
+		this.extension = data.extensionId;
+
+		this.label = data.label;
+		this.description = data.description;
+		this.detail = data.detail;
+		this.kind = data.kind;
+
+		this.connectionState = ProxyKernelState.Disconnected;
+	}
+
+	update(data: Partial<INotebookProxyKernel>) {
+		const event: INotebookProxyKernelChangeEvent = Object.create(null);
+		if (data.label !== undefined) {
+			this.label = data.label;
+			event.label = true;
+		}
+		if (data.description !== undefined) {
+			this.description = data.description;
+			event.description = true;
+		}
+		if (data.detail !== undefined) {
+			this.detail = data.detail;
+			event.detail = true;
+		}
+		if (data.kind !== undefined) {
+			this.kind = data.kind;
+			event.kind = true;
+		}
+
+		this._onDidChange.fire(event);
+	}
+
+	abstract resolveKernel(): Promise<string | null>;
+}
+
+@extHostNamedCustomer(MainContext.MainThreadNotebookProxyKernels)
+export class MainThreadNotebookProxyKernels implements MainThreadNotebookProxyKernelsShape {
+
+	private readonly _disposables = new DisposableStore();
+
+	private readonly _proxyKernels = new Map<number, [kernel: MainThreadProxyKernel, registraion: IDisposable]>();
+	private readonly _proxyKernelProxy: ExtHostNotebookProxyKernelsShape;
+
+	constructor(
+		extHostContext: IExtHostContext,
+		@INotebookKernelService private readonly _notebookKernelService: INotebookKernelService,
+	) {
+		this._proxyKernelProxy = extHostContext.getProxy(ExtHostContext.ExtHostNotebookProxyKernels);
+	}
+
+	dispose(): void {
+		this._disposables.dispose();
+
+		for (let [, registration] of this._proxyKernels.values()) {
+			registration.dispose();
+		}
+	}
+
+	// -- Proxy kernel
+
+	async $addProxyKernel(handle: number, data: INotebookProxyKernelDto): Promise<void> {
+		const that = this;
+		const proxyKernel = new class extends MainThreadProxyKernel {
+			async resolveKernel(): Promise<string | null> {
+				this.connectionState = ProxyKernelState.Initializing;
+				this._onDidChange.fire({ connectionState: true });
+				const delegateKernel = await that._proxyKernelProxy.$resolveKernel(handle);
+				this.connectionState = ProxyKernelState.Connected;
+				this._onDidChange.fire({ connectionState: true });
+				return delegateKernel;
+			}
+		}(data);
+
+		const registration = this._notebookKernelService.registerKernel(proxyKernel);
+		this._proxyKernels.set(handle, [proxyKernel, registration]);
+	}
+
+	$updateProxyKernel(handle: number, data: Partial<INotebookProxyKernelDto>): void {
+		const tuple = this._proxyKernels.get(handle);
+		if (tuple) {
+			tuple[0].update(data);
+		}
+	}
+
+	$removeProxyKernel(handle: number): void {
+		const tuple = this._proxyKernels.get(handle);
+		if (tuple) {
+			tuple[1].dispose();
+			this._proxyKernels.delete(handle);
+		}
+	}
+}

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

@@ -94,6 +94,7 @@ import { ExtHostInteractive } from 'vs/workbench/api/common/extHostInteractive';
 import { combinedDisposable } from 'vs/base/common/lifecycle';
 import { checkProposedApiEnabled, isProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
 import { DebugConfigurationProviderTriggerKind } from 'vs/workbench/contrib/debug/common/debug';
+import { ExtHostNotebookProxyKernels } from 'vs/workbench/api/common/extHostNotebookProxyKernels';
 
 export interface IExtensionApiFactory {
 	(extension: IExtensionDescription, registry: ExtensionDescriptionRegistry, configProvider: ExtHostConfigProvider): typeof vscode;
@@ -156,6 +157,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 	const extHostNotebookDocuments = rpcProtocol.set(ExtHostContext.ExtHostNotebookDocuments, new ExtHostNotebookDocuments(extHostNotebook));
 	const extHostNotebookEditors = rpcProtocol.set(ExtHostContext.ExtHostNotebookEditors, new ExtHostNotebookEditors(extHostLogService, rpcProtocol, extHostNotebook));
 	const extHostNotebookKernels = rpcProtocol.set(ExtHostContext.ExtHostNotebookKernels, new ExtHostNotebookKernels(rpcProtocol, initData, extHostNotebook, extHostCommands, extHostLogService));
+	const extHostNotebookProxyKernels = rpcProtocol.set(ExtHostContext.ExtHostNotebookProxyKernels, new ExtHostNotebookProxyKernels(rpcProtocol, extHostNotebookKernels, extHostLogService));
 	const extHostNotebookRenderers = rpcProtocol.set(ExtHostContext.ExtHostNotebookRenderers, new ExtHostNotebookRenderers(rpcProtocol, extHostNotebook));
 	const extHostEditors = rpcProtocol.set(ExtHostContext.ExtHostEditors, new ExtHostEditors(rpcProtocol, extHostDocumentsAndEditors));
 	const extHostTreeViews = rpcProtocol.set(ExtHostContext.ExtHostTreeViews, new ExtHostTreeViews(rpcProtocol.getProxy(MainContext.MainThreadTreeViews), extHostCommands, extHostLogService));
@@ -1135,7 +1137,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			},
 			createNotebookProxyController(id: string, notebookType: string, label: string, handler: () => vscode.NotebookController | string | Thenable<vscode.NotebookController | string>) {
 				checkProposedApiEnabled(extension, 'notebookProxyController');
-				return extHostNotebookKernels.createNotebookProxyController(extension, id, notebookType, label, handler);
+				return extHostNotebookProxyKernels.createNotebookProxyController(extension, id, notebookType, label, handler);
 			}
 		};
 

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

@@ -1018,8 +1018,6 @@ export interface MainThreadNotebookKernelsShape extends IDisposable {
 	$postMessage(handle: number, editorId: string | undefined, message: any): Promise<boolean>;
 	$addKernel(handle: number, data: INotebookKernelDto2): Promise<void>;
 	$updateKernel(handle: number, data: Partial<INotebookKernelDto2>): void;
-	$addProxyKernel(handle: number, data: INotebookProxyKernelDto): Promise<void>;
-	$updateProxyKernel(handle: number, data: Partial<INotebookProxyKernelDto>): void;
 	$removeKernel(handle: number): void;
 	$updateNotebookPriority(handle: number, uri: UriComponents, value: number | undefined): void;
 
@@ -1028,6 +1026,12 @@ export interface MainThreadNotebookKernelsShape extends IDisposable {
 	$completeExecution(handle: number, data: SerializableObjectWithBuffers<ICellExecutionCompleteDto>): void;
 }
 
+export interface MainThreadNotebookProxyKernelsShape extends IDisposable {
+	$addProxyKernel(handle: number, data: INotebookProxyKernelDto): Promise<void>;
+	$updateProxyKernel(handle: number, data: Partial<INotebookProxyKernelDto>): void;
+	$removeProxyKernel(handle: number): void;
+}
+
 export interface MainThreadNotebookRenderersShape extends IDisposable {
 	$postMessage(editorId: string | undefined, rendererId: string, message: unknown): Promise<boolean>;
 }
@@ -2114,6 +2118,9 @@ export interface ExtHostNotebookKernelsShape {
 	$cancelCells(handle: number, uri: UriComponents, handles: number[]): Promise<void>;
 	$acceptKernelMessageFromRenderer(handle: number, editorId: string, message: any): void;
 	$cellExecutionChanged(uri: UriComponents, cellHandle: number, state: notebookCommon.NotebookCellExecutionState | undefined): void;
+}
+
+export interface ExtHostNotebookProxyKernelsShape {
 	$resolveKernel(handle: number): Promise<string | null>;
 }
 
@@ -2293,6 +2300,7 @@ export const MainContext = {
 	MainThreadNotebookDocuments: createProxyIdentifier<MainThreadNotebookDocumentsShape>('MainThreadNotebookDocumentsShape'),
 	MainThreadNotebookEditors: createProxyIdentifier<MainThreadNotebookEditorsShape>('MainThreadNotebookEditorsShape'),
 	MainThreadNotebookKernels: createProxyIdentifier<MainThreadNotebookKernelsShape>('MainThreadNotebookKernels'),
+	MainThreadNotebookProxyKernels: createProxyIdentifier<MainThreadNotebookProxyKernelsShape>('MainThreadNotebookProxyKernels'),
 	MainThreadNotebookRenderers: createProxyIdentifier<MainThreadNotebookRenderersShape>('MainThreadNotebookRenderers'),
 	MainThreadInteractive: createProxyIdentifier<MainThreadInteractiveShape>('MainThreadInteractive'),
 	MainThreadTheming: createProxyIdentifier<MainThreadThemingShape>('MainThreadTheming'),
@@ -2345,6 +2353,7 @@ export const ExtHostContext = {
 	ExtHostNotebookDocuments: createProxyIdentifier<ExtHostNotebookDocumentsShape>('ExtHostNotebookDocuments'),
 	ExtHostNotebookEditors: createProxyIdentifier<ExtHostNotebookEditorsShape>('ExtHostNotebookEditors'),
 	ExtHostNotebookKernels: createProxyIdentifier<ExtHostNotebookKernelsShape>('ExtHostNotebookKernels'),
+	ExtHostNotebookProxyKernels: createProxyIdentifier<ExtHostNotebookProxyKernelsShape>('ExtHostNotebookProxyKernels'),
 	ExtHostNotebookRenderers: createProxyIdentifier<ExtHostNotebookRenderersShape>('ExtHostNotebookRenderers'),
 	ExtHostInteractive: createProxyIdentifier<ExtHostInteractiveShape>('ExtHostInteractive'),
 	ExtHostTheming: createProxyIdentifier<ExtHostThemingShape>('ExtHostTheming'),

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

@@ -12,7 +12,7 @@ import { ResourceMap } from 'vs/base/common/map';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { ExtensionIdentifier, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { ILogService } from 'vs/platform/log/common/log';
-import { ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, IMainContext, INotebookKernelDto2, INotebookProxyKernelDto, MainContext, MainThreadNotebookKernelsShape, NotebookOutputDto } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, IMainContext, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape, NotebookOutputDto } from 'vs/workbench/api/common/extHost.protocol';
 import { ApiCommand, ApiCommandArgument, ApiCommandResult, ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
 import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitDataService';
 import { ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
@@ -34,13 +34,6 @@ interface IKernelData {
 	associatedNotebooks: ResourceMap<boolean>;
 }
 
-interface IProxyKernelData {
-	extensionId: ExtensionIdentifier;
-	controller: vscode.NotebookProxyController;
-	onDidChangeSelection: Emitter<{ selected: boolean; notebook: vscode.NotebookDocument }>;
-	associatedNotebooks: ResourceMap<boolean>;
-}
-
 type ExtHostSelectKernelArgs = ControllerInfo | { notebookEditor: vscode.NotebookEditor } | ControllerInfo & { notebookEditor: vscode.NotebookEditor } | undefined;
 export type SelectKernelReturnArgs = ControllerInfo | { notebookEditorId: string } | ControllerInfo & { notebookEditorId: string } | undefined;
 type ControllerInfo = { id: string; extension: string };
@@ -52,7 +45,6 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 	private readonly _activeExecutions = new ResourceMap<NotebookCellExecutionTask>();
 
 	private readonly _kernelData = new Map<number, IKernelData>();
-	private readonly _proxyKernelData: Map<number, IProxyKernelData> = new Map<number, IProxyKernelData>();
 	private _handlePool: number = 0;
 
 	private readonly _onDidChangeCellExecutionState = new Emitter<vscode.NotebookCellExecutionStateChangeEvent>();
@@ -116,7 +108,7 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		const onDidReceiveMessage = new Emitter<{ editor: vscode.NotebookEditor; message: any }>();
 
 		const data: INotebookKernelDto2 = {
-			id: createKernelId(extension, id),
+			id: createKernelId(extension.identifier, id),
 			notebookType: viewType,
 			extensionId: extension.identifier,
 			extensionLocation: extension.extensionLocation,
@@ -226,7 +218,7 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 					that._logService.trace(`NotebookController[${handle}] NOT associated to notebook, associated to THESE notebooks:`, Array.from(associatedNotebooks.keys()).map(u => u.toString()));
 					throw new Error(`notebook controller is NOT associated to notebook: ${cell.notebook.uri.toString()}`);
 				}
-				return that._createNotebookCellExecution(cell, createKernelId(extension, this.id));
+				return that._createNotebookCellExecution(cell, createKernelId(extension.identifier, this.id));
 			},
 			dispose: () => {
 				if (!isDisposed) {
@@ -265,105 +257,18 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		return controller;
 	}
 
-	createNotebookProxyController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler: () => vscode.NotebookController | string | Thenable<vscode.NotebookController | string>): vscode.NotebookProxyController {
-		const handle = this._handlePool++;
-
-		let isDisposed = false;
-		const commandDisposables = new DisposableStore();
-		const onDidChangeSelection = new Emitter<{ selected: boolean; notebook: vscode.NotebookDocument }>();
-
-		const data: INotebookProxyKernelDto = {
-			id: createKernelId(extension, id),
-			notebookType: viewType,
-			extensionId: extension.identifier,
-			extensionLocation: extension.extensionLocation,
-			label: label || extension.identifier.value,
-		};
-
-		let _resolveHandler = handler;
-
-		this._proxy.$addProxyKernel(handle, data).catch(err => {
-			// this can happen when a kernel with that ID is already registered
-			console.log(err);
-			isDisposed = true;
-		});
-
-		let tokenPool = 0;
-		const _update = () => {
-			if (isDisposed) {
-				return;
+	getIdByController(controller: vscode.NotebookController) {
+		// return this._kernelData.(data => data.controller === controller);
+		for (const [_, candidate] of this._kernelData) {
+			if (candidate.controller === controller) {
+				return createKernelId(candidate.extensionId, controller.id);
 			}
-			const myToken = ++tokenPool;
-			Promise.resolve().then(() => {
-				if (myToken === tokenPool) {
-					this._proxy.$updateProxyKernel(handle, data);
-				}
-			});
-		};
-
-		// notebook documents that are associated to this controller
-		const associatedNotebooks = new ResourceMap<boolean>();
-
-		const controller: vscode.NotebookProxyController = {
-			get id() { return id; },
-			get notebookType() { return data.notebookType; },
-			onDidChangeSelectedNotebooks: onDidChangeSelection.event,
-			get label() {
-				return data.label;
-			},
-			set label(value) {
-				data.label = value ?? extension.displayName ?? extension.name;
-				_update();
-			},
-			get detail() {
-				return data.detail ?? '';
-			},
-			set detail(value) {
-				data.detail = value;
-				_update();
-			},
-			get description() {
-				return data.description ?? '';
-			},
-			set description(value) {
-				data.description = value;
-				_update();
-			},
-			get kind() {
-				checkProposedApiEnabled(extension, 'notebookControllerKind');
-				return data.kind ?? '';
-			},
-			set kind(value) {
-				checkProposedApiEnabled(extension, 'notebookControllerKind');
-				data.kind = value;
-				_update();
-			},
-			get resolveHandler() {
-				return _resolveHandler;
-			},
-			dispose: () => {
-				if (!isDisposed) {
-					this._logService.trace(`NotebookController[${handle}], DISPOSED`);
-					isDisposed = true;
-					this._kernelData.delete(handle);
-					commandDisposables.dispose();
-					onDidChangeSelection.dispose();
-					this._proxy.$removeKernel(handle);
-				}
-			}
-		};
-
-		this._proxyKernelData.set(handle, {
-			extensionId: extension.identifier,
-			controller,
-			onDidChangeSelection,
-			associatedNotebooks
-		});
-		return controller;
+		}
+		return null;
 	}
 
 	$acceptNotebookAssociation(handle: number, uri: UriComponents, value: boolean): void {
-		const obj = this._kernelData.get(handle) ?? this._proxyKernelData.get(handle);
+		const obj = this._kernelData.get(handle);
 		if (obj) {
 			// update data structure
 			const notebook = this._extHostNotebook.getNotebookDocument(URI.revive(uri))!;
@@ -451,34 +356,6 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		}
 	}
 
-	async $resolveKernel(handle: number): Promise<string | null> {
-		const obj = this._proxyKernelData.get(handle);
-		if (!obj) {
-			// extension can dispose kernels in the meantime
-			return null;
-		}
-
-		const controller = await obj.controller.resolveHandler();
-		let matchedKernelData: IKernelData | undefined;
-		this._kernelData.forEach(d => {
-			if (typeof controller === 'string') {
-				if (d.controller.id === controller) {
-					matchedKernelData = d;
-				}
-			} else {
-				if (d.controller.id === controller.id) {
-					matchedKernelData = d;
-				}
-			}
-		});
-
-		if (matchedKernelData) {
-			return `${matchedKernelData.extensionId.value}/${matchedKernelData.controller.id}`;
-		}
-
-		return null;
-	}
-
 	// ---
 
 	_createNotebookCellExecution(cell: vscode.NotebookCell, controllerId: string): vscode.NotebookCellExecution {
@@ -716,6 +593,6 @@ class TimeoutBasedCollector<T> {
 	}
 }
 
-function createKernelId(extension: IExtensionDescription, id: string): string {
-	return `${extension.identifier.value}/${id}`;
+export function createKernelId(extensionIdentifier: ExtensionIdentifier, id: string): string {
+	return `${extensionIdentifier.value}/${id}`;
 }

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

@@ -0,0 +1,157 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Emitter } from 'vs/base/common/event';
+import { DisposableStore } from 'vs/base/common/lifecycle';
+import { ResourceMap } from 'vs/base/common/map';
+import { ExtensionIdentifier, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
+import { ILogService } from 'vs/platform/log/common/log';
+import { ExtHostNotebookProxyKernelsShape, IMainContext, INotebookProxyKernelDto, MainContext, MainThreadNotebookProxyKernelsShape } from 'vs/workbench/api/common/extHost.protocol';
+import { createKernelId, ExtHostNotebookKernels } from 'vs/workbench/api/common/extHostNotebookKernels';
+import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
+import * as vscode from 'vscode';
+
+interface IProxyKernelData {
+	extensionId: ExtensionIdentifier;
+	controller: vscode.NotebookProxyController;
+	onDidChangeSelection: Emitter<{ selected: boolean; notebook: vscode.NotebookDocument }>;
+	associatedNotebooks: ResourceMap<boolean>;
+}
+
+export type SelectKernelReturnArgs = ControllerInfo | { notebookEditorId: string } | ControllerInfo & { notebookEditorId: string } | undefined;
+type ControllerInfo = { id: string; extension: string };
+
+
+export class ExtHostNotebookProxyKernels implements ExtHostNotebookProxyKernelsShape {
+
+	private readonly _proxy: MainThreadNotebookProxyKernelsShape;
+
+	private readonly _proxyKernelData: Map<number, IProxyKernelData> = new Map<number, IProxyKernelData>();
+	private _handlePool: number = 0;
+
+	private readonly _onDidChangeCellExecutionState = new Emitter<vscode.NotebookCellExecutionStateChangeEvent>();
+	readonly onDidChangeNotebookCellExecutionState = this._onDidChangeCellExecutionState.event;
+
+	constructor(
+		mainContext: IMainContext,
+		private readonly extHostNotebook: ExtHostNotebookKernels,
+		@ILogService private readonly _logService: ILogService
+	) {
+		this._proxy = mainContext.getProxy(MainContext.MainThreadNotebookProxyKernels);
+	}
+
+	createNotebookProxyController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler: () => vscode.NotebookController | string | Thenable<vscode.NotebookController | string>): vscode.NotebookProxyController {
+		const handle = this._handlePool++;
+
+		let isDisposed = false;
+		const commandDisposables = new DisposableStore();
+		const onDidChangeSelection = new Emitter<{ selected: boolean; notebook: vscode.NotebookDocument }>();
+
+		const data: INotebookProxyKernelDto = {
+			id: createKernelId(extension.identifier, id),
+			notebookType: viewType,
+			extensionId: extension.identifier,
+			extensionLocation: extension.extensionLocation,
+			label: label || extension.identifier.value,
+		};
+
+		let _resolveHandler = handler;
+
+		this._proxy.$addProxyKernel(handle, data).catch(err => {
+			// this can happen when a kernel with that ID is already registered
+			console.log(err);
+			isDisposed = true;
+		});
+
+		let tokenPool = 0;
+		const _update = () => {
+			if (isDisposed) {
+				return;
+			}
+			const myToken = ++tokenPool;
+			Promise.resolve().then(() => {
+				if (myToken === tokenPool) {
+					this._proxy.$updateProxyKernel(handle, data);
+				}
+			});
+		};
+
+		// notebook documents that are associated to this controller
+		const associatedNotebooks = new ResourceMap<boolean>();
+
+		const controller: vscode.NotebookProxyController = {
+			get id() { return id; },
+			get notebookType() { return data.notebookType; },
+			onDidChangeSelectedNotebooks: onDidChangeSelection.event,
+			get label() {
+				return data.label;
+			},
+			set label(value) {
+				data.label = value ?? extension.displayName ?? extension.name;
+				_update();
+			},
+			get detail() {
+				return data.detail ?? '';
+			},
+			set detail(value) {
+				data.detail = value;
+				_update();
+			},
+			get description() {
+				return data.description ?? '';
+			},
+			set description(value) {
+				data.description = value;
+				_update();
+			},
+			get kind() {
+				checkProposedApiEnabled(extension, 'notebookControllerKind');
+				return data.kind ?? '';
+			},
+			set kind(value) {
+				checkProposedApiEnabled(extension, 'notebookControllerKind');
+				data.kind = value;
+				_update();
+			},
+			get resolveHandler() {
+				return _resolveHandler;
+			},
+			dispose: () => {
+				if (!isDisposed) {
+					this._logService.trace(`NotebookProxyController[${handle}], DISPOSED`);
+					isDisposed = true;
+					this._proxyKernelData.delete(handle);
+					commandDisposables.dispose();
+					onDidChangeSelection.dispose();
+					this._proxy.$removeProxyKernel(handle);
+				}
+			}
+		};
+
+		this._proxyKernelData.set(handle, {
+			extensionId: extension.identifier,
+			controller,
+			onDidChangeSelection,
+			associatedNotebooks
+		});
+		return controller;
+	}
+
+	async $resolveKernel(handle: number): Promise<string | null> {
+		const obj = this._proxyKernelData.get(handle);
+		if (!obj) {
+			// extension can dispose kernels in the meantime
+			return null;
+		}
+
+		const controller = await obj.controller.resolveHandler();
+		if (typeof controller === 'string') {
+			return controller;
+		} else {
+			return this.extHostNotebook.getIdByController(controller);
+		}
+	}
+}
+