two kernel types

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

@@ -15,11 +15,12 @@ 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 { INotebookKernel, INotebookKernelChangeEvent, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
+import { IResolvedNotebookKernel, INotebookKernelChangeEvent, INotebookKernelService, INotebookProxyKernel, INotebookProxyKernelChangeEvent, ProxyKernelState, NotebookKernelType } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 import { SerializableObjectWithBuffers } from 'vs/workbench/services/extensions/common/proxyIdentifier';
-import { ExtHostContext, ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, ICellExecutionCompleteDto, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape } from '../common/extHost.protocol';
+import { ExtHostContext, ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, ICellExecutionCompleteDto, INotebookKernelDto2, INotebookProxyKernelDto, MainContext, MainThreadNotebookKernelsShape } from '../common/extHost.protocol';
 
-abstract class MainThreadKernel implements INotebookKernel {
+abstract class MainThreadKernel implements IResolvedNotebookKernel {
+	readonly type: NotebookKernelType.Resolved = NotebookKernelType.Resolved;
 
 	private readonly _onDidChange = new Emitter<INotebookKernelChangeEvent>();
 	private readonly preloads: { uri: URI; provides: string[] }[];
@@ -97,6 +98,58 @@ abstract class MainThreadKernel implements INotebookKernel {
 	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 {
 
@@ -104,6 +157,7 @@ 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>();
@@ -243,6 +297,40 @@ 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/common/extHost.api.impl.ts

@@ -1134,6 +1134,10 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 				extHostApiDeprecation.report('notebookConcatTextDocument', extension, 'This proposal is not on track for finalization and will be removed.');
 				return new ExtHostNotebookConcatDocument(extHostNotebookDocuments, extHostDocuments, notebook, selector);
 			},
+			createNotebookProxyController(id: string, notebookType: string, label: string, handler: () => vscode.NotebookController | Thenable<vscode.NotebookController>) {
+				checkProposedApiEnabled(extension, 'notebookProxyController');
+				return extHostNotebookKernels.createNotebookProxyController(extension, id, notebookType, label, handler);
+			}
 		};
 
 		return <typeof vscode>{

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

@@ -979,6 +979,17 @@ export interface INotebookKernelDto2 {
 	preloads?: { uri: UriComponents; provides: string[] }[];
 }
 
+export interface INotebookProxyKernelDto {
+	id: string;
+	notebookType: string;
+	extensionId: ExtensionIdentifier;
+	extensionLocation: UriComponents;
+	label: string;
+	detail?: string;
+	description?: string;
+	kind?: string;
+}
+
 export interface ICellExecuteOutputEditDto {
 	editType: CellExecutionUpdateType.Output;
 	append?: boolean;
@@ -1004,6 +1015,8 @@ 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;
 
@@ -2097,6 +2110,7 @@ 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;
+	$resolveKernel(handle: number): Promise<string | null>;
 }
 
 export interface ExtHostInteractiveShape {

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, MainContext, MainThreadNotebookKernelsShape, NotebookOutputDto } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, IMainContext, INotebookKernelDto2, INotebookProxyKernelDto, 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,6 +34,13 @@ 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 };
@@ -45,6 +52,7 @@ 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>();
@@ -257,8 +265,105 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		return controller;
 	}
 
+	createNotebookProxyController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler: () => vscode.NotebookController | Thenable<vscode.NotebookController>): 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;
+			}
+			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;
+	}
+
 	$acceptNotebookAssociation(handle: number, uri: UriComponents, value: boolean): void {
-		const obj = this._kernelData.get(handle);
+		const obj = this._kernelData.get(handle) ?? this._proxyKernelData.get(handle);
 		if (obj) {
 			// update data structure
 			const notebook = this._extHostNotebook.getNotebookDocument(URI.revive(uri))!;
@@ -346,6 +451,28 @@ 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 (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 {