Skip to content

Commit 27267ee

Browse files
roblourensroblourens
authored
Include session id with LanguageModelToolInvocationPrepareOptions (microsoft#252257)
* Include session id with LanguageModelToolInvocationPrepareOptions * Fix random token change from copilot
1 parent d967218 commit 27267ee

File tree

10 files changed

+46
-19
lines changed

10 files changed

+46
-19
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
8787
this._runningToolCalls.delete(dto.callId);
8888
}
8989
},
90-
prepareToolInvocation: (parameters, token) => this._proxy.$prepareToolInvocation(id, parameters, token),
90+
prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token),
9191
});
9292
this._tools.set(id, disposable);
9393
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ import { IChatContentInlineReference, IChatFollowup, IChatNotebookEdit, IChatPro
6161
import { IChatRequestVariableValue } from '../../contrib/chat/common/chatVariables.js';
6262
import { ChatAgentLocation } from '../../contrib/chat/common/constants.js';
6363
import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata, ILanguageModelChatSelector, ILanguageModelsChangeEvent } from '../../contrib/chat/common/languageModels.js';
64-
import { IPreparedToolInvocation, IToolInvocation, IToolProgressStep, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
64+
import { IPreparedToolInvocation, IToolInvocation, IToolInvocationPreparationContext, IToolProgressStep, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
6565
import { DebugConfigurationProviderTriggerKind, IAdapterDescriptor, IConfig, IDebugSessionReplMode, IDebugTestRunReference, IDebugVisualization, IDebugVisualizationContext, IDebugVisualizationTreeItem, MainThreadDebugVisualization } from '../../contrib/debug/common/debug.js';
6666
import { McpCollectionDefinition, McpConnectionState, McpServerDefinition, McpServerLaunch } from '../../contrib/mcp/common/mcpTypes.js';
6767
import * as notebookCommon from '../../contrib/notebook/common/notebookCommon.js';
@@ -1405,7 +1405,7 @@ export interface ExtHostLanguageModelToolsShape {
14051405
$invokeTool(dto: IToolInvocation, token: CancellationToken): Promise<Dto<IToolResult> | SerializableObjectWithBuffers<Dto<IToolResult>>>;
14061406
$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number>;
14071407

1408-
$prepareToolInvocation(toolId: string, parameters: unknown, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
1408+
$prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
14091409
}
14101410

14111411
export interface MainThreadUrlsShape extends IDisposable {

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
1111
import { revive } from '../../../base/common/marshalling.js';
1212
import { generateUuid } from '../../../base/common/uuid.js';
1313
import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
14-
import { IPreparedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
14+
import { IPreparedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolInvocationPreparationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
1515
import { ExtensionEditToolId, InternalEditToolId } from '../../contrib/chat/common/tools/editFileTool.js';
1616
import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/tools.js';
1717
import { checkProposedApiEnabled, isProposedApiEnabled } from '../../services/extensions/common/extensions.js';
@@ -220,13 +220,18 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
220220
return model;
221221
}
222222

223-
async $prepareToolInvocation(toolId: string, input: unknown, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
223+
async $prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
224224
const item = this._registeredTools.get(toolId);
225225
if (!item) {
226226
throw new Error(`Unknown tool ${toolId}`);
227227
}
228228

229-
const options: vscode.LanguageModelToolInvocationPrepareOptions<any> = { input };
229+
const options: vscode.LanguageModelToolInvocationPrepareOptions<any> = {
230+
input: context.parameters,
231+
chatRequestId: context.chatRequestId,
232+
chatSessionId: context.chatSessionId,
233+
chatInteractionId: context.chatInteractionId
234+
};
230235
if (isProposedApiEnabled(item.extension, 'chatParticipantPrivate') && item.tool.prepareInvocation2) {
231236
const result = await item.tool.prepareInvocation2(options, token);
232237
if (!result) {

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -355,7 +355,12 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
355355

356356
private async prepareToolInvocation(tool: IToolEntry, dto: IToolInvocation, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
357357
const prepared = tool.impl!.prepareToolInvocation ?
358-
await tool.impl!.prepareToolInvocation(dto.parameters, token)
358+
await tool.impl!.prepareToolInvocation({
359+
parameters: dto.parameters,
360+
chatRequestId: dto.chatRequestId,
361+
chatSessionId: dto.context?.sessionId,
362+
chatInteractionId: dto.chatInteractionId
363+
}, token)
359364
: undefined;
360365

361366
if (prepared?.confirmationMessages) {

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,13 @@ export function isToolInvocationContext(obj: any): obj is IToolInvocationContext
123123
return typeof obj === 'object' && typeof obj.sessionId === 'string';
124124
}
125125

126+
export interface IToolInvocationPreparationContext {
127+
parameters: any;
128+
chatRequestId?: string;
129+
chatSessionId?: string;
130+
chatInteractionId?: string;
131+
}
132+
126133
export interface IToolResultInputOutputDetails {
127134
readonly input: string;
128135
readonly output: ({
@@ -195,7 +202,7 @@ export interface IPreparedToolInvocation {
195202

196203
export interface IToolImpl {
197204
invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken): Promise<IToolResult>;
198-
prepareToolInvocation?(parameters: any, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
205+
prepareToolInvocation?(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
199206
}
200207

201208
export class ToolSet {

‎src/vs/workbench/contrib/chat/common/tools/editFileTool.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { INotebookService } from '../../../notebook/common/notebookService.js';
1313
import { ICodeMapperService } from '../../common/chatCodeMapperService.js';
1414
import { ChatModel } from '../../common/chatModel.js';
1515
import { IChatService } from '../../common/chatService.js';
16-
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolDataSource, ToolProgress } from '../../common/languageModelToolsService.js';
16+
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, ToolDataSource, ToolProgress } from '../../common/languageModelToolsService.js';
1717

1818
export const ExtensionEditToolId = 'vscode_editFile';
1919
export const InternalEditToolId = 'vscode_editFile_internal';
@@ -136,7 +136,7 @@ export class EditTool implements IToolImpl {
136136
};
137137
}
138138

139-
async prepareToolInvocation(parameters: any, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
139+
async prepareToolInvocation(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
140140
return {
141141
presentation: 'hidden'
142142
};

‎src/vs/workbench/contrib/chat/electron-browser/tools/fetchPageTool.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { ResourceSet } from '../../../../../base/common/map.js';
99
import { URI } from '../../../../../base/common/uri.js';
1010
import { localize } from '../../../../../nls.js';
1111
import { IWebContentExtractorService } from '../../../../../platform/webContentExtractor/common/webContentExtractor.js';
12-
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultTextPart, ToolDataSource, ToolProgress } from '../../common/languageModelToolsService.js';
12+
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, IToolResultTextPart, ToolDataSource, ToolProgress } from '../../common/languageModelToolsService.js';
1313
import { InternalFetchWebPageToolId } from '../../common/tools/tools.js';
1414

1515
export const FetchWebPageToolData: IToolData = {
@@ -75,8 +75,8 @@ export class FetchWebPageTool implements IToolImpl {
7575
};
7676
}
7777

78-
async prepareToolInvocation(parameters: any, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
79-
const map = this._parseUris(parameters.urls);
78+
async prepareToolInvocation(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
79+
const map = this._parseUris(context.parameters.urls);
8080
const invalid = new Array<string>();
8181
const valid = new Array<URI>();
8282
map.forEach((uri, url) => {

‎src/vs/workbench/contrib/extensions/common/installExtensionsTool.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { CancellationToken } from '../../../../base/common/cancellation.js';
77
import { MarkdownString } from '../../../../base/common/htmlContent.js';
88
import { localize } from '../../../../nls.js';
99
import { areSameExtensions } from '../../../../platform/extensionManagement/common/extensionManagementUtil.js';
10-
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolDataSource, ToolProgress } from '../../chat/common/languageModelToolsService.js';
10+
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, ToolDataSource, ToolProgress } from '../../chat/common/languageModelToolsService.js';
1111
import { IExtensionsWorkbenchService } from './extensions.js';
1212

1313
export const InstallExtensionsToolId = 'vscode_installExtensions';
@@ -44,7 +44,8 @@ export class InstallExtensionsTool implements IToolImpl {
4444
@IExtensionsWorkbenchService private readonly extensionsWorkbenchService: IExtensionsWorkbenchService,
4545
) { }
4646

47-
async prepareToolInvocation(parameters: InputParams, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
47+
async prepareToolInvocation(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
48+
const parameters = context.parameters as InputParams;
4849
return {
4950
confirmationMessages: {
5051
title: localize('installExtensionsTool.confirmationTitle', 'Install Extensions'),
@@ -68,4 +69,3 @@ export class InstallExtensionsTool implements IToolImpl {
6869
};
6970
}
7071
}
71-

‎src/vs/workbench/contrib/mcp/common/mcpService.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import { ILogService } from '../../../../platform/log/common/log.js';
1919
import { IProductService } from '../../../../platform/product/common/productService.js';
2020
import { StorageScope } from '../../../../platform/storage/common/storage.js';
2121
import { ChatResponseResource, getAttachableImageExtension } from '../../chat/common/chatModel.js';
22-
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultInputOutputDetails, ToolDataSource, ToolProgress, ToolSet } from '../../chat/common/languageModelToolsService.js';
22+
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, IToolResultInputOutputDetails, ToolDataSource, ToolProgress, ToolSet } from '../../chat/common/languageModelToolsService.js';
2323
import { McpCommandIds } from './mcpCommandIds.js';
2424
import { IMcpRegistry } from './mcpRegistryTypes.js';
2525
import { McpServer, McpServerMetadataCache } from './mcpServer.js';
@@ -246,7 +246,7 @@ class McpToolImplementation implements IToolImpl {
246246
@IProductService private readonly _productService: IProductService,
247247
) { }
248248

249-
async prepareToolInvocation(parameters: any): Promise<IPreparedToolInvocation> {
249+
async prepareToolInvocation(context: IToolInvocationPreparationContext): Promise<IPreparedToolInvocation> {
250250
const tool = this._tool;
251251
const server = this._server;
252252

@@ -277,7 +277,7 @@ class McpToolImplementation implements IToolImpl {
277277
}), { isTrusted: true }),
278278
toolSpecificData: {
279279
kind: 'input',
280-
rawInput: parameters
280+
rawInput: context.parameters
281281
}
282282
};
283283
}

‎src/vscode-dts/vscode.proposed.chatParticipantPrivate.d.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,16 @@ declare module 'vscode' {
190190
terminalCommand?: string;
191191
}
192192

193+
export interface LanguageModelToolInvocationPrepareOptions<T> {
194+
/**
195+
* The input that the tool is being invoked with.
196+
*/
197+
input: T;
198+
chatRequestId?: string;
199+
chatSessionId?: string;
200+
chatInteractionId?: string;
201+
}
202+
193203
export interface PreparedToolInvocation {
194204
pastTenseMessage?: string | MarkdownString;
195205
presentation?: 'hidden' | undefined;

0 commit comments

Comments
 (0)