feat(amazonq): Get repomap from Flare directly

package.json

@@ -74,7 +74,7 @@
         "webpack-merge": "^5.10.0"
     },
     "dependencies": {
-        "@aws/language-server-runtimes": "^0.2.125",
+        "@aws/language-server-runtimes": "^0.2.128",
         "@types/node": "^22.7.5",
         "jaro-winkler": "^0.2.8",
         "vscode-nls": "^5.2.0",

packages/amazonq/src/app/inline/activation.ts

@@ -22,14 +22,16 @@ import {
     vsCodeState,
 } from 'aws-core-vscode/codewhisperer'
 import { Commands, getLogger, globals, sleep } from 'aws-core-vscode/shared'
+import { LanguageClient } from 'vscode-languageclient'
 
-export async function activate() {
+export async function activate(languageClient: LanguageClient) {
     const codewhispererSettings = CodeWhispererSettings.instance
     const client = new DefaultCodeWhispererClient()
 
     if (isInlineCompletionEnabled()) {
         await setSubscriptionsforInlineCompletion()
         await AuthUtil.instance.setVscodeContextProps()
+        RecommendationHandler.instance.setLanguageClient(languageClient)
     }
 
     function getAutoTriggerStatus(): boolean {

packages/amazonq/src/lsp/client.ts

@@ -342,12 +342,35 @@ async function onLanguageServerReady(
     const enableInlineRollback = true
     if (enableInlineRollback) {
         // use VSC inline
-        await activateInline()
+        await activateInline(client)
     } else {
         // use language server for inline completion
         const inlineManager = new InlineCompletionManager(client, sessionManager, lineTracker, inlineTutorialAnnotation)
         inlineManager.registerInlineCompletion()
-        toDispose.push(inlineManager)
+        toDispose.push(
+            inlineManager,
+            Commands.register('aws.amazonq.showPrev', async () => {
+                await sessionManager.maybeRefreshSessionUx()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
+                sessionManager.onPrevSuggestion()
+            }),
+            Commands.register('aws.amazonq.showNext', async () => {
+                await sessionManager.maybeRefreshSessionUx()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
+                sessionManager.onNextSuggestion()
+            }),
+            // this is a workaround since handleDidShowCompletionItem is not public API
+            Commands.register('aws.amazonq.checkInlineSuggestionVisibility', async () => {
+                sessionManager.checkInlineSuggestionVisibility()
+            }),
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                vsCodeState.lastManualTriggerTime = performance.now()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
+            }),
+            vscode.workspace.onDidCloseTextDocument(async () => {
+                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
+            })
+        )
     }
 
     activateInlineChat(extensionContext, client, encryptionKey, inlineChatTutorialAnnotation)
@@ -364,24 +387,6 @@ async function onLanguageServerReady(
     await initializeLanguageServerConfiguration(client, 'startup')
 
     toDispose.push(
-        Commands.register('aws.amazonq.showPrev', async () => {
-            await sessionManager.maybeRefreshSessionUx()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
-            sessionManager.onPrevSuggestion()
-        }),
-        Commands.register('aws.amazonq.showNext', async () => {
-            await sessionManager.maybeRefreshSessionUx()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
-            sessionManager.onNextSuggestion()
-        }),
-        // this is a workaround since handleDidShowCompletionItem is not public API
-        Commands.register('aws.amazonq.checkInlineSuggestionVisibility', async () => {
-            sessionManager.checkInlineSuggestionVisibility()
-        }),
-        Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-            vsCodeState.lastManualTriggerTime = performance.now()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
-        }),
         Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => {
             telemetry.record({
                 traceId: TelemetryHelper.instance.traceId,
@@ -407,9 +412,6 @@ async function onLanguageServerReady(
                 getLogger().debug(`codewhisperer: user dismiss tutorial.`)
             }
         }),
-        vscode.workspace.onDidCloseTextDocument(async () => {
-            await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
-        }),
         AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
             await auth.refreshConnection()
         }),

packages/core/src/codewhisperer/service/recommendationHandler.ts

@@ -44,6 +44,7 @@ import { indent } from '../../shared/utilities/textUtilities'
 import path from 'path'
 import { isIamConnection } from '../../auth/connection'
 import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker'
+import { LanguageClient } from 'vscode-languageclient'
 
 /**
  * This class is for getRecommendation/listRecommendation API calls and its states
@@ -99,6 +100,7 @@ export class RecommendationHandler {
     private next: vscode.Disposable
     private prev: vscode.Disposable
     private _timer?: NodeJS.Timer
+    private languageClient?: LanguageClient
     documentUri: vscode.Uri | undefined = undefined
 
     constructor() {
@@ -121,6 +123,10 @@ export class RecommendationHandler {
         return session.recommendations.some((r) => r.content.trim() !== '')
     }
 
+    setLanguageClient(languageClient: LanguageClient) {
+        this.languageClient = languageClient
+    }
+
     async getServerResponse(
         triggerType: CodewhispererTriggerType,
         isManualTriggerOn: boolean,
@@ -204,7 +210,8 @@ export class RecommendationHandler {
                 session.requestContext = await EditorContext.buildListRecommendationRequest(
                     editor as vscode.TextEditor,
                     this.nextToken,
-                    config.isSuggestionsWithCodeReferencesEnabled
+                    config.isSuggestionsWithCodeReferencesEnabled,
+                    this.languageClient
                 )
             } else {
                 session.requestContext = {

packages/core/src/codewhisperer/util/editorContext.ts

@@ -22,6 +22,7 @@ import { indent } from '../../shared/utilities/textUtilities'
 import { isInDirectory } from '../../shared/filesystemUtilities'
 import { AuthUtil } from './authUtil'
 import { predictionTracker } from '../nextEditPrediction/activation'
+import { LanguageClient } from 'vscode-languageclient'
 
 let tabSize: number = getTabSizeSetting()
 
@@ -224,7 +225,8 @@ async function getWorkspaceId(editor: vscode.TextEditor): Promise<string | undef
 export async function buildListRecommendationRequest(
     editor: vscode.TextEditor,
     nextToken: string,
-    allowCodeWithReference: boolean
+    allowCodeWithReference: boolean,
+    languageClient?: LanguageClient
 ): Promise<{
     request: codewhispererClient.ListRecommendationsRequest
     supplementalMetadata: CodeWhispererSupplementalContext | undefined
@@ -236,7 +238,7 @@ export async function buildListRecommendationRequest(
         tokenSource.cancel()
     }, supplementalContextTimeoutInMs)
 
-    const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token)
+    const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token, languageClient)
 
     logSupplementalContext(supplementalContexts)
 

packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -23,7 +23,13 @@ import {
 import { waitUntil } from '../../../shared/utilities/timeoutUtils'
 import { FeatureConfigProvider } from '../../../shared/featureConfig'
 import fs from '../../../shared/fs/fs'
+import { LanguageClient } from 'vscode-languageclient'
 
+import {
+    GetSupplementalContextParams,
+    getSupplementalContextRequestType,
+    SupplementalContextItem,
+} from '@aws/language-server-runtimes/protocol'
 type CrossFileSupportedLanguage =
     | 'java'
     | 'python'
@@ -66,7 +72,8 @@ type SupplementalContextConfig = 'none' | 'opentabs' | 'codemap' | 'bm25' | 'def
 
 export async function fetchSupplementalContextForSrc(
     editor: vscode.TextEditor,
-    cancellationToken: vscode.CancellationToken
+    cancellationToken: vscode.CancellationToken,
+    languageClient?: LanguageClient
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     const supplementalContextConfig = getSupplementalContextConfig(editor.document.languageId)
 
@@ -101,7 +108,7 @@ export async function fetchSupplementalContextForSrc(
             async function () {
                 const result: CodeWhispererSupplementalContextItem[] = []
                 const opentabsContext = await fetchOpentabsContext(editor, cancellationToken)
-                const codemap = await fetchProjectContext(editor, 'codemap')
+                const codemap = await fetchProjectContext(editor, 'codemap', languageClient)
 
                 function addToResult(items: CodeWhispererSupplementalContextItem[]) {
                     for (const item of items) {
@@ -145,7 +152,7 @@ export async function fetchSupplementalContextForSrc(
     if (supplementalContextConfig === 'bm25') {
         const projectBM25Promise = waitUntil(
             async function () {
-                return await fetchProjectContext(editor, 'bm25')
+                return await fetchProjectContext(editor, 'bm25', languageClient)
             },
             { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
         )
@@ -168,7 +175,7 @@ export async function fetchSupplementalContextForSrc(
     // global bm25 with repomap
     const projectContextAndCodemapPromise = waitUntil(
         async function () {
-            return await fetchProjectContext(editor, 'default')
+            return await fetchProjectContext(editor, 'default', languageClient)
         },
         { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
     )
@@ -192,13 +199,20 @@ export async function fetchSupplementalContextForSrc(
 
 export async function fetchProjectContext(
     editor: vscode.TextEditor,
-    target: 'default' | 'codemap' | 'bm25'
+    target: 'default' | 'codemap' | 'bm25',
+    languageclient?: LanguageClient
 ): Promise<CodeWhispererSupplementalContextItem[]> {
-    // const inputChunkContent = getInputChunk(editor)
-    // TODO:
-    const inlineProjectContext: { content: string; score: number; filePath: string }[] = []
-
-    return inlineProjectContext
+    if (languageclient) {
+        const request: GetSupplementalContextParams = {
+            filePath: editor.document.uri.fsPath,
+        }
+        const response = await languageclient.sendRequest<SupplementalContextItem[]>(
+            getSupplementalContextRequestType.method,
+            request
+        )
+        return response as CodeWhispererSupplementalContextItem[]
+    }
+    return []
 }
 
 export async function fetchOpentabsContext(

packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts

@@ -13,10 +13,12 @@ import { getLogger } from '../../../shared/logger/logger'
 import { CodeWhispererSupplementalContext } from '../../models/model'
 import * as os from 'os'
 import { crossFileContextConfig } from '../../models/constants'
+import { LanguageClient } from 'vscode-languageclient'
 
 export async function fetchSupplementalContext(
     editor: vscode.TextEditor,
-    cancellationToken: vscode.CancellationToken
+    cancellationToken: vscode.CancellationToken,
+    languageClient?: LanguageClient
 ): Promise<CodeWhispererSupplementalContext | undefined> {
     const timesBeforeFetching = performance.now()
 
@@ -32,7 +34,7 @@ export async function fetchSupplementalContext(
     if (isUtg) {
         supplementalContextPromise = fetchSupplementalContextForTest(editor, cancellationToken)
     } else {
-        supplementalContextPromise = fetchSupplementalContextForSrc(editor, cancellationToken)
+        supplementalContextPromise = fetchSupplementalContextForSrc(editor, cancellationToken, languageClient)
     }
 
     return supplementalContextPromise