fix leak + better logic for hidden tool calls

src/vs/workbench/contrib/chat/browser/widget/chatContentParts/chatThinkingContentPart.ts

@@ -34,6 +34,7 @@ import { IChatMarkdownAnchorService } from './chatMarkdownAnchorService.js';
 import { ChatMessageRole, ILanguageModelsService } from '../../../common/languageModels.js';
 import './media/chatThinkingContent.css';
 import { IHoverService } from '../../../../../../platform/hover/browser/hover.js';
+import { shouldHideToolInvocationAfterCompletion } from './toolInvocationParts/chatToolInvocationVisibility.js';
 
 
 function extractTextFromPart(content: IChatThinkingPart): string {
@@ -1103,12 +1104,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 			const result = factory();
 			this.appendItemToDOM(result.domNode, toolInvocationId, toolInvocationOrMarkdown, originalParent);
 			if (result.disposable) {
-				const toolCallId = toolInvocationOrMarkdown && (toolInvocationOrMarkdown.kind === 'toolInvocation' || toolInvocationOrMarkdown.kind === 'toolInvocationSerialized') ? toolInvocationOrMarkdown.toolCallId : undefined;
-				if (toolCallId) {
-					this.toolDisposables.get(toolCallId)?.add(result.disposable);
-				} else {
-					this._register(result.disposable);
-				}
+				this.registerToolDisposable(result.disposable, toolInvocationOrMarkdown);
 			}
 		} else {
 			// Defer rendering until expanded
@@ -1155,6 +1151,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 			this.toolInvocations.splice(toolInvocationsIndex, 1);
 		}
 
+		this.refreshDisplayedTitle();
 		this.updateDropdownClickability();
 		this._onDidChangeHeight.fire();
 	}
@@ -1200,6 +1197,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 			this.toolInvocations.splice(toolInvocationsIndex, 1);
 		}
 
+		this.refreshDisplayedTitle();
 		this.updateDropdownClickability();
 		return true;
 	}
@@ -1273,10 +1271,52 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 		if (titleIndex !== -1) {
 			this.extractedTitles.splice(titleIndex, 1);
 		}
+		this.refreshDisplayedTitle();
 		this.updateDropdownClickability();
 		this._onDidChangeHeight.fire();
 	}
 
