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 { raceTimeoutAndCancellationError } from '../../../util/common/racePromise';
+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,32 @@ 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);
+
+		const results = await raceTimeoutAndCancellationError(
+			Promise.resolve(this.searchService.findFiles(pattern, undefined, searchCancellation.token)),
+			searchCancellation,
+			timeoutInMs,
+			'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

@@ -10,8 +10,9 @@ import { OffsetLineColumnConverter } from '../../../platform/editing/common/offs
 import { IPromptPathRepresentationService } from '../../../platform/prompts/common/promptPathRepresentationService';
 import { ISearchService } from '../../../platform/search/common/searchService';
 import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService';
+import { raceTimeoutAndCancellationError } from '../../../util/common/racePromise';
 import { asArray } from '../../../util/vs/base/common/arrays';
-import { CancellationToken } from '../../../util/vs/base/common/cancellation';
+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,41 @@ 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);
+
+		let results = await raceTimeoutAndCancellationError(
+			this.searchAndCollectResults(options.input.query, isRegExp, patterns, maxResults, searchCancellation.token),
+			searchCancellation,
+			timeoutInMs,
+			// 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),
+				searchCancellation,
+				timeoutInMs,
+				// 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));

src/util/common/racePromise.ts

@@ -0,0 +1,39 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { raceCancellation, raceTimeout } from '../vs/base/common/async';
+import { CancellationTokenSource } from '../vs/base/common/cancellation';
+import { CancellationError } from '../vs/base/common/errors';
+
+// sentinel value to indicate cancellation
+const CANCELLED = Symbol('cancelled');
+
+/**
+ * Races a promise against a cancellation token and a timeout.
+ * @param promise The promise to race.
+ * @param cancellationSource The cancellation token source to use.
+ * @param timeoutInMs The timeout in milliseconds.
+ * @param timeoutMessage The message to use for the timeout error.
+ * @returns The result of the promise if it completes before the timeout, or throws an error if it times out or is cancelled.
+ */
+export async function raceTimeoutAndCancellationError<T>(
+	promise: Promise<T>,
+	cancellationSource: CancellationTokenSource,
+	timeoutInMs: number,
+	timeoutMessage: string): Promise<T> {
+	const result = await raceTimeout(raceCancellation(promise, cancellationSource.token, CANCELLED as T), timeoutInMs);
+
+	if (result === CANCELLED) { // cancelled sentinel from raceCancellation
+		throw new CancellationError();
+	}
+
+	if (result === undefined) { // timeout sentinel from raceTimeout
+		// signal ongoing work to cancel in the promise
+		cancellationSource.cancel();
+		throw new Error(timeoutMessage);
+	}
+
+	return result;
+}