From 381c4ffde9b1c1eb5a6f8bc1fd3ad43333f3f1fe Mon Sep 17 00:00:00 2001
From: Osvaldo Ortega <48293249+osortega@users.noreply.github.com>
Date: Wed, 27 Nov 2024 22:27:20 +0000
Subject: [PATCH] Allowing selecting multiple files from the search widget into
 chat

---
 .../browser/actions/chatContextActions.ts     | 57 +++++++++++++++----
 .../contrib/search/browser/searchView.ts      | 14 ++++-
 2 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
index be164a6ec03..a38b8e77527 100644
--- a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
+++ b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
@@ -13,7 +13,7 @@ import { compare } from '../../../../../base/common/strings.js';
 import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { URI } from '../../../../../base/common/uri.js';
 import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions.js';
-import { IRange } from '../../../../../editor/common/core/range.js';
+import { IRange, Range } from '../../../../../editor/common/core/range.js';
 import { EditorType } from '../../../../../editor/common/editorCommon.js';
 import { Command } from '../../../../../editor/common/languages.js';
 import { AbstractGotoSymbolQuickAccessProvider, IGotoSymbolQuickPickItem } from '../../../../../editor/contrib/quickAccess/browser/gotoSymbolQuickAccess.js';
@@ -250,26 +250,59 @@ class AttachSelectionToChatAction extends Action2 {
 			title: localize2('workbench.action.chat.attachSelection.label', "Add Selection to Chat"),
 			category: CHAT_CATEGORY,
 			f1: false,
-			precondition: ContextKeyExpr.and(ChatContextKeys.enabled, ActiveEditorContext.isEqualTo('workbench.editors.files.textFileEditor')),
-			menu: {
+			precondition: ChatContextKeys.enabled,
+			menu: [{
 				id: MenuId.ChatCommandCenter,
 				group: 'b_chat_context',
+				when: ActiveEditorContext.isEqualTo('workbench.editors.files.textFileEditor'),
 				order: 15,
-			}
+			}, {
+				id: MenuId.SearchContext,
+				group: 'z_chat',
+				order: 2
+			}]
 		});
 	}
 
 	override async run(accessor: ServicesAccessor, ...args: any[]): Promise<void> {
 		const variablesService = accessor.get(IChatVariablesService);
 		const textEditorService = accessor.get(IEditorService);
-
-		const activeEditor = textEditorService.activeTextEditorControl;
-		const activeUri = textEditorService.activeEditor?.resource;
-		if (textEditorService.activeTextEditorControl?.getEditorType() === EditorType.ICodeEditor && activeUri && [Schemas.file, Schemas.vscodeRemote, Schemas.untitled].includes(activeUri.scheme)) {
-			const selection = activeEditor?.getSelection();
-			if (selection) {
-				(await showChatView(accessor.get(IViewsService)))?.focusInput();
-				variablesService.attachContext('file', { uri: activeUri, range: selection }, ChatAgentLocation.Panel);
+		const [_, matches] = args;
+		// It means this is coming from the search widget
+		if (matches && matches.length > 0) {
+			const uris = new Map<URI, Range | undefined>();
+			for (const match of matches) {
+				// If it's a file match, we want to attach the file
+				// If it's a tree, we want to make sure that at least has one file, if not just add the whole file
+				if (isSearchTreeFileMatch(match)) {
+					uris.set(match.resource, undefined);
+				} else {
+					const context = { uri: match._parent.resource, range: match._range };
+					const range = uris.get(context.uri);
+					// if the lines in the y axis are the same then contineu
+					if (!range ||
+						range.startLineNumber !== context.range.startLineNumber && range.endLineNumber !== context.range.endLineNumber) {
+						uris.set(context.uri, context.range);
+						variablesService.attachContext('file', context, ChatAgentLocation.Panel);
+					}
+				}
+			}
+			// Add the root files for all of the ones that didn't have a match
+			for (const uri of uris) {
+				const [resource, range] = uri;
+				if (!range) {
+					variablesService.attachContext('file', { uri: resource }, ChatAgentLocation.Panel);
+				}
+			}
+		} else {
+			const activeEditor = textEditorService.activeTextEditorControl;
+			const activeUri = textEditorService.activeEditor?.resource;
+			if (textEditorService.activeTextEditorControl?.getEditorType() === EditorType.ICodeEditor && activeUri && [Schemas.file, Schemas.vscodeRemote, Schemas.untitled].includes(activeUri.scheme)) {
+				const selection = activeEditor?.getSelection();
+				if (selection) {
+					(await showChatView(accessor.get(IViewsService)))?.focusInput();
+					variablesService.attachContext('file', { uri: activeUri, range: selection }, ChatAgentLocation.Panel);
+				}
 			}
 		}
 	}
diff --git a/src/vs/workbench/contrib/search/browser/searchView.ts b/src/vs/workbench/contrib/search/browser/searchView.ts
index 697012dbc5a..de360365c66 100644
--- a/src/vs/workbench/contrib/search/browser/searchView.ts
+++ b/src/vs/workbench/contrib/search/browser/searchView.ts
@@ -971,12 +971,22 @@ export class SearchView extends ViewPane {
 		e.browserEvent.preventDefault();
 		e.browserEvent.stopPropagation();
 
+		const selection = this.tree.getSelection();
+		let arg;
+		let context;
+		if (selection && selection.length > 0) {
+			arg = e.element;
+			context = selection;
+		} else {
+			context = e.element;
+		}
+
 		this.contextMenuService.showContextMenu({
 			menuId: MenuId.SearchContext,
-			menuActionOptions: { shouldForwardArgs: true },
+			menuActionOptions: { shouldForwardArgs: true, arg },
 			contextKeyService: this.contextKeyService,
 			getAnchor: () => e.anchor,
-			getActionsContext: () => e.element,
+			getActionsContext: () => context,
 		});
 	}