chat: allow tools to report progress (#246768)

* chat: allow tools to report progress

This adds a `toolProgress` proposed API that allows extensions to report
progress keyed on the `toolInvocationToken`. Internally, when given to a
tool, this now includes the call ID of the tool.

We can use that both from extensions and internally to report progress
for tools, and I hooked this up for MCP servers. Currently only the text
is updated.

Involved some changes in the progress service internally:

- Previously `viewId` was a naked string, I wrapped it in an object to
  make it more identifiable.
- The progress service is in `workbench/services` and directly calls
  into other services to effect progress. In leui of going for a full
  'contribution' model, I made a small `ILanguageModelToolProgressService`
	that it writes state into and that can be read back out. The state
	(an observable) is kept as long as progress is ongoing or this is
	is an observer.

* be less creative ;)

* pr comments

* pass in an options object instead

* Revert "pass in an options object instead"

This reverts commit c6c85b3e9d.

* update with api feedback
This commit is contained in:
Connor Peet
2025-04-23 11:23:45 -07:00
committed by GitHub
21 changed files with 194 additions and 67 deletions
@@ -395,6 +395,9 @@ const _allApiProposals = {
tokenInformation: {
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.tokenInformation.d.ts',
},
toolProgress: {
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.toolProgress.d.ts',
},
treeViewActiveItem: {
proposal: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.treeViewActiveItem.d.ts',
},
@@ -6,7 +6,7 @@
import { CancellationToken } from '../../../base/common/cancellation.js';
import { Disposable, DisposableMap } from '../../../base/common/lifecycle.js';
import { revive } from '../../../base/common/marshalling.js';
import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolInvocation, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolInvocation, IToolProgressStep, IToolResult, ToolProgress } from '../../contrib/chat/common/languageModelToolsService.js';
import { IExtHostContext, extHostNamedCustomer } from '../../services/extensions/common/extHostCustomers.js';
import { Dto } from '../../services/extensions/common/proxyIdentifier.js';
import { ExtHostContext, ExtHostLanguageModelToolsShape, MainContext, MainThreadLanguageModelToolsShape } from '../common/extHost.protocol.js';
@@ -16,7 +16,10 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
private readonly _proxy: ExtHostLanguageModelToolsShape;
private readonly _tools = this._register(new DisposableMap<string>());
private readonly _countTokenCallbacks = new Map</* call ID */string, CountTokensCallback>();
private readonly _runningToolCalls = new Map</* call ID */string, {
countTokens: CountTokensCallback;
progress: ToolProgress;
}>();
constructor(
extHostContext: IExtHostContext,
@@ -45,26 +48,30 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
};
}
$acceptToolProgress(callId: string, progress: IToolProgressStep): void {
this._runningToolCalls.get(callId)?.progress.report(progress);
}
$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number> {
const fn = this._countTokenCallbacks.get(callId);
const fn = this._runningToolCalls.get(callId);
if (!fn) {
throw new Error(`Tool invocation call ${callId} not found`);
}
return fn(input, token);
return fn.countTokens(input, token);
}
$registerTool(id: string): void {
const disposable = this._languageModelToolsService.registerToolImplementation(
id,
{
invoke: async (dto, countTokens, token) => {
invoke: async (dto, countTokens, progress, token) => {
try {
this._countTokenCallbacks.set(dto.callId, countTokens);
this._runningToolCalls.set(dto.callId, { countTokens, progress });
const resultDto = await this._proxy.$invokeTool(dto, token);
return revive(resultDto) as IToolResult;
} finally {
this._countTokenCallbacks.delete(dto.callId);
this._runningToolCalls.delete(dto.callId);
}
},
prepareToolInvocation: (parameters, token) => this._proxy.$prepareToolInvocation(id, parameters, token),
@@ -59,7 +59,7 @@ import { IChatContentInlineReference, IChatFollowup, IChatNotebookEdit, IChatPro
import { IChatRequestVariableValue } from '../../contrib/chat/common/chatVariables.js';
import { ChatAgentLocation } from '../../contrib/chat/common/constants.js';
import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata, ILanguageModelChatSelector, ILanguageModelsChangeEvent } from '../../contrib/chat/common/languageModels.js';
import { IPreparedToolInvocation, IToolData, IToolInvocation, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
import { IPreparedToolInvocation, IToolData, IToolInvocation, IToolProgressStep, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
import { DebugConfigurationProviderTriggerKind, IAdapterDescriptor, IConfig, IDebugSessionReplMode, IDebugTestRunReference, IDebugVisualization, IDebugVisualizationContext, IDebugVisualizationTreeItem, MainThreadDebugVisualization } from '../../contrib/debug/common/debug.js';
import { McpCollectionDefinition, McpConnectionState, McpServerDefinition, McpServerLaunch } from '../../contrib/mcp/common/mcpTypes.js';
import * as notebookCommon from '../../contrib/notebook/common/notebookCommon.js';
@@ -1380,6 +1380,7 @@ export type IToolDataDto = Omit<IToolData, 'when'>;
export interface MainThreadLanguageModelToolsShape extends IDisposable {
$getTools(): Promise<Dto<IToolDataDto>[]>;
$acceptToolProgress(callId: string, progress: IToolProgressStep): void;
$invokeTool(dto: IToolInvocation, token?: CancellationToken): Promise<Dto<IToolResult>>;
$countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number>;
$registerTool(id: string): void;
@@ -138,7 +138,21 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
};
}
const extensionResult = await raceCancellation(Promise.resolve(item.tool.invoke(options, token)), token);
let progress: vscode.Progress<{ message?: string | vscode.MarkdownString; increment?: number }> | undefined;
if (isProposedApiEnabled(item.extension, 'toolProgress')) {
progress = {
report: value => {
this._proxy.$acceptToolProgress(dto.callId, {
message: typeConvert.MarkdownString.fromStrict(value.message),
increment: value.increment,
total: 100,
});
}
};
}
// todo: 'any' cast because TS can't handle the overloads
const extensionResult = await raceCancellation(Promise.resolve((item.tool.invoke as any)(options, token, progress!)), token);
if (!extensionResult) {
throw new CancellationError();
}
@@ -9,6 +9,7 @@ import { Codicon } from '../../../../../base/common/codicons.js';
import { Emitter } from '../../../../../base/common/event.js';
import { IMarkdownString, MarkdownString } from '../../../../../base/common/htmlContent.js';
import { Disposable, DisposableStore, IDisposable, thenIfNotDisposed, toDisposable } from '../../../../../base/common/lifecycle.js';
import { autorunWithStore } from '../../../../../base/common/observable.js';
import { ThemeIcon } from '../../../../../base/common/themables.js';
import { URI } from '../../../../../base/common/uri.js';
import { generateUuid } from '../../../../../base/common/uuid.js';
@@ -458,27 +459,37 @@ class ChatToolInvocationSubPart extends Disposable {
}
private createProgressPart(): HTMLElement {
let content: IMarkdownString;
if (this.toolInvocation.isComplete && this.toolInvocation.isConfirmed !== false && this.toolInvocation.pastTenseMessage) {
content = typeof this.toolInvocation.pastTenseMessage === 'string' ?
new MarkdownString().appendText(this.toolInvocation.pastTenseMessage) :
this.toolInvocation.pastTenseMessage;
const part = this.renderProgressContent(this.toolInvocation.pastTenseMessage);
this._register(part);
return part.domNode;
} else {
content = typeof this.toolInvocation.invocationMessage === 'string' ?
new MarkdownString().appendText(this.toolInvocation.invocationMessage + '…') :
MarkdownString.lift(this.toolInvocation.invocationMessage).appendText('…');
const container = document.createElement('div');
const progressObservable = this.toolInvocation.kind === 'toolInvocation' ? this.toolInvocation.progress : undefined;
this._register(autorunWithStore((reader, store) => {
const progress = progressObservable?.read(reader);
const part = store.add(this.renderProgressContent(progress?.message || this.toolInvocation.invocationMessage));
dom.reset(container, part.domNode);
}));
return container;
}
}
private renderProgressContent(content: IMarkdownString | string) {
if (typeof content === 'string') {
content = new MarkdownString().appendText(content);
}
const progressMessage: IChatProgressMessage = {
kind: 'progressMessage',
content
};
const iconOverride = !this.toolInvocation.isConfirmed ?
Codicon.error :
this.toolInvocation.isComplete ?
Codicon.check : undefined;
const progressPart = this._register(this.instantiationService.createInstance(ChatProgressContentPart, progressMessage, this.renderer, this.context, undefined, true, iconOverride));
return progressPart.domNode;
return this.instantiationService.createInstance(ChatProgressContentPart, progressMessage, this.renderer, this.context, undefined, true, iconOverride);
}
private createTerminalMarkdownProgressPart(toolInvocation: IChatToolInvocation | IChatToolInvocationSerialized, terminalData: IChatTerminalToolInvocationData): HTMLElement {
@@ -7,6 +7,7 @@ import { $ } from '../../../../base/browser/dom.js';
import { Dialog } from '../../../../base/browser/ui/dialog/dialog.js';
import { toAction, WorkbenchActionExecutedClassification, WorkbenchActionExecutedEvent } from '../../../../base/common/actions.js';
import { timeout } from '../../../../base/common/async.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { Codicon } from '../../../../base/common/codicons.js';
import { toErrorMessage } from '../../../../base/common/errorMessage.js';
import { isCancellationError } from '../../../../base/common/errors.js';
@@ -52,11 +53,13 @@ import { IHostService } from '../../../services/host/browser/host.js';
import { IWorkbenchLayoutService, Parts } from '../../../services/layout/browser/layoutService.js';
import { ILifecycleService } from '../../../services/lifecycle/common/lifecycle.js';
import { IViewsService } from '../../../services/views/common/viewsService.js';
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolProgress } from '../../chat/common/languageModelToolsService.js';
import { IExtensionsWorkbenchService } from '../../extensions/common/extensions.js';
import { IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../common/chatAgents.js';
import { ChatContextKeys } from '../common/chatContextKeys.js';
import { ChatEntitlement, ChatEntitlementContext, ChatEntitlementRequests, ChatEntitlementService, IChatEntitlementService } from '../common/chatEntitlementService.js';
import { IChatRequestModel, ChatRequestModel, ChatModel, IChatRequestVariableData, IChatRequestToolEntry } from '../common/chatModel.js';
import { ChatModel, ChatRequestModel, IChatRequestModel, IChatRequestToolEntry, IChatRequestVariableData } from '../common/chatModel.js';
import { ChatRequestAgentPart, ChatRequestToolPart } from '../common/chatParserTypes.js';
import { IChatProgress, IChatService } from '../common/chatService.js';
import { ChatAgentLocation, ChatConfiguration, ChatMode, validateChatMode } from '../common/constants.js';
import { ILanguageModelsService } from '../common/languageModels.js';
@@ -64,9 +67,6 @@ import { CHAT_CATEGORY, CHAT_OPEN_ACTION_ID, CHAT_SETUP_ACTION_ID } from './acti
import { ChatViewId, IChatWidgetService, showCopilotView } from './chat.js';
import { CHAT_SIDEBAR_PANEL_ID } from './chatViewPane.js';
import './media/chatSetup.css';
import { ChatRequestAgentPart, ChatRequestToolPart } from '../common/chatParserTypes.js';
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../chat/common/languageModelToolsService.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
const defaultChat = {
extensionId: product.defaultChatAgent?.extensionId ?? '',
@@ -518,7 +518,7 @@ class SetupTool extends Disposable implements IToolImpl {
super();
}
invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken): Promise<IToolResult> {
const result: IToolResult = {
content: [
{
@@ -12,7 +12,7 @@ import { Emitter } from '../../../../base/common/event.js';
import { MarkdownString } from '../../../../base/common/htmlContent.js';
import { Iterable } from '../../../../base/common/iterator.js';
import { Lazy } from '../../../../base/common/lazy.js';
import { Disposable, DisposableStore, dispose, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
import { LRUCache } from '../../../../base/common/map.js';
import { localize } from '../../../../nls.js';
import { IAccessibilityService } from '../../../../platform/accessibility/common/accessibility.js';
@@ -40,6 +40,11 @@ interface IToolEntry {
impl?: IToolImpl;
}
interface ITrackedCall {
invocation?: ChatToolInvocation;
store: IDisposable;
}
export class LanguageModelToolsService extends Disposable implements ILanguageModelToolsService {
_serviceBrand: undefined;
@@ -53,7 +58,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
private _toolContextKeys = new Set<string>();
private readonly _ctxToolsCount: IContextKey<number>;
private _callsByRequestId = new Map<string, IDisposable[]>();
private _callsByRequestId = new Map<string, ITrackedCall[]>();
private _workspaceToolConfirmStore: Lazy<ToolConfirmStore>;
private _profileToolConfirmStore: Lazy<ToolConfirmStore>;
@@ -235,7 +240,8 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
if (!this._callsByRequestId.has(requestId)) {
this._callsByRequestId.set(requestId, []);
}
this._callsByRequestId.get(requestId)!.push(store);
const trackedCall: ITrackedCall = { store };
this._callsByRequestId.get(requestId)!.push(trackedCall);
const source = new CancellationTokenSource();
store.add(toDisposable(() => {
@@ -252,6 +258,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
const prepared = await this.prepareToolInvocation(tool, dto, token);
toolInvocation = new ChatToolInvocation(prepared, tool.data, dto.callId);
trackedCall.invocation = toolInvocation;
if (this.shouldAutoConfirm(tool.data.id, tool.data.runsInWorkspace)) {
toolInvocation.confirmed.complete(true);
}
@@ -285,7 +292,11 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
throw new CancellationError();
}
toolResult = await tool.impl.invoke(dto, countTokens, token);
toolResult = await tool.impl.invoke(dto, countTokens, {
report: step => {
toolInvocation?.acceptProgress(step);
}
}, token);
this.ensureToolDetails(dto, toolResult, tool.data);
this._telemetryService.publicLog2<LanguageModelToolInvokedEvent, LanguageModelToolInvokedClassification>(
@@ -403,7 +414,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
private cleanupCallDisposables(requestId: string, store: DisposableStore): void {
const disposables = this._callsByRequestId.get(requestId);
if (disposables) {
const index = disposables.indexOf(store);
const index = disposables.findIndex(d => d.store === store);
if (index > -1) {
disposables.splice(index, 1);
}
@@ -417,7 +428,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
cancelToolCallsForRequest(requestId: string): void {
const calls = this._callsByRequestId.get(requestId);
if (calls) {
calls.forEach(call => call.dispose());
calls.forEach(call => call.store.dispose());
this._callsByRequestId.delete(requestId);
}
}
@@ -425,7 +436,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
public override dispose(): void {
super.dispose();
this._callsByRequestId.forEach(calls => dispose(calls));
this._callsByRequestId.forEach(calls => calls.forEach(call => call.store.dispose()));
this._ctxToolsCount.reset();
}
}
@@ -5,9 +5,10 @@
import { DeferredPromise } from '../../../../../base/common/async.js';
import { IMarkdownString } from '../../../../../base/common/htmlContent.js';
import { observableValue } from '../../../../../base/common/observable.js';
import { localize } from '../../../../../nls.js';
import { IChatTerminalToolInvocationData, IChatToolInputInvocationData, IChatToolInvocation, IChatToolInvocationSerialized } from '../chatService.js';
import { IPreparedToolInvocation, IToolConfirmationMessages, IToolData, IToolResult } from '../languageModelToolsService.js';
import { IPreparedToolInvocation, IToolConfirmationMessages, IToolData, IToolProgressStep, IToolResult } from '../languageModelToolsService.js';
export class ChatToolInvocation implements IChatToolInvocation {
public readonly kind: 'toolInvocation' = 'toolInvocation';
@@ -45,6 +46,8 @@ export class ChatToolInvocation implements IChatToolInvocation {
public readonly toolSpecificData?: IChatTerminalToolInvocationData | IChatToolInputInvocationData;
public readonly progress = observableValue<{ message?: string | IMarkdownString; progress: number }>(this, { progress: 0 });
constructor(preparedInvocation: IPreparedToolInvocation | undefined, toolData: IToolData, public readonly toolCallId: string) {
const defaultMessage = localize('toolInvocationMessage', "Using {0}", `"${toolData.displayName}"`);
const invocationMessage = preparedInvocation?.invocationMessage ?? defaultMessage;
@@ -84,6 +87,14 @@ export class ChatToolInvocation implements IChatToolInvocation {
return this._confirmationMessages;
}
public acceptProgress(step: IToolProgressStep) {
const prev = this.progress.get();
this.progress.set({
progress: step.increment ? (prev.progress + step.increment) : prev.progress,
message: step.message,
}, undefined);
}
public toJSON(): IChatToolInvocationSerialized {
return {
kind: 'toolInvocationSerialized',
@@ -7,6 +7,7 @@ import { DeferredPromise } from '../../../../base/common/async.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { Event } from '../../../../base/common/event.js';
import { IMarkdownString } from '../../../../base/common/htmlContent.js';
import { IObservable } from '../../../../base/common/observable.js';
import { ThemeIcon } from '../../../../base/common/themables.js';
import { URI } from '../../../../base/common/uri.js';
import { IRange, Range } from '../../../../editor/common/core/range.js';
@@ -234,6 +235,7 @@ export interface IChatToolInvocation {
invocationMessage: string | IMarkdownString;
pastTenseMessage: string | IMarkdownString | undefined;
resultDetails: IToolResult['toolResultDetails'];
progress: IObservable<{ message?: string | IMarkdownString; progress: number }>;
readonly toolId: string;
readonly toolCallId: string;
@@ -8,14 +8,15 @@ import { Event } from '../../../../base/common/event.js';
import { IMarkdownString } from '../../../../base/common/htmlContent.js';
import { IJSONSchema } from '../../../../base/common/jsonSchema.js';
import { IDisposable } from '../../../../base/common/lifecycle.js';
import { Schemas } from '../../../../base/common/network.js';
import { ThemeIcon } from '../../../../base/common/themables.js';
import { URI } from '../../../../base/common/uri.js';
import { Location } from '../../../../editor/common/languages.js';
import { ContextKeyExpression } from '../../../../platform/contextkey/common/contextkey.js';
import { ExtensionIdentifier } from '../../../../platform/extensions/common/extensions.js';
import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
import { Location } from '../../../../editor/common/languages.js';
import { IProgress } from '../../../../platform/progress/common/progress.js';
import { IChatTerminalToolInvocationData, IChatToolInputInvocationData } from './chatService.js';
import { Schemas } from '../../../../base/common/network.js';
import { PromptElementJSON, stringifyPromptElementJSON } from './tools/promptTsxTypes.js';
export interface IToolData {
@@ -40,6 +41,14 @@ export interface IToolData {
supportsToolPicker?: boolean;
}
export interface IToolProgressStep {
readonly message: string | IMarkdownString | undefined;
readonly increment: number | undefined;
readonly total: number | undefined;
}
export type ToolProgress = IProgress<IToolProgressStep>;
export type ToolDataSource =
| {
type: 'extension';
@@ -134,7 +143,7 @@ export interface IPreparedToolInvocation {
}
export interface IToolImpl {
invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult>;
invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken): Promise<IToolResult>;
prepareToolInvocation?(parameters: any, token: CancellationToken): Promise<IPreparedToolInvocation | undefined>;
}
@@ -16,7 +16,7 @@ import { INotebookService } from '../../../notebook/common/notebookService.js';
import { ICodeMapperService } from '../../common/chatCodeMapperService.js';
import { ChatModel } from '../../common/chatModel.js';
import { IChatService } from '../../common/chatService.js';
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../common/languageModelToolsService.js';
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolProgress } from '../../common/languageModelToolsService.js';
export const ExtensionEditToolId = 'vscode_editFile';
export const InternalEditToolId = 'vscode_editFile_internal';
@@ -42,7 +42,7 @@ export class EditTool implements IToolImpl {
@INotebookService private readonly notebookService: INotebookService,
) { }
async invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
async invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, _progress: ToolProgress, token: CancellationToken): Promise<IToolResult> {
if (!invocation.context) {
throw new Error('toolInvocationToken is required for this tool');
}
@@ -3,13 +3,13 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import { localize } from '../../../../../nls.js';
import { CancellationToken } from '../../../../../base/common/cancellation.js';
import { MarkdownString } from '../../../../../base/common/htmlContent.js';
import { URI } from '../../../../../base/common/uri.js';
import { localize } from '../../../../../nls.js';
import { IWebContentExtractorService } from '../../../../../platform/webContentExtractor/common/webContentExtractor.js';
import { ITrustedDomainService } from '../../../url/browser/trustedDomainService.js';
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultTextPart } from '../../common/languageModelToolsService.js';
import { MarkdownString } from '../../../../../base/common/htmlContent.js';
import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultTextPart, ToolProgress } from '../../common/languageModelToolsService.js';
import { InternalFetchWebPageToolId } from '../../common/tools/tools.js';
export const FetchWebPageToolData: IToolData = {
@@ -41,7 +41,7 @@ export class FetchWebPageTool implements IToolImpl {
@ITrustedDomainService private readonly _trustedDomainService: ITrustedDomainService,
) { }
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _token: CancellationToken): Promise<IToolResult> {
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _progress: ToolProgress, _token: CancellationToken): Promise<IToolResult> {
const parsedUriResults = this._parseUris((invocation.parameters as { urls?: string[] }).urls);
const validUris = Array.from(parsedUriResults.values()).filter((uri): uri is URI => !!uri);
if (!validUris.length) {
@@ -149,7 +149,7 @@ suite('LanguageModelToolsService', () => {
const toolBarrier = new Barrier();
const toolImpl: IToolImpl = {
invoke: async (invocation, countTokens, cancelToken) => {
invoke: async (invocation, countTokens, progress, cancelToken) => {
assert.strictEqual(invocation.callId, '1');
assert.strictEqual(invocation.toolId, 'testTool');
assert.deepStrictEqual(invocation.parameters, { a: 1 });
@@ -6,6 +6,7 @@
import { CancellationToken } from '../../../../../base/common/cancellation.js';
import { Event } from '../../../../../base/common/event.js';
import { Disposable, IDisposable } from '../../../../../base/common/lifecycle.js';
import { IProgressStep } from '../../../../../platform/progress/common/progress.js';
import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../common/languageModelToolsService.js';
export class MockLanguageModelToolsService implements ILanguageModelToolsService {
@@ -46,6 +47,10 @@ export class MockLanguageModelToolsService implements ILanguageModelToolsService
return undefined;
}
acceptProgress(sessionId: string | undefined, callId: string, progress: IProgressStep): void {
}
async invokeTool(dto: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
return {
content: [{ kind: 'text', value: 'result' }]
@@ -9,8 +9,8 @@ import { ThemeIcon } from '../../../../base/common/themables.js';
import { localize } from '../../../../nls.js';
import { SortBy } from '../../../../platform/extensionManagement/common/extensionManagement.js';
import { EXTENSION_CATEGORIES } from '../../../../platform/extensions/common/extensions.js';
import { CountTokensCallback, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../chat/common/languageModelToolsService.js';
import { ExtensionState, IExtensionsWorkbenchService, IExtension } from '../common/extensions.js';
import { CountTokensCallback, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolProgress } from '../../chat/common/languageModelToolsService.js';
import { ExtensionState, IExtension, IExtensionsWorkbenchService } from '../common/extensions.js';
export const SearchExtensionsToolId = 'vscode_searchExtensions_internal';
@@ -72,7 +72,7 @@ export class SearchExtensionsTool implements IToolImpl {
@IExtensionsWorkbenchService private readonly extensionWorkbenchService: IExtensionsWorkbenchService,
) { }
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _progress: ToolProgress, token: CancellationToken): Promise<IToolResult> {
const params = invocation.parameters as InputParams;
if (!params.keywords?.length && !params.category && !params.ids?.length) {
return {
@@ -16,12 +16,6 @@ const mcpActivationEventPrefix = 'onMcpCollection:';
export const mcpActivationEvent = (collectionId: string) => mcpActivationEventPrefix + collectionId;
const mcpSchemaExampleServer = {
command: 'node',
args: ['my-mcp-server.js'],
env: {},
};
export const enum DiscoverySource {
ClaudeDesktop = 'claude-desktop',
Windsurf = 'windsurf',
@@ -55,15 +49,17 @@ export const mcpSchemaExampleServers = {
}
};
const httpSchemaExample = {
url: 'http://localhost:3001/mcp',
headers: {},
const httpSchemaExamples = {
'my-mcp-server': {
url: 'http://localhost:3001/mcp',
headers: {},
}
};
export const mcpStdioServerSchema: IJSONSchema = {
type: 'object',
additionalProperties: false,
examples: [mcpSchemaExampleServer],
examples: [mcpSchemaExampleServers['mcp-server-time']],
properties: {
type: {
type: 'string',
@@ -110,14 +106,14 @@ export const mcpServerSchema: IJSONSchema = {
servers: {
examples: [
mcpSchemaExampleServers,
httpSchemaExample,
httpSchemaExamples,
],
additionalProperties: {
oneOf: [mcpStdioServerSchema, {
type: 'object',
additionalProperties: false,
required: ['url'],
examples: [httpSchemaExample],
examples: [httpSchemaExamples['my-mcp-server']],
properties: {
type: {
type: 'string',
@@ -4,29 +4,31 @@
*--------------------------------------------------------------------------------------------*/
import { raceCancellationError, Sequencer } from '../../../../base/common/async.js';
import * as json from '../../../../base/common/json.js';
import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
import * as json from '../../../../base/common/json.js';
import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
import { LRUCache } from '../../../../base/common/map.js';
import { autorun, autorunWithStore, derived, disposableObservableValue, IObservable, ITransaction, observableFromEvent, ObservablePromise, observableValue, transaction } from '../../../../base/common/observable.js';
import { basename } from '../../../../base/common/resources.js';
import { URI } from '../../../../base/common/uri.js';
import { generateUuid } from '../../../../base/common/uuid.js';
import { localize } from '../../../../nls.js';
import { ICommandService } from '../../../../platform/commands/common/commands.js';
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { ILogger, ILoggerService } from '../../../../platform/log/common/log.js';
import { INotificationService, Severity } from '../../../../platform/notification/common/notification.js';
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry.js';
import { IWorkspaceContextService } from '../../../../platform/workspace/common/workspace.js';
import { IEditorService } from '../../../services/editor/common/editorService.js';
import { IExtensionService } from '../../../services/extensions/common/extensions.js';
import { IOutputService } from '../../../services/output/common/output.js';
import { ToolProgress } from '../../chat/common/languageModelToolsService.js';
import { mcpActivationEvent } from './mcpConfiguration.js';
import { IMcpRegistry } from './mcpRegistryTypes.js';
import { McpServerRequestHandler } from './mcpServerRequestHandler.js';
import { extensionMcpCollectionPrefix, IMcpServer, IMcpServerConnection, IMcpTool, McpCollectionReference, McpConnectionFailedError, McpConnectionState, McpDefinitionReference, McpServerDefinition, McpServerToolsState } from './mcpTypes.js';
import { MCP } from './modelContextProtocol.js';
import { INotificationService, Severity } from '../../../../platform/notification/common/notification.js';
import { localize } from '../../../../nls.js';
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { IEditorService } from '../../../services/editor/common/editorService.js';
type ServerBootData = {
supportsLogging: boolean;
@@ -497,7 +499,31 @@ export class McpTool implements IMcpTool {
call(params: Record<string, unknown>, token?: CancellationToken): Promise<MCP.CallToolResult> {
// serverToolName is always set now, but older cache entries (from 1.99-Insiders) may not have it.
const name = this._definition.serverToolName ?? this._definition.name;
return this._server.callOn(h => h.callTool({ name, arguments: params }), token);
return this._server.callOn(h => h.callTool({ name, arguments: params }, token), token);
}
callWithProgress(params: Record<string, unknown>, progress: ToolProgress, token?: CancellationToken): Promise<MCP.CallToolResult> {
// serverToolName is always set now, but older cache entries (from 1.99-Insiders) may not have it.
const name = this._definition.serverToolName ?? this._definition.name;
const progressToken = generateUuid();
return this._server.callOn(h => {
let lastProgressN = 0;
const listener = h.onDidReceiveProgressNotification((e) => {
if (e.params.progressToken === progressToken) {
progress.report({
message: e.params.message,
increment: e.params.progress - lastProgressN,
total: e.params.total,
});
lastProgressN = e.params.progress;
}
});
return h.callTool({ name, arguments: params, _meta: { progressToken } }, token).finally(() => {
listener.dispose();
});
}, token);
}
compare(other: IMcpTool): number {
@@ -466,7 +466,7 @@ export class McpServerRequestHandler extends Disposable {
/**
* Call a specific tool
*/
callTool(params: MCP.CallToolRequest['params'], token?: CancellationToken): Promise<MCP.CallToolResult> {
callTool(params: MCP.CallToolRequest['params'] & MCP.Request['params'], token?: CancellationToken): Promise<MCP.CallToolResult> {
return this.sendRequest<MCP.CallToolRequest, MCP.CallToolResult>({ method: 'tools/call', params }, token);
}
@@ -15,7 +15,7 @@ import { IInstantiationService } from '../../../../platform/instantiation/common
import { ILogService } from '../../../../platform/log/common/log.js';
import { IProductService } from '../../../../platform/product/common/productService.js';
import { StorageScope } from '../../../../platform/storage/common/storage.js';
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../../chat/common/languageModelToolsService.js';
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, ToolProgress } from '../../chat/common/languageModelToolsService.js';
import { IMcpRegistry } from './mcpRegistryTypes.js';
import { McpServer, McpServerMetadataCache } from './mcpServer.js';
import { IMcpServer, IMcpService, IMcpTool, McpCollectionDefinition, McpServerDefinition, McpServerToolsState } from './mcpTypes.js';
@@ -249,7 +249,7 @@ class McpToolImplementation implements IToolImpl {
};
}
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, token: CancellationToken) {
async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken) {
const result: IToolResult = {
content: []
@@ -257,7 +257,7 @@ class McpToolImplementation implements IToolImpl {
const outputParts: string[] = [];
const callResult = await this._tool.call(invocation.parameters as Record<string, any>, token);
const callResult = await this._tool.callWithProgress(invocation.parameters as Record<string, any>, progress, token);
for (const item of callResult.content) {
if (item.type === 'text') {
result.content.push({
@@ -3,11 +3,11 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import { equals as arraysEqual } from '../../../../base/common/arrays.js';
import { assertNever } from '../../../../base/common/assert.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { IDisposable } from '../../../../base/common/lifecycle.js';
import { equals as objectsEqual } from '../../../../base/common/objects.js';
import { equals as arraysEqual } from '../../../../base/common/arrays.js';
import { IObservable } from '../../../../base/common/observable.js';
import { URI, UriComponents } from '../../../../base/common/uri.js';
import { Location } from '../../../../editor/common/languages.js';
@@ -17,6 +17,7 @@ import { ExtensionIdentifier } from '../../../../platform/extensions/common/exte
import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
import { StorageScope } from '../../../../platform/storage/common/storage.js';
import { IWorkspaceFolderData } from '../../../../platform/workspace/common/workspace.js';
import { ToolProgress } from '../../chat/common/languageModelToolsService.js';
import { McpServerRequestHandler } from './mcpServerRequestHandler.js';
import { MCP } from './modelContextProtocol.js';
@@ -259,6 +260,11 @@ export interface IMcpTool {
* @throws {@link McpConnectionFailedError} if the connection to the server fails
*/
call(params: Record<string, unknown>, token?: CancellationToken): Promise<MCP.CallToolResult>;
/**
* Identical to {@link call}, but reports progress.
*/
callWithProgress(params: Record<string, unknown>, progress: ToolProgress, token?: CancellationToken): Promise<MCP.CallToolResult>;
}
export const enum McpServerTransportType {
+25
View File
@@ -0,0 +1,25 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
declare module 'vscode' {
/**
* A progress update during an {@link LanguageModelTool.invoke} call.
*/
export interface ToolProgressStep {
/**
* A progress message that represents a chunk of work
*/
message?: string | MarkdownString;
/**
* An increment for discrete progress. Increments will be summed up until 100 (100%) is reached
*/
increment?: number;
}
export interface LanguageModelTool<T> {
invoke(options: LanguageModelToolInvocationOptions<T>, token: CancellationToken, progress: Progress<ToolProgressStep>): ProviderResult<LanguageModelToolResult>;
}
}