refactor: clean up chat context types (microsoft#213264)

PIC1G55 · May 26, 2024 · 97e8321 · 97e8321
1 parent eb2f661
commit 97e8321
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 15 deletions.
diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
@@ -23,6 +23,7 @@ import { ChatContextAttachments } from 'vs/workbench/contrib/chat/browser/contri
 import { SelectAndInsertFileAction } from 'vs/workbench/contrib/chat/browser/contrib/chatDynamicVariables';
 import { ChatAgentLocation, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { CONTEXT_CHAT_LOCATION, CONTEXT_IN_CHAT_INPUT } from 'vs/workbench/contrib/chat/common/chatContextKeys';
+import { IChatRequestVariableEntry } from 'vs/workbench/contrib/chat/common/chatModel';
 import { ChatRequestAgentPart } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
 import { AnythingQuickAccessProvider } from 'vs/workbench/contrib/search/browser/anythingQuickAccess';
@@ -31,6 +32,36 @@ export function registerChatContextActions() {
 	registerAction2(AttachContextAction);
 }
 
+export type IChatContextQuickPickItem = IFileQuickPickItem | IDynamicVariableQuickPickItem | IStaticVariableQuickPickItem;
+
+export interface IFileQuickPickItem extends IQuickPickItem {
+	id: string;
+	name: string;
+	value: URI;
+
+	resource: URI;
+	isDynamic: true;
+}
+
+export interface IDynamicVariableQuickPickItem extends IQuickPickItem {
+	id: string;
+	name?: string;
+	value: unknown;
+	icon?: ThemeIcon;
+
+	command?: Command;
+	isDynamic: true;
+}
+
+export interface IStaticVariableQuickPickItem extends IQuickPickItem {
+	id: string;
+	name: string;
+	value: unknown;
+	icon?: ThemeIcon;
+
+	isDynamic?: false;
+}
+
 class AttachContextAction extends Action2 {
 
 	static readonly ID = 'workbench.action.chat.attachContext';
@@ -61,20 +92,43 @@ class AttachContextAction extends Action2 {
 		});
 	}
 
-	private async _attachContext(widget: IChatWidget, commandService: ICommandService, ...picks: any[]) {
-		const toAttach = [];
+	private async _attachContext(widget: IChatWidget, commandService: ICommandService, ...picks: IChatContextQuickPickItem[]) {
+		const toAttach: IChatRequestVariableEntry[] = [];
 		for (const pick of picks) {
-			if (pick && typeof pick === 'object' && 'command' in pick) {
+			if (pick && typeof pick === 'object' && 'command' in pick && pick.command) {
+				// Dynamic variable with a followup command
 				const selection = await commandService.executeCommand(pick.command.id, ...(pick.command.arguments ?? []));
-				if (selection) {
-					const qualifiedName = `${typeof pick.value === 'string' && pick.value.startsWith('#') ? pick.value.slice(1) : ''}${selection}`;
-
-					toAttach.push({ ...pick, isDynamic: pick.isDynamic, value: pick.value, name: qualifiedName, fullName: `$(${pick.icon.id}) ${selection}` });
+				if (!selection) {
+					// User made no selection, skip this variable
+					continue;
 				}
+				toAttach.push({
+					...pick,
+					isDynamic: pick.isDynamic,
+					value: pick.value,
+					name: `${typeof pick.value === 'string' && pick.value.startsWith('#') ? pick.value.slice(1) : ''}${selection}`,
+					// Apply the original icon with the new name
+					fullName: `${pick.icon ? `$(${pick.icon.id}) ` : ''}${selection}`
+				});
 			} else if (pick && typeof pick === 'object' && 'resource' in pick) {
-				toAttach.push({ ...pick, value: pick.resource, name: pick.label, id: pick.resource.toString(), isDynamic: true });
+				// #file variable
+				toAttach.push({
+					...pick,
+					id: pick.resource.toString(),
+					value: pick.resource,
+					name: pick.label,
+					isDynamic: true
+				});
 			} else {
-				toAttach.push({ ...pick, fullName: pick.label, name: 'name' in pick && typeof pick.name === 'string' ? pick.name : pick.label, icon: 'icon' in pick && ThemeIcon.isThemeIcon(pick.icon) ? pick.icon : undefined });
+				// All other dynamic variables and static variables
+				toAttach.push({
+					...pick,
+					id: pick.id,
+					value: pick.value,
+					fullName: pick.label,
+					name: 'name' in pick && typeof pick.name === 'string' ? pick.name : pick.label,
+					icon: 'icon' in pick && ThemeIcon.isThemeIcon(pick.icon) ? pick.icon : undefined
+				});
 			}
 		}
 
@@ -95,10 +149,15 @@ class AttachContextAction extends Action2 {
 
 		const usedAgent = widget.parsedInput.parts.find(p => p instanceof ChatRequestAgentPart);
 		const slowSupported = usedAgent ? usedAgent.agent.metadata.supportsSlowVariables : true;
-		const quickPickItems: (QuickPickItem & { name?: string; icon?: ThemeIcon })[] = [];
+		const quickPickItems: (IChatContextQuickPickItem | QuickPickItem)[] = [];
 		for (const variable of chatVariablesService.getVariables()) {
 			if (variable.fullName && (!variable.isSlow || slowSupported)) {
-				quickPickItems.push({ label: `${variable.icon ? `$(${variable.icon.id}) ` : ''}${variable.fullName}`, name: variable.name, id: variable.id, icon: variable.icon });
+				quickPickItems.push({
+					label: `${variable.icon ? `$(${variable.icon.id}) ` : ''}${variable.fullName}`,
+					name: variable.name,
+					id: variable.id,
+					icon: variable.icon
+				});
 			}
 		}
 
@@ -108,7 +167,15 @@ class AttachContextAction extends Action2 {
 				const completions = await chatAgentService.getAgentCompletionItems(agentPart.agent.id, '', CancellationToken.None);
 				for (const variable of completions) {
 					if (variable.fullName) {
-						quickPickItems.push({ label: `${variable.icon ? `$(${variable.icon.id}) ` : ''}${variable.fullName}`, id: variable.id, command: variable.command, icon: variable.icon, value: variable.value, isDynamic: true });
+						quickPickItems.push({
+							label: `${variable.icon ? `$(${variable.icon.id}) ` : ''}${variable.fullName}`,
+							id: variable.id,
+							command: variable.command,
+							icon: variable.icon,
+							value: variable.value,
+							isDynamic: true,
+							name: variable.name
+						});
 					}
 				}
 			}
@@ -123,7 +190,7 @@ class AttachContextAction extends Action2 {
 			enabledProviderPrefixes: [AnythingQuickAccessProvider.PREFIX],
 			placeholder: localize('chatContext.attach.placeholder', 'Search attachments'),
 			providerOptions: <AnythingQuickAccessProviderRunOptions>{
-				handleAccept: (item: IQuickPickItem) => {
+				handleAccept: (item: IChatContextQuickPickItem) => {
 					this._attachContext(widget, commandService, item);
 				},
 				additionPicks: quickPickItems,

diff --git a/src/vs/workbench/contrib/chat/browser/chatVariables.ts b/src/vs/workbench/contrib/chat/browser/chatVariables.ts
@@ -63,6 +63,7 @@ export class ChatVariablesService implements IChatVariablesService {
 				}
 			});
 
+		const resolvedAttachedContext: IChatRequestVariableEntry[] = [];
 		attachedContextVariables
 			?.forEach((attachment, i) => {
 				const data = this._resolver.get(attachment.name?.toLowerCase());
@@ -77,11 +78,11 @@ export class ChatVariablesService implements IChatVariablesService {
 					};
 					jobs.push(data.resolver(prompt.text, '', model, variableProgressCallback, token).then(value => {
 						if (value) {
-							resolvedVariables[i] = { id: data.data.id, modelDescription: data.data.modelDescription, name: attachment.name, range: attachment.range, value, references };
+							resolvedAttachedContext[i] = { id: data.data.id, modelDescription: data.data.modelDescription, name: attachment.name, range: attachment.range, value, references };
 						}
 					}).catch(onUnexpectedExternalError));
 				} else if (attachment.isDynamic) {
-					resolvedVariables[i] = { id: attachment.id, name: attachment.name, value: attachment.value };
+					resolvedAttachedContext[i] = { id: attachment.id, name: attachment.name, value: attachment.value };
 				}
 			});
 
@@ -91,6 +92,7 @@ export class ChatVariablesService implements IChatVariablesService {
 
 		// "reverse", high index first so that replacement is simple
 		resolvedVariables.sort((a, b) => b.range!.start - a.range!.start);
+		resolvedVariables.push(...resolvedAttachedContext);
 
 		return {
 			variables: resolvedVariables,

diff --git a/src/vs/workbench/contrib/chat/common/chatAgents.ts b/src/vs/workbench/contrib/chat/common/chatAgents.ts
@@ -148,6 +148,7 @@ interface IChatAgentEntry {
 
 export interface IChatAgentCompletionItem {
 	id: string;
+	name?: string;
 	fullName?: string;
 	icon?: ThemeIcon;
 	value: unknown;