Implement chat agent variable provider

src/vs/editor/common/core/wordHelper.ts

@@ -96,6 +96,8 @@ export function setDefaultGetWordAtTextConfig(value: IGetWordAtTextConfig) {
 }
 
 export function getWordAtText(column: number, wordDefinition: RegExp, text: string, textOffset: number, config?: IGetWordAtTextConfig): IWordAtPosition | null {
+	// Ensure the regex has the 'g' flag, otherwise this will loop forever
+	wordDefinition = ensureValidWordDefinition(wordDefinition);
 
 	if (!config) {
 		config = Iterable.first(_defaultConfig)!;

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

@@ -4,11 +4,25 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { DeferredPromise } from 'vs/base/common/async';
+import { CancellationToken } from 'vs/base/common/cancellation';
 import { IMarkdownString } from 'vs/base/common/htmlContent';
-import { Disposable, DisposableMap } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableMap, IDisposable } from 'vs/base/common/lifecycle';
 import { revive } from 'vs/base/common/marshalling';
+import { escapeRegExpCharacters } from 'vs/base/common/strings';
+import { Position } from 'vs/editor/common/core/position';
+import { Range } from 'vs/editor/common/core/range';
+import { getWordAtText } from 'vs/editor/common/core/wordHelper';
+import { CompletionContext, CompletionItem, CompletionItemKind, CompletionList } from 'vs/editor/common/languages';
+import { ITextModel } from 'vs/editor/common/model';
+import { ILanguageFeaturesService } from 'vs/editor/common/services/languageFeatures';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ExtHostChatAgentsShape2, ExtHostContext, IChatProgressDto, IExtensionChatAgentMetadata, MainContext, MainThreadChatAgentsShape2 } from 'vs/workbench/api/common/extHost.protocol';
+import { IChatWidgetService } from 'vs/workbench/contrib/chat/browser/chat';
+import { ChatInputPart } from 'vs/workbench/contrib/chat/browser/chatInputPart';
+import { AddDynamicVariableAction, IAddDynamicVariableContext } from 'vs/workbench/contrib/chat/browser/contrib/chatDynamicVariables';
 import { IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
+import { ChatRequestAgentPart } from 'vs/workbench/contrib/chat/common/chatParserTypes';
+import { ChatRequestParser } from 'vs/workbench/contrib/chat/common/chatRequestParser';
 import { IChatFollowup, IChatProgress, IChatService, IChatTreeData } from 'vs/workbench/contrib/chat/common/chatService';
 import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
 
@@ -23,6 +37,8 @@ type AgentData = {
 export class MainThreadChatAgents2 extends Disposable implements MainThreadChatAgentsShape2 {
 
 	private readonly _agents = this._register(new DisposableMap<number, AgentData>());
+	private readonly _agentCompletionProviders = this._register(new DisposableMap<number, IDisposable>());
+
 	private readonly _pendingProgress = new Map<string, (part: IChatProgress) => void>();
 	private readonly _proxy: ExtHostChatAgentsShape2;
 
@@ -33,6 +49,9 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
 		extHostContext: IExtHostContext,
 		@IChatAgentService private readonly _chatAgentService: IChatAgentService,
 		@IChatService private readonly _chatService: IChatService,
+		@ILanguageFeaturesService private readonly _languageFeaturesService: ILanguageFeaturesService,
+		@IChatWidgetService private readonly _chatWidgetService: IChatWidgetService,
+		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 	) {
 		super();
 		this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostChatAgents2);
@@ -135,4 +154,80 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
 		const revivedProgress = revive(progress);
 		this._pendingProgress.get(requestId)?.(revivedProgress as IChatProgress);
 	}
+
+	$registerAgentCompletionsProvider(handle: number, triggerCharacters: string[]): void {
+		this._agentCompletionProviders.set(handle, this._languageFeaturesService.completionProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, hasAccessToAllModels: true }, {
+			_debugDisplayName: 'chatAgentCompletions:' + handle,
+			triggerCharacters,
+			provideCompletionItems: async (model: ITextModel, position: Position, _context: CompletionContext, token: CancellationToken) => {
+				const widget = this._chatWidgetService.getWidgetByInputUri(model.uri);
+				if (!widget || !widget.viewModel) {
+					return;
+				}
+
+				const triggerCharsPart = triggerCharacters.map(c => escapeRegExpCharacters(c)).join('');
+				const wordRegex = new RegExp(`[${triggerCharsPart}]\\S*`, 'g');
+				const query = getWordAtText(position.column, wordRegex, model.getLineContent(position.lineNumber), 0)?.word ?? '';
+
+				if (query && !triggerCharacters.some(c => query.startsWith(c))) {
+					return;
+				}
+
+				const parsedRequest = (await this._instantiationService.createInstance(ChatRequestParser).parseChatRequest(widget.viewModel.sessionId, model.getValue())).parts;
+				const agentPart = parsedRequest.find((part): part is ChatRequestAgentPart => part instanceof ChatRequestAgentPart);
+				const thisAgentName = this._agents.get(handle)?.name;
+				if (agentPart?.agent.id !== thisAgentName) {
+					return;
+				}
+
+				const range = computeCompletionRanges(model, position, wordRegex);
+				if (!range) {
+					return null;
+				}
+
+				const result = await this._proxy.$invokeCompletionProvider(handle, query, token);
+				const variableItems = result.map(v => {
+					const insertText = v.insertText ?? (typeof v.label === 'string' ? v.label : v.label.label);
+					const rangeAfterInsert = new Range(range.insert.startLineNumber, range.insert.startColumn, range.insert.endLineNumber, range.insert.startColumn + insertText.length);
+					return {
+						label: v.label,
+						range,
+						insertText: insertText + ' ',
+						kind: CompletionItemKind.Text,
+						detail: v.detail,
+						documentation: v.documentation,
+						command: { id: AddDynamicVariableAction.ID, title: '', arguments: [{ widget, range: rangeAfterInsert, variableData: v.values } satisfies IAddDynamicVariableContext] }
+					} satisfies CompletionItem;
+				});
+
+				return {
+					suggestions: variableItems
+				} satisfies CompletionList;
+			}
+		}));
+	}
+
+	$unregisterAgentCompletionsProvider(handle: number): void {
+		this._agentCompletionProviders.deleteAndDispose(handle);
+	}
+}
+
+
+function computeCompletionRanges(model: ITextModel, position: Position, reg: RegExp): { insert: Range; replace: Range } | undefined {
+	const varWord = getWordAtText(position.column, reg, model.getLineContent(position.lineNumber), 0);
+	if (!varWord && model.getWordUntilPosition(position).word) {
+		// inside a "normal" word
+		return;
+	}
+
+	let insert: Range;
+	let replace: Range;
+	if (!varWord) {
+		insert = replace = Range.fromPositions(position);
+	} else {
+		insert = new Range(position.lineNumber, varWord.startColumn, position.lineNumber, position.column);
+		replace = new Range(position.lineNumber, varWord.startColumn, position.lineNumber, varWord.endColumn);
+	}
+
+	return { insert, replace };
 }

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

