diff --git a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts index 2a9423e6159..d2d7709f265 100644 --- a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts +++ b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts @@ -73,94 +73,93 @@ export async function fetchSupplementalContextForSrc( return undefined } - // opentabs context will use bm25 and users' open tabs to fetch supplemental context - if (supplementalContextConfig === 'opentabs') { - return { - supplementalContextItems: (await fetchOpentabsContext(editor, cancellationToken)) ?? [], - strategy: 'opentabs', - } - } + const opentabsContextPromise = waitUntil( + async function () { + return await fetchOpentabsContext(editor, cancellationToken) + }, + { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } + ) - // codemap will use opentabs context plus repomap if it's present - if (supplementalContextConfig === 'codemap') { - const opentabsContextAndCodemap = await waitUntil( - async function () { - const result: CodeWhispererSupplementalContextItem[] = [] - const opentabsContext = await fetchOpentabsContext(editor, cancellationToken) - const codemap = await fetchProjectContext(editor, 'codemap') + let supContextPromise: Promise | undefined = new Promise( + (resolve) => resolve(undefined) + ) - if (codemap && codemap.length > 0) { - result.push(...codemap) - } + let target: 'default' | 'codemap' | 'bm25' | undefined = undefined + switch (supplementalContextConfig) { + case 'codemap': + target = 'codemap' + break + case 'bm25': + target = 'bm25' + break + case 'default': + target = 'default' + break + default: + break + } - if (opentabsContext && opentabsContext.length > 0) { - result.push(...opentabsContext) + if (target) { + supContextPromise = waitUntil( + async function () { + if (target) { + return await fetchProjectContext(editor, target) } - return result + return undefined }, { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } ) + } + + const [nullableOpentabsContext, nullableSupContext] = await Promise.all([opentabsContextPromise, supContextPromise]) + const opentabsContext = nullableOpentabsContext ?? [] + const supContext = nullableSupContext ?? [] + // opentabs context will use bm25 and users' open tabs to fetch supplemental context + if (supplementalContextConfig === 'opentabs') { return { - supplementalContextItems: opentabsContextAndCodemap ?? [], - strategy: 'codemap', + supplementalContextItems: opentabsContext, + strategy: opentabsContext.length === 0 ? 'Empty' : 'opentabs', } } - // fallback to opentabs if projectContext timeout for 'default' | 'bm25' - const opentabsContextPromise = waitUntil( - async function () { - return await fetchOpentabsContext(editor, cancellationToken) - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) + // repomap + opentabs + if (supplementalContextConfig === 'codemap') { + if (supContext.length > 0) { + // TODO: truncation + const res = [...supContext, ...opentabsContext] + return { + supplementalContextItems: res, + strategy: 'codemap', + } + } + } // global bm25 without repomap if (supplementalContextConfig === 'bm25') { - const projectBM25Promise = waitUntil( - async function () { - return await fetchProjectContext(editor, 'bm25') - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - const [projectContext, opentabsContext] = await Promise.all([projectBM25Promise, opentabsContextPromise]) - if (projectContext && projectContext.length > 0) { + if (supContext.length > 0) { return { - supplementalContextItems: projectContext, + supplementalContextItems: supContext, strategy: 'bm25', } } - - return { - supplementalContextItems: opentabsContext ?? [], - strategy: 'opentabs', - } } - // global bm25 with repomap - const projectContextAndCodemapPromise = waitUntil( - async function () { - return await fetchProjectContext(editor, 'default') - }, - { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } - ) - - const [projectContext, opentabsContext] = await Promise.all([ - projectContextAndCodemapPromise, - opentabsContextPromise, - ]) - if (projectContext && projectContext.length > 0) { - return { - supplementalContextItems: projectContext, - strategy: 'default', + // repomap + global bm25 + if (supplementalContextConfig === 'default') { + if (supContext.length > 0) { + return { + supplementalContextItems: supContext, + strategy: 'default', + } } } + // fallback to use opentabs if other srategy return empty or timeout return { - supplementalContextItems: opentabsContext ?? [], - strategy: 'opentabs', + supplementalContextItems: opentabsContext, + strategy: opentabsContext.length === 0 ? 'Empty' : 'opentabs', } }