refactor(inline): recommendationHandler

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)
             })

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')

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)
             })
         })

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<GetRecommendationsResponse> {
+        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
                     )
                 }
             }

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)
 

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<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | 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<CodeWhispererSupplementalContextItem[] | undefined> {
     const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
 

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<CodeWhispererSupplementalContext | undefined> {
     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