Update kernels dropdown after installing a notebook extension (#20759)

src/sql/workbench/services/notebook/browser/interface.ts

@@ -23,6 +23,7 @@ export abstract class INotebookInput extends EditorInput {
 	readonly standardKernels: IStandardKernelWithProvider[];
 	readonly providersLoaded: Promise<void>;
 	readonly showActions: boolean;
+	readonly languageMode: string;
 }
 
 export function isINotebookInput(value: any): value is INotebookInput {

src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts

@@ -622,6 +622,7 @@ export interface INotebookModelOptions {
 	notificationService: INotificationService;
 	connectionService: IConnectionManagementService;
 	capabilitiesService: ICapabilitiesService;
+	getInputLanguageMode: () => string;
 	editorLoadedTimestamp?: number;
 }
 

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -23,7 +23,7 @@ import { INotebookEditOperation, NotebookEditOperationType } from 'sql/workbench
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { uriPrefixes } from 'sql/platform/connection/common/utils';
 import { ILogService } from 'vs/platform/log/common/log';
-import { getErrorMessage } from 'vs/base/common/errors';
+import { getErrorMessage, onUnexpectedError } from 'vs/base/common/errors';
 import { notebookConstants } from 'sql/workbench/services/notebook/browser/interfaces';
 import { IAdsTelemetryService, ITelemetryEvent, ITelemetryEventProperties } from 'sql/platform/telemetry/common/telemetry';
 import { Deferred } from 'sql/base/common/promise';
@@ -40,6 +40,8 @@ import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
 import { deepClone } from 'vs/base/common/objects';
 import { DotnetInteractiveDisplayName } from 'sql/workbench/api/common/notebooks/notebookUtils';
 import { IPYKERNEL_DISPLAY_NAME } from 'sql/workbench/common/constants';
+import * as path from 'vs/base/common/path';
+import { IModeService } from 'vs/editor/common/services/modeService';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -134,7 +136,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		@IConnectionManagementService private connectionManagementService: IConnectionManagementService,
 		@IConfigurationService private configurationService: IConfigurationService,
 		@IUndoRedoService private undoService: IUndoRedoService,
-		@ICapabilitiesService private _capabilitiesService?: ICapabilitiesService,
+		@INotebookService private _notebookService: INotebookService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
+		@IModeService private _modeService: IModeService,
 	) {
 		super();
 		if (!_notebookOptions || !_notebookOptions.notebookUri || !_notebookOptions.executeManagers) {
@@ -147,6 +151,45 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			this._notebookOptions.layoutChanged(() => this._layoutChanged.fire());
 		}
 		this._defaultKernel = _notebookOptions.defaultKernel;
+
+		this._register(this._notebookService.onNotebookKernelsAdded(async kernels => this.handleNewKernelsAdded(kernels).catch(error => onUnexpectedError(error))));
+	}
+
+	// Add new kernels to the model's list as they're registered so that we don't
+	// need to restart the notebook to select them in the kernel dropdown.
+	private async handleNewKernelsAdded(kernels: notebookUtils.IStandardKernelWithProvider[]): Promise<void> {
+		// Kernels are file-specific, so we need to check the file extension
+		// to see if the kernel is supported for this notebook.
+		let extensions: string[];
+		let fileExt = path.extname(this._notebookOptions.notebookUri.path);
+		if (!fileExt) {
+			let languageMode = this._notebookOptions.getInputLanguageMode();
+			if (languageMode) {
+				let languageName = this._modeService.getLanguageName(languageMode);
+				let fileExtensions = this._modeService.getExtensions(languageName);
+				if (fileExtensions?.length > 0) {
+					extensions = fileExtensions;
+				} else {
+					this.logService.warn(`Could not retrieve file extensions for language mode '${languageMode}' in notebook '${this._notebookOptions.notebookUri.toString()}'`);
+				}
+			} else {
+				this.logService.warn(`Could not determine language mode for notebook '${this._notebookOptions.notebookUri.toString()}'`);
+			}
+		} else {
+			extensions = [fileExt];
+		}
+		// All kernels from the same provider share the same supported file extensions,
+		// so we only need to check the first one here.
+		if (extensions?.some(ext => kernels[0]?.supportedFileExtensions?.includes(ext))) {
+			this._standardKernels.push(...kernels);
+			this.setDisplayNameMapsForKernels(kernels);
+
+			// Also add corresponding execute manager so that we can change to the new kernels
+			let manager = await this._notebookService.getOrCreateExecuteManager(kernels[0].notebookProvider, this.notebookUri);
+			this._notebookOptions.executeManagers.push(manager);
+
+			this._kernelsChangedEmitter.fire(this._activeClientSession?.kernel);
+		}
 	}
 
 	private get serializationManagers(): ISerializationManager[] {
@@ -401,7 +444,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public set standardKernels(kernels: notebookUtils.IStandardKernelWithProvider[]) {
 		this._standardKernels = kernels;
-		this.setKernelDisplayNameMapsWithStandardKernels();
+		this.setDisplayNameMapsForKernels(kernels);
 	}
 
 	public getApplicableConnectionProviderIds(kernelDisplayName: string): string[] {
@@ -1567,8 +1610,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	 * Set maps with values to have a way to determine the connection
 	 * provider and notebook provider ids from a kernel display name
 	 */
-	private setKernelDisplayNameMapsWithStandardKernels(): void {
-		this._standardKernels.forEach(kernel => {
+	private setDisplayNameMapsForKernels(kernels: notebookUtils.IStandardKernelWithProvider[]): void {
+		kernels.forEach(kernel => {
 			let displayName = kernel.displayName;
 			if (!displayName) {
 				displayName = kernel.name;

src/sql/workbench/services/notebook/browser/models/notebookUtils.ts

@@ -65,6 +65,7 @@ export interface IStandardKernelWithProvider {
 	readonly connectionProviderIds: string[];
 	readonly notebookProvider: string;
 	readonly supportedLanguages: string[];
+	readonly supportedFileExtensions?: string[];
 }
 
 export interface IEndpoint {

src/sql/workbench/services/notebook/browser/notebookService.ts

@@ -23,6 +23,7 @@ import { INotebookShowOptions } from 'sql/workbench/api/common/sqlExtHost.protoc
 import { NotebookViewsExtension } from 'sql/workbench/services/notebook/browser/notebookViews/notebookViewsExtension';
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
 import { ICodeEditorViewState } from 'vs/editor/common/editorCommon';
+import { IStandardKernelWithProvider } from 'sql/workbench/services/notebook/browser/models/notebookUtils';
 
 export const SERVICE_ID = 'sqlNotebookService';
 export const INotebookService = createDecorator<INotebookService>(SERVICE_ID);
@@ -55,6 +56,7 @@ export interface INotebookService {
 	readonly onNotebookEditorAdd: Event<INotebookEditor>;
 	readonly onNotebookEditorRemove: Event<INotebookEditor>;
 	onNotebookEditorRename: Event<INotebookEditor>;
+	readonly onNotebookKernelsAdded: Event<IStandardKernelWithProvider[]>;
 
 	readonly isRegistrationComplete: boolean;
 	readonly registrationComplete: Promise<void>;

src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts

@@ -54,6 +54,7 @@ import { DEFAULT_NB_LANGUAGE_MODE, INTERACTIVE_LANGUAGE_MODE, INTERACTIVE_PROVID
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { SqlSerializationProvider } from 'sql/workbench/services/notebook/browser/sql/sqlSerializationProvider';
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
+import { IStandardKernelWithProvider } from 'sql/workbench/services/notebook/browser/models/notebookUtils';
 
 const languageAssociationRegistry = Registry.as<ILanguageAssociationRegistry>(LanguageAssociationExtensions.LanguageAssociations);
 
@@ -178,6 +179,7 @@ export class NotebookService extends Disposable implements INotebookService {
 	private _onNotebookEditorAdd = new Emitter<INotebookEditor>();
 	private _onNotebookEditorRemove = new Emitter<INotebookEditor>();
 	private _onNotebookEditorRename = new Emitter<INotebookEditor>();
+	private _onNotebookKernelsAdded = new Emitter<IStandardKernelWithProvider[]>();
 	private _editors = new Map<string, INotebookEditor>();
 	private _fileToProviderDescriptions = new Map<string, ProviderDescriptionRegistration[]>();
 	private _providerToStandardKernels = new Map<string, StandardKernelsDescriptor>(); // Note: providerId key here should be in upper case
@@ -426,7 +428,7 @@ export class NotebookService extends Disposable implements INotebookService {
 			if (!this._executeProviders.has(p.id)) {
 				this._executeProviders.set(p.id, new ExecuteProviderDescriptor(p.id));
 			}
-			this.addStandardKernels(registration);
+			this.addStandardKernels(registration, registration.fileExtensions);
 		} else {
 			// Standard kernels might get registered later for VSCode notebooks, so add a descriptor to wait on
 			if (!this._providerToStandardKernels.has(p.id)) {
@@ -506,7 +508,7 @@ export class NotebookService extends Disposable implements INotebookService {
 	// in the kernels dropdown list before a SessionManager has been started; this way,
 	// every NotebookProvider doesn't need to have an active SessionManager in order to contribute
 	// kernels to the dropdown
-	private addStandardKernels(provider: ProviderDescriptionRegistration) {
+	private addStandardKernels(provider: ProviderDescriptionRegistration, supportedFileExtensions?: string[]) {
 		let providerUpperCase = provider.provider.toUpperCase();
 		let descriptor = this._providerToStandardKernels.get(providerUpperCase);
 		if (!descriptor) {
@@ -526,6 +528,20 @@ export class NotebookService extends Disposable implements INotebookService {
 		}
 		descriptor.instance = standardKernels;
 		this._providerToStandardKernels.set(providerUpperCase, descriptor);
+
+		// Emit update event if the provider is not one of the default options
+		if (provider.provider !== SQL_NOTEBOOK_PROVIDER && provider.provider !== JUPYTER_PROVIDER_ID && standardKernels.length > 0) {
+			this._onNotebookKernelsAdded.fire(standardKernels.map(kernel => {
+				return {
+					name: kernel.name,
+					displayName: kernel.displayName,
+					connectionProviderIds: kernel.connectionProviderIds,
+					notebookProvider: provider.provider,
+					supportedLanguages: kernel.supportedLanguages,
+					supportedFileExtensions: supportedFileExtensions
+				};
+			}));
+		}
 	}
 
 	getSupportedFileExtensions(): string[] {
@@ -632,6 +648,10 @@ export class NotebookService extends Disposable implements INotebookService {
 		return this._onNotebookEditorRename.event;
 	}
 
+	get onNotebookKernelsAdded(): Event<IStandardKernelWithProvider[]> {
+		return this._onNotebookKernelsAdded.event;
+	}
+
 	addNotebookEditor(editor: INotebookEditor): void {
 		this._editors.set(editor.id, editor);
 		this._onNotebookEditorAdd.fire(editor);