feat: add timeout in searching text in files

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<IFindFilesToolParams> {
 		// 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<T>(promise: Promise<T>, timeoutMessage: string): Promise<T> {
+			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`) :

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<IFindTextInFilesToolPar
 		const maxResults = Math.min(options.input.maxResults ?? 20, MaxResultsCap);
 		const isRegExp = options.input.isRegexp ?? true;
 		const queryIsValidRegex = this.isValidRegex(options.input.query);
-		let results = await this.searchAndCollectResults(options.input.query, isRegExp, patterns, maxResults, 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<T>(promise: Promise<T>, timeoutMessage: string): Promise<T> {
+			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 <match> 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<IFindMatchProps> {
 	constructor(
 		props: PromptElementProps<IFindMatchProps>,