fix: correct the implementation of gathering open tabs in cross file context (aws#2040)

* fix: pass correct open tab files for inline completion

* fix: add open tab in the cross file context of inline completion

Author: leigaol

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

@@ -338,7 +338,8 @@ export const CodewhispererServerFactory =
                                   workspace,
                                   logging,
                                   token,
-                                  amazonQServiceManager
+                                  amazonQServiceManager,
+                                  params.openTabFilepaths
                               )
                             : Promise.resolve(undefined)
 

server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.ts

@@ -67,7 +67,8 @@ export async function fetchSupplementalContextForSrc(
     position: Position,
     workspace: Workspace,
     cancellationToken: CancellationToken,
-    amazonQServiceManager?: AmazonQBaseServiceManager
+    amazonQServiceManager?: AmazonQBaseServiceManager,
+    openTabFiles?: string[]
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     const supplementalContextConfig = getSupplementalContextConfig(document.languageId)
 
@@ -76,7 +77,14 @@ export async function fetchSupplementalContextForSrc(
     }
     //TODO: add logic for other strategies once available
     if (supplementalContextConfig === 'codemap') {
-        return await codemapContext(document, position, workspace, cancellationToken, amazonQServiceManager)
+        return await codemapContext(
+            document,
+            position,
+            workspace,
+            cancellationToken,
+            amazonQServiceManager,
+            openTabFiles
+        )
     }
     return { supplementalContextItems: [], strategy: 'Empty' }
 }
@@ -86,13 +94,14 @@ export async function codemapContext(
     position: Position,
     workspace: Workspace,
     cancellationToken: CancellationToken,
-    amazonQServiceManager?: AmazonQBaseServiceManager
+    amazonQServiceManager?: AmazonQBaseServiceManager,
+    openTabFiles?: string[]
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     let strategy: SupplementalContextStrategy = 'Empty'
 
     const openTabsContextPromise = waitUntil(
         async function () {
-            return await fetchOpenTabsContext(document, position, workspace, cancellationToken)
+            return await fetchOpenTabsContext(document, position, workspace, cancellationToken, openTabFiles)
         },
         { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
     )
@@ -160,12 +169,13 @@ export async function fetchOpenTabsContext(
     document: TextDocument,
     position: Position,
     workspace: Workspace,
-    cancellationToken: CancellationToken
+    cancellationToken: CancellationToken,
+    openTabFiles?: string[]
 ): Promise<CodeWhispererSupplementalContextItem[]> {
     const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
 
     // Step 1: Get relevant cross files to refer
-    const relevantCrossFileCandidates = await getCrossFileCandidates(document, workspace)
+    const relevantCrossFileCandidates = await getCrossFileCandidates(document, workspace, openTabFiles)
 
     throwIfCancelled(cancellationToken)
 
@@ -318,11 +328,102 @@ function createFileUrl(uri: string): URL {
     return nodeUrl.pathToFileURL(resolvedPath)
 }
 
+/**
+ * Returns the language id for construction a TextDocument
+ * @param filepath
+ * @returns
+ */
+
+function guessLanguageId(filepath: string): string {
+    const ext = path.extname(filepath).toLowerCase()
+    switch (ext) {
+        case '.abap':
+            return 'abap'
+        case '.c':
+            return 'c'
+        case '.cpp':
+        case '.cc':
+        case '.cxx':
+        case '.hpp':
+        case '.h':
+            return 'cpp'
+        case '.cs':
+            return 'csharp'
+        case '.dart':
+            return 'dart'
+        case '.go':
+            return 'go'
+        case '.java':
+            return 'java'
+        case '.js':
+            return 'javascript'
+        case '.json':
+            return 'json'
+        case '.jsx':
+            return 'jsx'
+        case '.kt':
+        case '.kts':
+            return 'kotlin'
+        case '.lua':
+            return 'lua'
+        case '.php':
+            return 'php'
+        case '.txt':
+            return 'plaintext'
+        case '.ps1':
+        case '.psm1':
+        case '.psd1':
+            return 'powershell'
+        case '.py':
+            return 'python'
+        case '.r':
+        case '.R':
+            return 'r'
+        case '.rb':
+        case '.rbw':
+            return 'ruby'
+        case '.rs':
+            return 'rust'
+        case '.scala':
+        case '.sc':
+            return 'scala'
+        case '.sh':
+        case '.bash':
+            return 'shell'
+        case '.sql':
+            return 'sql'
+        case '.swift':
+            return 'swift'
+        case '.sv':
+        case '.svh':
+        case '.v':
+            return 'systemverilog'
+        case '.tf':
+        case '.tfvars':
+            return 'tf'
+        case '.tsx':
+            return 'tsx'
+        case '.ts':
+            return 'typescript'
+        case '.vue':
+            return 'vue'
+        case '.yml':
+        case '.yaml':
+            return 'yaml'
+        default:
+            return ''
+    }
+}
+
 /**
  * This function will return relevant cross files sorted by file distance for the given editor file
  * by referencing open files, imported files and same package files.
  */
-export async function getCrossFileCandidates(document: TextDocument, workspace: Workspace): Promise<TextDocument[]> {
+export async function getCrossFileCandidates(
+    document: TextDocument,
+    workspace: Workspace,
+    openTabFiles?: string[]
+): Promise<TextDocument[]> {
     const targetFile = document.uri
     const language = document.languageId as CrossFileSupportedLanguage
     const dialects = supportedLanguageToDialects[language]
@@ -335,8 +436,34 @@ export async function getCrossFileCandidates(document: TextDocument, workspace:
      *
      * Porting note: this function relies of Workspace feature to get all documents,
      * managed by this language server, instead of VSCode `vscode.window` API as VSCode toolkit does.
+     * this function only gets the user opened tab in this IDE session
+     * for a resumed IDE session, opened tabs are restored but this getAllTextDocuments function returns empty
+     * in that case we manually create TextDocuments from it
      */
-    const unsortedCandidates = await workspace.getAllTextDocuments()
+    let unsortedCandidates: TextDocument[] = await workspace.getAllTextDocuments()
+    if (openTabFiles && openTabFiles.length > 0) {
+        for (const openTabFile of openTabFiles) {
+            try {
+                const openTabFilesUri = URI.file(openTabFile)
+                if (!unsortedCandidates.some(x => x.uri === openTabFilesUri.toString())) {
+                    const content = await workspace.fs.readFile(openTabFilesUri.fsPath)
+                    if (content) {
+                        unsortedCandidates.push(
+                            TextDocument.create(
+                                URI.file(openTabFile).toString(),
+                                guessLanguageId(openTabFilesUri.fsPath),
+                                1,
+                                content
+                            )
+                        )
+                    }
+                }
+            } catch (e) {
+                // do not throw here.
+            }
+        }
+    }
+
     return unsortedCandidates
         .filter((candidateFile: TextDocument) => {
             let candidateFileURL

server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.ts

@@ -29,7 +29,8 @@ export async function fetchSupplementalContext(
     workspace: Workspace,
     logging: Logging,
     cancellationToken: CancellationToken,
-    amazonQServiceManager?: AmazonQBaseServiceManager
+    amazonQServiceManager?: AmazonQBaseServiceManager,
+    openTabFiles?: string[]
 ): Promise<CodeWhispererSupplementalContext | undefined> {
     const timesBeforeFetching = performance.now()
 
@@ -73,7 +74,8 @@ export async function fetchSupplementalContext(
                 position,
                 workspace,
                 cancellationToken,
-                amazonQServiceManager
+                amazonQServiceManager,
+                openTabFiles
             )
         }