Handle invalid uris better and move out internal tool name (#244049)

Follow up to @connor4312's comment here #243922 (comment)

Author: TylerLeonhardt

src/vs/workbench/contrib/chat/common/tools/tools.ts

@@ -28,3 +28,5 @@ export class BuiltinToolsContribution extends Disposable implements IWorkbenchCo
 export interface IToolInputProcessor {
 	processInput(input: any): any;
 }
+
+export const InternalFetchWebPageToolId = 'vscode_fetchWebPage_internal';

src/vs/workbench/contrib/chat/electron-sandbox/tools/fetchPageTool.ts

@@ -10,8 +10,8 @@ import { IWebContentExtractorService } from '../../../../../platform/webContentE
 import { ITrustedDomainService } from '../../../url/browser/trustedDomainService.js';
 import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultTextPart } from '../../common/languageModelToolsService.js';
 import { MarkdownString } from '../../../../../base/common/htmlContent.js';
+import { InternalFetchWebPageToolId } from '../../common/tools/tools.js';
 
-export const InternalFetchWebPageToolId = 'vscode_fetchWebPage_internal';
 export const FetchWebPageToolData: IToolData = {
 	id: InternalFetchWebPageToolId,
 	displayName: 'Fetch Web Page',
@@ -41,26 +41,50 @@ export class FetchWebPageTool implements IToolImpl {
 	) { }
 
 	async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, _token: CancellationToken): Promise<IToolResult> {
-		const { valid } = this._parseUris((invocation.parameters as { urls?: string[] }).urls);
-		if (!valid.length) {
+		const parsedUriResults = this._parseUris((invocation.parameters as { urls?: string[] }).urls);
+		const validUris = Array.from(parsedUriResults.values()).filter((uri): uri is URI => !!uri);
+		if (!validUris.length) {
 			return {
 				content: [{ kind: 'text', value: localize('fetchWebPage.noValidUrls', 'No valid URLs provided.') }]
 			};
 		}
 
-		for (const uri of valid) {
+		// We approved these via confirmation, so mark them as "approved" in this session
+		// if they are not approved via the trusted domain service.
+		for (const uri of validUris) {
 			if (!this._trustedDomainService.isValid(uri)) {
 				this._alreadyApprovedDomains.add(uri.toString(true));
 			}
 		}
 
-		const contents = await this._readerModeService.extract(valid);
+		const contents = await this._readerModeService.extract(validUris);
+		// Make an array that conatains either the content or undefined for invalid URLs
+		const contentsWithUndefined = new Map<string, string | undefined>();
+		let indexInContents = 0;
+		parsedUriResults.forEach((uri, url) => {
+			if (uri) {
+				contentsWithUndefined.set(url, contents[indexInContents]);
+				indexInContents++;
+			} else {
+				contentsWithUndefined.set(url, undefined);
+			}
+		});
+
 		// TODO: Should we return a content for invalid URLs so that the indexes are aligned?
-		return { content: contents.map((content, index) => this._getPromptPartForWebPageContents(content, valid[index])) };
+		return { content: this._getPromptPartsForResults(contentsWithUndefined) };
 	}
 
 	async prepareToolInvocation(parameters: any, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
-		const { invalid, valid } = this._parseUris(parameters.urls);
+		const map = this._parseUris(parameters.urls);
+		const invalid = new Array<string>();
+		const valid = new Array<URI>();
+		map.forEach((uri, url) => {
+			if (!uri) {
+				invalid.push(url);
+			} else {
+				valid.push(uri);
+			}
+		});
 		const urlsNeedingConfirmation = valid.filter(url => !this._trustedDomainService.isValid(url) && !this._alreadyApprovedDomains.has(url.toString(true)));
 
 		const pastTenseMessage = invalid.length
@@ -118,25 +142,34 @@ export class FetchWebPageTool implements IToolImpl {
 		return result;
 	}
 
-	private _parseUris(urls?: string[]): { invalid: string[]; valid: URI[] } {
-		const invalidUrls: string[] = [];
-		const validUrls: URI[] = [];
+	private _parseUris(urls?: string[]): Map<string, URI | undefined> {
+		const results = new Map<string, URI | undefined>();
 		urls?.forEach(uri => {
 			try {
 				const uriObj = URI.parse(uri);
-				validUrls.push(uriObj);
+				results.set(uri, uriObj);
 			} catch (e) {
-				invalidUrls.push(uri);
+				results.set(uri, undefined);
 			}
 		});
-
-		return { invalid: invalidUrls, valid: validUrls };
+		return results;
 	}
 
-	private _getPromptPartForWebPageContents(webPageContents: string, uri: URI): IToolResultTextPart {
-		return {
-			kind: 'text',
-			value: `<!-- ${uri.toString(true)} -->\n\n` + webPageContents
-		};
+	private _getPromptPartsForResults(results: Map<string, string | undefined>): IToolResultTextPart[] {
+		const arr = new Array<IToolResultTextPart>();
+		for (const [url, content] of results.entries()) {
+			if (content) {
+				arr.push({
+					kind: 'text',
+					value: `<!-- ${url} -->\n\n` + content
+				});
+			} else {
+				arr.push({
+					kind: 'text',
+					value: `<!-- ${url} -->\n\n` + localize('fetchWebPage.invalidUrl', 'Invalid URL')
+				});
+			}
+		}
+		return arr;
 	}
 }