From 877af02b7fae00a1a20f5acca92af321b57f6f30 Mon Sep 17 00:00:00 2001
From: Johannes Rieken <johannes.rieken@gmail.com>
Date: Fri, 14 Mar 2025 14:16:31 +0100
Subject: [PATCH] polish tool picker (#243560)

* reset selected tools for new chat
* honor selected tools on subsequent pick
---
 .../chat/browser/actions/chatToolActions.ts   | 22 +++++++++++--------
 .../contrib/chat/browser/chatInputPart.ts     |  2 ++
 .../contrib/chat/browser/chatSelectedTools.ts |  6 ++++-
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
index b020d03d83e..5d71ac109d1 100644
--- a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
+++ b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
@@ -137,6 +137,7 @@ export class AttachToolsAction extends Action2 {
 			picked: true,
 		};
 
+		const nowSelectedTools = new Set(widget.input.selectedToolsModel.tools.get());
 		const toolBuckets = new Map<string, BucketPick>();
 
 		for (const tool of toolsService.getTools()) {
@@ -156,7 +157,7 @@ export class AttachToolsAction extends Action2 {
 					// description: mcpServer.definition.,
 					status: localize('desc', "MCP - {0} ({1})", mcpServer.collection.label, McpConnectionState.toString(mcpServer.connectionState.get())),
 					ordinal: BucketOrdinal.Mcp,
-					picked: true,
+					picked: false,
 					children: []
 				};
 				toolBuckets.set(mcpServer.definition.id, bucket);
@@ -165,7 +166,7 @@ export class AttachToolsAction extends Action2 {
 					type: 'item',
 					label: ext.displayName ?? ext.name,
 					ordinal: BucketOrdinal.Extension,
-					picked: true,
+					picked: false,
 					children: []
 				};
 				toolBuckets.set(ExtensionIdentifier.toKey(ext.identifier), bucket);
@@ -173,15 +174,21 @@ export class AttachToolsAction extends Action2 {
 				bucket = defaultBucket;
 			}
 
+			const picked = nowSelectedTools.has(tool);
+
 			bucket.children.push({
 				tool,
 				parent: bucket,
 				type: 'item',
 				label: `$(tools) ${tool.displayName}`,
 				description: tool.userDescription,
-				picked: true,
+				picked,
 				iconClasses: ['tool-pick']
 			});
+
+			if (picked) {
+				bucket.picked = true;
+			}
 		}
 
 		function isBucketPick(obj: any): obj is BucketPick {
@@ -220,6 +227,9 @@ export class AttachToolsAction extends Action2 {
 				lastSelectedItems = new Set(items);
 				picker.items = picks;
 				picker.selectedItems = items;
+
+				widget.input.selectedToolsModel.update(items.filter(isToolPick).map(tool => tool.tool));
+
 			} finally {
 				ignoreEvent = false;
 			}
@@ -266,14 +276,8 @@ export class AttachToolsAction extends Action2 {
 			_update();
 		}));
 
-
 		await Promise.race([Event.toPromise(Event.any(picker.onDidAccept, picker.onDidHide))]);
-
-		const selectedTools = picker.selectedItems.filter(isToolPick);
-
 		store.dispose();
-
-		widget.input.selectedToolsModel.update(selectedTools.map(tool => tool.tool));
 	}
 }
 
diff --git a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts
index e910bcc9652..77c919afc17 100644
--- a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts
+++ b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts
@@ -565,6 +565,8 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		} else if (this.location === ChatAgentLocation.EditingSession) {
 			this.setChatMode(ChatMode.Edit);
 		}
+
+		this.selectedToolsModel.reset();
 	}
 
 	logInputHistory(): void {
diff --git a/src/vs/workbench/contrib/chat/browser/chatSelectedTools.ts b/src/vs/workbench/contrib/chat/browser/chatSelectedTools.ts
index f3635c1a603..57ae4a9faf2 100644
--- a/src/vs/workbench/contrib/chat/browser/chatSelectedTools.ts
+++ b/src/vs/workbench/contrib/chat/browser/chatSelectedTools.ts
@@ -84,7 +84,11 @@ export class ChatSelectedTools extends Disposable {
 		}, Event.fromObservable(toolsCount)));
 	}
 
-	update(tools: IToolData[] | undefined): void {
+	update(tools: IToolData[]): void {
 		this._selectedTools.set(tools, undefined);
 	}
+
+	reset(): void {
+		this._selectedTools.set(undefined, undefined);
+	}
 }