+	private refreshDisplayedTitle(): void {
+		const extractedThinkingTitle = extractTitleFromThinkingContent(this.currentThinkingValue);
+		const nextTitle = this.extractedTitles[this.extractedTitles.length - 1]
+			?? (extractedThinkingTitle && extractedThinkingTitle !== this.defaultTitle ? extractedThinkingTitle : undefined);
+
+		this.lastExtractedTitle = nextTitle;
+
+		if (this.fixedScrollingMode || this._isExpanded.get() || this.element.isComplete || this.streamingCompleted) {
+			return;
+		}
+
+		if (nextTitle) {
+			this.setTitle(nextTitle);
+		} else {
+			this.setTitle(this.defaultTitle, true);
+			this.currentTitle = this.defaultTitle;
+		}
+	}
+
+	private registerToolDisposable(
+		disposable: IDisposable,
+		toolInvocationOrMarkdown?: IChatToolInvocation | IChatToolInvocationSerialized | IChatMarkdownContent
+	): void {
+		const toolCallId = toolInvocationOrMarkdown && (toolInvocationOrMarkdown.kind === 'toolInvocation' || toolInvocationOrMarkdown.kind === 'toolInvocationSerialized')
+			? toolInvocationOrMarkdown.toolCallId
+			: undefined;
+
+		if (!toolCallId) {
+			this._register(disposable);
+			return;
+		}
+
+		let toolStore = this.toolDisposables.get(toolCallId);
+		if (!toolStore) {
+			toolStore = new DisposableStore();
+			this.toolDisposables.set(toolCallId, toolStore);
+		}
+
+		toolStore.add(disposable);
+	}
+
 	private trackToolMetadata(
 		toolInvocationId?: string,
 		toolInvocationOrMarkdown?: IChatToolInvocation | IChatToolInvocationSerialized | IChatMarkdownContent
@@ -1357,7 +1397,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 					if (currentState.type === IChatToolInvocation.StateKind.Completed ||
 						currentState.type === IChatToolInvocation.StateKind.Cancelled) {
 						// Remove tools that should be hidden now or after completion.
-						if (toolInvocationOrMarkdown.presentation === 'hidden' || toolInvocationOrMarkdown.presentation === 'hiddenAfterComplete') {
+						if (toolInvocationOrMarkdown.presentation === 'hidden' || shouldHideToolInvocationAfterCompletion(toolInvocationOrMarkdown)) {
 							this.pendingRemovals.push({ toolCallId: toolInvocationOrMarkdown.toolCallId, toolLabel: currentToolLabel });
 							this.schedulePendingRemovalsFlush();
 						}
@@ -1510,12 +1550,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 		this.appendItemToDOM(result.domNode, item.toolInvocationId, item.toolInvocationOrMarkdown, item.originalParent);
 
 		if (result.disposable) {
-			const toolCallId = item.toolInvocationOrMarkdown && (item.toolInvocationOrMarkdown.kind === 'toolInvocation' || item.toolInvocationOrMarkdown.kind === 'toolInvocationSerialized') ? item.toolInvocationOrMarkdown.toolCallId : undefined;
-			if (toolCallId) {
-				this.toolDisposables.get(toolCallId)?.add(result.disposable);
-			} else {
-				this._register(result.disposable);
-			}
+			this.registerToolDisposable(result.disposable, item.toolInvocationOrMarkdown);
 		}
 	}
 

src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolInvocationPart.ts

@@ -13,7 +13,7 @@ import { IChatToolInvocation, IChatToolInvocationSerialized, ToolConfirmKind } f
 import { IChatRendererContent } from '../../../../common/model/chatViewModel.js';
 import { IChatTodoListService } from '../../../../common/tools/chatTodoListService.js';
 import { CodeBlockModelCollection } from '../../../../common/widget/codeBlockModelCollection.js';
-import { isToolResultInputOutputDetails, isToolResultOutputDetails, ToolInvocationPresentation } from '../../../../common/tools/languageModelToolsService.js';
+import { isToolResultInputOutputDetails, isToolResultOutputDetails } from '../../../../common/tools/languageModelToolsService.js';
 import { ChatTreeItem, IChatCodeBlockInfo } from '../../../chat.js';
 import { EditorPool } from '../chatContentCodePools.js';
 import { IChatContentPart, IChatContentPartRenderContext } from '../chatContentParts.js';
@@ -32,6 +32,7 @@ import { ChatToolOutputSubPart } from './chatToolOutputPart.js';
 import { ChatToolPostExecuteConfirmationPart } from './chatToolPostExecuteConfirmationPart.js';
 import { ChatToolProgressSubPart } from './chatToolProgressPart.js';
 import { ChatToolStreamingSubPart } from './chatToolStreamingSubPart.js';
+import { shouldHideToolInvocationAfterCompletion } from './chatToolInvocationVisibility.js';
 
 export class ChatToolInvocationPart extends Disposable implements IChatContentPart {
 	public readonly domNode: HTMLElement;
@@ -108,7 +109,7 @@ export class ChatToolInvocationPart extends Disposable implements IChatContentPa
 		const render = () => {
 			partStore.clear();
 
-			if (toolInvocation.presentation === ToolInvocationPresentation.HiddenAfterComplete && IChatToolInvocation.isComplete(toolInvocation)) {
+			if (shouldHideToolInvocationAfterCompletion(toolInvocation)) {
 				dom.hide(this.domNode);
 				return;
 			}

src/vs/workbench/contrib/chat/browser/widget/chatContentParts/toolInvocationParts/chatToolInvocationVisibility.ts

@@ -0,0 +1,52 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IChatToolInvocation, IChatToolInvocationSerialized } from '../../../../common/chatService/chatService.js';
+import { isToolResultInputOutputDetails, isToolResultOutputDetails, ToolInvocationPresentation } from '../../../../common/tools/languageModelToolsService.js';
+
+export function shouldKeepToolInvocationVisibleAfterCompletion(invocation: IChatToolInvocation | IChatToolInvocationSerialized): boolean {
+	if (invocation.presentation !== ToolInvocationPresentation.HiddenAfterComplete || !IChatToolInvocation.isComplete(invocation)) {
+		return false;
+	}
+
+	if (hasVisibleToolResult(invocation)) {
+		return true;
+	}
+
+	return !!invocation.isAttachedToThinking && !IChatToolInvocation.getConfirmationMessages(invocation);
+}
+
+export function shouldHideToolInvocationAfterCompletion(invocation: IChatToolInvocation | IChatToolInvocationSerialized): boolean {
+	if (invocation.presentation !== ToolInvocationPresentation.HiddenAfterComplete || !IChatToolInvocation.isComplete(invocation)) {
+		return false;
+	}
+
+	return !shouldKeepToolInvocationVisibleAfterCompletion(invocation);
+}
+
+function hasVisibleToolResult(invocation: IChatToolInvocation | IChatToolInvocationSerialized): boolean {
+	const resultDetails = IChatToolInvocation.resultDetails(invocation);
+	if (Array.isArray(resultDetails) && resultDetails.length > 0) {
+		return true;
+	}
+
+	if (isToolResultOutputDetails(resultDetails) || isToolResultInputOutputDetails(resultDetails)) {
+		return true;
+	}
+
+	if (invocation.toolSpecificData?.kind === 'resources' && invocation.toolSpecificData.values.length > 0) {
+		return true;
+	}
+
+	if (invocation.toolSpecificData?.kind === 'simpleToolInvocation') {
+		return true;
+	}
+
+	if (invocation.toolSpecificData?.kind === 'input' && !!invocation.toolSpecificData.mcpAppData) {
+		return true;
+	}
+
+	return false;
+}

src/vs/workbench/contrib/chat/test/browser/widget/chatContentParts/chatThinkingContentPart.test.ts

@@ -14,14 +14,15 @@ import { workbenchInstantiationService } from '../../../../../../test/browser/wo
 import { IConfigurationService } from '../../../../../../../platform/configuration/common/configuration.js';
 import { TestConfigurationService } from '../../../../../../../platform/configuration/test/common/testConfigurationService.js';
 import { ChatThinkingContentPart } from '../../../../browser/widget/chatContentParts/chatThinkingContentPart.js';
-import { IChatMarkdownContent, IChatThinkingPart } from '../../../../common/chatService/chatService.js';
+import { IChatMarkdownContent, IChatThinkingPart, IChatToolInvocation } from '../../../../common/chatService/chatService.js';
 import { IChatContentPartRenderContext, InlineTextModelCollection } from '../../../../browser/widget/chatContentParts/chatContentParts.js';
 import { IChatRendererContent, IChatResponseViewModel } from '../../../../common/model/chatViewModel.js';
 import { IChatMarkdownAnchorService } from '../../../../browser/widget/chatContentParts/chatMarkdownAnchorService.js';
 import { IMarkdownRenderer } from '../../../../../../../platform/markdown/browser/markdownRenderer.js';
 import { IRenderedMarkdown, MarkdownRenderOptions } from '../../../../../../../base/browser/markdownRenderer.js';
 import { IMarkdownString } from '../../../../../../../base/common/htmlContent.js';
 import { ThinkingDisplayMode } from '../../../../common/constants.js';
+import { ToolInvocationPresentation } from '../../../../common/tools/languageModelToolsService.js';
 import { CodeBlockModelCollection } from '../../../../common/widget/codeBlockModelCollection.js';
 import { EditorPool, DiffEditorPool } from '../../../../browser/widget/chatContentParts/chatContentCodePools.js';
 import { IHoverService } from '../../../../../../../platform/hover/browser/hover.js';
@@ -72,6 +73,56 @@ suite('ChatThinkingContentPart', () => {
 		};
 	}
 
+	function createMockToolInvocation(options?: {
+		toolCallId?: string;
+		toolId?: string;
+		invocationMessage?: string;
+		presentation?: ToolInvocationPresentation;
+		isAttachedToThinking?: boolean;
+		resultDetails?: unknown;
+		confirmationMessages?: unknown;
+	}) {
+		const progress = observableValue('toolProgress', { progress: undefined as number | undefined });
+		type MockToolState =
+			| {
+				type: IChatToolInvocation.StateKind.Executing;
+				progress: typeof progress;
+			}
+			| {
+				type: IChatToolInvocation.StateKind.Completed;
+				resultDetails: unknown;
+				confirmationMessages: unknown;
+				confirmed: undefined;
+				postConfirmed: undefined;
+				contentForModel: [];
+			};
+
+		const state = observableValue<MockToolState, void>('toolState', {
+			type: IChatToolInvocation.StateKind.Executing,
+			progress
+		});
+
+		const invocation = {
+			kind: 'toolInvocation' as const,
+			toolCallId: options?.toolCallId ?? 'test-tool-call',
+			toolId: options?.toolId ?? 'test_tool',
+			invocationMessage: options?.invocationMessage ?? 'Test tool invocation',
+			originMessage: undefined,
+			pastTenseMessage: options?.invocationMessage ?? 'Test tool invocation',
+			toolSpecificData: undefined,
+			presentation: options?.presentation ?? ToolInvocationPresentation.HiddenAfterComplete,
+			state,
+			source: undefined,
+			isAttachedToThinking: options?.isAttachedToThinking ?? true,
+		} as unknown as IChatToolInvocation;
+
+		return { invocation, state };
+	}
+
+	async function waitForAnimationFrame(): Promise<void> {
+		await new Promise<void>(resolve => mainWindow.requestAnimationFrame(() => resolve()));
+	}
+
 	setup(() => {
 		disposables = store.add(new DisposableStore());
 		instantiationService = workbenchInstantiationService(undefined, store);
@@ -593,6 +644,87 @@ suite('ChatThinkingContentPart', () => {
 			assert.strictEqual(factoryCalled, false, 'Factory should never have been called');
 		});
 
+		test('hiddenAfterComplete tools with visible results stay visible after completion', async () => {
+			const content = createThinkingPart('**Working**');
+			const context = createMockRenderContext(false);
+			const part = store.add(instantiationService.createInstance(
+				ChatThinkingContentPart,
+				content,
+				context,
+				mockMarkdownRenderer,
+				false
+			));
+
+			mainWindow.document.body.appendChild(part.domNode);
+			disposables.add(toDisposable(() => part.domNode.remove()));
+
+			const { invocation, state } = createMockToolInvocation({
+				invocationMessage: 'Captured DOM node screenshot',
+				resultDetails: { input: 'node', output: 'image bytes', isError: false }
+			});
+
+			part.appendItem(() => ({ domNode: $('div.test-tool-item', undefined, 'Result-bearing tool') }), invocation.toolId, invocation);
+
+			const button = part.domNode.querySelector('.monaco-button') as HTMLElement;
+			button.click();
+			assert.ok(part.domNode.querySelector('.test-tool-item'), 'Tool should render after expanding the thinking container');
+
+			state.set({
+				type: IChatToolInvocation.StateKind.Completed,
+				resultDetails: { input: 'node', output: 'image bytes', isError: false },
+				confirmationMessages: undefined,
+				confirmed: undefined,
+				postConfirmed: undefined,
+				contentForModel: []
+			}, undefined, undefined);
+
+			await waitForAnimationFrame();
+
+			assert.ok(part.domNode.querySelector('.test-tool-item'), 'Completed tool with visible result should remain rendered inside thinking');
+		});
+
+		test('confirmation-only hiddenAfterComplete tools clear stale titles after completion', async () => {
+			const content = createThinkingPart('**Working**');
+			const context = createMockRenderContext(false);
+			const part = store.add(instantiationService.createInstance(
+				ChatThinkingContentPart,
+				content,
+				context,
+				mockMarkdownRenderer,
+				false
+			));
+
+			mainWindow.document.body.appendChild(part.domNode);
+			disposables.add(toDisposable(() => part.domNode.remove()));
+
+			let factoryCalled = false;
+			const { invocation, state } = createMockToolInvocation({
+				invocationMessage: 'Captured DOM node screenshot',
+				confirmationMessages: { title: 'Confirm', message: { value: 'Confirm' } }
+			});
+
+			part.appendItem(() => {
+				factoryCalled = true;
+				return { domNode: $('div.test-tool-item', undefined, 'Confirmation-only tool') };
+			}, invocation.toolId, invocation);
+
+			assert.ok(part.domNode.textContent?.includes('Captured DOM node screenshot'), 'Collapsed thinking title should reflect the pending tool');
+
+			state.set({
+				type: IChatToolInvocation.StateKind.Completed,
+				resultDetails: undefined,
+				confirmationMessages: { title: 'Confirm', message: { value: 'Confirm' } },
+				confirmed: undefined,
+				postConfirmed: undefined,
+				contentForModel: []
+			}, undefined, undefined);
+
+			await waitForAnimationFrame();
+
+			assert.strictEqual(factoryCalled, false, 'Confirmation-only tool should be removed before it materializes');
+			assert.ok(!part.domNode.textContent?.includes('Captured DOM node screenshot'), 'Collapsed thinking title should no longer show the removed tool label');
+		});
+
 		test('lazy items should preserve append order when mixing tool and markdown items', () => {
 			// This test verifies that when tool invocations and markdown items are appended
 			// in a specific order while collapsed, the DOM order matches the append order