fix: some performance issues via abort controller

src/runtime/generator.ts

@@ -17,11 +17,12 @@ const getGeneratorContext = (cwd: string): Fig.GeneratorContext => {
 };
 
 // TODO: add support for caching, trigger, & getQueryTerm
-export const runGenerator = async (generator: Fig.Generator, tokens: string[], cwd: string): Promise<Fig.Suggestion[]> => {
+export const runGenerator = async (generator: Fig.Generator, tokens: string[], cwd: string, signal?: AbortSignal): Promise<Fig.Suggestion[]> => {
   // TODO: support trigger
+  signal?.throwIfAborted();
   const { script, postProcess, scriptTimeout, splitOn, custom, template, filterTemplateSuggestions } = generator;
 
-  const executeShellCommand = await buildExecuteShellCommand(scriptTimeout ?? 5000);
+  const executeShellCommand = await buildExecuteShellCommand(scriptTimeout ?? 5000, signal);
   const suggestions = [];
   try {
     if (script) {
@@ -43,7 +44,7 @@ export const runGenerator = async (generator: Fig.Generator, tokens: string[], c
     }
 
     if (template != null) {
-      const templateSuggestions = await runTemplates(template, cwd);
+      const templateSuggestions = await runTemplates(template, cwd, signal);
       if (filterTemplateSuggestions) {
         suggestions.push(...filterTemplateSuggestions(templateSuggestions));
       } else {

src/runtime/runtime.ts

@@ -47,7 +47,7 @@ loadSpecsSet(speclist as string[], versionedSpeclist, specResourcesPath);
 
 const loadedSpecs: { [key: string]: Fig.Spec } = {};
 
-const loadSpec = async (cmd: CommandToken[]): Promise<Fig.Spec | undefined> => {
+const loadSpec = async (cmd: CommandToken[], signal?: AbortSignal): Promise<Fig.Spec | undefined> => {
   const rootToken = cmd.at(0);
   if (!rootToken?.complete) {
     return;
@@ -57,6 +57,7 @@ const loadSpec = async (cmd: CommandToken[]): Promise<Fig.Spec | undefined> => {
     return loadedSpecs[rootToken.token];
   }
   if (specSet[rootToken.token]) {
+    signal?.throwIfAborted();
     const specPath = specSet[rootToken.token];
     const importPath = path.isAbsolute(specPath) ? pathToFileURL(specPath).href : specPath;
     const spec = (await import(importPath)).default;
@@ -66,14 +67,15 @@ const loadSpec = async (cmd: CommandToken[]): Promise<Fig.Spec | undefined> => {
 };
 
 // this load spec function should only be used for `loadSpec` on the fly as it is cacheless
-const lazyLoadSpec = async (key: string): Promise<Fig.Spec | undefined> => {
+const lazyLoadSpec = async (key: string, signal?: AbortSignal): Promise<Fig.Spec | undefined> => {
+  signal?.throwIfAborted();
   const specPath = path.join(specResourcesPath, `${key}.js`);
   const importPath = path.isAbsolute(specPath) ? pathToFileURL(specPath).href : specPath;
   return (await import(importPath)).default;
 };
 
 // eslint-disable-next-line @typescript-eslint/no-unused-vars -- will be implemented in below TODO
-const lazyLoadSpecLocation = async (location: Fig.SpecLocation): Promise<Fig.Spec | undefined> => {
+const lazyLoadSpecLocation = async (location: Fig.SpecLocation, signal?: AbortSignal): Promise<Fig.Spec | undefined> => {
   return; //TODO: implement spec location loading
 };
 
@@ -95,7 +97,7 @@ export const loadLocalSpecsSet = async () => {
   );
 };
 
-export const getSuggestions = async (cmd: string, cwd: string, shell: Shell): Promise<SuggestionBlob | undefined> => {
+export const getSuggestions = async (cmd: string, cwd: string, shell: Shell, signal?: AbortSignal): Promise<SuggestionBlob | undefined> => {
   let activeCmd = parseCommand(cmd, shell);
   const rootToken = activeCmd.at(0);
   if (activeCmd.length === 0) {
@@ -107,18 +109,20 @@ export const getSuggestions = async (cmd: string, cwd: string, shell: Shell): Pr
 
   activeCmd = aliasExpand(activeCmd);
 
-  const spec = await loadSpec(activeCmd);
+  signal?.throwIfAborted();
+  const spec = await loadSpec(activeCmd, signal);
   if (spec == null) return;
   const subcommand = getSubcommand(spec);
   if (subcommand == null) return;
 
+  signal?.throwIfAborted();
   const lastCommand = activeCmd.at(-1);
-  const { cwd: resolvedCwd, pathy, complete: pathyComplete } = await resolveCwd(lastCommand, cwd, shell);
+  const { cwd: resolvedCwd, pathy, complete: pathyComplete } = await resolveCwd(lastCommand, cwd, shell, signal);
   if (pathy && lastCommand) {
     lastCommand.isPath = true;
     lastCommand.isPathComplete = pathyComplete;
   }
-  const result = await runSubcommand(activeCmd.slice(1), activeCmd, subcommand, resolvedCwd, shell);
+  const result = await runSubcommand(activeCmd.slice(1), activeCmd, subcommand, resolvedCwd, shell, undefined, undefined, undefined, undefined, signal);
   if (result == null) return;
   if (result.suggestions.length == 0 && !result.argumentDescription) return;
 
@@ -152,9 +156,7 @@ const getSubcommand = (spec?: Fig.Spec): Fig.Subcommand | undefined => {
   return spec;
 };
 
-const executeShellCommand = buildExecuteShellCommand(5000);
-
-const genSubcommand = async (command: string, parentCommand: Fig.Subcommand): Promise<Fig.Subcommand | undefined> => {
+const genSubcommand = async (command: string, parentCommand: Fig.Subcommand, signal?: AbortSignal): Promise<Fig.Subcommand | undefined> => {
   if (!parentCommand.subcommands || parentCommand.subcommands.length === 0) return;
 
   const subcommandIdx = parentCommand.subcommands.findIndex((s) => (Array.isArray(s.name) ? s.name.includes(command) : s.name === command));
@@ -164,11 +166,13 @@ const genSubcommand = async (command: string, parentCommand: Fig.Subcommand): Pr
 
   // this pulls in the spec from the load spec and overwrites the subcommand in the parent with the loaded spec.
   // then it returns the subcommand and clears the loadSpec field so that it doesn't get called again
+  signal?.throwIfAborted();
+  const executeShellCommand = buildExecuteShellCommand(5000, signal);
   switch (typeof subcommand.loadSpec) {
     case "function": {
       const partSpec = await subcommand.loadSpec(command, executeShellCommand);
       if (partSpec instanceof Array) {
-        const locationSpecs = (await Promise.all(partSpec.map((s) => lazyLoadSpecLocation(s)))).filter((s) => s != null) as Fig.Spec[];
+        const locationSpecs = (await Promise.all(partSpec.map((s) => lazyLoadSpecLocation(s, signal)))).filter((s) => s != null) as Fig.Spec[];
         const subcommands = locationSpecs.map((s) => getSubcommand(s)).filter((s) => s != null) as Fig.Subcommand[];
         (parentCommand.subcommands as Fig.Subcommand[])[subcommandIdx] = {
           ...subcommand,
@@ -177,7 +181,7 @@ const genSubcommand = async (command: string, parentCommand: Fig.Subcommand): Pr
         };
         return (parentCommand.subcommands as Fig.Subcommand[])[subcommandIdx];
       } else if (Object.prototype.hasOwnProperty.call(partSpec, "type")) {
-        const locationSingleSpec = await lazyLoadSpecLocation(partSpec as Fig.SpecLocation);
+        const locationSingleSpec = await lazyLoadSpecLocation(partSpec as Fig.SpecLocation, signal);
         (parentCommand.subcommands as Fig.Subcommand[])[subcommandIdx] = {
           ...subcommand,
           ...(getSubcommand(locationSingleSpec) ?? []),
@@ -194,7 +198,7 @@ const genSubcommand = async (command: string, parentCommand: Fig.Subcommand): Pr
       }
     }
     case "string": {
-      const spec = await lazyLoadSpec(subcommand.loadSpec as string);
+      const spec = await lazyLoadSpec(subcommand.loadSpec as string, signal);
       (parentCommand.subcommands as Fig.Subcommand[])[subcommandIdx] = {
         ...subcommand,
         ...(getSubcommand(spec) ?? []),
@@ -237,6 +241,7 @@ const runOption = async (
   shell: Shell,
   persistentOptions: Fig.Option[],
   acceptedTokens: CommandToken[],
+  signal?: AbortSignal,
 ): Promise<SuggestionBlob | undefined> => {
   if (tokens.length === 0) {
     throw new Error("invalid state reached, option expected but no tokens found");
@@ -245,7 +250,7 @@ const runOption = async (
   const isPersistent = persistentOptions.some((o) => (typeof o.name === "string" ? o.name === activeToken.token : o.name.includes(activeToken.token)));
   if ((option.args instanceof Array && option.args.length > 0) || option.args != null) {
     const args = option.args instanceof Array ? option.args : [option.args];
-    return runArg(tokens.slice(1), allTokens, args, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), true, false);
+    return runArg(tokens.slice(1), allTokens, args, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), true, false, signal);
   }
   return runSubcommand(
     tokens.slice(1),
@@ -258,6 +263,9 @@ const runOption = async (
       ...activeToken,
       isPersistent,
     }),
+    undefined,
+    undefined,
+    signal,
   );
 };
 
@@ -272,45 +280,70 @@ const runArg = async (
   acceptedTokens: CommandToken[],
   fromOption: boolean,
   fromVariadic: boolean,
+  signal?: AbortSignal,
 ): Promise<SuggestionBlob | undefined> => {
+  signal?.throwIfAborted();
   if (args.length === 0) {
-    return runSubcommand(tokens, allTokens, subcommand, cwd, shell, persistentOptions, acceptedTokens, true, !fromOption);
+    return runSubcommand(tokens, allTokens, subcommand, cwd, shell, persistentOptions, acceptedTokens, true, !fromOption, signal);
   } else if (tokens.length === 0) {
-    return await getArgDrivenRecommendation(args, subcommand, persistentOptions, undefined, acceptedTokens, allTokens, fromVariadic, cwd, shell);
+    return await getArgDrivenRecommendation(args, subcommand, persistentOptions, undefined, acceptedTokens, allTokens, fromVariadic, cwd, shell, signal);
   } else if (!tokens.at(0)?.complete) {
-    return await getArgDrivenRecommendation(args, subcommand, persistentOptions, tokens[0], acceptedTokens, allTokens, fromVariadic, cwd, shell);
+    return await getArgDrivenRecommendation(args, subcommand, persistentOptions, tokens[0], acceptedTokens, allTokens, fromVariadic, cwd, shell, signal);
   }
 
   const activeToken = tokens[0];
   if (args.every((a) => a.isOptional)) {
     if (activeToken.isOption) {
       const option = getOption(activeToken, persistentOptions.concat(subcommand.options ?? []));
       if (option != null) {
-        return runOption(tokens, allTokens, option, subcommand, cwd, shell, persistentOptions, acceptedTokens);
+        return runOption(tokens, allTokens, option, subcommand, cwd, shell, persistentOptions, acceptedTokens, signal);
       }
       return;
     }
 
-    const nextSubcommand = await genSubcommand(activeToken.token, subcommand);
+    const nextSubcommand = await genSubcommand(activeToken.token, subcommand, signal);
     if (nextSubcommand != null) {
-      return runSubcommand(tokens.slice(1), allTokens, nextSubcommand, cwd, shell, persistentOptions, getPersistentTokens(acceptedTokens.concat(activeToken)));
+      return runSubcommand(
+        tokens.slice(1),
+        allTokens,
+        nextSubcommand,
+        cwd,
+        shell,
+        persistentOptions,
+        getPersistentTokens(acceptedTokens.concat(activeToken)),
+        undefined,
+        undefined,
+        signal,
+      );
     }
   }
 
   const activeArg = args[0];
   if (activeArg.isVariadic) {
-    return runArg(tokens.slice(1), allTokens, args, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), fromOption, true);
+    return runArg(tokens.slice(1), allTokens, args, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), fromOption, true, signal);
   } else if (activeArg.isCommand) {
     if (tokens.length <= 0) {
       return;
     }
-    const spec = await loadSpec(tokens);
+    const spec = await loadSpec(tokens, signal);
     if (spec == null) return;
     const subcommand = getSubcommand(spec);
     if (subcommand == null) return;
-    return runSubcommand(tokens.slice(1), allTokens, subcommand, cwd, shell);
+    return runSubcommand(tokens.slice(1), allTokens, subcommand, cwd, shell, undefined, undefined, undefined, undefined, signal);
   }
-  return runArg(tokens.slice(1), allTokens, args.slice(1), subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), fromOption, false);
+  return runArg(
+    tokens.slice(1),
+    allTokens,
+    args.slice(1),
+    subcommand,
+    cwd,
+    shell,
+    persistentOptions,
+    acceptedTokens.concat(activeToken),
+    fromOption,
+    false,
+    signal,
+  );
 };
 
 const runSubcommand = async (
@@ -323,11 +356,13 @@ const runSubcommand = async (
   acceptedTokens: CommandToken[] = [],
   argsDepleted = false,
   argsUsed = false,
+  signal?: AbortSignal,
 ): Promise<SuggestionBlob | undefined> => {
+  signal?.throwIfAborted();
   if (tokens.length === 0) {
-    return getSubcommandDrivenRecommendation(subcommand, persistentOptions, undefined, argsDepleted, argsUsed, acceptedTokens, allTokens, cwd, shell);
+    return getSubcommandDrivenRecommendation(subcommand, persistentOptions, undefined, argsDepleted, argsUsed, acceptedTokens, allTokens, cwd, shell, signal);
   } else if (!tokens.at(0)?.complete) {
-    return getSubcommandDrivenRecommendation(subcommand, persistentOptions, tokens[0], argsDepleted, argsUsed, acceptedTokens, allTokens, cwd, shell);
+    return getSubcommandDrivenRecommendation(subcommand, persistentOptions, tokens[0], argsDepleted, argsUsed, acceptedTokens, allTokens, cwd, shell, signal);
   }
 
   const activeToken = tokens[0];
@@ -337,12 +372,12 @@ const runSubcommand = async (
   if (activeToken.isOption) {
     const option = getOption(activeToken, allOptions);
     if (option != null) {
-      return runOption(tokens, allTokens, option, subcommand, cwd, shell, persistentOptions, acceptedTokens);
+      return runOption(tokens, allTokens, option, subcommand, cwd, shell, persistentOptions, acceptedTokens, signal);
     }
     return;
   }
 
-  const nextSubcommand = await genSubcommand(activeToken.token, subcommand);
+  const nextSubcommand = await genSubcommand(activeToken.token, subcommand, signal);
   if (nextSubcommand != null) {
     return runSubcommand(
       tokens.slice(1),
@@ -352,6 +387,9 @@ const runSubcommand = async (
       shell,
       getPersistentOptions(persistentOptions, subcommand.options),
       getPersistentTokens(acceptedTokens.concat(activeToken)),
+      undefined,
+      undefined,
+      signal,
     );
   }
 
@@ -361,10 +399,10 @@ const runSubcommand = async (
 
   const args = getArgs(subcommand.args);
   if (args.length != 0) {
-    return runArg(tokens, allTokens, args, subcommand, cwd, shell, allOptions, acceptedTokens, false, false);
+    return runArg(tokens, allTokens, args, subcommand, cwd, shell, allOptions, acceptedTokens, false, false, signal);
   }
   // if the subcommand has no args specified, fallback to the subcommand and ignore this item
-  return runSubcommand(tokens.slice(1), allTokens, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken));
+  return runSubcommand(tokens.slice(1), allTokens, subcommand, cwd, shell, persistentOptions, acceptedTokens.concat(activeToken), undefined, undefined, signal);
 };
 
 const runCommand = async (token: CommandToken): Promise<SuggestionBlob | undefined> => {

src/runtime/suggestion.ts

@@ -108,13 +108,23 @@ const toSuggestion = (suggestion: Fig.Suggestion, name?: string, type?: Fig.Sugg
     description: suggestion.description,
     icon: getIcon(suggestion.icon, type ?? suggestion.type),
     allNames: suggestion.name instanceof Array ? suggestion.name : [suggestion.name],
-    priority: suggestion.priority ?? 50,
+    priority: getSuggestionPriority(suggestion),
     insertValue: suggestion.insertValue,
     type: suggestion.type,
     hidden: suggestion.hidden,
   };
 };
 
+const getSuggestionPriority = (suggestion: Fig.Suggestion): number => {
+  const isOption =
+    suggestion.type === "option" ||
+    (suggestion.name instanceof Array
+      ? suggestion.name.some((n) => n.startsWith("--") || n.startsWith("-"))
+      : suggestion?.name?.startsWith("--") || suggestion?.name?.startsWith("-"));
+  if (isOption) return 45;
+  return suggestion.priority ?? 50;
+};
+
 function filter<T extends Fig.BaseSuggestion & { name?: Fig.SingleOrArray<string>; type?: Fig.SuggestionType | undefined }>(
   suggestions: T[],
   filterStrategy: FilterStrategy | undefined,
@@ -201,11 +211,13 @@ const generatorSuggestions = async (
   filterStrategy: FilterStrategy | undefined,
   partialCmd: string | undefined,
   cwd: string,
+  signal?: AbortSignal,
 ): Promise<Suggestion[]> => {
   const generators = generator instanceof Array ? generator : generator ? [generator] : [];
   const tokens = allTokens.map((t) => t.token);
   if (partialCmd) tokens.push(partialCmd);
-  const suggestions = (await Promise.all(generators.map((gen) => runGenerator(gen, tokens, cwd)))).flat();
+  signal?.throwIfAborted();
+  const suggestions = (await Promise.all(generators.map((gen) => runGenerator(gen, tokens, cwd, signal)))).flat();
   return filter<Fig.Suggestion>(
     suggestions.map((suggestion) => ({ ...suggestion, priority: suggestion.priority ?? 60 })),
     filterStrategy,
@@ -219,8 +231,9 @@ const templateSuggestions = async (
   filterStrategy: FilterStrategy | undefined,
   partialCmd: string | undefined,
   cwd: string,
+  signal?: AbortSignal,
 ): Promise<Suggestion[]> => {
-  return filter<Fig.Suggestion>(await runTemplates(templates ?? [], cwd), filterStrategy, partialCmd, undefined);
+  return filter<Fig.Suggestion>(await runTemplates(templates ?? [], cwd, signal), filterStrategy, partialCmd, undefined);
 };
 
 const suggestionSuggestions = (
@@ -310,6 +323,7 @@ export const getSubcommandDrivenRecommendation = async (
   allTokens: CommandToken[],
   cwd: string,
   shell: Shell,
+  signal?: AbortSignal,
 ): Promise<SuggestionBlob | undefined> => {
   log.debug({ msg: "suggestion point", subcommand, persistentOptions, partialToken, argsDepleted, argsFromSubcommand, acceptedTokens, cwd });
   if (argsDepleted && argsFromSubcommand) {
@@ -331,9 +345,9 @@ export const getSubcommandDrivenRecommendation = async (
   }
   if (argLength != 0) {
     const activeArg = subcommand.args instanceof Array ? subcommand.args[0] : subcommand.args;
-    suggestions.push(...(await generatorSuggestions(activeArg?.generators, allTokens, activeArg?.filterStrategy, partialCmd, cwd)));
+    suggestions.push(...(await generatorSuggestions(activeArg?.generators, allTokens, activeArg?.filterStrategy, partialCmd, cwd, signal)));
     suggestions.push(...suggestionSuggestions(activeArg?.suggestions, activeArg?.filterStrategy, partialCmd));
-    suggestions.push(...(await templateSuggestions(activeArg?.template, activeArg?.filterStrategy, partialCmd, cwd)));
+    suggestions.push(...(await templateSuggestions(activeArg?.template, activeArg?.filterStrategy, partialCmd, cwd, signal)));
   }
 
   return {
@@ -369,6 +383,7 @@ export const getArgDrivenRecommendation = async (
   variadicArgBound: boolean,
   cwd: string,
   shell: Shell,
+  signal?: AbortSignal,
 ): Promise<SuggestionBlob | undefined> => {
   let partialCmd = partialToken?.token;
   if (partialToken?.isPath) {
@@ -379,9 +394,9 @@ export const getArgDrivenRecommendation = async (
   const activeArg = args[0];
   const allOptions = persistentOptions.concat(subcommand.options ?? []);
   const suggestions = [
-    ...(await generatorSuggestions(args[0].generators, allTokens, activeArg?.filterStrategy, partialCmd, cwd)),
+    ...(await generatorSuggestions(args[0].generators, allTokens, activeArg?.filterStrategy, partialCmd, cwd, signal)),
     ...suggestionSuggestions(args[0].suggestions, activeArg?.filterStrategy, partialCmd),
-    ...(await templateSuggestions(args[0].template, activeArg?.filterStrategy, partialCmd, cwd)),
+    ...(await templateSuggestions(args[0].template, activeArg?.filterStrategy, partialCmd, cwd, signal)),
   ];
 
   if (activeArg.isOptional || (activeArg.isVariadic && variadicArgBound)) {