feat: add timeout in searching text in files (#1233)

* feat: add timeout handling for file and text search operations

* feat: add timeout handling for search and rendering operations in FindTextInFilesTool

* use raceTimeout and raceCancellation

* Updated per comments

* extract a raceTimeoutAndCancellationError()

* enhance timeout handling and cancellation for file and text search operations

* enhance raceTimeoutAndCancellationError

* removed unnecessary check

* refactor: streamline cancellation token usage in file and text search operations

Author: bryanchen-d

src/extension/tools/node/findFilesTool.tsx

@@ -11,6 +11,7 @@ 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 { raceTimeoutAndCancellationError } from '../../../util/common/racePromise';
 import { CancellationToken } from '../../../util/vs/base/common/cancellation';
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
 import { ExtendedLanguageModelToolResult, LanguageModelPromptTsxPart, MarkdownString } from '../../../vscodeTypes';
@@ -40,14 +41,24 @@ 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;
+
+
+		const results = await raceTimeoutAndCancellationError(
+			(searchToken) => Promise.resolve(this.searchService.findFiles(pattern, undefined, searchToken)),
+			token,
+			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);
+		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,6 +10,7 @@ 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 { count } from '../../../util/vs/base/common/strings';
@@ -51,14 +52,36 @@ 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;
+
+		let results = await raceTimeoutAndCancellationError(
+			(searchToken) => this.searchAndCollectResults(options.input.query, isRegExp, patterns, maxResults, searchToken),
+			token,
+			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(
+				(searchToken) => this.searchAndCollectResults(options.input.query, !isRegExp, patterns, maxResults, searchToken),
+				token,
+				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,45 @@
+/*---------------------------------------------------------------------------------------------
+ *  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 { CancellationToken, 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 promiseGenerator A function that generates the promise to race against cancellation and timeout.
+ * @param parentToken The cancellation token 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>(
+	promiseGenerator: (cancellationToken: CancellationToken) => Promise<T>,
+	parentToken: CancellationToken,
+	timeoutInMs: number,
+	timeoutMessage: string): Promise<T> {
+	const cancellationSource = new CancellationTokenSource(parentToken);
+	try {
+		const result = await raceTimeout(raceCancellation(promiseGenerator(cancellationSource.token), 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;
+	}
+	finally {
+		cancellationSource.dispose();
+	}
+}