perf: Fix leak in rendering markdown/edits in thinking/subagent parts (#308939)

Co-authored-by: Copilot <copilot@github.com>

Author: roblourens

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

@@ -58,6 +58,12 @@ interface ILazyToolItem {
 interface ILazyMarkdownItem {
 	kind: 'markdown';
 	lazy: Lazy<{ domNode: HTMLElement; disposable?: IDisposable }>;
+	/**
+	 * True when the caller passed an eagerDisposable that has already been registered on this
+	 * subagent part. In that case, materializeLazyItem must not register the factory's returned
+	 * disposable again.
+	 */
+	eagerlyRegistered?: boolean;
 }
 
 /**
@@ -864,25 +870,39 @@ export class ChatSubagentContentPart extends ChatCollapsibleContentPart implemen
 	/**
 	 * Appends a markdown item (e.g., an edit pill) to the subagent content part.
 	 * This is used to route codeblockUri parts with subAgentInvocationId to this subagent's container.
+	 *
+	 * When the caller has already created the content part eagerly (for example, a
+	 * pre-built `ChatMarkdownContentPart` wrapped in a factory), the caller MUST pass
+	 * that part as `eagerDisposable` so it is registered on this subagent part
+	 * immediately. Otherwise, if the subagent section is collapsed and the lazy item
+	 * is never materialized, the eagerly-created part would leak.
 	 */
 	public appendMarkdownItem(
 		factory: () => { domNode: HTMLElement; disposable?: IDisposable },
 		_codeblocksPartId: string | undefined,
 		_markdown: IChatMarkdownContent,
-		_originalParent?: HTMLElement
+		_originalParent?: HTMLElement,
+		eagerDisposable?: IDisposable,
 	): void {
+		// Register any caller-owned disposable up-front so it is always cleaned up
+		// with this subagent part, even if the lazy item is never materialized.
+		if (eagerDisposable) {
+			this._register(eagerDisposable);
+		}
+
 		// If expanded or has been expanded once, render immediately
 		if (this.isExpanded() || this.hasExpandedOnce) {
 			const result = factory();
 			this.appendMarkdownItemToDOM(result.domNode);
-			if (result.disposable) {
+			if (result.disposable && result.disposable !== eagerDisposable) {
 				this._register(result.disposable);
 			}
 		} else {
 			// Defer rendering until expanded
 			const item: ILazyMarkdownItem = {
 				kind: 'markdown',
 				lazy: new Lazy(factory),
+				eagerlyRegistered: !!eagerDisposable,
 			};
 			this.lazyItems.push(item);
 		}
@@ -1079,7 +1099,7 @@ export class ChatSubagentContentPart extends ChatCollapsibleContentPart implemen
 		} else if (item.kind === 'markdown') {
 			const result = item.lazy.value;
 			this.appendMarkdownItemToDOM(result.domNode);
-			if (result.disposable) {
+			if (result.disposable && !item.eagerlyRegistered) {
 				this._register(result.disposable);
 			}
 		} else if (item.kind === 'hook') {

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

@@ -1330,13 +1330,22 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 	 * Appends a tool invocation or content item to the thinking group.
 	 * The factory is called lazily - only when the thinking section is expanded.
 	 * If already expanded, the factory is called immediately.
+	 *
+	 * When the caller has already created the content part eagerly (for example, a
+	 * pre-built `ChatMarkdownContentPart` wrapped in a factory), the caller MUST pass
+	 * that part as `eagerDisposable` so it is registered on this thinking part
+	 * immediately. Otherwise, if the thinking section is collapsed and the lazy item
+	 * is never materialized (because the user never expands it), the eagerly-created
+	 * part would leak: its disposable is only referenced from inside the factory's
+	 * closure, which nothing ever calls.
 	 */
 	public appendItem(
 		factory: () => { domNode: HTMLElement; disposable?: IDisposable },
 		toolInvocationId?: string,
 		toolInvocationOrMarkdown?: IChatToolInvocation | IChatToolInvocationSerialized | IChatMarkdownContent,
 		originalParent?: HTMLElement,
-		onDidChangeDiff?: Event<IEditSessionDiffStats>
+		onDidChangeDiff?: Event<IEditSessionDiffStats>,
+		eagerDisposable?: IDisposable,
 	): void {
 		this.processPendingRemovals();
 
@@ -1352,6 +1361,12 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 			}));
 		}
 
+		// Register any caller-owned disposable up-front so it is always cleaned up
+		// with this thinking part, even if the lazy item is never materialized.
+		if (eagerDisposable) {
+			this._register(eagerDisposable);
+		}
+
 		// get random message based on tool type
 		if (this.workingSpinnerLabel) {
 			const isTerminalTool = toolInvocationOrMarkdown && (toolInvocationOrMarkdown.kind === 'toolInvocation' || toolInvocationOrMarkdown.kind === 'toolInvocationSerialized') && toolInvocationOrMarkdown.toolSpecificData?.kind === 'terminal';
@@ -1379,7 +1394,7 @@ ${this.hookCount > 0 ? `EXAMPLES WITH BLOCKED CONTENT (from hooks):
 				toolInvocationId,
 				toolInvocationOrMarkdown,
 				originalParent,
-				isHook: !toolInvocationOrMarkdown && !!toolInvocationId
+				isHook: !toolInvocationOrMarkdown && !!toolInvocationId,
 			};
 			this.lazyItems.push(item);
 		}

src/vs/workbench/contrib/chat/browser/widget/chatListRenderer.ts

@@ -2689,7 +2689,8 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 						() => ({ domNode: markdownPart.domNode, disposable: markdownPart }),
 						markdownPart.codeblocksPartId,
 						markdown,
-						templateData.value
+						templateData.value,
+						markdownPart,
 					);
 					return subagentPart;
 				}
@@ -2709,7 +2710,8 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 						markdownPart.codeblocksPartId,
 						markdown,
 						templateData.value,
-						markdownPart.onDidChangeDiff
+						markdownPart.onDidChangeDiff,
+						markdownPart,
 					);
 				}
 
@@ -2718,7 +2720,10 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 
 			if (this.shouldPinPart(markdown, context.element) && isComplete) {
 				if (lastThinking && markdownPart?.domNode) {
-					// Factory wrapping already-created markdown part
+					// Factory wrapping already-created markdown part.
+					// No eagerDisposable needed here because the markdownPart is returned
+					// from this method and tracked directly in renderedParts, so it will
+					// be disposed by clearRenderedParts.
 					lastThinking.appendItem(
 						() => ({ domNode: markdownPart.domNode, disposable: markdownPart }),
 						markdownPart.codeblocksPartId,

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

@@ -1606,4 +1606,114 @@ suite('ChatThinkingContentPart', () => {
 			assert.strictEqual(diffContainer, null, 'Should not render diff container when no diffs exist');
 		});
 	});
+
+	suite('eagerDisposable lifecycle', () => {
+		setup(() => {
+			mockConfigurationService.setUserConfiguration('chat.agent.thinkingStyle', ThinkingDisplayMode.Collapsed);
+		});
+
+		test('eagerDisposable is disposed when thinking part is disposed even if factory was never called', () => {
+			const content = createThinkingPart('**Working**');
+			const context = createMockRenderContext(false);
+
+			const part = instantiationService.createInstance(
+				ChatThinkingContentPart,
+				content,
+				context,
+				mockMarkdownRenderer,
+				false
+			);
+
+			mainWindow.document.body.appendChild(part.domNode);
+
+			let disposed = false;
+			const eagerDisposable = toDisposable(() => { disposed = true; });
+			const factory = () => ({
+				domNode: $('div.test-item'),
+				disposable: eagerDisposable,
+			});
+
+			// Append while collapsed — factory is NOT called
+			part.appendItem(factory, 'test-tool', undefined, undefined, undefined, eagerDisposable);
+
+			assert.strictEqual(disposed, false, 'Should not be disposed yet');
+
+			// Dispose the thinking part without ever expanding
+			part.domNode.remove();
+			part.dispose();
+
+			assert.strictEqual(disposed, true, 'eagerDisposable should be disposed with the thinking part');
+		});
+
+		test('eagerDisposable is disposed when thinking part is disposed after factory was called', () => {
+			const content = createThinkingPart('**Working**\nSome detailed analysis');
+			const context = createMockRenderContext(false);
+
+			const part = instantiationService.createInstance(
+				ChatThinkingContentPart,
+				content,
+				context,
+				mockMarkdownRenderer,
+				false
+			);
+
+			mainWindow.document.body.appendChild(part.domNode);
+
+			let disposed = false;
+			const eagerDisposable = toDisposable(() => { disposed = true; });
+			const factory = () => ({
+				domNode: $('div.test-item'),
+				disposable: eagerDisposable,
+			});
+
+			// Append while collapsed
+			part.appendItem(factory, 'test-tool', undefined, undefined, undefined, eagerDisposable);
+
+			// Expand to trigger factory call
+			const button = part.domNode.querySelector('.monaco-button') as HTMLElement;
+			button?.click();
+
+			assert.strictEqual(disposed, false, 'Should not be disposed yet');
+
+			// Dispose
+			part.domNode.remove();
+			part.dispose();
+
+			assert.strictEqual(disposed, true, 'eagerDisposable should be disposed even after being materialized');
+		});
+
+		test('appendItem without eagerDisposable disposes factory result on thinking part disposal', () => {
+			const content = createThinkingPart('**Working**\nSome detailed analysis');
+			const context = createMockRenderContext(false);
+
+			const part = instantiationService.createInstance(
+				ChatThinkingContentPart,
+				content,
+				context,
+				mockMarkdownRenderer,
+				false
+			);
+
+			mainWindow.document.body.appendChild(part.domNode);
+
+			// Expand first so factory is called immediately
+			const button = part.domNode.querySelector('.monaco-button') as HTMLElement;
+			button?.click();
+
+			let disposed = false;
+			const factory = () => ({
+				domNode: $('div.test-item'),
+				disposable: toDisposable(() => { disposed = true; }),
+			});
+
+			part.appendItem(factory, 'test-tool');
+
+			assert.strictEqual(disposed, false, 'Should not be disposed yet');
+
+			part.domNode.remove();
+			part.dispose();
+
+			assert.strictEqual(disposed, true, 'Factory disposable should be disposed with thinking part');
+		});
+	});
 });