feat: allow extensions to attach chat context (#213172)

* feat: allow extensions to attach chat context * fix: allow attaching different ranges from same file * fix: restore accepting files from picker
microsoft · May 21, 2024 · f2f4adc · f2f4adc
1 parent 6743fa3
commit f2f4adc
Show file tree

Hide file tree

Showing 12 changed files with 83 additions and 1 deletion.
diff --git a/src/vs/workbench/api/browser/mainThreadChatVariables.ts b/src/vs/workbench/api/browser/mainThreadChatVariables.ts
@@ -5,7 +5,10 @@
 
 import { DisposableMap } from 'vs/base/common/lifecycle';
 import { revive } from 'vs/base/common/marshalling';
+import { URI } from 'vs/base/common/uri';
+import { Location } from 'vs/editor/common/languages';
 import { ExtHostChatVariablesShape, ExtHostContext, IChatVariableResolverProgressDto, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
+import { ChatAgentLocation } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolverProgress, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
 import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
 
@@ -47,4 +50,8 @@ export class MainThreadChatVariables implements MainThreadChatVariablesShape {
 	$unregisterVariable(handle: number): void {
 		this._variables.deleteAndDispose(handle);
 	}
+
+	$attachContext(name: string, value: string | URI | Location | unknown, location: ChatAgentLocation.Panel): void {
+		this._chatVariablesService.attachContext(name, revive(value), location);
+	}
 }
diff --git a/src/vs/workbench/api/common/extHost.api.impl.ts b/src/vs/workbench/api/common/extHost.api.impl.ts
@@ -1425,6 +1425,10 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			createDynamicChatParticipant(id: string, dynamicProps: vscode.DynamicChatParticipantProps, handler: vscode.ChatExtendedRequestHandler): vscode.ChatParticipant {
 				checkProposedApiEnabled(extension, 'chatParticipantAdditions');
 				return extHostChatAgents2.createDynamicChatAgent(extension, id, dynamicProps, handler);
+			},
+			attachContext(name: string, value: string | vscode.Uri | vscode.Location | unknown, location: vscode.ChatLocation.Panel) {
+				checkProposedApiEnabled(extension, 'chatVariableResolver');
+				return extHostChatVariables.attachContext(name, value, location);
 			}
 		};
 

diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts
@@ -6,6 +6,7 @@
 import { VSBuffer } from 'vs/base/common/buffer';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { IRemoteConsoleLog } from 'vs/base/common/console';
+import { Location } from 'vs/editor/common/languages';
 import { SerializedError } from 'vs/base/common/errors';
 import { IRelativePattern } from 'vs/base/common/glob';
 import { IMarkdownString } from 'vs/base/common/htmlContent';
@@ -1290,6 +1291,7 @@ export interface MainThreadChatVariablesShape extends IDisposable {
 	$registerVariable(handle: number, data: IChatVariableData): void;
 	$handleProgressChunk(requestId: string, progress: IChatVariableResolverProgressDto): Promise<number | void>;
 	$unregisterVariable(handle: number): void;
+	$attachContext(name: string, value: string | Dto<Location> | URI | unknown, location: ChatAgentLocation): void;
 }
 
 export type IChatRequestVariableValueDto = Dto<IChatRequestVariableValue>;

diff --git a/src/vs/workbench/api/common/extHostChatVariables.ts b/src/vs/workbench/api/common/extHostChatVariables.ts
@@ -64,6 +64,10 @@ export class ExtHostChatVariables implements ExtHostChatVariablesShape {
 			this._proxy.$unregisterVariable(handle);
 		});
 	}
+
+	attachContext(name: string, value: string | vscode.Location | vscode.Uri | unknown, location: vscode.ChatLocation.Panel) {
+		this._proxy.$attachContext(name, extHostTypes.Location.isLocation(value) ? typeConvert.Location.from(value) : value, typeConvert.ChatLocation.from(location));
+	}
 }
 
 class ChatVariableResolverResponseStream {

diff --git a/src/vs/workbench/api/common/extHostTypeConverters.ts b/src/vs/workbench/api/common/extHostTypeConverters.ts
@@ -2566,6 +2566,15 @@ export namespace ChatLocation {
 			case ChatAgentLocation.Editor: return types.ChatLocation.Editor;
 		}
 	}
