diff --git a/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts index 91e26e36111..e197cbeedca 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts @@ -30,11 +30,6 @@ import { LspController } from 'aws-core-vscode/amazonq' let tempFolder: string describe('crossFileContextUtil', function () { - const fakeCancellationToken: vscode.CancellationToken = { - isCancellationRequested: false, - onCancellationRequested: sinon.spy(), - } - let mockEditor: vscode.TextEditor let clock: FakeTimers.InstalledClock @@ -68,7 +63,7 @@ describe('crossFileContextUtil', function () { await assertTabCount(2) - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor) assert.ok(actual) assert.strictEqual(actual.supplementalContextItems.length, 3) assert.strictEqual(actual.supplementalContextItems[0].content.split('\n').length, 50) @@ -96,7 +91,7 @@ describe('crossFileContextUtil', function () { }, ]) - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor) assert.ok(actual) assert.strictEqual(actual.supplementalContextItems.length, 3) assert.strictEqual(actual?.strategy, 'codemap') @@ -149,7 +144,7 @@ describe('crossFileContextUtil', function () { }, ]) - const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor) assert.ok(actual) assert.strictEqual(actual.supplementalContextItems.length, 5) assert.strictEqual(actual?.strategy, 'bm25') @@ -188,14 +183,14 @@ describe('crossFileContextUtil', function () { describe('non supported language should return undefined', function () { it('c++', async function () { mockEditor = createMockTextEditor('content', 'fileName', 'cpp') - const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor) assert.strictEqual(actual, undefined) }) it('ruby', async function () { mockEditor = createMockTextEditor('content', 'fileName', 'ruby') - const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(mockEditor) assert.strictEqual(actual, undefined) }) @@ -286,7 +281,7 @@ describe('crossFileContextUtil', function () { await assertTabCount(4) - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(editor) assert.ok(actual && actual.supplementalContextItems.length === 0) }) @@ -317,7 +312,7 @@ describe('crossFileContextUtil', function () { await assertTabCount(4) - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(editor) assert.ok(actual && actual.supplementalContextItems.length !== 0) }) @@ -360,7 +355,7 @@ describe('crossFileContextUtil', function () { await assertTabCount(4) - const actual = await crossFile.fetchSupplementalContextForSrc(editor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContextForSrc(editor) assert.ok(actual && actual.supplementalContextItems.length !== 0) }) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index d5085e4db0c..b151da71332 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -160,7 +160,7 @@ describe('editorContext', function () { const optOutPreference = false await globals.telemetry.setTelemetryEnabled(false) const editor = createMockTextEditor('import math\ndef two_sum(nums, target):\n', 'test.py', 'python', 1, 17) - const actual = await EditorContext.buildListRecommendationRequest(editor, nextToken, optOutPreference) + const actual = await EditorContext.buildGenerateCompletionRequest(editor, nextToken, optOutPreference) assert.strictEqual(actual.request.nextToken, nextToken) assert.strictEqual((actual.request as GenerateCompletionsRequest).optOutPreference, 'OPTOUT') diff --git a/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts index 051ac65bee1..e771b7a8aef 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts @@ -5,7 +5,6 @@ import assert from 'assert' import * as FakeTimers from '@sinonjs/fake-timers' -import * as vscode from 'vscode' import * as sinon from 'sinon' import * as crossFile from 'aws-core-vscode/codewhisperer' import { TestFolder, assertTabCount, installFakeClock } from 'aws-core-vscode/test' @@ -17,11 +16,6 @@ describe('supplementalContextUtil', function () { let testFolder: TestFolder let clock: FakeTimers.InstalledClock - const fakeCancellationToken: vscode.CancellationToken = { - isCancellationRequested: false, - onCancellationRequested: sinon.spy(), - } - before(function () { clock = installFakeClock() }) @@ -62,7 +56,7 @@ describe('supplementalContextUtil', function () { await assertTabCount(4) - const actual = await crossFile.fetchSupplementalContext(editor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContext(editor) assert.ok(actual?.supplementalContextItems.length === 4) }) @@ -78,7 +72,7 @@ describe('supplementalContextUtil', function () { await assertTabCount(4) - const actual = await crossFile.fetchSupplementalContext(editor, fakeCancellationToken) + const actual = await crossFile.fetchSupplementalContext(editor) assert.ok(actual?.supplementalContextItems.length === 0) }) }) diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/packages/core/src/codewhisperer/service/recommendationHandler.ts index 1f5096ad1cc..6fd9fe773aa 100644 --- a/packages/core/src/codewhisperer/service/recommendationHandler.ts +++ b/packages/core/src/codewhisperer/service/recommendationHandler.ts @@ -156,9 +156,9 @@ export class RecommendationHandler { config: ConfigurationEntry, autoTriggerType?: CodewhispererAutomatedTriggerType, pagination: boolean = true, - page: number = 0, - generate: boolean = isIamConnection(AuthUtil.instance.conn) + page: number = 0 ): Promise { + const isIamConn: boolean = isIamConnection(AuthUtil.instance.conn) let invocationResult: 'Succeeded' | 'Failed' = 'Failed' let errorMessage: string | undefined = undefined let errorCode: string | undefined = undefined @@ -184,9 +184,9 @@ export class RecommendationHandler { ).language session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName) - if (pagination && !generate) { + if (pagination && !isIamConn) { if (page === 0) { - session.requestContext = await EditorContext.buildListRecommendationRequest( + session.requestContext = await EditorContext.buildGenerateCompletionRequest( editor as vscode.TextEditor, this.nextToken, config.isSuggestionsWithCodeReferencesEnabled @@ -238,7 +238,7 @@ export class RecommendationHandler { this.lastInvocationTime = startTime const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request) const codewhispererPromise = - pagination && !generate + pagination && !isIamConn ? client.listRecommendations(mappedReq) : client.generateRecommendations(mappedReq) const resp = await this.getServerResponse(triggerType, config.isManualTriggerEnabled, codewhispererPromise) @@ -333,8 +333,7 @@ export class RecommendationHandler { config, autoTriggerType, pagination, - page, - true + page ) } } diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index 99a15fd1f02..cba7b92d56d 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -11,7 +11,6 @@ import { getTabSizeSetting } from '../../shared/utilities/editorUtilities' import { getLogger } from '../../shared/logger/logger' import { runtimeLanguageContext } from './runtimeLanguageContext' import { fetchSupplementalContext } from './supplementalContext/supplementalContextUtil' -import { supplementalContextTimeoutInMs } from '../models/constants' import { getSelectedCustomization } from './customizationUtil' import { selectFrom } from '../../shared/utilities/tsUtils' import { checkLeftContextKeywordsForJson } from './commonUtil' @@ -82,7 +81,7 @@ export function getFileRelativePath(editor: vscode.TextEditor): string { return relativePath.substring(0, CodeWhispererConstants.filenameCharsLimit) } -export async function buildListRecommendationRequest( +export async function buildGenerateCompletionRequest( editor: vscode.TextEditor, nextToken: string, allowCodeWithReference: boolean @@ -91,13 +90,7 @@ export async function buildListRecommendationRequest( supplementalMetadata: CodeWhispererSupplementalContext | undefined }> { const fileContext = extractContextForCodeWhisperer(editor) - - const tokenSource = new vscode.CancellationTokenSource() - setTimeout(() => { - tokenSource.cancel() - }, supplementalContextTimeoutInMs) - - const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token) + const supplementalContexts = await fetchSupplementalContext(editor) logSupplementalContext(supplementalContexts) @@ -128,14 +121,7 @@ export async function buildGenerateRecommendationRequest(editor: vscode.TextEdit supplementalMetadata: CodeWhispererSupplementalContext | undefined }> { const fileContext = extractContextForCodeWhisperer(editor) - - const tokenSource = new vscode.CancellationTokenSource() - // the supplement context fetch mechanisms each has a timeout of supplementalContextTimeoutInMs - // adding 10 ms for overall timeout as buffer - setTimeout(() => { - tokenSource.cancel() - }, supplementalContextTimeoutInMs + 10) - const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token) + const supplementalContexts = await fetchSupplementalContext(editor) logSupplementalContext(supplementalContexts) diff --git a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts index f4688e2b5a9..2066ef0e16a 100644 --- a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts +++ b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts @@ -65,8 +65,7 @@ interface Chunk { type SupplementalContextConfig = 'none' | 'opentabs' | 'codemap' | 'bm25' | 'default' export async function fetchSupplementalContextForSrc( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken + editor: vscode.TextEditor ): Promise | undefined> { const supplementalContextConfig = getSupplementalContextConfig(editor.document.languageId) @@ -78,7 +77,7 @@ export async function fetchSupplementalContextForSrc( // fallback to opentabs if projectContext timeout const opentabsContextPromise = waitUntil( async function () { - return await fetchOpentabsContext(editor, cancellationToken) + return await fetchOpentabsContext(editor) }, { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } ) @@ -100,7 +99,7 @@ export async function fetchSupplementalContextForSrc( const opentabsContextAndCodemap = await waitUntil( async function () { const result: CodeWhispererSupplementalContextItem[] = [] - const opentabsContext = await fetchOpentabsContext(editor, cancellationToken) + const opentabsContext = await fetchOpentabsContext(editor) const codemap = await fetchProjectContext(editor, 'codemap') function addToResult(items: CodeWhispererSupplementalContextItem[]) { @@ -207,8 +206,7 @@ export async function fetchProjectContext( } export async function fetchOpentabsContext( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken + editor: vscode.TextEditor ): Promise { const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch diff --git a/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts index 03f9d59b3f2..88f7a83c838 100644 --- a/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts +++ b/packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts @@ -13,8 +13,7 @@ import { getLogger } from '../../../shared/logger/logger' import { CodeWhispererSupplementalContext } from '../../models/model' export async function fetchSupplementalContext( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken + editor: vscode.TextEditor ): Promise { const timesBeforeFetching = performance.now() @@ -28,9 +27,9 @@ export async function fetchSupplementalContext( > if (isUtg) { - supplementalContextPromise = fetchSupplementalContextForTest(editor, cancellationToken) + supplementalContextPromise = fetchSupplementalContextForTest(editor) } else { - supplementalContextPromise = fetchSupplementalContextForSrc(editor, cancellationToken) + supplementalContextPromise = fetchSupplementalContextForSrc(editor) } return supplementalContextPromise diff --git a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts b/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts index a39a48183b0..8e5a3a6cadb 100644 --- a/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts +++ b/packages/core/src/codewhisperer/util/supplementalContext/utgUtils.ts @@ -4,7 +4,7 @@ */ import * as path from 'path' -import { fs } from '../../../shared' +import { fs, waitUntil } from '../../../shared' import * as vscode from 'vscode' import { countSubstringMatches, @@ -14,10 +14,7 @@ import { utgLanguageConfig, utgLanguageConfigs, } from './codeParsingUtil' -import { ToolkitError } from '../../../shared/errors' -import { supplemetalContextFetchingTimeoutMsg } from '../../models/constants' -import { CancellationError } from '../../../shared/utilities/timeoutUtils' -import { utgConfig } from '../../models/constants' +import { supplementalContextTimeoutInMs, utgConfig } from '../../models/constants' import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities' import { getLogger } from '../../../shared/logger/logger' import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem, UtgStrategy } from '../../models/model' @@ -48,8 +45,7 @@ export function shouldFetchUtgContext(languageId: vscode.TextDocument['languageI * @returns */ export async function fetchSupplementalContextForTest( - editor: vscode.TextEditor, - cancellationToken: vscode.CancellationToken + editor: vscode.TextEditor ): Promise | undefined> { const shouldProceed = shouldFetchUtgContext(editor.document.languageId) @@ -59,50 +55,53 @@ export async function fetchSupplementalContextForTest( const languageConfig = utgLanguageConfigs[editor.document.languageId] - // TODO (Metrics): 1. Total number of calls to fetchSupplementalContextForTest - throwIfCancelled(cancellationToken) - - let crossSourceFile = await findSourceFileByName(editor, languageConfig, cancellationToken) - if (crossSourceFile) { - // TODO (Metrics): 2. Success count for fetchSourceFileByName (find source file by name) - getLogger().debug(`CodeWhisperer finished fetching utg context by file name`) - return { - supplementalContextItems: await generateSupplementalContextFromFocalFile( - crossSourceFile, - 'byName', - cancellationToken - ), - strategy: 'byName', - } - } - throwIfCancelled(cancellationToken) - - crossSourceFile = await findSourceFileByContent(editor, languageConfig, cancellationToken) - if (crossSourceFile) { - // TODO (Metrics): 3. Success count for fetchSourceFileByContent (find source file by content) - getLogger().debug(`CodeWhisperer finished fetching utg context by file content`) - return { - supplementalContextItems: await generateSupplementalContextFromFocalFile( - crossSourceFile, - 'byContent', - cancellationToken - ), - strategy: 'byContent', - } - } + const utgContext: Pick | undefined = + await waitUntil( + async function () { + let crossSourceFile = await findSourceFileByName(editor) + if (crossSourceFile) { + getLogger().debug(`CodeWhisperer finished fetching utg context by file name`) + return { + supplementalContextItems: await generateSupplementalContextFromFocalFile( + crossSourceFile, + 'byName' + ), + strategy: 'byName', + } + } + + crossSourceFile = await findSourceFileByContent(editor, languageConfig) + if (crossSourceFile) { + getLogger().debug(`CodeWhisperer finished fetching utg context by file content`) + return { + supplementalContextItems: await generateSupplementalContextFromFocalFile( + crossSourceFile, + 'byContent' + ), + strategy: 'byContent', + } + } + + return undefined + }, + { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false } + ) - // TODO (Metrics): 4. Failure count - when unable to find focal file (supplemental context empty) - getLogger().debug(`CodeWhisperer failed to fetch utg context`) - return { - supplementalContextItems: [], - strategy: 'empty', + if (!utgContext) { + getLogger().debug(`CodeWhisperer failed to fetch utg context`) } + + return ( + utgContext ?? { + supplementalContextItems: [], + strategy: 'empty', + } + ) } async function generateSupplementalContextFromFocalFile( filePath: string, - strategy: UtgStrategy, - cancellationToken: vscode.CancellationToken + strategy: UtgStrategy ): Promise { const fileContent = await fs.readFileText(vscode.Uri.parse(filePath!).fsPath) @@ -121,34 +120,23 @@ async function generateSupplementalContextFromFocalFile( async function findSourceFileByContent( editor: vscode.TextEditor, - languageConfig: utgLanguageConfig, - cancellationToken: vscode.CancellationToken + languageConfig: utgLanguageConfig ): Promise { const testFileContent = await fs.readFileText(editor.document.fileName) const testElementList = extractFunctions(testFileContent, languageConfig.functionExtractionPattern) - throwIfCancelled(cancellationToken) - testElementList.push(...extractClasses(testFileContent, languageConfig.classExtractionPattern)) - throwIfCancelled(cancellationToken) - let sourceFilePath: string | undefined = undefined let maxMatchCount = 0 if (testElementList.length === 0) { - // TODO: Add metrics here, as unable to parse test file using Regex. return sourceFilePath } const relevantFilePaths = await getRelevantUtgFiles(editor) - throwIfCancelled(cancellationToken) - - // TODO (Metrics):Add metrics for relevantFilePaths length for (const filePath of relevantFilePaths) { - throwIfCancelled(cancellationToken) - const fileContent = await fs.readFileText(filePath) const elementList = extractFunctions(fileContent, languageConfig.functionExtractionPattern) elementList.push(...extractClasses(fileContent, languageConfig.classExtractionPattern)) @@ -201,11 +189,7 @@ export function guessSrcFileName( return undefined } -async function findSourceFileByName( - editor: vscode.TextEditor, - languageConfig: utgLanguageConfig, - cancellationToken: vscode.CancellationToken -): Promise { +async function findSourceFileByName(editor: vscode.TextEditor): Promise { const testFileName = path.basename(editor.document.fileName) const assumedSrcFileName = guessSrcFileName(testFileName, editor.document.languageId) if (!assumedSrcFileName) { @@ -214,16 +198,8 @@ async function findSourceFileByName( const sourceFiles = await vscode.workspace.findFiles(`**/${assumedSrcFileName}`) - throwIfCancelled(cancellationToken) - if (sourceFiles.length > 0) { return sourceFiles[0].toString() } return undefined } - -function throwIfCancelled(token: vscode.CancellationToken): void | never { - if (token.isCancellationRequested) { - throw new ToolkitError(supplemetalContextFetchingTimeoutMsg, { cause: new CancellationError('timeout') }) - } -}