feat(amazonq): Add project context to inline completion (#5729)

## Problem

1. Inline Completion of Q is not context aware.
2. The `@workspace` index does not update on file add/remove.



## Solution

1. We are adding new global BM25 index and class/function digest from
imported files to provide project context for inline completion. This
feature will be added as AB testing. In this PR, there will be no
customer facing changes since it is hidden behind a AB testing flag. We
will toggle it a few days later. For this, we added new LSP APIs.

2. One customer facing change, when user add or remove a file in the VS
Code IDE, its `@workspace` index will be updated

The net-browserify is added to prevent npm run package failure for
browser.

---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: Lei Gao <[email protected]>
Co-authored-by: Lei Gao <[email protected]>
Co-authored-by: Maxim Hayes <[email protected]>

Author: 3 people

packages/amazonq/.changes/next-release/Bug Fix-59d88db7-cfd5-41da-9a63-2a6391b27add.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Update `@workspace` index when adding or deleting a file"
+}

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

@@ -25,7 +25,11 @@ export async function activate(context: ExtensionContext) {
     await amazonq.TryChatCodeLensProvider.register(appInitContext.onDidChangeAmazonQVisibility.event)
 
     const setupLsp = funcUtil.debounce(async () => {
-        void amazonq.LspController.instance.trySetupLsp(context)
+        void amazonq.LspController.instance.trySetupLsp(context, {
+            startUrl: AuthUtil.instance.startUrl,
+            maxIndexSize: CodeWhispererSettings.instance.getMaxIndexSize(),
+            isVectorIndexEnabled: CodeWhispererSettings.instance.isLocalIndexEnabled(),
+        })
     }, 5000)
 
     context.subscriptions.push(

packages/amazonq/test/unit/amazonq/lsp/lspClient.test.ts

@@ -18,7 +18,7 @@ describe('Amazon Q LSP client', function () {
     })
 
     it('encrypts payload of query ', async () => {
-        await lspClient.query('mock_input')
+        await lspClient.queryVectorIndex('mock_input')
         assert.ok(encryptFunc.calledOnce)
         assert.ok(encryptFunc.calledWith(JSON.stringify({ query: 'mock_input' })))
         const value = await encryptFunc.returnValues[0]
@@ -27,14 +27,15 @@ describe('Amazon Q LSP client', function () {
     })
 
     it('encrypts payload of index files ', async () => {
-        await lspClient.indexFiles(['fileA'], 'path', false)
+        await lspClient.buildIndex(['fileA'], 'path', 'all')
         assert.ok(encryptFunc.calledOnce)
         assert.ok(
             encryptFunc.calledWith(
                 JSON.stringify({
                     filePaths: ['fileA'],
-                    rootPath: 'path',
-                    refresh: false,
+                    projectRoot: 'path',
+                    config: 'all',
+                    language: '',
                 })
             )
         )

packages/core/src/amazonq/lsp/lspClient.ts

@@ -17,7 +17,17 @@ import * as jose from 'jose'
 import { Disposable, ExtensionContext } from 'vscode'
 
 import { LanguageClient, LanguageClientOptions, ServerOptions, TransportKind } from 'vscode-languageclient'
-import { GetUsageRequestType, IndexRequestType, QueryRequestType, UpdateIndexRequestType, Usage } from './types'
+import {
+    BuildIndexRequestPayload,
+    BuildIndexRequestType,
+    GetUsageRequestType,
+    IndexConfig,
+    QueryInlineProjectContextRequestType,
+    QueryVectorIndexRequestType,
+    UpdateIndexV2RequestPayload,
+    UpdateIndexV2RequestType,
+    Usage,
+} from './types'
 import { Writable } from 'stream'
 import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { fs, getLogger } from '../../shared'
@@ -61,52 +71,67 @@ export class LspClient {
             .encrypt(key)
     }
 
-    async indexFiles(request: string[], rootPath: string, refresh: boolean) {
+    async buildIndex(paths: string[], rootPath: string, config: IndexConfig) {
+        const payload: BuildIndexRequestPayload = {
+            filePaths: paths,
+            projectRoot: rootPath,
+            config: config,
+            language: '',
+        }
         try {
-            const encryptedRequest = await this.encrypt(
-                JSON.stringify({
-                    filePaths: request,
-                    rootPath: rootPath,
-                    refresh: refresh,
-                })
-            )
-            const resp = await this.client?.sendRequest(IndexRequestType, encryptedRequest)
+            const encryptedRequest = await this.encrypt(JSON.stringify(payload))
+            const resp = await this.client?.sendRequest(BuildIndexRequestType, encryptedRequest)
             return resp
         } catch (e) {
-            getLogger().error(`LspClient: indexFiles error: ${e}`)
+            getLogger().error(`LspClient: buildIndex error: ${e}`)
             return undefined
         }
     }
 
-    async query(request: string) {
+    async queryVectorIndex(request: string) {
         try {
             const encryptedRequest = await this.encrypt(
                 JSON.stringify({
                     query: request,
                 })
             )
-            const resp = await this.client?.sendRequest(QueryRequestType, encryptedRequest)
+            const resp = await this.client?.sendRequest(QueryVectorIndexRequestType, encryptedRequest)
             return resp
         } catch (e) {
-            getLogger().error(`LspClient: query error: ${e}`)
+            getLogger().error(`LspClient: queryVectorIndex error: ${e}`)
             return []
         }
     }
 
+    async queryInlineProjectContext(query: string, path: string) {
+        try {
+            const request = JSON.stringify({
+                query: query,
+                filePath: path,
+            })
+            const encrypted = await this.encrypt(request)
+            const resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrypted)
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: queryInlineProjectContext error: ${e}`)
+            throw e
+        }
+    }
+
     async getLspServerUsage(): Promise<Usage | undefined> {
         if (this.client) {
             return (await this.client.sendRequest(GetUsageRequestType, '')) as Usage
         }
     }
 
-    async updateIndex(filePath: string) {
+    async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add') {
+        const payload: UpdateIndexV2RequestPayload = {
+            filePaths: filePath,
+            updateMode: mode,
+        }
         try {
-            const encryptedRequest = await this.encrypt(
-                JSON.stringify({
-                    filePath: filePath,
-                })
-            )
-            const resp = await this.client?.sendRequest(UpdateIndexRequestType, encryptedRequest)
+            const encryptedRequest = await this.encrypt(JSON.stringify(payload))
+            const resp = await this.client?.sendRequest(UpdateIndexV2RequestType, encryptedRequest)
             return resp
         } catch (e) {
             getLogger().error(`LspClient: updateIndex error: ${e}`)
@@ -197,15 +222,26 @@ export async function activate(extensionContext: ExtensionContext) {
                 return
             }
             savedDocument = document.uri
-        })
-    )
-    toDispose.push(
+        }),
         vscode.window.onDidChangeActiveTextEditor((editor) => {
             if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
-                void LspClient.instance.updateIndex(savedDocument.fsPath)
+                void LspClient.instance.updateIndex([savedDocument.fsPath], 'update')
             }
+        }),
+        vscode.workspace.onDidCreateFiles((e) => {
+            void LspClient.instance.updateIndex(
+                e.files.map((f) => f.fsPath),
+                'add'
+            )
+        }),
+        vscode.workspace.onDidDeleteFiles((e) => {
+            void LspClient.instance.updateIndex(
+                e.files.map((f) => f.fsPath),
+                'remove'
+            )
         })
     )
+
     return LspClient.instance.client.onReady().then(() => {
         const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
         toDispose.push(disposableFunc)

packages/core/src/amazonq/lsp/lspController.ts

@@ -15,12 +15,10 @@ import { LspClient } from './lspClient'
 import AdmZip from 'adm-zip'
 import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
 import { makeTemporaryToolkitFolder, tryRemoveFolder } from '../../shared/filesystemUtilities'
-import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { activate as activateLsp } from './lspClient'
 import { telemetry } from '../../shared/telemetry'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { fs, globals, ToolkitError } from '../../shared'
-import { AuthUtil } from '../../codewhisperer'
 import { isWeb } from '../../shared/extensionGlobals'
 import { getUserAgent } from '../../shared/telemetry/util'
 import { isAmazonInternalOs } from '../../shared/vscode/env'
@@ -68,9 +66,16 @@ export interface Manifest {
 }
 const manifestUrl = 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json'
 // this LSP client in Q extension is only going to work with these LSP server versions
-const supportedLspServerVersions = ['0.1.13']
+const supportedLspServerVersions = ['0.1.22', '0.1.19']
 
 const nodeBinName = process.platform === 'win32' ? 'node.exe' : 'node'
+
+export interface BuildIndexConfig {
+    startUrl?: string
+    maxIndexSize: number
+    isVectorIndexEnabled: boolean
+}
+
 /*
  * LSP Controller manages the status of Amazon Q LSP:
  * 1. Downloading, verifying and installing LSP using DEXP LSP manifest and CDN.
@@ -281,7 +286,7 @@ export class LspController {
     }
 
     async query(s: string): Promise<RelevantTextDocument[]> {
-        const chunks: Chunk[] | undefined = await LspClient.instance.query(s)
+        const chunks: Chunk[] | undefined = await LspClient.instance.queryVectorIndex(s)
         const resp: RelevantTextDocument[] = []
         chunks?.forEach((chunk) => {
             const text = chunk.context ? chunk.context : chunk.content
@@ -303,7 +308,15 @@ export class LspController {
         return resp
     }
 
-    async buildIndex() {
+    async queryInlineProjectContext(query: string, path: string) {
+        try {
+            return await LspClient.instance.queryInlineProjectContext(query, path)
+        } catch (e) {
+            return []
+        }
+    }
+
+    async buildIndex(buildIndexConfig: BuildIndexConfig) {
         getLogger().info(`LspController: Starting to build index of project`)
         const start = performance.now()
         const projPaths = getProjectPaths()
@@ -318,18 +331,16 @@ export class LspController {
                 projPaths,
                 vscode.workspace.workspaceFolders as CurrentWsFolders,
                 true,
-                CodeWhispererSettings.instance.getMaxIndexSize() * 1024 * 1024
+                buildIndexConfig.maxIndexSize * 1024 * 1024
             )
             const totalSizeBytes = files.reduce(
                 (accumulator, currentFile) => accumulator + currentFile.fileSizeBytes,
                 0
             )
             getLogger().info(`LspController: Found ${files.length} files in current project ${getProjectPaths()}`)
-            const resp = await LspClient.instance.indexFiles(
-                files.map((f) => f.fileUri.fsPath),
-                projRoot,
-                false
-            )
+            const config = buildIndexConfig.isVectorIndexEnabled ? 'all' : 'default'
+            const r = files.map((f) => f.fileUri.fsPath)
+            const resp = await LspClient.instance.buildIndex(r, projRoot, config)
             if (resp) {
                 getLogger().debug(`LspController: Finish building index of project`)
                 const usage = await LspClient.instance.getLspServerUsage()
@@ -340,31 +351,36 @@ export class LspController {
                     amazonqIndexMemoryUsageInMB: usage ? usage.memoryUsage / (1024 * 1024) : undefined,
                     amazonqIndexCpuUsagePercentage: usage ? usage.cpuUsage : undefined,
                     amazonqIndexFileSizeInMB: totalSizeBytes / (1024 * 1024),
-                    credentialStartUrl: AuthUtil.instance.startUrl,
+                    credentialStartUrl: buildIndexConfig.startUrl,
                 })
             } else {
-                getLogger().error(`LspController: Failed to build index of project`)
-                telemetry.amazonq_indexWorkspace.emit({
-                    duration: performance.now() - start,
-                    result: 'Failed',
-                    amazonqIndexFileCount: 0,
-                    amazonqIndexFileSizeInMB: 0,
-                })
+                // TODO: Re-enable this code path for LSP 0.1.20+
+                // getLogger().error(`LspController: Failed to build index of project`)
+                // telemetry.amazonq_indexWorkspace.emit({
+                //     duration: performance.now() - start,
+                //     result: 'Failed',
+                //     amazonqIndexFileCount: 0,
+                //     amazonqIndexFileSizeInMB: 0,
+                //     reason: `Unknown`,
+                // })
             }
-        } catch (e) {
+        } catch (error) {
+            //TODO: use telemetry.run()
             getLogger().error(`LspController: Failed to build index of project`)
             telemetry.amazonq_indexWorkspace.emit({
                 duration: performance.now() - start,
                 result: 'Failed',
                 amazonqIndexFileCount: 0,
                 amazonqIndexFileSizeInMB: 0,
+                reason: `${error instanceof Error ? error.name : 'Unknown'}`,
+                reasonDesc: `Error when building index. ${error instanceof Error ? error.message : error}`,
             })
         } finally {
             this._isIndexingInProgress = false
         }
     }
 
-    async trySetupLsp(context: vscode.ExtensionContext) {
+    async trySetupLsp(context: vscode.ExtensionContext, buildIndexConfig: BuildIndexConfig) {
         if (isCloud9() || isWeb() || isAmazonInternalOs()) {
             getLogger().warn('LspController: Skipping LSP setup. LSP is not compatible with the current environment. ')
             // do not do anything if in Cloud9 or Web mode or in AL2 (AL2 does not support node v18+)
@@ -378,9 +394,7 @@ export class LspController {
             try {
                 await activateLsp(context)
                 getLogger().info('LspController: LSP activated')
-                if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
-                    void LspController.instance.buildIndex()
-                }
+                void LspController.instance.buildIndex(buildIndexConfig)
                 // log the LSP server CPU and Memory usage per 30 minutes.
                 globals.clock.setInterval(
                     async () => {

packages/core/src/amazonq/lsp/types.ts

@@ -11,10 +11,6 @@ export type IndexRequestPayload = {
     refresh: boolean
 }
 
-export type IndexRequest = string
-
-export const IndexRequestType: RequestType<IndexRequest, any, any> = new RequestType('lsp/index')
-
 export type ClearRequest = string
 
 export const ClearRequestType: RequestType<ClearRequest, any, any> = new RequestType('lsp/clear')
@@ -23,10 +19,6 @@ export type QueryRequest = string
 
 export const QueryRequestType: RequestType<QueryRequest, any, any> = new RequestType('lsp/query')
 
-export type UpdateIndexRequest = string
-
-export const UpdateIndexRequestType: RequestType<UpdateIndexRequest, any, any> = new RequestType('lsp/updateIndex')
-
 export type GetUsageRequest = string
 
 export const GetUsageRequestType: RequestType<GetUsageRequest, any, any> = new RequestType('lsp/getUsage')
@@ -35,3 +27,40 @@ export interface Usage {
     memoryUsage: number
     cpuUsage: number
 }
+
+export type BuildIndexRequestPayload = {
+    filePaths: string[]
+    projectRoot: string
+    config: string
+    language: string
+}
+
+export type BuildIndexRequest = string
+
+export const BuildIndexRequestType: RequestType<BuildIndexRequest, any, any> = new RequestType('lsp/buildIndex')
+
+export type UpdateIndexV2Request = string
+
+export type UpdateIndexV2RequestPayload = { filePaths: string[]; updateMode: string }
+
+export const UpdateIndexV2RequestType: RequestType<UpdateIndexV2Request, any, any> = new RequestType(
+    'lsp/updateIndexV2'
+)
+
+export type QueryInlineProjectContextRequest = string
+export type QueryInlineProjectContextRequestPayload = {
+    query: string
+    filePath: string
+}
+export const QueryInlineProjectContextRequestType: RequestType<QueryInlineProjectContextRequest, any, any> =
+    new RequestType('lsp/queryInlineProjectContext')
+
+export type QueryVectorIndexRequestPayload = { query: string }
+
+export type QueryVectorIndexRequest = string
+
+export const QueryVectorIndexRequestType: RequestType<QueryVectorIndexRequest, any, any> = new RequestType(
+    'lsp/queryVectorIndex'
+)
+
+export type IndexConfig = 'all' | 'default'

packages/core/src/codewhisperer/models/model.ts

@@ -48,7 +48,7 @@ export type UtgStrategy = 'ByName' | 'ByContent'
 
 export type CrossFileStrategy = 'OpenTabs_BM25'
 
-export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
+export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty' | 'LSP'
 
 export interface CodeWhispererSupplementalContext {
     isUtg: boolean