Add api command for fetching notebook variables (microsoft#205224)

Author: rebornix

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, MainContext, MainThreadNotebookKernelsShape, NotebookOutputDto, VariablesResult } 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';
@@ -90,7 +90,24 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 				})
 			],
 			ApiCommandResult.Void);
+
+		const requestKernelVariablesApiCommand = new ApiCommand(
+			'vscode.executeNotebookVariableProvider',
+			'_executeNotebookVariableProvider',
+			'Execute notebook variable provider',
+			[ApiCommandArgument.Uri],
+			new ApiCommandResult<VariablesResult[], vscode.Variable[]>('A promise that resolves to an array of variables', (value, apiArgs) => {
+				return value.map(variable => {
+					return {
+						name: variable.name,
+						value: variable.value,
+						editable: false
+					};
+				});
+			})
+		);
 		this._commands.registerApiCommand(selectKernelApiCommand);
+		this._commands.registerApiCommand(requestKernelVariablesApiCommand);
 	}
 
 	createNotebookController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable<void>, preloads?: vscode.NotebookRendererScript[]): vscode.NotebookController {

src/vs/workbench/contrib/notebook/browser/contrib/notebookVariables/notebookVariableCommands.ts

@@ -3,11 +3,15 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { CancellationToken } from 'vs/base/common/cancellation';
+import { URI, UriComponents } from 'vs/base/common/uri';
 import { localize } from 'vs/nls';
 import { Action2, registerAction2 } from 'vs/platform/actions/common/actions';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
 import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
 import { contextMenuArg } from 'vs/workbench/contrib/notebook/browser/contrib/notebookVariables/notebookVariablesView';
+import { INotebookKernelService, VariablesResult } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
+import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 
 export const COPY_NOTEBOOK_VARIABLE_VALUE_ID = 'workbench.debug.viewlet.action.copyWorkspaceVariableValue';
 export const COPY_NOTEBOOK_VARIABLE_VALUE_LABEL = localize('copyWorkspaceVariableValue', "Copy Value");
@@ -28,3 +32,39 @@ registerAction2(class extends Action2 {
 		}
 	}
 });
+
+
+registerAction2(class extends Action2 {
+	constructor() {
+		super({
+			id: '_executeNotebookVariableProvider',
+			title: localize('executeNotebookVariableProvider', "Execute Notebook Variable Provider"),
+			f1: false,
+		});
+	}
+
+	async run(accessor: ServicesAccessor, resource: UriComponents | undefined): Promise<VariablesResult[]> {
+		if (!resource) {
+			return [];
+		}
+
+		const uri = URI.revive(resource);
+		const notebookKernelService = accessor.get(INotebookKernelService);
+		const notebookService = accessor.get(INotebookService);
+		const notebookTextModel = notebookService.getNotebookTextModel(uri);
+
+		if (!notebookTextModel) {
+			return [];
+		}
+
+		const selectedKernel = notebookKernelService.getMatchingKernel(notebookTextModel).selected;
+		if (selectedKernel && selectedKernel.hasVariableProvider) {
+			const variables = selectedKernel.provideVariables(notebookTextModel.uri, undefined, 'named', 0, CancellationToken.None);
+			return await variables
+				.map(variable => { return variable; })
+				.toPromise();
+		}
+
+		return [];
+	}
+});

src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts

@@ -176,6 +176,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	private _strategy: EditStrategy | undefined;
 	private _sessionCtor: CancelablePromise<void> | undefined;
 	private _activeSession?: Session;
+	private _warmupRequestCts?: CancellationTokenSource;
 	private readonly _ctxHasActiveRequest: IContextKey<boolean>;
 	private readonly _ctxCellWidgetFocused: IContextKey<boolean>;
 	private readonly _ctxUserDidEdit: IContextKey<boolean>;
@@ -275,6 +276,10 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		inlineChatWidget.placeholder = localize('default.placeholder', "Ask a question");
 		inlineChatWidget.updateInfo(localize('welcome.1', "AI-generated code may be incorrect"));
 		widgetContainer.appendChild(inlineChatWidget.domNode);
+		this._widgetDisposableStore.add(inlineChatWidget.onDidChangeInput(() => {
+			this._warmupRequestCts?.dispose(true);
+			this._warmupRequestCts = undefined;
+		}));
 
 		this._notebookEditor.changeViewZones(accessor => {
 			const notebookViewZone = {
@@ -307,6 +312,8 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 
 				if (fakeParentEditor.hasModel()) {
 					await this._startSession(fakeParentEditor, token);
+					this._warmupRequestCts = new CancellationTokenSource();
+					this._startInitialFolowups(fakeParentEditor, this._warmupRequestCts.token);
 
 					if (this._widget) {
 						this._widget.inlineChatWidget.placeholder = this._activeSession?.session.placeholder ?? localize('default.placeholder', "Ask a question");
@@ -368,6 +375,8 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	async acceptInput() {
 		assertType(this._activeSession);
 		assertType(this._widget);
+		this._warmupRequestCts?.dispose(true);
+		this._warmupRequestCts = undefined;
 		this._activeSession.addInput(new SessionPrompt(this._widget.inlineChatWidget.value));
 
 		assertType(this._activeSession.lastInput);
@@ -559,6 +568,50 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		this._strategy = new EditStrategy(session);
 	}
 
+	private async _startInitialFolowups(editor: IActiveCodeEditor, token: CancellationToken) {
+		if (!this._activeSession || !this._activeSession.provider.provideFollowups) {
+			return;
+		}
+
+		const request: IInlineChatRequest = {
+			requestId: generateUuid(),
+			prompt: '',
+			attempt: 0,
+			selection: { selectionStartLineNumber: 1, selectionStartColumn: 1, positionLineNumber: 1, positionColumn: 1 },
+			wholeRange: { startLineNumber: 1, startColumn: 1, endLineNumber: 1, endColumn: 1 },
+			live: true,
+			previewDocument: editor.getModel().uri,
+			withIntentDetection: true
+		};
+
+		const progress = new AsyncProgress<IInlineChatProgressItem>(async data => { });
+		const task = this._activeSession.provider.provideResponse(this._activeSession.session, request, progress, token);
+		const reply = await raceCancellationError(Promise.resolve(task), token);
+		if (token.isCancellationRequested) {
+			return;
+		}
+
+		if (!reply) {
+			return;
+		}
+
+		const markdownContents = new MarkdownString('', { supportThemeIcons: true, supportHtml: true, isTrusted: false });
+		const response = this._instantiationService.createInstance(ReplyResponse, reply, markdownContents, this._activeSession.textModelN.uri, this._activeSession.textModelN.getAlternativeVersionId(), [], request.requestId);
+		const followups = await this._activeSession.provider.provideFollowups(this._activeSession.session, response.raw, token);
+		if (followups && this._widget) {
+			const widget = this._widget;
+			widget.inlineChatWidget.updateFollowUps(followups, async followup => {
+				if (followup.kind === 'reply') {
+					widget.inlineChatWidget.value = followup.message;
+					this.acceptInput();
+				} else {
+					await this.acceptSession();
+					this._commandService.executeCommand(followup.commandId, ...(followup.args ?? []));
+				}
+			});
+		}
+	}
+
 	private async _makeChanges(edits: TextEdit[], opts: ProgressingEditsOptions | undefined) {
 		assertType(this._activeSession);
 		assertType(this._strategy);