+
+	export function from(loc: types.ChatLocation): ChatAgentLocation {
+		switch (loc) {
+			case types.ChatLocation.Notebook: return ChatAgentLocation.Notebook;
+			case types.ChatLocation.Terminal: return ChatAgentLocation.Terminal;
+			case types.ChatLocation.Panel: return ChatAgentLocation.Panel;
+			case types.ChatLocation.Editor: return ChatAgentLocation.Editor;
+		}
+	}
 }
 
 export namespace ChatAgentValueReference {
@@ -2579,7 +2588,9 @@ export namespace ChatAgentValueReference {
 			id: variable.id,
 			name: variable.name,
 			range: variable.range && [variable.range.start, variable.range.endExclusive],
-			value: isUriComponents(value) ? URI.revive(value) : value,
+			value: isUriComponents(value) ? URI.revive(value) :
+				value && typeof value === 'object' && 'uri' in value && 'range' in value && isUriComponents(value.uri) ?
+					Location.to(revive(value)) : value,
 			modelDescription: variable.modelDescription
 		};
 	}

diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
@@ -70,6 +70,8 @@ class AttachContextAction extends Action2 {
 
 					toAttach.push({ ...pick, isDynamic: pick.isDynamic, value: pick.value, name: qualifiedName, fullName: `$(${pick.icon.id}) ${selection}` });
 				}
+			} else if (pick && typeof pick === 'object' && 'resource' in pick) {
+				toAttach.push({ ...pick, value: pick.resource });
 			} else {
 				toAttach.push({ ...pick, fullName: pick.label });
 			}

diff --git a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts
@@ -7,6 +7,7 @@ import * as dom from 'vs/base/browser/dom';
 import { DEFAULT_FONT_FAMILY } from 'vs/base/browser/fonts';
 import { IHistoryNavigationWidget } from 'vs/base/browser/history';
 import * as aria from 'vs/base/browser/ui/aria/aria';
+import { Range } from 'vs/editor/common/core/range';
 import { Button } from 'vs/base/browser/ui/button/button';
 import { IAction } from 'vs/base/common/actions';
 import { Codicon } from 'vs/base/common/codicons';
@@ -442,6 +443,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 				label.setFile(file, {
 					fileKind: FileKind.FILE,
 					hidePath: true,
+					range: attachment.value && typeof attachment.value === 'object' && 'range' in attachment.value && Range.isIRange(attachment.value.range) ? attachment.value.range : undefined,
 				});
 			} else {
 				label.setLabel(attachment.fullName ?? attachment.name);

diff --git a/src/vs/workbench/contrib/chat/browser/chatVariables.ts b/src/vs/workbench/contrib/chat/browser/chatVariables.ts
@@ -3,13 +3,17 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { basename } from 'vs/base/common/path';
 import { coalesce } from 'vs/base/common/arrays';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { onUnexpectedExternalError } from 'vs/base/common/errors';
 import { Iterable } from 'vs/base/common/iterator';
 import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
+import { URI } from 'vs/base/common/uri';
+import { Location } from 'vs/editor/common/languages';
 import { IChatWidgetService } from 'vs/workbench/contrib/chat/browser/chat';
 import { ChatDynamicVariableModel } from 'vs/workbench/contrib/chat/browser/contrib/chatDynamicVariables';
+import { ChatAgentLocation } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { IChatModel, IChatRequestVariableData, IChatRequestVariableEntry } from 'vs/workbench/contrib/chat/common/chatModel';
 import { ChatRequestDynamicVariablePart, ChatRequestVariablePart, IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { IChatContentReference } from 'vs/workbench/contrib/chat/common/chatService';
@@ -142,4 +146,30 @@ export class ChatVariablesService implements IChatVariablesService {
 			this._resolver.delete(key);
 		});
 	}
