fix: run custom generator commands in the generator cwd and clean their output

src/runtime/generator.ts

@@ -16,6 +16,18 @@ const getGeneratorContext = (cwd: string): Fig.GeneratorContext => {
   };
 };
 
+// Matches Fig's cleanOutput: CRLF becomes LF, the cursor-show sequence and leading/trailing newlines are removed.
+const cleanOutput = (output: string): string => output.replaceAll("\r\n", "\n").replaceAll("\u001b[?25h", "").replace(/^\n+/, "").replace(/\n+$/, "");
+
+// Custom generators call executeShellCommand themselves. Like Fig, run their commands in the generator's
+// working directory unless they pass one, and give them cleaned output (e.g. `npm prefix` without a trailing newline).
+const buildCustomExecuteCommand =
+  (executeShellCommand: Fig.ExecuteCommandFunction, cwd: string): Fig.ExecuteCommandFunction =>
+  async (input) => {
+    const output = await executeShellCommand({ ...input, cwd: input.cwd ?? cwd });
+    return { ...output, stdout: cleanOutput(output.stdout), stderr: cleanOutput(output.stderr) };
+  };
+
 export const executeGenerator = async (generator: Fig.Generator, tokens: string[], cwd: string, signal?: AbortSignal): Promise<Fig.Suggestion[]> => {
   const generatorTiming = startTiming();
   try {
@@ -39,7 +51,7 @@ export const executeGenerator = async (generator: Fig.Generator, tokens: string[
     }
 
     if (custom) {
-      suggestions.push(...(await custom(tokens, executeShellCommand, getGeneratorContext(cwd))));
+      suggestions.push(...(await custom(tokens, buildCustomExecuteCommand(executeShellCommand, cwd), getGeneratorContext(cwd))));
     }
 
     if (template != null) {

src/tests/runtime/generator.test.ts

@@ -4,7 +4,10 @@
 import { jest } from "@jest/globals";
 
 import { clearGeneratorState, clearImplicitGeneratorState, getGeneratorSuggestions } from "../../runtime/generatorCache.js";
-import { getGeneratorQueryTerm } from "../../runtime/generator.js";
+import { executeGenerator, getGeneratorQueryTerm } from "../../runtime/generator.js";
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
 
 const deferred = <T>() => {
   let resolve!: (value: T) => void;
@@ -210,3 +213,38 @@ test("reuses a generator when a path query moves past a delimiter", async () =>
 
   expect(calls).toBe(1);
 });
+
+describe("custom generator executeShellCommand", () => {
+  const nodeGenerator = (script: string, cwd?: string): Fig.Generator => ({
+    custom: async (_tokens, executeShellCommand) => {
+      const { stdout } = await executeShellCommand({ command: process.execPath, args: ["-e", script], cwd });
+      return [{ name: stdout }];
+    },
+  });
+
+  test("runs in the generator's cwd when the generator does not pass one", async () => {
+    const dir = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "generator-cwd-")));
+    try {
+      const [suggestion] = await executeGenerator(nodeGenerator("process.stdout.write(process.cwd())"), [], dir);
+      expect(suggestion.name).toBe(dir);
+    } finally {
+      fs.rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  test("keeps a cwd passed by the generator", async () => {
+    const dir = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "generator-cwd-")));
+    try {
+      const [suggestion] = await executeGenerator(nodeGenerator("process.stdout.write(process.cwd())", dir), [], os.tmpdir());
+      expect(suggestion.name).toBe(dir);
+    } finally {
+      fs.rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  test("cleans output like Fig: CRLF to LF, no leading or trailing newlines", async () => {
+    const script = String.raw`process.stdout.write("\r\n\n  a\r\nb  \r\n\n")`;
+    const [suggestion] = await executeGenerator(nodeGenerator(script), [], process.cwd());
+    expect(suggestion.name).toBe("  a\nb  ");
+  });
+});