Integrate with fig's generator code

extensions/terminal-suggest/src/fig/api-binding-wrappers/executeCommand.ts

@@ -0,0 +1,56 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { osIsWindows } from '../../helpers/os';
+import { spawnHelper2 } from '../../shell/common';
+import { withTimeout } from '../shared/utils';
+
+export const cleanOutput = (output: string) =>
+	output
+		.replace(/\r\n/g, '\n') // Replace carriage returns with just a normal return
+		.replace(/\x1b\[\?25h/g, '') // removes cursor character if present
+		.replace(/^\n+/, '') // strips new lines from start of output
+		.replace(/\n+$/, ''); // strips new lines from end of output
+
+export const executeCommandTimeout = async (
+	input: Fig.ExecuteCommandInput,
+	timeout = osIsWindows() ? 20000 : 5000,
+): Promise<Fig.ExecuteCommandOutput> => {
+	const command = [input.command, ...input.args].join(' ');
+	try {
+		console.info(`About to run shell command '${command}'`);
+		const start = performance.now();
+		const result = await withTimeout(
+			Math.max(timeout, input.timeout ?? 0),
+			spawnHelper2(input.command, input.args, {
+				env: input.env,
+				cwd: input.cwd,
+				timeout: input.timeout,
+			})
+		);
+		const end = performance.now();
+		console.info(`Result of shell command '${command}'`, {
+			result,
+			time: end - start,
+		});
+
+		const cleanStdout = cleanOutput(result.stdout);
+		const cleanStderr = cleanOutput(result.stderr);
+
+		if (result.exitCode !== 0) {
+			console.warn(
+				`Command ${command} exited with exit code ${result.exitCode}: ${cleanStderr}`,
+			);
+		}
+		return {
+			status: result.exitCode,
+			stdout: cleanStdout,
+			stderr: cleanStderr,
+		};
+	} catch (err) {
+		console.error(`Error running shell command '${command}'`, { err });
+		throw err;
+	}
+};

extensions/terminal-suggest/src/fig/api-binding-wrappers/executeCommandWrappers.ts

@@ -0,0 +1,9 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { executeCommandTimeout } from './executeCommand';
+
+export const executeCommand: Fig.ExecuteCommandFunction = (args) =>
+	executeCommandTimeout(args);

extensions/terminal-suggest/src/fig/api-bindings/types.ts

@@ -0,0 +1,34 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export interface EnvironmentVariable {
+	key: string;
+	value?: string;
+}
+
+export interface ShellContext {
+	/** The current PID of the shell process */
+	pid?: number;
+	/** /dev/ttys## of terminal session */
+	ttys?: string;
+	/** the name of the process */
+	processName?: string;
+	/** the directory where the user ran the command */
+	currentWorkingDirectory?: string;
+	/** the value of $TERM_SESSION_ID */
+	sessionId?: string;
+	/** the integration version of figterm */
+	integrationVersion?: number;
+	/** the parent terminal of figterm */
+	terminal?: string;
+	/** the hostname of the computer figterm is running on */
+	hostname?: string;
+	/** path to the shell being used in the terminal */
+	shellPath?: string;
+	/** the environment variables of the shell, note that only exported variables are included */
+	environmentVariables?: EnvironmentVariable[];
+	/** the raw output of alias */
+	alias?: string;
+}

extensions/terminal-suggest/src/fig/autocomplete/fig/hooks.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as Types from '../../api-bindings/types';
+import type { AliasMap } from '../../shell-parser';
+
+export type FigState = {
+	buffer: string;
+	cursorLocation: number;
+	cwd: string | null;
+	processUserIsIn: string | null;
+	sshContextString: string | null;
+	aliases: AliasMap;
+	environmentVariables: Record<string, string>;
+	shellContext?: Types.ShellContext | undefined;
+};

extensions/terminal-suggest/src/fig/autocomplete/generators/cache.ts

@@ -0,0 +1,92 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+class CacheEntry<T> {
+	private lastFetch = 0;
+
+	private promise: Promise<T> | undefined = undefined;
+
+	value: T | undefined = undefined;
+
+	private isInitialized = false;
+
+	private get isFetching() {
+		return !!this.promise;
+	}
+
+	private async fetchSync(run: () => Promise<T>) {
+		this.lastFetch = Date.now();
+		this.promise = run();
+		this.value = await this.promise;
+		if (!this.isInitialized) {
+			this.isInitialized = true;
+		}
+		this.promise = undefined;
+		return this.value;
+	}
+
+	private async fetchAsync(run: () => Promise<T>): Promise<T | undefined> {
+		if (this.isFetching) {
+			await this.promise;
+			return this.value;
+		}
+		return this.fetchSync(run);
+	}
+
+	private async maxAgeCache(run: () => Promise<T>, maxAge: number) {
+		if (Date.now() > maxAge + this.lastFetch) {
+			return this.fetchAsync(run);
+		}
+		return this.value;
+	}
+
+	private async swrCache(run: () => Promise<T>, maxAge = 0) {
+		if (!this.isFetching && Date.now() > this.lastFetch + maxAge) {
+			return this.fetchAsync(run);
+		}
+		return this.value as T;
+	}
+
+	async entry(run: () => Promise<T>, cache: Fig.Cache): Promise<T | undefined> {
+		if (!this.isInitialized) {
+			return this.fetchAsync(run);
+		}
+		switch (cache.strategy || 'stale-while-revalidate') {
+			case 'max-age':
+				return this.maxAgeCache(run, cache.ttl!);
+			case 'stale-while-revalidate':
+				// cache.ttl must be defined when no strategy is specified
+				return this.swrCache(run, cache.ttl!);
+			default:
+				return this.fetchAsync(run);
+		}
+	}
+}
+
+export class Cache {
+	private cache = new Map<string, CacheEntry<unknown>>();
+
+	async entry<T>(
+		key: string,
+		run: () => Promise<T>,
+		cache: Fig.Cache,
+	): Promise<T> {
+		if (!this.cache.has(key)) {
+			this.cache.set(key, new CacheEntry());
+		}
+		return this.cache.get(key)!.entry(run, cache) as Promise<T>;
+	}
+
+	currentValue<T>(key: string, _cache: Fig.Cache): T | undefined {
+		if (!this.cache.has(key)) {
+			this.cache.set(key, new CacheEntry());
+		}
+		return this.cache.get(key)!.value as T | undefined;
+	}
+
+	clear() {
+		this.cache.clear();
+	}
+}

extensions/terminal-suggest/src/fig/autocomplete/generators/customSuggestionsGenerator.ts

@@ -0,0 +1,57 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { executeCommand } from '../../api-binding-wrappers/executeCommandWrappers';
+import {
+	runCachedGenerator,
+	GeneratorContext,
+	haveContextForGenerator,
+} from './helpers';
+
+export async function getCustomSuggestions(
+	generator: Fig.Generator,
+	context: GeneratorContext,
+): Promise<Fig.Suggestion[] | undefined> {
+	if (!generator.custom) {
+		return [];
+	}
+
+	if (!haveContextForGenerator(context)) {
+		console.info('Don\'t have context for custom generator');
+		return [];
+	}
+
+	const {
+		tokenArray,
+		currentWorkingDirectory,
+		currentProcess,
+		isDangerous,
+		searchTerm,
+		environmentVariables,
+	} = context;
+
+	try {
+		const result = await runCachedGenerator(
+			generator,
+			context,
+			() =>
+				generator.custom!(tokenArray, executeCommand, {
+					currentWorkingDirectory,
+					currentProcess,
+					sshPrefix: '',
+					searchTerm,
+					environmentVariables,
+					isDangerous,
+				}),
+			generator.cache?.cacheKey,
+		);
+
+		return result?.map((name) => ({ ...name, type: name?.type || 'arg' }));
+	} catch (e) {
+		console.error('we had an error with the custom function generator', e);
+
+		return [];
+	}
+}

extensions/terminal-suggest/src/fig/autocomplete/generators/helpers.ts

@@ -0,0 +1,58 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Cache } from './cache';
+import type { Annotation } from '../../autocomplete-parser/parseArguments';
+import type { Suggestion } from '../../shared/internal';
+import { getCWDForFilesAndFolders } from '../../shared/utils';
+
+export type GeneratorContext = Fig.ShellContext & {
+	annotations: Annotation[];
+	tokenArray: string[];
+	isDangerous?: boolean;
+	searchTerm: string;
+};
+
+export type GeneratorState = {
+	generator: Fig.Generator;
+	context: GeneratorContext;
+	loading: boolean;
+	result: Suggestion[];
+	request?: Promise<Fig.Suggestion[] | undefined>;
+};
+
+export const haveContextForGenerator = (context: GeneratorContext): boolean =>
+	Boolean(context.currentWorkingDirectory);
+
+export const generatorCache = new Cache();
+
+export async function runCachedGenerator<T>(
+	generator: Fig.Generator,
+	context: GeneratorContext,
+	initialRun: () => Promise<T>,
+	cacheKey?: string /* This is generator.script or generator.script(...) */,
+): Promise<T> {
+	const cacheDefault = false; // getSetting<boolean>(SETTINGS.CACHE_ALL_GENERATORS) ?? false;
+	let { cache } = generator;
+	if (!cache && cacheDefault) {
+		cache = { strategy: 'stale-while-revalidate', ttl: 1_000 };
+	}
+	if (!cache) {
+		return initialRun();
+	}
+	const { tokenArray, currentWorkingDirectory, searchTerm } = context;
+
+	const directory = generator.template
+		? getCWDForFilesAndFolders(currentWorkingDirectory, searchTerm)
+		: currentWorkingDirectory;
+
+	// we cache generator results by script, if no script was provided we use the tokens instead
+	const key = [
+		cache.cacheByDirectory ? directory : undefined,
+		cacheKey || tokenArray.join(' '),
+	].toString();
+
+	return generatorCache.entry(key, initialRun, cache);
+}