+
+	async attachContext(name: string, value: string | URI | Location, location: ChatAgentLocation) {
+		if (location !== ChatAgentLocation.Panel) {
+			return;
+		}
+
+		const widget = this.chatWidgetService.lastFocusedWidget;
+		if (!widget || !widget.viewModel) {
+			return;
+		}
+
+		const key = name.toLowerCase();
+		if (key === 'file' && typeof value !== 'string') {
+			const uri = URI.isUri(value) ? value : value.uri;
+			const range = 'range' in value ? value.range : undefined;
+			widget.attachContext({ value, id: uri.toString() + (range?.toString() ?? ''), name: basename(uri.path), isDynamic: true });
+			return;
+		}
+
+		const resolved = this._resolver.get(key);
+		if (!resolved) {
+			return;
+		}
+
+		widget.attachContext({ ...resolved.data, value });
+	}
 }
diff --git a/src/vs/workbench/contrib/chat/browser/media/chat.css b/src/vs/workbench/contrib/chat/browser/media/chat.css
@@ -493,6 +493,10 @@
 	align-items: center;
 }
 
+.interactive-session .chat-attached-context .chat-attached-context-attachment .monaco-icon-label-container {
+	display: flex;
+}
+
 .interactive-session .chat-attached-context .chat-attached-context-attachment .monaco-icon-label-container .monaco-highlighted-label {
 	display: flex !important;
 	align-items: center !important;

diff --git a/src/vs/workbench/contrib/chat/common/chatVariables.ts b/src/vs/workbench/contrib/chat/common/chatVariables.ts
@@ -10,6 +10,7 @@ import { URI } from 'vs/base/common/uri';
 import { IRange } from 'vs/editor/common/core/range';
 import { Location } from 'vs/editor/common/languages';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
+import { ChatAgentLocation } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { IChatModel, IChatRequestVariableData, IChatRequestVariableEntry } from 'vs/workbench/contrib/chat/common/chatModel';
 import { IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { IChatContentReference, IChatProgressMessage } from 'vs/workbench/contrib/chat/common/chatService';
@@ -45,6 +46,7 @@ export interface IChatVariablesService {
 	getVariable(name: string): IChatVariableData | undefined;
 	getVariables(): Iterable<Readonly<IChatVariableData>>;
 	getDynamicVariables(sessionId: string): ReadonlyArray<IDynamicVariable>; // should be its own service?
+	attachContext(name: string, value: string | URI | Location | unknown, location: ChatAgentLocation): void;
 
 	/**
 	 * Resolves all variables that occur in `prompt`

diff --git a/src/vs/workbench/contrib/chat/test/common/mockChatVariables.ts b/src/vs/workbench/contrib/chat/test/common/mockChatVariables.ts
@@ -5,6 +5,7 @@
 
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { IDisposable } from 'vs/base/common/lifecycle';
+import { ChatAgentLocation } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { IChatModel, IChatRequestVariableData, IChatRequestVariableEntry } from 'vs/workbench/contrib/chat/common/chatModel';
 import { IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolver, IChatVariableResolverProgress, IChatVariablesService, IDynamicVariable } from 'vs/workbench/contrib/chat/common/chatVariables';
@@ -37,6 +38,10 @@ export class MockChatVariablesService implements IChatVariablesService {
 		};
 	}
 
+	attachContext(name: string, value: unknown, location: ChatAgentLocation): void {
+		throw new Error('Method not implemented.');
+	}
+
 	resolveVariable(variableName: string, promptText: string, model: IChatModel, progress: (part: IChatVariableResolverProgress) => void, token: CancellationToken): Promise<IChatRequestVariableValue> {
 		throw new Error('Method not implemented.');
 	}

diff --git a/src/vscode-dts/vscode.proposed.chatVariableResolver.d.ts b/src/vscode-dts/vscode.proposed.chatVariableResolver.d.ts
@@ -19,6 +19,15 @@ declare module 'vscode' {
 		 * @param icon An icon to display when selecting context in the picker UI.
 		 */
 		export function registerChatVariableResolver(id: string, name: string, userDescription: string, modelDescription: string | undefined, isSlow: boolean | undefined, resolver: ChatVariableResolver, fullName?: string, icon?: ThemeIcon): Disposable;
+
+		/**
+		 * Attaches a chat context with the specified name, value, and location.
+		 *
+		 * @param name - The name of the chat context.
+		 * @param value - The value of the chat context.
+		 * @param location - The location of the chat context.
+		 */
+		export function attachContext(name: string, value: string | Uri | Location | unknown, location: ChatLocation.Panel): void;
 	}
 
 	export interface ChatVariableValue {