Skip to content

Commit e9352a7

Browse files
roblourensroblourens
authored
Enable progress messages and references from variable resolvers (microsoft#205174)
* Enable progress messages and references from variable resolvers For microsoft#204539 * Fixes
1 parent bbd14fc commit e9352a7

File tree

11 files changed

+162
-27
lines changed

11 files changed

+162
-27
lines changed

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

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,16 @@
55

66
import { DisposableMap } from 'vs/base/common/lifecycle';
77
import { revive } from 'vs/base/common/marshalling';
8-
import { ExtHostChatVariablesShape, ExtHostContext, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
9-
import { IChatRequestVariableValue, IChatVariableData, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
8+
import { ExtHostChatVariablesShape, ExtHostContext, IChatVariableResolverProgressDto, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
9+
import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolverProgress, IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
1010
import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
1111

1212
@extHostNamedCustomer(MainContext.MainThreadChatVariables)
1313
export class MainThreadChatVariables implements MainThreadChatVariablesShape {
1414

1515
private readonly _proxy: ExtHostChatVariablesShape;
1616
private readonly _variables = new DisposableMap<number>();
17+
private readonly _pendingProgress = new Map<string, (part: IChatVariableResolverProgress) => void>();
1718

1819
constructor(
1920
extHostContext: IExtHostContext,
@@ -27,12 +28,22 @@ export class MainThreadChatVariables implements MainThreadChatVariablesShape {
2728
}
2829

2930
$registerVariable(handle: number, data: IChatVariableData): void {
30-
const registration = this._chatVariablesService.registerVariable(data, async (messageText, _arg, _model, token) => {
31-
return revive<IChatRequestVariableValue[]>(await this._proxy.$resolveVariable(handle, messageText, token));
31+
const registration = this._chatVariablesService.registerVariable(data, async (messageText, _arg, model, progress, token) => {
32+
const varRequestId = `${model.sessionId}-${handle}`;
33+
this._pendingProgress.set(varRequestId, progress);
34+
const result = revive<IChatRequestVariableValue[]>(await this._proxy.$resolveVariable(handle, varRequestId, messageText, token));
35+
36+
this._pendingProgress.delete(varRequestId);
37+
return result;
3238
});
3339
this._variables.set(handle, registration);
3440
}
3541

42+
async $handleProgressChunk(requestId: string, progress: IChatVariableResolverProgressDto): Promise<number | void> {
43+
const revivedProgress = revive(progress);
44+
this._pendingProgress.get(requestId)?.(revivedProgress as IChatVariableResolverProgress);
45+
}
46+
3647
$unregisterVariable(handle: number): void {
3748
this._variables.deleteAndDispose(handle);
3849
}

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ import { IChatAgentCommand, IChatAgentMetadata, IChatAgentRequest, IChatAgentRes
5454
import { IChatProgressResponseContent } from 'vs/workbench/contrib/chat/common/chatModel';
5555
import { IChatMessage, IChatResponseFragment, IChatResponseProviderMetadata } from 'vs/workbench/contrib/chat/common/chatProvider';
5656
import { IChatDynamicRequest, IChatProgress, IChatResponseErrorDetails, IChatUserActionEvent, InteractiveSessionVoteDirection, IChatFollowup } from 'vs/workbench/contrib/chat/common/chatService';
57-
import { IChatRequestVariableValue, IChatVariableData } from 'vs/workbench/contrib/chat/common/chatVariables';
57+
import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolverProgress } from 'vs/workbench/contrib/chat/common/chatVariables';
5858
import { DebugConfigurationProviderTriggerKind, MainThreadDebugVisualization, IAdapterDescriptor, IConfig, IDebugSessionReplMode, IDebugVisualization, IDebugVisualizationContext, IDebugVisualizationTreeItem } from 'vs/workbench/contrib/debug/common/debug';
5959
import { IInlineChatBulkEditResponse, IInlineChatEditResponse, IInlineChatFollowup, IInlineChatProgressItem, IInlineChatRequest, IInlineChatSession, InlineChatResponseFeedbackKind } from 'vs/workbench/contrib/inlineChat/common/inlineChat';
6060
import * as notebookCommon from 'vs/workbench/contrib/notebook/common/notebookCommon';
@@ -1231,15 +1231,19 @@ export interface ExtHostChatAgentsShape2 {
12311231
$releaseSession(sessionId: string): void;
12321232
}
12331233

1234+
export type IChatVariableResolverProgressDto =
1235+
| Dto<IChatVariableResolverProgress>;
1236+
12341237
export interface MainThreadChatVariablesShape extends IDisposable {
12351238
$registerVariable(handle: number, data: IChatVariableData): void;
1239+
$handleProgressChunk(requestId: string, progress: IChatVariableResolverProgressDto): Promise<number | void>;
12361240
$unregisterVariable(handle: number): void;
12371241
}
12381242

12391243
export type IChatRequestVariableValueDto = Dto<IChatRequestVariableValue>;
12401244

12411245
export interface ExtHostChatVariablesShape {
1242-
$resolveVariable(handle: number, messageText: string, token: CancellationToken): Promise<IChatRequestVariableValueDto[] | undefined>;
1246+
$resolveVariable(handle: number, requestId: string, messageText: string, token: CancellationToken): Promise<IChatRequestVariableValueDto[] | undefined>;
12431247
}
12441248

12451249
export interface MainThreadInlineChatShape extends IDisposable {

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

Lines changed: 80 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,35 +3,46 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import type * as vscode from 'vscode';
76
import { CancellationToken } from 'vs/base/common/cancellation';
7+
import { onUnexpectedExternalError } from 'vs/base/common/errors';
88
import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
9-
import { ExtensionIdentifier, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
10-
import { ExtHostChatVariablesShape, IMainContext, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
9+
import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
10+
import { ExtHostChatVariablesShape, IChatVariableResolverProgressDto, IMainContext, MainContext, MainThreadChatVariablesShape } from 'vs/workbench/api/common/extHost.protocol';
11+
import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters';
12+
import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
1113
import { IChatRequestVariableValue, IChatVariableData } from 'vs/workbench/contrib/chat/common/chatVariables';
12-
import { onUnexpectedExternalError } from 'vs/base/common/errors';
13-
import { ChatVariable } from 'vs/workbench/api/common/extHostTypeConverters';
14+
import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
15+
import type * as vscode from 'vscode';
1416

1517
export class ExtHostChatVariables implements ExtHostChatVariablesShape {
1618

1719
private static _idPool = 0;
1820

19-
private readonly _resolver = new Map<number, { extension: ExtensionIdentifier; data: IChatVariableData; resolver: vscode.ChatVariableResolver }>();
21+
private readonly _resolver = new Map<number, { extension: IExtensionDescription; data: IChatVariableData; resolver: vscode.ChatVariableResolver }>();
2022
private readonly _proxy: MainThreadChatVariablesShape;
2123

2224
constructor(mainContext: IMainContext) {
2325
this._proxy = mainContext.getProxy(MainContext.MainThreadChatVariables);
2426
}
2527

26-
async $resolveVariable(handle: number, messageText: string, token: CancellationToken): Promise<IChatRequestVariableValue[] | undefined> {
28+
async $resolveVariable(handle: number, requestId: string, messageText: string, token: CancellationToken): Promise<IChatRequestVariableValue[] | undefined> {
2729
const item = this._resolver.get(handle);
2830
if (!item) {
2931
return undefined;
3032
}
3133
try {
32-
const value = await item.resolver.resolve(item.data.name, { prompt: messageText }, token);
33-
if (value) {
34-
return value.map(ChatVariable.from);
34+
if (item.resolver.resolve2) {
35+
checkProposedApiEnabled(item.extension, 'chatAgents2Additions');
36+
const stream = new ChatVariableResolverResponseStream(requestId, this._proxy);
37+
const value = await item.resolver.resolve2(item.data.name, { prompt: messageText }, stream.apiObject, token);
38+
if (value) {
39+
return value.map(typeConvert.ChatVariable.from);
40+
}
41+
} else {
42+
const value = await item.resolver.resolve(item.data.name, { prompt: messageText }, token);
43+
if (value) {
44+
return value.map(typeConvert.ChatVariable.from);
45+
}
3546
}
3647
} catch (err) {
3748
onUnexpectedExternalError(err);
@@ -41,7 +52,7 @@ export class ExtHostChatVariables implements ExtHostChatVariablesShape {
4152

4253
registerVariableResolver(extension: IExtensionDescription, name: string, description: string, resolver: vscode.ChatVariableResolver): IDisposable {
4354
const handle = ExtHostChatVariables._idPool++;
44-
this._resolver.set(handle, { extension: extension.identifier, data: { name, description }, resolver: resolver });
55+
this._resolver.set(handle, { extension, data: { name, description }, resolver: resolver });
4556
this._proxy.$registerVariable(handle, { name, description });
4657

4758
return toDisposable(() => {
@@ -50,3 +61,61 @@ export class ExtHostChatVariables implements ExtHostChatVariablesShape {
5061
});
5162
}
5263
}
64+
65+
class ChatVariableResolverResponseStream {
66+
67+
private _isClosed: boolean = false;
68+
private _apiObject: vscode.ChatVariableResolverResponseStream | undefined;
69+
70+
constructor(
71+
private readonly _requestId: string,
72+
private readonly _proxy: MainThreadChatVariablesShape,
73+
) { }
74+
75+
close() {
76+
this._isClosed = true;
77+
}
78+
79+
get apiObject() {
80+
if (!this._apiObject) {
81+
const that = this;
82+
83+
function throwIfDone(source: Function | undefined) {
84+
if (that._isClosed) {
85+
const err = new Error('Response stream has been closed');
86+
Error.captureStackTrace(err, source);
87+
throw err;
88+
}
89+
}
90+
91+
const _report = (progress: IChatVariableResolverProgressDto) => {
92+
this._proxy.$handleProgressChunk(this._requestId, progress);
93+
};
94+
95+
this._apiObject = {
96+
progress(value) {
97+
throwIfDone(this.progress);
98+
const part = new extHostTypes.ChatResponseProgressPart(value);
99+
const dto = typeConvert.ChatResponseProgressPart.to(part);
100+
_report(dto);
101+
return this;
102+
},
103+
reference(value) {
104+
throwIfDone(this.reference);
105+
const part = new extHostTypes.ChatResponseReferencePart(value);
106+
const dto = typeConvert.ChatResponseReferencePart.to(part);
107+
_report(dto);
108+
return this;
109+
},
110+
push(part) {
111+
throwIfDone(this.push);
112+
const dto = typeConvert.ChatResponsePart.to(part);
113+
_report(dto as IChatVariableResolverProgressDto);
114+
return this;
115+
}
116+
};
117+
}
118+
119+
return this._apiObject;
120+
}
121+
}

‎src/vs/workbench/contrib/chat/browser/chatFollowups.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ export class ChatFollowups<T extends IChatFollowup | IInlineChatFollowup> extend
3737
return;
3838
}
3939

40+
if (!this.chatAgentService.getDefaultAgent()) {
41+
// No default agent yet, which affects how followups are rendered, so can't render this yet
42+
return;
43+
}
44+
4045
const tooltip = 'tooltip' in followup ? followup.tooltip : undefined;
4146
const button = this._register(new Button(container, { ...this.options, supportIcons: true, title: tooltip }));
4247
if (followup.kind === 'reply') {

‎src/vs/workbench/contrib/chat/browser/chatVariables.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { IChatWidgetService } from 'vs/workbench/contrib/chat/browser/chat';
1313
import { ChatDynamicVariableModel } from 'vs/workbench/contrib/chat/browser/contrib/chatDynamicVariables';
1414
import { IChatModel, IChatRequestVariableData } from 'vs/workbench/contrib/chat/common/chatModel';
1515
import { IParsedChatRequest, ChatRequestVariablePart, ChatRequestDynamicVariablePart } from 'vs/workbench/contrib/chat/common/chatParserTypes';
16-
import { IChatVariablesService, IChatRequestVariableValue, IChatVariableData, IChatVariableResolver, IDynamicVariable } from 'vs/workbench/contrib/chat/common/chatVariables';
16+
import { IChatVariablesService, IChatRequestVariableValue, IChatVariableData, IChatVariableResolver, IDynamicVariable, IChatVariableResolverProgress } from 'vs/workbench/contrib/chat/common/chatVariables';
1717

1818
interface IChatData {
1919
data: IChatVariableData;
@@ -30,7 +30,7 @@ export class ChatVariablesService implements IChatVariablesService {
3030
) {
3131
}
3232

33-
async resolveVariables(prompt: IParsedChatRequest, model: IChatModel, token: CancellationToken): Promise<IChatRequestVariableData> {
33+
async resolveVariables(prompt: IParsedChatRequest, model: IChatModel, progress: (part: IChatVariableResolverProgress) => void, token: CancellationToken): Promise<IChatRequestVariableData> {
3434
let resolvedVariables: { name: string; range: IOffsetRange; values: IChatRequestVariableValue[] }[] = [];
3535
const jobs: Promise<any>[] = [];
3636

@@ -39,7 +39,7 @@ export class ChatVariablesService implements IChatVariablesService {
3939
if (part instanceof ChatRequestVariablePart) {
4040
const data = this._resolver.get(part.variableName.toLowerCase());
4141
if (data) {
42-
jobs.push(data.resolver(prompt.text, part.variableArg, model, token).then(values => {
42+
jobs.push(data.resolver(prompt.text, part.variableArg, model, progress, token).then(values => {
4343
if (values?.length) {
4444
resolvedVariables[i] = { name: part.variableName, range: part.range, values };
4545
}

‎src/vs/workbench/contrib/chat/browser/contrib/chatHistoryVariables.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ class ChatHistoryVariables extends Disposable {
1515
) {
1616
super();
1717

18-
this._register(chatVariablesService.registerVariable({ name: 'response', description: '', canTakeArgument: true, hidden: true }, async (message, arg, model, token) => {
18+
this._register(chatVariablesService.registerVariable({ name: 'response', description: '', canTakeArgument: true, hidden: true }, async (message, arg, model, progress, token) => {
1919
if (!arg) {
2020
return undefined;
2121
}

‎src/vs/workbench/contrib/chat/common/chatServiceImpl.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -551,7 +551,7 @@ export class ChatService extends Disposable implements IChatService {
551551

552552
const initVariableData: IChatRequestVariableData = { variables: [] };
553553
request = model.addRequest(parsedRequest, initVariableData, agent, agentSlashCommandPart?.command);
554-
const variableData = await this.chatVariablesService.resolveVariables(parsedRequest, model, token);
554+
const variableData = await this.chatVariablesService.resolveVariables(parsedRequest, model, progressCallback, token);
555555
request.variableData = variableData;
556556

557557
const promptTextResult = getPromptText(request.message);

‎src/vs/workbench/contrib/chat/common/chatVariables.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { IRange } from 'vs/editor/common/core/range';
1010
import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
1111
import { IChatModel, IChatRequestVariableData } from 'vs/workbench/contrib/chat/common/chatModel';
1212
import { IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
13+
import { IChatContentReference, IChatProgressMessage } from 'vs/workbench/contrib/chat/common/chatService';
1314

1415
export interface IChatVariableData {
1516
name: string;
@@ -25,9 +26,13 @@ export interface IChatRequestVariableValue {
2526
description?: string;
2627
}
2728

29+
export type IChatVariableResolverProgress =
30+
| IChatContentReference
31+
| IChatProgressMessage;
32+
2833
export interface IChatVariableResolver {
2934
// TODO should we spec "zoom level"
30-
(messageText: string, arg: string | undefined, model: IChatModel, token: CancellationToken): Promise<IChatRequestVariableValue[] | undefined>;
35+
(messageText: string, arg: string | undefined, model: IChatModel, progress: (part: IChatVariableResolverProgress) => void, token: CancellationToken): Promise<IChatRequestVariableValue[] | undefined>;
3136
}
3237

3338
export const IChatVariablesService = createDecorator<IChatVariablesService>('IChatVariablesService');
@@ -42,7 +47,7 @@ export interface IChatVariablesService {
4247
/**
4348
* Resolves all variables that occur in `prompt`
4449
*/
45-
resolveVariables(prompt: IParsedChatRequest, model: IChatModel, token: CancellationToken): Promise<IChatRequestVariableData>;
50+
resolveVariables(prompt: IParsedChatRequest, model: IChatModel, progress: (part: IChatVariableResolverProgress) => void, token: CancellationToken): Promise<IChatRequestVariableData>;
4651
}
4752

4853
export interface IDynamicVariable {

‎src/vs/workbench/contrib/chat/test/browser/chatVariables.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ suite('ChatVariables', function () {
4242

4343
const resolveVariables = async (text: string) => {
4444
const result = parser.parseChatRequest('1', text);
45-
return await service.resolveVariables(result, null!, CancellationToken.None);
45+
return await service.resolveVariables(result, null!, () => { }, CancellationToken.None);
4646
};
4747

4848
{

‎src/vs/workbench/contrib/chat/test/common/mockChatVariables.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { CancellationToken } from 'vs/base/common/cancellation';
77
import { IDisposable } from 'vs/base/common/lifecycle';
88
import { IChatModel, IChatRequestVariableData } from 'vs/workbench/contrib/chat/common/chatModel';
99
import { IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
10-
import { IChatVariableData, IChatVariableResolver, IChatVariablesService, IDynamicVariable } from 'vs/workbench/contrib/chat/common/chatVariables';
10+
import { IChatVariableData, IChatVariableResolver, IChatVariableResolverProgress, IChatVariablesService, IDynamicVariable } from 'vs/workbench/contrib/chat/common/chatVariables';
1111

1212
export class MockChatVariablesService implements IChatVariablesService {
1313
_serviceBrand: undefined;
@@ -27,7 +27,7 @@ export class MockChatVariablesService implements IChatVariablesService {
2727
return [];
2828
}
2929

30-
async resolveVariables(prompt: IParsedChatRequest, model: IChatModel, token: CancellationToken): Promise<IChatRequestVariableData> {
30+
async resolveVariables(prompt: IParsedChatRequest, model: IChatModel, progress: (part: IChatVariableResolverProgress) => void, token: CancellationToken): Promise<IChatRequestVariableData> {
3131
return {
3232
variables: []
3333
};

0 commit comments

Comments
 (0)