@@ -4,8 +4,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { DisposableMap } from 'vs/base/common/lifecycle';
+import { revive } from 'vs/base/common/marshalling';
 import { ExtHostChatVariablesShape, ExtHostContext, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
-import { IChatVariableData, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
+import { IChatRequestVariableValue, IChatVariableData, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
 import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
 
 @extHostNamedCustomer(MainContext.MainThreadChatVariables)
@@ -26,8 +27,8 @@ export class MainThreadChatVariables implements MainThreadChatVariablesShape {
 	}
 
 	$registerVariable(handle: number, data: IChatVariableData): void {
-		const registration = this._chatVariablesService.registerVariable(data, (messageText, _arg, _model, token) => {
-			return this._proxy.$resolveVariable(handle, messageText, token);
+		const registration = this._chatVariablesService.registerVariable(data, async (messageText, _arg, _model, token) => {
+			return revive<IChatRequestVariableValue[]>(await this._proxy.$resolveVariable(handle, messageText, token));
 		});
 		this._variables.set(handle, registration);
 	}

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

@@ -1424,6 +1424,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			ChatMessage: extHostTypes.ChatMessage,
 			ChatMessageRole: extHostTypes.ChatMessageRole,
 			ChatVariableLevel: extHostTypes.ChatVariableLevel,
+			ChatAgentCompletionItem: extHostTypes.ChatAgentCompletionItem,
 			CallHierarchyIncomingCall: extHostTypes.CallHierarchyIncomingCall,
 			CallHierarchyItem: extHostTypes.CallHierarchyItem,
 			CallHierarchyOutgoingCall: extHostTypes.CallHierarchyOutgoingCall,

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

@@ -1177,17 +1177,28 @@ export interface IExtensionChatAgentMetadata extends Dto<IChatAgentMetadata> {
 
 export interface MainThreadChatAgentsShape2 extends IDisposable {
 	$registerAgent(handle: number, name: string, metadata: IExtensionChatAgentMetadata): void;
+	$registerAgentCompletionsProvider(handle: number, triggerCharacters: string[]): void;
+	$unregisterAgentCompletionsProvider(handle: number): void;
 	$updateAgent(handle: number, metadataUpdate: IExtensionChatAgentMetadata): void;
 	$unregisterAgent(handle: number): void;
 	$handleProgressChunk(requestId: string, chunk: IChatProgressDto, responsePartHandle?: number): Promise<number | void>;
 }
 
+export interface IChatAgentCompletionItem {
+	insertText?: string;
+	label: string | languages.CompletionItemLabel;
+	values: IChatRequestVariableValue[];
+	detail?: string;
+	documentation?: string | IMarkdownString;
+}
+
 export interface ExtHostChatAgentsShape2 {
 	$invokeAgent(handle: number, sessionId: string, requestId: string, request: IChatAgentRequest, context: { history: IChatMessage[] }, token: CancellationToken): Promise<IChatAgentResult | undefined>;
 	$provideSlashCommands(handle: number, token: CancellationToken): Promise<IChatAgentCommand[]>;
 	$provideFollowups(handle: number, sessionId: string, token: CancellationToken): Promise<IChatFollowup[]>;
 	$acceptFeedback(handle: number, sessionId: string, requestId: string, vote: InteractiveSessionVoteDirection, reportIssue?: boolean): void;
 	$acceptAction(handle: number, sessionId: string, requestId: string, action: IChatUserActionEvent): void;
+	$invokeCompletionProvider(handle: number, query: string, token: CancellationToken): Promise<IChatAgentCompletionItem[]>;
 	$releaseSession(sessionId: string): void;
 }
 

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

@@ -14,7 +14,7 @@ import { localize } from 'vs/nls';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { ILogService } from 'vs/platform/log/common/log';
 import { Progress } from 'vs/platform/progress/common/progress';
-import { ExtHostChatAgentsShape2, IMainContext, MainContext, MainThreadChatAgentsShape2 } from 'vs/workbench/api/common/extHost.protocol';
+import { ExtHostChatAgentsShape2, IChatAgentCompletionItem, IMainContext, MainContext, MainThreadChatAgentsShape2 } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostChatProvider } from 'vs/workbench/api/common/extHostChatProvider';
 import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
@@ -217,6 +217,16 @@ export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 		}
 		agent.acceptAction(Object.freeze({ action: action.action, result }));
 	}
+
+	async $invokeCompletionProvider(handle: number, query: string, token: CancellationToken): Promise<IChatAgentCompletionItem[]> {
+		const agent = this._agents.get(handle);
+		if (!agent) {
+			return [];
+		}
+
+		const items = await agent.invokeCompletionProvider(query, token);
+		return items.map(typeConvert.ChatAgentCompletionItem.from);
+	}
 }
 
 class ExtHostChatAgent {
@@ -235,6 +245,7 @@ class ExtHostChatAgent {
 	private _onDidReceiveFeedback = new Emitter<vscode.ChatAgentResult2Feedback>();
 	private _onDidPerformAction = new Emitter<vscode.ChatAgentUserActionEvent>();
 	private _supportIssueReporting: boolean | undefined;
+	private _agentVariableProvider?: { provider: vscode.ChatAgentCompletionItemProvider; triggerCharacters: string[] };
 
 	constructor(
 		public readonly extension: IExtensionDescription,
@@ -252,6 +263,14 @@ class ExtHostChatAgent {
 		this._onDidPerformAction.fire(event);
 	}
 
+	async invokeCompletionProvider(query: string, token: CancellationToken): Promise<vscode.ChatAgentCompletionItem[]> {
+		if (!this._agentVariableProvider) {
+			return [];
+		}
+
+		return await this._agentVariableProvider.provider.provideCompletionItems(query, token) ?? [];
+	}
+
 	async validateSlashCommand(command: string) {
 		if (!this._lastSlashCommands) {
 			await this.provideSlashCommand(CancellationToken.None);
@@ -434,6 +453,21 @@ class ExtHostChatAgent {
 			get onDidReceiveFeedback() {
 				return that._onDidReceiveFeedback.event;
 			},
+			set agentVariableProvider(v) {
+				that._agentVariableProvider = v;
+				if (v) {
+					if (!v.triggerCharacters.length) {
+						throw new Error('triggerCharacters are required');
+					}
+
+					that._proxy.$registerAgentCompletionsProvider(that._handle, v.triggerCharacters);
+				} else {
+					that._proxy.$unregisterAgentCompletionsProvider(that._handle);
+				}
+			},
+			get agentVariableProvider() {
+				return that._agentVariableProvider;
+			},
 			onDidPerformAction: !isProposedApiEnabled(this.extension, 'chatAgents2Additions')
 				? undefined!
 				: this._onDidPerformAction.event

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

@@ -11,7 +11,7 @@ import * as htmlContent from 'vs/base/common/htmlContent';
 import { DisposableStore } from 'vs/base/common/lifecycle';
 import { ResourceMap, ResourceSet } from 'vs/base/common/map';
 import { marked } from 'vs/base/common/marked/marked';
-import { parse } from 'vs/base/common/marshalling';
+import { parse, revive } from 'vs/base/common/marshalling';
 import { Mimes } from 'vs/base/common/mime';
 import { cloneAndChange } from 'vs/base/common/objects';
 import { isEmptyObject, isNumber, isString, isUndefinedOrNull } from 'vs/base/common/types';
@@ -2253,14 +2253,16 @@ export namespace ChatVariable {
 	export function to(variable: IChatRequestVariableValue): vscode.ChatVariableValue {
 		return {
 			level: ChatVariableLevel.to(variable.level),
-			value: variable.value,
+			kind: variable.kind,
+			value: revive(variable.value),
 			description: variable.description
 		};
 	}
 
 	export function from(variable: vscode.ChatVariableValue): IChatRequestVariableValue {
 		return {
 			level: ChatVariableLevel.from(variable.level),
+			kind: variable.kind,
 			value: variable.value,
 			description: variable.description
 		};
@@ -2368,6 +2370,18 @@ export namespace ChatResponseProgress {
 	}
 }
 
+export namespace ChatAgentCompletionItem {
+	export function from(item: vscode.ChatAgentCompletionItem): extHostProtocol.IChatAgentCompletionItem {
+		return {
+			label: item.label,
+			values: item.values.map(ChatVariable.from),
+			insertText: item.insertText,
+			detail: item.detail,
+			documentation: item.documentation,
+		};
+	}
+}
+
 
 export namespace TerminalQuickFix {
 	export function from(quickFix: vscode.TerminalQuickFixTerminalCommand | vscode.TerminalQuickFixOpener | vscode.Command, converter: Command.ICommandsConverter, disposables: DisposableStore): extHostProtocol.ITerminalQuickFixTerminalCommandDto | extHostProtocol.ITerminalQuickFixOpenerDto | extHostProtocol.ICommandDto | undefined {

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

@@ -4104,7 +4104,7 @@ export class ChatEditorTabInput {
 }
 //#endregion
 
-//#region Interactive session
+//#region Chat
 
 export enum InteractiveSessionVoteDirection {
 	Down = 0,
@@ -4116,6 +4116,25 @@ export enum InteractiveSessionCopyKind {
 	Toolbar = 2
 }
 
+export enum ChatVariableLevel {
+	Short = 1,
+	Medium = 2,
+	Full = 3
+}
+
+export class ChatAgentCompletionItem implements vscode.ChatAgentCompletionItem {
+	label: string | CompletionItemLabel;
+	insertText?: string;
+	values: vscode.ChatVariableValue[];
+	detail?: string;
+	documentation?: string | MarkdownString;
+
+	constructor(label: string | CompletionItemLabel, values: vscode.ChatVariableValue[]) {
+		this.label = label;
+		this.values = values;
+	}
+}
+
 //#endregion
 
 //#region Interactive Editor
@@ -4135,12 +4154,6 @@ export enum ChatMessageRole {
 	Function = 3,
 }
 
-export enum ChatVariableLevel {
-	Short = 1,
-	Medium = 2,
-	Full = 3
-}
-
 export class ChatMessage implements vscode.ChatMessage {
 
 	role: ChatMessageRole;