Skip to content

Commit 50d370b

Browse files
mjbvzmjbvz
committed
Add basic API for custom rendering inside of chat output
First sketch for a simple API that lets extensions render content in chat using a webview Right now this is targeting results from tool calls but we could potentially extend this to work with a more generic version of our chat response image part
1 parent 34d2ec4 commit 50d370b

13 files changed

+332
-13
lines changed

‎src/vs/platform/extensions/common/extensionsApiProposals.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ const _allApiProposals = {
3737
chatEditing: {
3838
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.chatEditing.d.ts',
3939
},
40+
chatOutputRenderer: {
41+
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.chatOutputRenderer.d.ts',
42+
},
4043
chatParticipantAdditions: {
4144
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts',
4245
},

‎src/vs/workbench/api/browser/extensionHost.contribution.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ import './mainThreadAiEmbeddingVector.js';
9191
import './mainThreadAiSettingsSearch.js';
9292
import './mainThreadMcp.js';
9393
import './mainThreadChatStatus.js';
94+
import './mainThreadChatOutputRenderer.js';
9495

9596
export class ExtensionPoints implements IWorkbenchContribution {
9697

‎src/vs/workbench/api/browser/mainThreadChatOutputRenderer.ts

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { VSBuffer } from '../../../base/common/buffer.js';
7+
import { Disposable, IDisposable } from '../../../base/common/lifecycle.js';
8+
import { IChatOutputItemRendererService } from '../../contrib/chat/browser/chatOutputItemRenderer.js';
9+
import { IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
10+
import { ExtHostChatOutputRendererShape, ExtHostContext, MainThreadChatOutputRendererShape } from '../common/extHost.protocol.js';
11+
import { MainThreadWebviews } from './mainThreadWebviews.js';
12+
13+
export class MainThreadChatOutputRenderer extends Disposable implements MainThreadChatOutputRendererShape {
14+
15+
private readonly _proxy: ExtHostChatOutputRendererShape;
16+
17+
private _webviewHandlePool = 0;
18+
19+
private readonly registeredRenderers = new Map<string, IDisposable>();
20+
21+
constructor(
22+
extHostContext: IExtHostContext,
23+
private readonly _mainThreadWebview: MainThreadWebviews,
24+
@IChatOutputItemRendererService private readonly _rendererService: IChatOutputItemRendererService,
25+
) {
26+
super();
27+
this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostChatOutputRenderer);
28+
}
29+
30+
override dispose(): void {
31+
super.dispose();
32+
33+
this.registeredRenderers.forEach(disposable => disposable.dispose());
34+
this.registeredRenderers.clear();
35+
}
36+
37+
$registerChatOutputRenderer(mime: string): void {
38+
this._rendererService.registerRenderer(mime, {
39+
renderOutputPart: async (mime, data, webview, token) => {
40+
const webviewHandle = `chat-output-${++this._webviewHandlePool}`;
41+
42+
this._mainThreadWebview.addWebview(webviewHandle, webview, {
43+
serializeBuffersForPostMessage: true,
44+
});
45+
46+
this._proxy.$renderChatPart(mime, VSBuffer.wrap(data), webviewHandle, token);
47+
},
48+
});
49+
}
50+
51+
$unregisterChatOutputRenderer(mime: string): void {
52+
this.registeredRenderers.get(mime)?.dispose();
53+
}
54+
}

‎src/vs/workbench/api/browser/mainThreadWebviewManager.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { MainThreadWebviews } from './mainThreadWebviews.js';
1111
import { MainThreadWebviewsViews } from './mainThreadWebviewViews.js';
1212
import * as extHostProtocol from '../common/extHost.protocol.js';
1313
import { extHostCustomer, IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
14+
import { MainThreadChatOutputRenderer } from './mainThreadChatOutputRenderer.js';
1415

1516
@extHostCustomer
1617
export class MainThreadWebviewManager extends Disposable {
@@ -31,5 +32,8 @@ export class MainThreadWebviewManager extends Disposable {
3132

3233
const webviewViews = this._register(instantiationService.createInstance(MainThreadWebviewsViews, context, webviews));
3334
context.set(extHostProtocol.MainContext.MainThreadWebviewViews, webviewViews);
35+
36+
const chatOutputRenderers = this._register(instantiationService.createInstance(MainThreadChatOutputRenderer, context, webviews));
37+
context.set(extHostProtocol.MainContext.MainThreadChatOutputRenderer, chatOutputRenderers);
3438
}
3539
}

‎src/vs/workbench/api/browser/mainThreadWebviews.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@ import { localize } from '../../../nls.js';
1313
import { ExtensionIdentifier } from '../../../platform/extensions/common/extensions.js';
1414
import { IOpenerService } from '../../../platform/opener/common/opener.js';
1515
import { IProductService } from '../../../platform/product/common/productService.js';
16-
import * as extHostProtocol from '../common/extHost.protocol.js';
17-
import { deserializeWebviewMessage, serializeWebviewMessage } from '../common/extHostWebviewMessaging.js';
18-
import { IOverlayWebview, IWebview, WebviewContentOptions, WebviewExtensionDescription } from '../../contrib/webview/browser/webview.js';
16+
import { IWebview, WebviewContentOptions, WebviewExtensionDescription } from '../../contrib/webview/browser/webview.js';
1917
import { IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
2018
import { SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
19+
import * as extHostProtocol from '../common/extHost.protocol.js';
20+
import { deserializeWebviewMessage, serializeWebviewMessage } from '../common/extHostWebviewMessaging.js';
2121

2222
export class MainThreadWebviews extends Disposable implements extHostProtocol.MainThreadWebviewsShape {
2323

@@ -43,7 +43,7 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
4343
this._proxy = context.getProxy(extHostProtocol.ExtHostContext.ExtHostWebviews);
4444
}
4545

46-
public addWebview(handle: extHostProtocol.WebviewHandle, webview: IOverlayWebview, options: { serializeBuffersForPostMessage: boolean }): void {
46+
public addWebview(handle: extHostProtocol.WebviewHandle, webview: IWebview, options: { serializeBuffersForPostMessage: boolean }): void {
4747
if (this._webviews.has(handle)) {
4848
throw new Error('Webview already registered');
4949
}
@@ -72,7 +72,7 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
7272
return webview.postMessage(message, arrayBuffers);
7373
}
7474

75-
private hookupWebviewEventDelegate(handle: extHostProtocol.WebviewHandle, webview: IOverlayWebview, options: { serializeBuffersForPostMessage: boolean }) {
75+
private hookupWebviewEventDelegate(handle: extHostProtocol.WebviewHandle, webview: IWebview, options: { serializeBuffersForPostMessage: boolean }) {
7676
const disposables = new DisposableStore();
7777

7878
disposables.add(webview.onDidClickLink((uri) => this.onDidClickLink(handle, uri)));

‎src/vs/workbench/api/common/extHost.api.impl.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ import { ExtHostWebviewViews } from './extHostWebviewView.js';
111111
import { IExtHostWindow } from './extHostWindow.js';
112112
import { IExtHostWorkspace } from './extHostWorkspace.js';
113113
import { ExtHostAiSettingsSearch } from './extHostAiSettingsSearch.js';
114+
import { ExtHostChatOutputRenderer } from './extHostChatOutputRenderer.js';
114115

115116
export interface IExtensionRegistries {
116117
mine: ExtensionDescriptionRegistry;
@@ -216,6 +217,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
216217
const extHostTesting = rpcProtocol.set(ExtHostContext.ExtHostTesting, accessor.get(IExtHostTesting));
217218
const extHostUriOpeners = rpcProtocol.set(ExtHostContext.ExtHostUriOpeners, new ExtHostUriOpeners(rpcProtocol));
218219
const extHostProfileContentHandlers = rpcProtocol.set(ExtHostContext.ExtHostProfileContentHandlers, new ExtHostProfileContentHandlers(rpcProtocol));
220+
const extHostChatOutputRenderer = rpcProtocol.set(ExtHostContext.ExtHostChatOutputRenderer, new ExtHostChatOutputRenderer(rpcProtocol, extHostWebviews));
219221
rpcProtocol.set(ExtHostContext.ExtHostInteractive, new ExtHostInteractive(rpcProtocol, extHostNotebook, extHostDocumentsAndEditors, extHostCommands, extHostLogService));
220222
const extHostLanguageModelTools = rpcProtocol.set(ExtHostContext.ExtHostLanguageModelTools, new ExtHostLanguageModelTools(rpcProtocol, extHostLanguageModels));
221223
const extHostChatAgents2 = rpcProtocol.set(ExtHostContext.ExtHostChatAgents2, new ExtHostChatAgents2(rpcProtocol, extHostLogService, extHostCommands, extHostDocuments, extHostLanguageModels, extHostDiagnostics, extHostLanguageModelTools));
@@ -1482,6 +1484,10 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
14821484
onDidDisposeChatSession: (listeners, thisArgs?, disposables?) => {
14831485
checkProposedApiEnabled(extension, 'chatParticipantPrivate');
14841486
return _asExtensionEvent(extHostChatAgents2.onDidDisposeChatSession)(listeners, thisArgs, disposables);
1487+
},
1488+
registerChatOutputRenderer: (mime: string, renderer: vscode.ChatOutputRenderer) => {
1489+
checkProposedApiEnabled(extension, 'chatOutputRenderer');
1490+
return extHostChatOutputRenderer.registerChatOutputRenderer(extension, mime, renderer);
14851491
}
14861492
};
14871493

‎src/vs/workbench/api/common/extHost.protocol.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1463,6 +1463,15 @@ export interface ExtHostUriOpenersShape {
14631463
$openUri(id: string, context: { resolvedUri: UriComponents; sourceUri: UriComponents }, token: CancellationToken): Promise<void>;
14641464
}
14651465

1466+
export interface MainThreadChatOutputRendererShape extends IDisposable {
1467+
$registerChatOutputRenderer(mime: string): void;
1468+
$unregisterChatOutputRenderer(mime: string): void;
1469+
}
1470+
1471+
export interface ExtHostChatOutputRendererShape {
1472+
$renderChatPart(mime: string, valueData: VSBuffer, webviewHandle: string, token: CancellationToken): Promise<void>;
1473+
}
1474+
14661475
export interface MainThreadProfileContentHandlersShape {
14671476
$registerProfileContentHandler(id: string, name: string, description: string | undefined, extensionId: string): Promise<void>;
14681477
$unregisterProfileContentHandler(id: string): Promise<void>;
@@ -3180,6 +3189,7 @@ export const MainContext = {
31803189
MainThreadAiEmbeddingVector: createProxyIdentifier<MainThreadAiEmbeddingVectorShape>('MainThreadAiEmbeddingVector'),
31813190
MainThreadChatStatus: createProxyIdentifier<MainThreadChatStatusShape>('MainThreadChatStatus'),
31823191
MainThreadAiSettingsSearch: createProxyIdentifier<MainThreadAiSettingsSearchShape>('MainThreadAiSettingsSearch'),
3192+
MainThreadChatOutputRenderer: createProxyIdentifier<MainThreadChatOutputRendererShape>('MainThreadChatOutputRenderer'),
31833193
};
31843194

31853195
export const ExtHostContext = {
@@ -3225,6 +3235,7 @@ export const ExtHostContext = {
32253235
ExtHostStorage: createProxyIdentifier<ExtHostStorageShape>('ExtHostStorage'),
32263236
ExtHostUrls: createProxyIdentifier<ExtHostUrlsShape>('ExtHostUrls'),
32273237
ExtHostUriOpeners: createProxyIdentifier<ExtHostUriOpenersShape>('ExtHostUriOpeners'),
3238+
ExtHostChatOutputRenderer: createProxyIdentifier<ExtHostChatOutputRendererShape>('ExtHostChatOutputRenderer'),
32283239
ExtHostProfileContentHandlers: createProxyIdentifier<ExtHostProfileContentHandlersShape>('ExtHostProfileContentHandlers'),
32293240
ExtHostOutputService: createProxyIdentifier<ExtHostOutputServiceShape>('ExtHostOutputService'),
32303241
ExtHostLabelService: createProxyIdentifier<ExtHostLabelServiceShape>('ExtHostLabelService'),

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import type * as vscode from 'vscode';
7+
import { CancellationToken } from '../../../base/common/cancellation.js';
8+
import { ExtHostChatOutputRendererShape, IMainContext, MainContext, MainThreadChatOutputRendererShape } from './extHost.protocol.js';
9+
import { Disposable } from './extHostTypes.js';
10+
import { ExtHostWebviews } from './extHostWebview.js';
11+
import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
12+
import { VSBuffer } from '../../../base/common/buffer.js';
13+
14+
export class ExtHostChatOutputRenderer implements ExtHostChatOutputRendererShape {
15+
16+
private readonly _proxy: MainThreadChatOutputRendererShape;
17+
18+
private readonly _renderers = new Map</*mime*/ string, { renderer: vscode.ChatOutputRenderer; extension: IExtensionDescription }>();
19+
20+
constructor(
21+
mainContext: IMainContext,
22+
private readonly webviews: ExtHostWebviews,
23+
) {
24+
this._proxy = mainContext.getProxy(MainContext.MainThreadChatOutputRenderer);
25+
}
26+
27+
registerChatOutputRenderer(extension: IExtensionDescription, mime: string, renderer: vscode.ChatOutputRenderer): vscode.Disposable {
28+
if (this._renderers.has(mime)) {
29+
throw new Error(`Chat response output renderer already registered for mime type: ${mime}`);
30+
}
31+
32+
this._renderers.set(mime, { extension, renderer });
33+
this._proxy.$registerChatOutputRenderer(mime);
34+
35+
return new Disposable(() => {
36+
this._renderers.delete(mime);
37+
this._proxy.$unregisterChatOutputRenderer(mime);
38+
});
39+
}
40+
41+
async $renderChatPart(mime: string, valueData: VSBuffer, webviewHandle: string, token: CancellationToken): Promise<void> {
42+
const entry = this._renderers.get(mime);
43+
if (!entry) {
44+
throw new Error(`No chat response output renderer registered for mime type: ${mime}`);
45+
}
46+
47+
const webview = this.webviews.createNewWebview(webviewHandle, {}, entry.extension);
48+
49+
const part = Object.freeze<vscode.ToolResultDataOutput>({ mime, value: valueData.buffer });
50+
return entry.renderer.renderChatOutput(part, webview, token);
51+
}
52+
}

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

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import { IChatAgentRequest, IChatAgentResult } from '../../contrib/chat/common/c
4242
import { IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
4343
import { IChatRequestVariableEntry, isImageVariableEntry } from '../../contrib/chat/common/chatVariableEntries.js';
4444
import { IChatAgentMarkdownContentWithVulnerability, IChatCodeCitation, IChatCommandButton, IChatConfirmation, IChatContentInlineReference, IChatContentReference, IChatExtensionsContent, IChatFollowup, IChatMarkdownContent, IChatMoveMessage, IChatPrepareToolInvocationPart, IChatProgressMessage, IChatResponseCodeblockUriPart, IChatTaskDto, IChatTaskResult, IChatTextEdit, IChatTreeData, IChatUserActionEvent, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
45-
import { IToolData, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
45+
import { IToolData, IToolResult, IToolResultInputOutputDetails, IToolResultOutputDetails } from '../../contrib/chat/common/languageModelToolsService.js';
4646
import * as chatProvider from '../../contrib/chat/common/languageModels.js';
4747
import { IChatMessageDataPart, IChatResponseDataPart, IChatResponsePromptTsxPart, IChatResponseTextPart } from '../../contrib/chat/common/languageModels.js';
4848
import { DebugTreeItemCollapsibleState, IDebugVisualizationTreeItem } from '../../contrib/debug/common/debug.js';
@@ -3335,11 +3335,12 @@ export namespace LanguageModelToolResult2 {
33353335
if (item.kind === 'text') {
33363336
return new types.LanguageModelTextPart(item.value);
33373337
} else if (item.kind === 'data') {
3338-
const mimeType = Object.values(types.ChatImageMimeType).includes(item.value.mimeType as types.ChatImageMimeType) ? item.value.mimeType as types.ChatImageMimeType : undefined;
3339-
if (!mimeType) {
3340-
throw new Error('Invalid MIME type');
3341-
}
3342-
return new types.LanguageModelDataPart(item.value.data.buffer, mimeType);
3338+
// TODO: make sure we can remove this
3339+
// const mimeType = Object.values(types.ChatImageMimeType).includes(item.value.mimeType as types.ChatImageMimeType) ? item.value.mimeType as types.ChatImageMimeType : undefined;
3340+
// // if (!mimeType) {
3341+
// // throw new Error('Invalid MIME type');
3342+
// // }
3343+
return new types.LanguageModelDataPart(item.value.data.buffer, item.value.mimeType);
33433344
} else {
33443345
return new types.LanguageModelPromptTsxPart(item.value);
33453346
}
@@ -3352,6 +3353,24 @@ export namespace LanguageModelToolResult2 {
33523353
}
33533354

33543355
let hasBuffers = false;
3356+
let detailsDto: Dto<Array<URI | types.Location> | IToolResultInputOutputDetails | IToolResultOutputDetails | undefined> = undefined;
3357+
if (Array.isArray(result.toolResultDetails)) {
3358+
detailsDto = result.toolResultDetails?.map(detail => {
3359+
return URI.isUri(detail) ? detail : Location.from(detail as vscode.Location);
3360+
});
3361+
} else {
3362+
if (result.toolResultDetails) {
3363+
detailsDto = {
3364+
output: {
3365+
type: 'data',
3366+
mimeType: (result.toolResultDetails as { mime: string; value: Uint8Array }).mime,
3367+
value: VSBuffer.wrap((result.toolResultDetails as { mime: string; value: Uint8Array }).value),
3368+
}
3369+
} satisfies IToolResultOutputDetails;
3370+
hasBuffers = true;
3371+
}
3372+
}
3373+
33553374
const dto: Dto<IToolResult> = {
33563375
content: result.content.map(item => {
33573376
if (item instanceof types.LanguageModelTextPart) {
@@ -3378,7 +3397,7 @@ export namespace LanguageModelToolResult2 {
33783397
}
33793398
}),
33803399
toolResultMessage: MarkdownString.fromStrict(result.toolResultMessage),
3381-
toolResultDetails: result.toolResultDetails?.map(detail => URI.isUri(detail) ? detail : Location.from(detail as vscode.Location)),
3400+
toolResultDetails: detailsDto,
33823401
};
33833402

33843403
return hasBuffers ? new SerializableObjectWithBuffers(dto) : dto;

0 commit comments

Comments
 (0)