PR comment follow up

Adding a few notes and small fixes

Author: mjbvz

src/vs/workbench/api/common/extHostChatOutputRenderer.ts

@@ -48,6 +48,6 @@ export class ExtHostChatOutputRenderer implements ExtHostChatOutputRendererShape
 		}
 
 		const webview = this.webviews.createNewWebview(webviewHandle, {}, entry.extension);
-		return entry.renderer.renderChatOutput(valueData.buffer, webview, token);
+		return entry.renderer.renderChatOutput(valueData.buffer, webview, {}, token);
 	}
 }

src/vs/workbench/api/common/extHostTypeConverters.ts

@@ -3335,12 +3335,11 @@ export namespace LanguageModelToolResult2 {
 			if (item.kind === 'text') {
 				return new types.LanguageModelTextPart(item.value);
 			} else if (item.kind === 'data') {
-				// TODO: make sure we can remove this
-				// const mimeType = Object.values(types.ChatImageMimeType).includes(item.value.mimeType as types.ChatImageMimeType) ? item.value.mimeType as types.ChatImageMimeType : undefined;
-				// // if (!mimeType) {
-				// // 	throw new Error('Invalid MIME type');
-				// // }
-				return new types.LanguageModelDataPart(item.value.data.buffer, item.value.mimeType);
+				const mimeType = Object.values(types.ChatImageMimeType).includes(item.value.mimeType as types.ChatImageMimeType) ? item.value.mimeType as types.ChatImageMimeType : undefined;
+				if (!mimeType) {
+					throw new Error('Invalid MIME type');
+				}
+				return new types.LanguageModelDataPart(item.value.data.buffer, mimeType);
 			} else {
 				return new types.LanguageModelPromptTsxPart(item.value);
 			}

src/vs/workbench/contrib/chat/browser/chatContentParts/toolInvocationParts/chatToolOutputPart.ts

@@ -5,7 +5,7 @@
 
 import * as dom from '../../../../../../base/browser/dom.js';
 import { decodeBase64 } from '../../../../../../base/common/buffer.js';
-import { CancellationToken } from '../../../../../../base/common/cancellation.js';
+import { CancellationTokenSource } from '../../../../../../base/common/cancellation.js';
 import { Codicon } from '../../../../../../base/common/codicons.js';
 import { ThemeIcon } from '../../../../../../base/common/themables.js';
 import { localize } from '../../../../../../nls.js';
@@ -18,11 +18,14 @@ import { IChatContentPartRenderContext } from '../chatContentParts.js';
 import { ChatCustomProgressPart } from '../chatProgressContentPart.js';
 import { BaseChatToolInvocationSubPart } from './chatToolInvocationSubPart.js';
 
+// TODO: see if we can reuse existing types instead of adding ChatToolOutputSubPart
 export class ChatToolOutputSubPart extends BaseChatToolInvocationSubPart {
 	public readonly domNode: HTMLElement;
 
 	public override readonly codeblocks: IChatCodeBlockInfo[] = [];
 
+	private readonly _disposeCts = this._register(new CancellationTokenSource());
+
 	constructor(
 		toolInvocation: IChatToolInvocation | IChatToolInvocationSerialized,
 		_context: IChatContentPartRenderContext,
@@ -44,6 +47,11 @@ export class ChatToolOutputSubPart extends BaseChatToolInvocationSubPart {
 		this.domNode = this.createOutputPart(details);
 	}
 
+	public override dispose(): void {
+		this._disposeCts.dispose(true);
+		super.dispose();
+	}
+
 	private createOutputPart(details: IToolResultOutputDetails): HTMLElement {
 		const parent = dom.$('div.webview-output');
 		parent.style.maxHeight = '80vh';
@@ -53,7 +61,12 @@ export class ChatToolOutputSubPart extends BaseChatToolInvocationSubPart {
 		const progressPart = this.instantiationService.createInstance(ChatCustomProgressPart, progressMessage, ThemeIcon.modify(Codicon.loading, 'spin'));
 		parent.appendChild(progressPart.domNode);
 
-		this.chatOutputItemRendererService.renderOutputPart(details.output.mimeType, details.output.value.buffer, parent, CancellationToken.None).then((renderedItem) => {
+		// TODO: we also need to show the tool output in the UI
+		this.chatOutputItemRendererService.renderOutputPart(details.output.mimeType, details.output.value.buffer, parent, this._disposeCts.token).then((renderedItem) => {
+			if (this._disposeCts.token.isCancellationRequested) {
+				return;
+			}
+
 			this._register(renderedItem);
 
 			progressPart.domNode.remove();
@@ -62,6 +75,9 @@ export class ChatToolOutputSubPart extends BaseChatToolInvocationSubPart {
 			this._register(renderedItem.onDidChangeHeight(() => {
 				this._onDidChangeHeight.fire();
 			}));
+		}, (error) => {
+			// TODO: show error in UI too
+			console.error('Error rendering tool output:', error);
 		});
 
 		return parent;

src/vscode-dts/vscode.proposed.chatOutputRenderer.d.ts

@@ -24,6 +24,7 @@ declare module 'vscode' {
 
 	export interface ExtendedLanguageModelToolResult2 extends ExtendedLanguageModelToolResult {
 		// Temporary to allow `toolResultDetails` to return a ToolResultDataOutput
+		// TODO: Should this live here? Or should we be able to mark each `content` items as user/lm specific?
 		// TODO: Should we allow multiple per tool result?
 		toolResultDetails2?: Array<Uri | Location> | ToolResultDataOutput;
 	}
@@ -32,8 +33,10 @@ declare module 'vscode' {
 		/**
 		 * Given an output, render it into the provided webview.
 		 *
-		 * TODO: Should we make this more generic so that we could support inputs besides tool outputs?
-		 * For example, a generic `ChatResponseDataPart` type.
+		 * TODO:Should this take an object instead of Uint8Array? That would let you get the original mime. Useful
+		 * if we ever support registering for multiple mime types or using image/*.
+		 *
+		 * TODO: Figure out what to pass as context?
 		 *
 		 * @param data The data to render.
 		 * @param webview The webview to render the data into.
@@ -42,7 +45,7 @@ declare module 'vscode' {
 		 *
 		 * @returns A promise that resolves when the webview has been initialized and is ready to be presented to the user.
 		 */
-		renderChatOutput(data: Uint8Array, webview: Webview, token: CancellationToken): Thenable<void>;
+		renderChatOutput(data: Uint8Array, webview: Webview, ctx: {}, token: CancellationToken): Thenable<void>;
 	}
 
 	export namespace chat {