diff --git a/src/extension/tools/node/findFilesTool.tsx b/src/extension/tools/node/findFilesTool.tsx index e01d53363..471253632 100644 --- a/src/extension/tools/node/findFilesTool.tsx +++ b/src/extension/tools/node/findFilesTool.tsx @@ -11,7 +11,8 @@ import { URI } from '../../../util/vs/base/common/uri'; import * as l10n from '@vscode/l10n'; import { ISearchService } from '../../../platform/search/common/searchService'; import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService'; -import { CancellationToken } from '../../../util/vs/base/common/cancellation'; +import { raceCancellationError, raceTimeout } from '../../../util/vs/base/common/async'; +import { CancellationToken, CancellationTokenSource } from '../../../util/vs/base/common/cancellation'; import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation'; import { ExtendedLanguageModelToolResult, LanguageModelPromptTsxPart, MarkdownString } from '../../../vscodeTypes'; import { IBuildPromptContext } from '../../prompt/common/intents'; @@ -40,14 +41,41 @@ export class FindFilesTool implements ICopilotTool { // The input _should_ be a pattern matching inside a workspace, folder, but sometimes we get absolute paths, so try to resolve them const pattern = inputGlobToPattern(options.input.query, this.workspaceService); - const results = await this.searchService.findFiles(pattern, undefined, token); + // try find text with a timeout of 20s + const timeoutInMs = 20_000; + // create a new cancellation token to be used in search + // so in the case of timeout, we can cancel the search + // also in the case of the parent token being cancelled, it will cancel this one too + const searchCancellation = new CancellationTokenSource(token); + + async function raceTimeoutAndCancellationError(promise: Promise, timeoutMessage: string): Promise { + const result = await raceTimeout(raceCancellationError(promise, token), timeoutInMs); + if (result === undefined) { + // we have timed out, so cancel the search + searchCancellation.cancel(); + throw new Error(timeoutMessage); + } + + return result; + } + + const results = await raceTimeoutAndCancellationError( + Promise.resolve(this.searchService.findFiles(pattern, undefined, searchCancellation.token)), + 'Timeout in searching files, try a more specific search pattern' + ); + checkCancellation(token); const maxResults = options.input.maxResults ?? 20; const resultsToShow = results.slice(0, maxResults); - const result = new ExtendedLanguageModelToolResult([ - new LanguageModelPromptTsxPart( - await renderPromptElementJSON(this.instantiationService, FindFilesResult, { fileResults: resultsToShow, totalResults: results.length }, options.tokenizationOptions, token))]); + // Render the prompt element with a timeout + const prompt = await renderPromptElementJSON(this.instantiationService, FindFilesResult, { fileResults: resultsToShow, totalResults: results.length }, options.tokenizationOptions, token); + + if (prompt === undefined) { + throw new Error('Timeout in rendering prompt'); + } + + const result = new ExtendedLanguageModelToolResult([new LanguageModelPromptTsxPart(prompt)]); const query = `\`${options.input.query}\``; result.toolResultMessage = resultsToShow.length === 0 ? new MarkdownString(l10n.t`Searched for files matching ${query}, no matches`) : diff --git a/src/extension/tools/node/findTextInFilesTool.tsx b/src/extension/tools/node/findTextInFilesTool.tsx index 964ee9b22..13482e78e 100644 --- a/src/extension/tools/node/findTextInFilesTool.tsx +++ b/src/extension/tools/node/findTextInFilesTool.tsx @@ -11,7 +11,8 @@ import { IPromptPathRepresentationService } from '../../../platform/prompts/comm import { ISearchService } from '../../../platform/search/common/searchService'; import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService'; import { asArray } from '../../../util/vs/base/common/arrays'; -import { CancellationToken } from '../../../util/vs/base/common/cancellation'; +import { raceCancellationError, raceTimeout } from '../../../util/vs/base/common/async'; +import { CancellationToken, CancellationTokenSource } from '../../../util/vs/base/common/cancellation'; import { count } from '../../../util/vs/base/common/strings'; import { URI } from '../../../util/vs/base/common/uri'; import { Position as EditorPosition } from '../../../util/vs/editor/common/core/position'; @@ -51,14 +52,48 @@ export class FindTextInFilesTool implements ICopilotTool(promise: Promise, timeoutMessage: string): Promise { + const result = await raceTimeout(raceCancellationError(promise, token), timeoutInMs); + if (result === undefined) { + // we have timed out, so cancel the search + searchCancellation.cancel(); + throw new Error(timeoutMessage); + } + + return result; + } + + let results = await raceTimeoutAndCancellationError( + this.searchAndCollectResults(options.input.query, isRegExp, patterns, maxResults, searchCancellation.token), + // embed message to give LLM hint about what to do next + `Timeout in searching text in files with ${isRegExp ? 'regex' : 'literal'} search, try a more specific search pattern or change regex/literal mode` + ); + + // If we still have no results, we need to try the opposite regex mode if (!results.length && queryIsValidRegex) { - results = await this.searchAndCollectResults(options.input.query, !isRegExp, patterns, maxResults, token); + results = await raceTimeoutAndCancellationError( + this.searchAndCollectResults(options.input.query, !isRegExp, patterns, maxResults, searchCancellation.token), + // embed message to give LLM hint about what to do next + `Find ${results.length} results in searching text in files with ${isRegExp ? 'regex' : 'literal'} search, and then another searching hits timeout in with ${!isRegExp ? 'regex' : 'literal'} search, try a more specific search pattern` + ); } - const result = new ExtendedLanguageModelToolResult([ - new LanguageModelPromptTsxPart( - await renderPromptElementJSON(this.instantiationService, FindTextInFilesResult, { textResults: results, maxResults, askedForTooManyResults: Boolean(askedForTooManyResults) }, options.tokenizationOptions, token))]); + const prompt = await renderPromptElementJSON(this.instantiationService, + FindTextInFilesResult, + { textResults: results, maxResults, askedForTooManyResults: Boolean(askedForTooManyResults) }, + options.tokenizationOptions, + token); + + const result = new ExtendedLanguageModelToolResult([new LanguageModelPromptTsxPart(prompt)]); const textMatches = results.flatMap(r => { if ('ranges' in r) { return asArray(r.ranges).map(rangeInfo => new Location(r.uri, rangeInfo.sourceRange)); @@ -210,8 +245,8 @@ interface IFindMatchProps extends BasePromptElementProps { * 1. Removes excessive extra character data from the match, e.g. avoiding * giant minified lines * 2. Wraps the match in a tag - * 3. Prioritizes lines in the middle of the match where the range lies - */ + * 3. Prioritizes lines in the middle of the match where the range lies + */ export class FindMatch extends PromptElement { constructor( props: PromptElementProps,