feat: add timeout in searching text in files

src/extension/tools/node/findFilesTool.tsx

@@ -11,7 +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 { raceCancellationError, raceTimeout } from '../../../util/vs/base/common/async';
+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';
@@ -48,19 +48,10 @@ export class FindFilesTool implements ICopilotTool<IFindFilesToolParams> {
 		// 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)),
+			searchCancellation,
+			timeoutInMs,
 			'Timeout in searching files, try a more specific search pattern'
 		);
 

src/extension/tools/node/findTextInFilesTool.tsx

@@ -10,8 +10,8 @@ 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 { 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';
@@ -61,19 +61,10 @@ export class FindTextInFilesTool implements ICopilotTool<IFindTextInFilesToolPar
 		// 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),
+			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`
 		);
@@ -82,6 +73,8 @@ export class FindTextInFilesTool implements ICopilotTool<IFindTextInFilesToolPar
 		if (!results.length && queryIsValidRegex) {
 			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`
 			);
@@ -245,8 +238,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>,

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;
+}