diff --git a/src/runtime/generator.ts b/src/runtime/generator.ts index 0f7598e..7d4041e 100644 --- a/src/runtime/generator.ts +++ b/src/runtime/generator.ts @@ -5,18 +5,24 @@ import { runTemplates } from "./template.js"; import { buildExecuteShellCommand } from "./utils.js"; import { endTiming, startTiming } from "../utils/performance.js"; -const getGeneratorContext = (cwd: string): Fig.GeneratorContext => { +const getGeneratorContext = (cwd: string, searchTerm: string): Fig.GeneratorContext => { return { environmentVariables: Object.fromEntries(Object.entries(process.env).filter((entry): entry is [string, string] => entry[1] != null)), currentWorkingDirectory: cwd, currentProcess: "", // TODO: define current process sshPrefix: "", // deprecated, should be empty isDangerous: false, - searchTerm: "", // TODO: define search term + searchTerm, }; }; -export const executeGenerator = async (generator: Fig.Generator, tokens: string[], cwd: string, signal?: AbortSignal): Promise => { +export const executeGenerator = async ( + generator: Fig.Generator, + tokens: string[], + searchTerm: string, + cwd: string, + signal?: AbortSignal, +): Promise => { const generatorTiming = startTiming(); try { signal?.throwIfAborted(); @@ -39,7 +45,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, executeShellCommand, getGeneratorContext(cwd, searchTerm)))); } if (template != null) { diff --git a/src/runtime/generatorCache.ts b/src/runtime/generatorCache.ts index ad8fa8e..0f56139 100644 --- a/src/runtime/generatorCache.ts +++ b/src/runtime/generatorCache.ts @@ -103,9 +103,15 @@ const waitForDebounce = async (signal?: AbortSignal): Promise => { }); }; -const executeGeneratorSafely = async (generator: Fig.Generator, tokens: string[], cwd: string, signal?: AbortSignal): Promise => { +const executeGeneratorSafely = async ( + generator: Fig.Generator, + tokens: string[], + activeToken: string, + cwd: string, + signal?: AbortSignal, +): Promise => { try { - return { suggestions: await executeGenerator(generator, tokens, cwd, signal), cacheable: true }; + return { suggestions: await executeGenerator(generator, tokens, activeToken, cwd, signal), cacheable: true }; } catch (error) { if (signal?.aborted) signal.throwIfAborted(); const err = typeof error === "string" ? error : error instanceof Error ? error.message : error; @@ -129,7 +135,7 @@ const generate = async ( const generationVersion = ++state.generationVersion; const generation = (async (): Promise => { if (debounce) await waitForDebounce(signal); - const result = await executeGeneratorSafely(generator, tokens, cwd, signal); + const result = await executeGeneratorSafely(generator, tokens, activeToken, cwd, signal); return { ...result, current: state.generationVersion === generationVersion }; })(); state.inFlight = { key: generationKey, promise: generation, signal }; diff --git a/src/tests/runtime/generator.test.ts b/src/tests/runtime/generator.test.ts index e240fbc..c2e48f8 100644 --- a/src/tests/runtime/generator.test.ts +++ b/src/tests/runtime/generator.test.ts @@ -210,3 +210,16 @@ test("reuses a generator when a path query moves past a delimiter", async () => expect(calls).toBe(1); }); + +test("passes searchTerm correctly in generator context", async () => { + let passedSearchTerm = ""; + const generator: Fig.Generator = { + custom: async (_tokens, _executeShellCommand, context) => { + passedSearchTerm = context.searchTerm; + return []; + }, + }; + + await getGeneratorSuggestions(generator, ["tool", "mysearch"], "mysearch", process.cwd(), false); + expect(passedSearchTerm).toBe("mysearch"); +});