From 4d38837a83ecc9ab7444eac34626259736b57c49 Mon Sep 17 00:00:00 2001 From: Johannes Rieken Date: Fri, 13 Feb 2026 12:27:05 +0100 Subject: [PATCH] Move usages tool (#295139) * feat: add UsagesTool for finding code symbol usages and implement related tests * feat: handle late-registered tools in LanguageModelToolsExtensionPointHandler * feat: enhance UsagesTool to utilize ISearchService for symbol searches and improve reference classification --- .../editor/common/languageFeatureRegistry.ts | 11 +- src/vs/editor/common/languageSelector.ts | 15 + .../common/modes/languageSelector.test.ts | 25 +- .../contrib/chat/browser/chat.contribution.ts | 2 + .../contrib/chat/browser/tools/usagesTool.ts | 371 ++++++++++++++++++ .../tools/languageModelToolsContribution.ts | 27 +- .../test/browser/tools/usagesTool.test.ts | 320 +++++++++++++++ 7 files changed, 767 insertions(+), 4 deletions(-) create mode 100644 src/vs/workbench/contrib/chat/browser/tools/usagesTool.ts create mode 100644 src/vs/workbench/contrib/chat/test/browser/tools/usagesTool.test.ts diff --git a/src/vs/editor/common/languageFeatureRegistry.ts b/src/vs/editor/common/languageFeatureRegistry.ts index d76b0906419..c0a0c07d2f6 100644 --- a/src/vs/editor/common/languageFeatureRegistry.ts +++ b/src/vs/editor/common/languageFeatureRegistry.ts @@ -6,7 +6,7 @@ import { Emitter } from '../../base/common/event.js'; import { IDisposable, toDisposable } from '../../base/common/lifecycle.js'; import { ITextModel, shouldSynchronizeModel } from './model.js'; -import { LanguageFilter, LanguageSelector, score } from './languageSelector.js'; +import { LanguageFilter, LanguageSelector, score, selectLanguageIds } from './languageSelector.js'; import { URI } from '../../base/common/uri.js'; interface Entry { @@ -115,6 +115,14 @@ export class LanguageFeatureRegistry { return this._entries.map(entry => entry.provider); } + get registeredLanguageIds(): ReadonlySet { + const result = new Set(); + for (const entry of this._entries) { + selectLanguageIds(entry.selector, result); + } + return result; + } + ordered(model: ITextModel, recursive = false): T[] { const result: T[] = []; this._orderedForEach(model, recursive, entry => result.push(entry.provider)); @@ -226,4 +234,3 @@ function isBuiltinSelector(selector: LanguageSelector): boolean { return Boolean((selector as LanguageFilter).isBuiltin); } - diff --git a/src/vs/editor/common/languageSelector.ts b/src/vs/editor/common/languageSelector.ts index 6374d380f48..80ffb5450d1 100644 --- a/src/vs/editor/common/languageSelector.ts +++ b/src/vs/editor/common/languageSelector.ts @@ -142,3 +142,18 @@ export function targetsNotebooks(selector: LanguageSelector): boolean { return !!(selector).notebookType; } } + +export function selectLanguageIds(selector: LanguageSelector, into: Set): void { + if (typeof selector === 'string') { + into.add(selector); + } else if (Array.isArray(selector)) { + for (const item of selector) { + selectLanguageIds(item, into); + } + } else { + const language = (selector).language; + if (language) { + into.add(language); + } + } +} diff --git a/src/vs/editor/test/common/modes/languageSelector.test.ts b/src/vs/editor/test/common/modes/languageSelector.test.ts index 31f6c051af4..be9c597f98f 100644 --- a/src/vs/editor/test/common/modes/languageSelector.test.ts +++ b/src/vs/editor/test/common/modes/languageSelector.test.ts @@ -6,7 +6,7 @@ import assert from 'assert'; import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; -import { LanguageSelector, score } from '../../../common/languageSelector.js'; +import { LanguageSelector, score, selectLanguageIds } from '../../../common/languageSelector.js'; suite('LanguageSelector', function () { @@ -173,4 +173,27 @@ suite('LanguageSelector', function () { }, obj.uri, obj.langId, true, undefined, undefined); assert.strictEqual(value, 0); }); + + test('selectLanguageIds', function () { + const result = new Set(); + + selectLanguageIds('typescript', result); + assert.deepStrictEqual([...result], ['typescript']); + + result.clear(); + selectLanguageIds({ language: 'python', scheme: 'file' }, result); + assert.deepStrictEqual([...result], ['python']); + + result.clear(); + selectLanguageIds({ scheme: 'file' }, result); + assert.deepStrictEqual([...result], []); + + result.clear(); + selectLanguageIds(['javascript', { language: 'css' }, { scheme: 'untitled' }], result); + assert.deepStrictEqual([...result].sort(), ['css', 'javascript']); + + result.clear(); + selectLanguageIds('*', result); + assert.deepStrictEqual([...result], ['*']); + }); }); diff --git a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts index b0d814cbf57..aad7733c40f 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts @@ -60,6 +60,7 @@ import { IPromptsService } from '../common/promptSyntax/service/promptsService.j import { PromptsService } from '../common/promptSyntax/service/promptsServiceImpl.js'; import { LanguageModelToolsExtensionPointHandler } from '../common/tools/languageModelToolsContribution.js'; import { BuiltinToolsContribution } from '../common/tools/builtinTools/tools.js'; +import { UsagesToolContribution } from './tools/usagesTool.js'; import { IVoiceChatService, VoiceChatService } from '../common/voiceChatService.js'; import { registerChatAccessibilityActions } from './actions/chatAccessibilityActions.js'; import { AgentChatAccessibilityHelp, EditsChatAccessibilityHelp, PanelChatAccessibilityHelp, QuickChatAccessibilityHelp } from './actions/chatAccessibilityHelp.js'; @@ -1422,6 +1423,7 @@ registerWorkbenchContribution2(ChatSetupContribution.ID, ChatSetupContribution, registerWorkbenchContribution2(ChatTeardownContribution.ID, ChatTeardownContribution, WorkbenchPhase.AfterRestored); registerWorkbenchContribution2(ChatStatusBarEntry.ID, ChatStatusBarEntry, WorkbenchPhase.BlockRestore); registerWorkbenchContribution2(BuiltinToolsContribution.ID, BuiltinToolsContribution, WorkbenchPhase.Eventually); +registerWorkbenchContribution2(UsagesToolContribution.ID, UsagesToolContribution, WorkbenchPhase.BlockRestore); registerWorkbenchContribution2(ChatAgentSettingContribution.ID, ChatAgentSettingContribution, WorkbenchPhase.AfterRestored); registerWorkbenchContribution2(ChatAgentActionsContribution.ID, ChatAgentActionsContribution, WorkbenchPhase.Eventually); registerWorkbenchContribution2(ToolReferenceNamesContribution.ID, ToolReferenceNamesContribution, WorkbenchPhase.AfterRestored); diff --git a/src/vs/workbench/contrib/chat/browser/tools/usagesTool.ts b/src/vs/workbench/contrib/chat/browser/tools/usagesTool.ts new file mode 100644 index 00000000000..aa8e6731069 --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/tools/usagesTool.ts @@ -0,0 +1,371 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { CancellationToken } from '../../../../../base/common/cancellation.js'; +import { Codicon } from '../../../../../base/common/codicons.js'; +import { Emitter, Event } from '../../../../../base/common/event.js'; +import { MarkdownString } from '../../../../../base/common/htmlContent.js'; +import { escapeRegExpCharacters } from '../../../../../base/common/strings.js'; +import { Disposable, IDisposable } from '../../../../../base/common/lifecycle.js'; +import { ResourceSet } from '../../../../../base/common/map.js'; +import { ThemeIcon } from '../../../../../base/common/themables.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { relativePath } from '../../../../../base/common/resources.js'; +import { Position } from '../../../../../editor/common/core/position.js'; +import { Range } from '../../../../../editor/common/core/range.js'; +import { Location, LocationLink } from '../../../../../editor/common/languages.js'; +import { IModelService } from '../../../../../editor/common/services/model.js'; +import { ILanguageFeaturesService } from '../../../../../editor/common/services/languageFeatures.js'; +import { ITextModelService } from '../../../../../editor/common/services/resolverService.js'; +import { getDefinitionsAtPosition, getImplementationsAtPosition, getReferencesAtPosition } from '../../../../../editor/contrib/gotoSymbol/browser/goToSymbol.js'; +import { localize } from '../../../../../nls.js'; +import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; +import { IWorkspaceContextService } from '../../../../../platform/workspace/common/workspace.js'; +import { IWorkbenchContribution } from '../../../../common/contributions.js'; +import { ISearchService, QueryType, resultIsMatch } from '../../../../services/search/common/search.js'; +import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, ToolDataSource, ToolProgress, } from '../../common/tools/languageModelToolsService.js'; +import { createToolSimpleTextResult } from '../../common/tools/builtinTools/toolHelpers.js'; + +export const UsagesToolId = 'vscode_listCodeUsages'; + +interface IUsagesToolInput { + symbol: string; + uri?: string; + filePath?: string; + lineContent: string; +} + +const BaseModelDescription = `Find all usages (references, definitions, and implementations) of a code symbol across the workspace. This tool locates where a symbol is referenced, defined, or implemented. + +Input: +- "symbol": The exact name of the symbol to search for (function, class, method, variable, type, etc.). +- "uri": A full URI (e.g. "file:///path/to/file.ts") of a file where the symbol appears. Provide either "uri" or "filePath". +- "filePath": A workspace-relative file path (e.g. "src/utils/helpers.ts") of a file where the symbol appears. Provide either "uri" or "filePath". +- "lineContent": A substring of the line of code where the symbol appears. This is used to locate the exact position in the file. Must be the actual text from the file - do NOT fabricate it. + +IMPORTANT: The file and line do NOT need to be the definition of the symbol. Any occurrence works - a usage, an import, a call site, etc. You can pick whichever occurrence is most convenient. + +If the tool returns an error, retry with corrected input - ensure the file path is correct, the line content matches the actual file content, and the symbol name appears in that line.`; + +export class UsagesTool extends Disposable implements IToolImpl { + + private readonly _onDidUpdateToolData = this._store.add(new Emitter()); + readonly onDidUpdateToolData = this._onDidUpdateToolData.event; + + constructor( + @ILanguageFeaturesService private readonly _languageFeaturesService: ILanguageFeaturesService, + @IModelService private readonly _modelService: IModelService, + @ISearchService private readonly _searchService: ISearchService, + @ITextModelService private readonly _textModelService: ITextModelService, + @IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService, + ) { + super(); + + this._store.add(Event.debounce( + this._languageFeaturesService.referenceProvider.onDidChange, + () => { }, + 2000 + )((() => this._onDidUpdateToolData.fire()))); + } + + getToolData(): IToolData { + const languageIds = this._languageFeaturesService.referenceProvider.registeredLanguageIds; + + let modelDescription = BaseModelDescription; + if (languageIds.has('*')) { + modelDescription += '\n\nSupported for all languages.'; + } else if (languageIds.size > 0) { + const sorted = [...languageIds].sort(); + modelDescription += `\n\nCurrently supported for: ${sorted.join(', ')}.`; + } else { + modelDescription += '\n\nNo languages currently have reference providers registered.'; + } + + return { + id: UsagesToolId, + toolReferenceName: 'usages', + canBeReferencedInPrompt: false, + icon: ThemeIcon.fromId(Codicon.references.id), + displayName: localize('tool.usages.displayName', 'List Code Usages'), + userDescription: localize('tool.usages.userDescription', 'Find references, definitions, and implementations of a symbol'), + modelDescription, + source: ToolDataSource.Internal, + inputSchema: { + type: 'object', + properties: { + symbol: { + type: 'string', + description: 'The exact name of the symbol (function, class, method, variable, type, etc.) to find usages of.' + }, + uri: { + type: 'string', + description: 'A full URI of a file where the symbol appears (e.g. "file:///path/to/file.ts"). Provide either "uri" or "filePath".' + }, + filePath: { + type: 'string', + description: 'A workspace-relative file path where the symbol appears (e.g. "src/utils/helpers.ts"). Provide either "uri" or "filePath".' + }, + lineContent: { + type: 'string', + description: 'A substring of the line of code where the symbol appears. Used to locate the exact position. Must be actual text from the file.' + } + }, + required: ['symbol', 'lineContent'] + } + }; + } + + async prepareToolInvocation(context: IToolInvocationPreparationContext, _token: CancellationToken): Promise { + const input = context.parameters as IUsagesToolInput; + return { + invocationMessage: localize('tool.usages.invocationMessage', 'Analyzing usages of `{0}`', input.symbol), + }; + } + + async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _progress: ToolProgress, token: CancellationToken): Promise { + const input = invocation.parameters as IUsagesToolInput; + + // --- resolve URI --- + const uri = this._resolveUri(input); + if (!uri) { + return this._errorResult('Provide either "uri" (a full URI) or "filePath" (a workspace-relative path) to identify the file.'); + } + + // --- open text model --- + const ref = await this._textModelService.createModelReference(uri); + try { + const model = ref.object.textEditorModel; + + if (!this._languageFeaturesService.referenceProvider.has(model)) { + return this._errorResult(`No reference provider available for this file's language. The usages tool may not support this language.`); + } + + // --- find line containing lineContent --- + const parts = input.lineContent.trim().split(/\s+/); + const lineContent = parts.map(escapeRegExpCharacters).join('\\s+'); + const matches = model.findMatches(lineContent, false, true, false, null, false, 1); + if (matches.length === 0) { + return this._errorResult(`Could not find line content "${input.lineContent}" in ${uri.toString()}. Provide the exact text from the line where the symbol appears.`); + } + const lineNumber = matches[0].range.startLineNumber; + + // --- find symbol in that line --- + const lineText = model.getLineContent(lineNumber); + const column = this._findSymbolColumn(lineText, input.symbol); + if (column === undefined) { + return this._errorResult(`Could not find symbol "${input.symbol}" in the matched line. Ensure the symbol name is correct and appears in the provided line content.`); + } + + const position = new Position(lineNumber, column); + + // --- query references, definitions, implementations in parallel --- + const [definitions, references, implementations] = await Promise.all([ + getDefinitionsAtPosition(this._languageFeaturesService.definitionProvider, model, position, false, token), + getReferencesAtPosition(this._languageFeaturesService.referenceProvider, model, position, false, false, token), + getImplementationsAtPosition(this._languageFeaturesService.implementationProvider, model, position, false, token), + ]); + + if (references.length === 0) { + const result = createToolSimpleTextResult(`No usages found for \`${input.symbol}\`.`); + result.toolResultMessage = new MarkdownString(localize('tool.usages.noResults', 'Analyzed usages of `{0}`, no results', input.symbol)); + return result; + } + + // --- classify and format results with previews --- + const previews = await this._getLinePreviews(input.symbol, references, token); + + const lines: string[] = []; + lines.push(`${references.length} usages of \`${input.symbol}\`:\n`); + + for (let i = 0; i < references.length; i++) { + const ref = references[i]; + const kind = this._classifyReference(ref, definitions, implementations); + const startLine = Range.lift(ref.range).startLineNumber; + const preview = previews[i]; + if (preview) { + lines.push(``); + lines.push(`\t${preview}`); + lines.push(``); + } else { + lines.push(``); + } + } + + const text = lines.join('\n'); + const result = createToolSimpleTextResult(text); + + result.toolResultMessage = references.length === 1 + ? new MarkdownString(localize('tool.usages.oneResult', 'Analyzed usages of `{0}`, 1 result', input.symbol)) + : new MarkdownString(localize('tool.usages.results', 'Analyzed usages of `{0}`, {1} results', input.symbol, references.length)); + + result.toolResultDetails = references.map((r): Location => ({ uri: r.uri, range: r.range })); + + return result; + } finally { + ref.dispose(); + } + } + + private async _getLinePreviews(symbol: string, references: LocationLink[], token: CancellationToken): Promise<(string | undefined)[]> { + const previews: (string | undefined)[] = new Array(references.length); + + // Build a lookup: (uriString, lineNumber) → index in references array + const lookup = new Map(); + const needSearch = new ResourceSet(); + + for (let i = 0; i < references.length; i++) { + const ref = references[i]; + const lineNumber = Range.lift(ref.range).startLineNumber; + + // Try already-open models first + const existingModel = this._modelService.getModel(ref.uri); + if (existingModel) { + previews[i] = existingModel.getLineContent(lineNumber).trim(); + } else { + lookup.set(`${ref.uri.toString()}:${lineNumber}`, i); + needSearch.add(ref.uri); + } + } + + if (needSearch.size === 0 || token.isCancellationRequested) { + return previews; + } + + // Use ISearchService to search for the symbol name, restricted to the + // referenced files. This is backed by ripgrep for file:// URIs. + try { + // Build includePattern from workspace-relative paths + const folders = this._workspaceContextService.getWorkspace().folders; + const relativePaths: string[] = []; + for (const uri of needSearch) { + const folder = this._workspaceContextService.getWorkspaceFolder(uri); + if (folder) { + const rel = relativePath(folder.uri, uri); + if (rel) { + relativePaths.push(rel); + } + } + } + + if (relativePaths.length > 0) { + const includePattern: Record = {}; + if (relativePaths.length === 1) { + includePattern[relativePaths[0]] = true; + } else { + includePattern[`{${relativePaths.join(',')}}`] = true; + } + + const searchResult = await this._searchService.textSearch( + { + type: QueryType.Text, + contentPattern: { pattern: escapeRegExpCharacters(symbol), isRegExp: true, isWordMatch: true }, + folderQueries: folders.map(f => ({ folder: f.uri })), + includePattern, + }, + token, + ); + + for (const fileMatch of searchResult.results) { + if (!fileMatch.results) { + continue; + } + for (const textMatch of fileMatch.results) { + if (!resultIsMatch(textMatch)) { + continue; + } + for (const range of textMatch.rangeLocations) { + const lineNumber = range.source.startLineNumber + 1; // 0-based → 1-based + const key = `${fileMatch.resource.toString()}:${lineNumber}`; + const idx = lookup.get(key); + if (idx !== undefined) { + previews[idx] = textMatch.previewText.trim(); + lookup.delete(key); + } + } + } + } + } + } catch { + // search might fail, leave remaining previews as undefined + } + + return previews; + } + + private _resolveUri(input: IUsagesToolInput): URI | undefined { + if (input.uri) { + return URI.parse(input.uri); + } + if (input.filePath) { + const folders = this._workspaceContextService.getWorkspace().folders; + if (folders.length === 1) { + return folders[0].toResource(input.filePath); + } + // try each folder, return the first + for (const folder of folders) { + return folder.toResource(input.filePath); + } + } + return undefined; + } + + private _findSymbolColumn(lineText: string, symbol: string): number | undefined { + // use word boundary matching to avoid partial matches + const pattern = new RegExp(`\\b${escapeRegExpCharacters(symbol)}\\b`); + const match = pattern.exec(lineText); + if (match) { + return match.index + 1; // 1-based column + } + return undefined; + } + + private _classifyReference(ref: LocationLink, definitions: LocationLink[], implementations: LocationLink[]): string { + if (definitions.some(d => this._overlaps(ref, d))) { + return 'definition'; + } + if (implementations.some(d => this._overlaps(ref, d))) { + return 'implementation'; + } + return 'reference'; + } + + private _overlaps(a: LocationLink, b: LocationLink): boolean { + if (a.uri.toString() !== b.uri.toString()) { + return false; + } + return Range.areIntersectingOrTouching(a.range, b.range); + } + + private _errorResult(message: string): IToolResult { + const result = createToolSimpleTextResult(message); + result.toolResultMessage = new MarkdownString(message); + return result; + } +} + +export class UsagesToolContribution extends Disposable implements IWorkbenchContribution { + + static readonly ID = 'chat.usagesTool'; + + constructor( + @ILanguageModelToolsService toolsService: ILanguageModelToolsService, + @IInstantiationService instantiationService: IInstantiationService, + ) { + super(); + + const usagesTool = this._store.add(instantiationService.createInstance(UsagesTool)); + + let registration: IDisposable | undefined; + const registerUsagesTool = () => { + registration?.dispose(); + toolsService.flushToolUpdates(); + const toolData = usagesTool.getToolData(); + registration = toolsService.registerTool(toolData, usagesTool); + }; + registerUsagesTool(); + this._store.add(usagesTool.onDidUpdateToolData(registerUsagesTool)); + this._store.add({ dispose: () => registration?.dispose() }); + } +} diff --git a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts index 96fb13cb669..a17eb174f38 100644 --- a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts +++ b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts @@ -327,6 +327,7 @@ export class LanguageModelToolsExtensionPointHandler implements IWorkbenchContri const tools: IToolData[] = []; const toolSets: IToolSet[] = []; + const missingToolNames: string[] = []; for (const toolName of toolSet.tools) { const toolObj = languageModelToolsService.getToolByName(toolName); @@ -339,7 +340,7 @@ export class LanguageModelToolsExtensionPointHandler implements IWorkbenchContri toolSets.push(toolSetObj); continue; } - extension.collector.warn(`Tool set '${toolSet.name}' CANNOT find tool or tool set by name: ${toolName}`); + missingToolNames.push(toolName); } if (toolSets.length === 0 && tools.length === 0) { @@ -373,6 +374,30 @@ export class LanguageModelToolsExtensionPointHandler implements IWorkbenchContri toolSets.forEach(toolSet => store.add(obj.addToolSet(toolSet, tx))); }); + // Listen for late-registered tools that weren't available at contribution time + if (missingToolNames.length > 0) { + const pending = new Set(missingToolNames); + const listener = store.add(languageModelToolsService.onDidChangeTools(() => { + for (const toolName of pending) { + const toolObj = languageModelToolsService.getToolByName(toolName); + if (toolObj) { + store.add(obj.addTool(toolObj)); + pending.delete(toolName); + } else { + const toolSetObj = languageModelToolsService.getToolSetByName(toolName); + if (toolSetObj) { + store.add(obj.addToolSet(toolSetObj)); + pending.delete(toolName); + } + } + } + if (pending.size === 0) { + // done + store.delete(listener); + } + })); + } + this._registrationDisposables.set(toToolSetKey(extension.description.identifier, toolSet.name), store); } } diff --git a/src/vs/workbench/contrib/chat/test/browser/tools/usagesTool.test.ts b/src/vs/workbench/contrib/chat/test/browser/tools/usagesTool.test.ts new file mode 100644 index 00000000000..e0e20ec03f6 --- /dev/null +++ b/src/vs/workbench/contrib/chat/test/browser/tools/usagesTool.test.ts @@ -0,0 +1,320 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { CancellationToken } from '../../../../../../base/common/cancellation.js'; +import { DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { Range } from '../../../../../../editor/common/core/range.js'; +import { DefinitionProvider, ImplementationProvider, Location, ReferenceProvider } from '../../../../../../editor/common/languages.js'; +import { ITextModel } from '../../../../../../editor/common/model.js'; +import { LanguageFeaturesService } from '../../../../../../editor/common/services/languageFeaturesService.js'; +import { IModelService } from '../../../../../../editor/common/services/model.js'; +import { ITextModelService } from '../../../../../../editor/common/services/resolverService.js'; +import { createTextModel } from '../../../../../../editor/test/common/testTextModel.js'; +import { IWorkspaceContextService, IWorkspaceFolder } from '../../../../../../platform/workspace/common/workspace.js'; +import { FileMatch, ISearchComplete, ISearchService, ITextQuery, OneLineRange, TextSearchMatch } from '../../../../../services/search/common/search.js'; +import { UsagesTool, UsagesToolId } from '../../../browser/tools/usagesTool.js'; +import { IToolInvocation, IToolResult, IToolResultTextPart, ToolProgress } from '../../../common/tools/languageModelToolsService.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; + +function getTextContent(result: IToolResult): string { + const part = result.content.find((p): p is IToolResultTextPart => p.kind === 'text'); + return part?.value ?? ''; +} + +suite('UsagesTool', () => { + + const disposables = new DisposableStore(); + let langFeatures: LanguageFeaturesService; + + const testUri = URI.parse('file:///test/file.ts'); + const testContent = [ + 'import { MyClass } from "./myClass";', + '', + 'function doSomething() {', + '\tconst instance = new MyClass();', + '\tinstance.run();', + '}', + ].join('\n'); + + function createMockModelService(models?: ITextModel[]): IModelService { + return { + _serviceBrand: undefined, + getModel: (uri: URI) => models?.find(m => m.uri.toString() === uri.toString()) ?? null, + } as unknown as IModelService; + } + + function createMockSearchService(searchImpl?: (query: ITextQuery) => ISearchComplete): ISearchService { + return { + _serviceBrand: undefined, + textSearch: async (query: ITextQuery) => searchImpl?.(query) ?? { results: [], messages: [] }, + } as unknown as ISearchService; + } + + function createMockTextModelService(model: ITextModel): ITextModelService { + return { + _serviceBrand: undefined, + createModelReference: async () => ({ + object: { textEditorModel: model }, + dispose: () => { }, + }), + registerTextModelContentProvider: () => ({ dispose: () => { } }), + canHandleResource: () => false, + } as unknown as ITextModelService; + } + + function createMockWorkspaceService(): IWorkspaceContextService { + const folderUri = URI.parse('file:///test'); + const folder = { + uri: folderUri, + toResource: (relativePath: string) => URI.parse(`file:///test/${relativePath}`), + } as unknown as IWorkspaceFolder; + return { + _serviceBrand: undefined, + getWorkspace: () => ({ folders: [folder] }), + getWorkspaceFolder: (uri: URI) => { + if (uri.toString().startsWith(folderUri.toString())) { + return folder; + } + return null; + }, + } as unknown as IWorkspaceContextService; + } + + function createInvocation(parameters: Record): IToolInvocation { + return { parameters } as unknown as IToolInvocation; + } + + const noopCountTokens = async () => 0; + const noopProgress: ToolProgress = { report() { } }; + + function createTool(textModelService: ITextModelService, workspaceService: IWorkspaceContextService, options?: { modelService?: IModelService; searchService?: ISearchService }): UsagesTool { + return new UsagesTool(langFeatures, options?.modelService ?? createMockModelService(), options?.searchService ?? createMockSearchService(), textModelService, workspaceService); + } + + setup(() => { + langFeatures = new LanguageFeaturesService(); + }); + + teardown(() => { + disposables.clear(); + }); + + ensureNoDisposablesAreLeakedInTestSuite(); + + suite('getToolData', () => { + + test('reports no providers when none registered', () => { + const tool = disposables.add(createTool(createMockTextModelService(null!), createMockWorkspaceService())); + const data = tool.getToolData(); + assert.strictEqual(data.id, UsagesToolId); + assert.ok(data.modelDescription.includes('No languages currently have reference providers')); + }); + + test('lists registered language ids', () => { + const model = disposables.add(createTextModel('', 'typescript', undefined, testUri)); + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService())); + disposables.add(langFeatures.referenceProvider.register('typescript', { provideReferences: () => [] })); + const data = tool.getToolData(); + assert.ok(data.modelDescription.includes('typescript')); + }); + + test('reports all languages for wildcard', () => { + const tool = disposables.add(createTool(createMockTextModelService(null!), createMockWorkspaceService())); + disposables.add(langFeatures.referenceProvider.register('*', { provideReferences: () => [] })); + const data = tool.getToolData(); + assert.ok(data.modelDescription.includes('all languages')); + }); + }); + + suite('invoke', () => { + + test('returns error when no uri or filePath provided', async () => { + const tool = disposables.add(createTool(createMockTextModelService(null!), createMockWorkspaceService())); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', lineContent: 'MyClass' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + assert.ok(getTextContent(result).includes('Provide either')); + }); + + test('returns error when line content not found', async () => { + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, testUri)); + disposables.add(langFeatures.referenceProvider.register('typescript', { provideReferences: () => [] })); + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService())); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', uri: testUri.toString(), lineContent: 'nonexistent line' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + assert.ok(getTextContent(result).includes('Could not find line content')); + }); + + test('returns error when symbol not found in line', async () => { + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, testUri)); + disposables.add(langFeatures.referenceProvider.register('typescript', { provideReferences: () => [] })); + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService())); + const result = await tool.invoke( + createInvocation({ symbol: 'NotHere', uri: testUri.toString(), lineContent: 'function doSomething' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + assert.ok(getTextContent(result).includes('Could not find symbol')); + }); + + test('finds references and classifies them with usage tags', async () => { + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, testUri)); + const otherUri = URI.parse('file:///test/other.ts'); + + const refProvider: ReferenceProvider = { + provideReferences: (_model: ITextModel): Location[] => [ + { uri: testUri, range: new Range(1, 10, 1, 17) }, + { uri: testUri, range: new Range(4, 23, 4, 30) }, + { uri: otherUri, range: new Range(5, 1, 5, 8) }, + ] + }; + const defProvider: DefinitionProvider = { + provideDefinition: () => [{ uri: testUri, range: new Range(1, 10, 1, 17) }] + }; + const implProvider: ImplementationProvider = { + provideImplementation: () => [{ uri: otherUri, range: new Range(5, 1, 5, 8) }] + }; + + disposables.add(langFeatures.referenceProvider.register('typescript', refProvider)); + disposables.add(langFeatures.definitionProvider.register('typescript', defProvider)); + disposables.add(langFeatures.implementationProvider.register('typescript', implProvider)); + + // Model is open for testUri so IModelService returns it; otherUri needs search + const searchCalled: ITextQuery[] = []; + const searchService = createMockSearchService(query => { + searchCalled.push(query); + const fileMatch = new FileMatch(otherUri); + fileMatch.results = [new TextSearchMatch( + 'export class MyClass implements IMyClass {', + new OneLineRange(4, 0, 7) // 0-based line 4 = 1-based line 5 + )]; + return { results: [fileMatch], messages: [] }; + }); + const modelService = createMockModelService([model]); + + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService(), { modelService, searchService })); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', uri: testUri.toString(), lineContent: 'import { MyClass }' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + + const text = getTextContent(result); + + // Check overall structure + assert.ok(text.includes('3 usages of `MyClass`')); + + // Check usage tag format + assert.ok(text.includes(``)); + assert.ok(text.includes(``)); + assert.ok(text.includes(``)); + + // Check that previews from open model are included (testUri lines) + assert.ok(text.includes('import { MyClass } from "./myClass"')); + assert.ok(text.includes('const instance = new MyClass()')); + + // Check that preview from search service is included (otherUri) + assert.ok(text.includes('export class MyClass implements IMyClass {')); + + // Check closing tags + assert.ok(text.includes('')); + + // Verify search service was called for the non-open file + assert.strictEqual(searchCalled.length, 1); + assert.ok(searchCalled[0].contentPattern.pattern.includes('MyClass')); + assert.ok(searchCalled[0].contentPattern.isWordMatch); + }); + + test('uses self-closing tag when no preview available', async () => { + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, testUri)); + const otherUri = URI.parse('file:///test/other.ts'); + + disposables.add(langFeatures.referenceProvider.register('typescript', { + provideReferences: (): Location[] => [ + { uri: otherUri, range: new Range(10, 5, 10, 12) }, + ] + })); + + // Search returns no results for this file (symbol renamed/aliased) + const searchService = createMockSearchService(() => ({ results: [], messages: [] })); + + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService(), { searchService })); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', uri: testUri.toString(), lineContent: 'import { MyClass }' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + + const text = getTextContent(result); + assert.ok(text.includes(``)); + }); + + test('does not call search service for files already open in model service', async () => { + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, testUri)); + + disposables.add(langFeatures.referenceProvider.register('typescript', { + provideReferences: (): Location[] => [ + { uri: testUri, range: new Range(1, 10, 1, 17) }, + ] + })); + + let searchCalled = false; + const searchService = createMockSearchService(() => { + searchCalled = true; + return { results: [], messages: [] }; + }); + const modelService = createMockModelService([model]); + + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService(), { modelService, searchService })); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', uri: testUri.toString(), lineContent: 'import { MyClass }' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + + assert.ok(getTextContent(result).includes('1 usages')); + assert.strictEqual(searchCalled, false, 'search service should not be called when all files are open'); + }); + + test('handles whitespace normalization in lineContent', async () => { + const content = 'function doSomething(x: number) {}'; + const model = disposables.add(createTextModel(content, 'typescript', undefined, testUri)); + + disposables.add(langFeatures.referenceProvider.register('typescript', { + provideReferences: (): Location[] => [ + { uri: testUri, range: new Range(1, 12, 1, 23) }, + ] + })); + + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService())); + const result = await tool.invoke( + createInvocation({ symbol: 'doSomething', uri: testUri.toString(), lineContent: 'function doSomething(x: number)' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + + assert.ok(getTextContent(result).includes('1 usages')); + }); + + test('resolves filePath via workspace folders', async () => { + const fileUri = URI.parse('file:///test/src/file.ts'); + const model = disposables.add(createTextModel(testContent, 'typescript', undefined, fileUri)); + + disposables.add(langFeatures.referenceProvider.register('typescript', { + provideReferences: (): Location[] => [ + { uri: fileUri, range: new Range(1, 10, 1, 17) }, + ] + })); + + const tool = disposables.add(createTool(createMockTextModelService(model), createMockWorkspaceService())); + const result = await tool.invoke( + createInvocation({ symbol: 'MyClass', filePath: 'src/file.ts', lineContent: 'import { MyClass }' }), + noopCountTokens, noopProgress, CancellationToken.None + ); + + assert.ok(getTextContent(result).includes('1 usages')); + }); + }); +});