fix(amazonq): avoid workspace context server missing historical dependency events (aws#1946)

LiGaCu · Jiatong Li · web-flow · commit 3362956ded75 · 2025-07-22T11:29:16.000-07:00
Co-authored-by: Jiatong Li &lt;jiatonli@amazon.com&gt;
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyDiscoverer.ts
@@ -10,13 +10,15 @@ import {
 } from './dependencyHandler/LanguageDependencyHandler'
 import { ArtifactManager } from '../artifactManager'
 import { supportedWorkspaceContextLanguages } from '../../../shared/languageDetection'
+import { DependencyEventBundler } from './dependencyEventBundler'
 
 export class DependencyDiscoverer {
     private logging: Logging
     private workspaceFolders: WorkspaceFolder[]
     public dependencyHandlerRegistry: LanguageDependencyHandler<BaseDependencyInfo>[] = []
     private initializedWorkspaceFolder = new Map<WorkspaceFolder, boolean>()
     private sharedState: DependencyHandlerSharedState = { isDisposed: false, dependencyUploadedSizeSum: 0 }
+    private dependencyEventsIngestedFolderUris = new Set<string>()
 
     constructor(
         workspace: Workspace,
@@ -70,6 +72,9 @@ export class DependencyDiscoverer {
     async searchDependencies(folders: WorkspaceFolder[]): Promise<void> {
         this.logging.log('Starting dependency search across workspace folders')
 
+        // ingest recorded dependency events to corresponding dependency maps first
+        this.ingestRecordedDependencyEvents(folders)
+
         for (const workspaceFolder of folders) {
             if (
                 this.initializedWorkspaceFolder.has(workspaceFolder) &&
@@ -144,12 +149,45 @@ export class DependencyDiscoverer {
         }
     }
 
-    async handleDependencyUpdateFromLSP(language: string, paths: string[], workspaceRoot?: WorkspaceFolder) {
+    public isDependencyEventsIngested(workspaceFolderUri: string): boolean {
+        return this.dependencyEventsIngestedFolderUris.has(workspaceFolderUri)
+    }
+
+    private ingestRecordedDependencyEvents(workspaceFolders: WorkspaceFolder[]): void {
+        let ingestedDependencyCount = 0
+        for (const workspaceFolder of workspaceFolders) {
+            for (const dependencyHandler of this.dependencyHandlerRegistry) {
+                try {
+                    const recordedPaths = DependencyEventBundler.getRecordedDependencyPaths(
+                        dependencyHandler.language,
+                        workspaceFolder.uri
+                    )
+                    if (!recordedPaths) {
+                        continue
+                    }
+                    dependencyHandler.updateDependencyMapBasedOnLSP(recordedPaths, workspaceFolder)
+                    ingestedDependencyCount += recordedPaths.length
+                } catch (error) {
+                    this.logging.debug(`Error ingesting dependency events for ${workspaceFolder.uri}: ${error}`)
+                }
+            }
+            this.dependencyEventsIngestedFolderUris.add(workspaceFolder.uri)
+        }
+        if (ingestedDependencyCount > 0) {
+            this.logging.log(`Ingested ${ingestedDependencyCount} dependencies from didChangeDependencyPaths events`)
+        }
+    }
+
+    async handleDependencyUpdateFromLSP(language: string, paths: string[], folder?: WorkspaceFolder) {
+        if (folder === undefined) {
+            return
+        }
         for (const dependencyHandler of this.dependencyHandlerRegistry) {
             if (dependencyHandler.language != language) {
                 continue
             }
-            await dependencyHandler.updateDependencyMapBasedOnLSP(paths, workspaceRoot)
+            const changedDependencyList = dependencyHandler.updateDependencyMapBasedOnLSP(paths, folder)
+            await dependencyHandler.zipAndUploadDependenciesByChunk(changedDependencyList, folder)
         }
     }
 
@@ -160,6 +198,7 @@ export class DependencyDiscoverer {
 
     public dispose(): void {
         this.initializedWorkspaceFolder.clear()
+        this.dependencyEventsIngestedFolderUris.clear()
         this.dependencyHandlerRegistry.forEach(dependencyHandler => {
             dependencyHandler.dispose()
         })
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyEventBundler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyEventBundler.ts
@@ -1,23 +1,34 @@
-import { Logging, WorkspaceFolder } from '@aws/language-server-runtimes/server-interface'
+import { Logging } from '@aws/language-server-runtimes/server-interface'
 import { DependencyDiscoverer } from './dependencyDiscoverer'
+import { WorkspaceFolderManager } from '../workspaceFolderManager'
 
 export interface DependencyEvent {
     language: string
     paths: string[]
-    workspaceFolder: WorkspaceFolder | undefined
+    workspaceFolderUri: string
 }
 
 export class DependencyEventBundler {
+    // Map storing historically received dependency events from extension
+    // Key is <language>-<workspaceFolderUri> and value is a set of paths
+    private static readonly recordedDependencies = new Map<string, Set<string>>()
+
     private readonly logging: Logging
     private readonly dependencyDiscoverer: DependencyDiscoverer
+    private readonly workspaceFolderManager: WorkspaceFolderManager
     private readonly BUNDLER_PROCESS_INTERVAL: number = 500 // 500 milliseconds
-    public eventQueue: DependencyEvent[] = []
+    private eventSendingQueue: DependencyEvent[] = []
     private eventBundlerInterval: NodeJS.Timeout | undefined
     private isBundlerWorking: boolean = false
 
-    constructor(logging: Logging, dependencyDiscoverer: DependencyDiscoverer) {
+    constructor(
+        logging: Logging,
+        dependencyDiscoverer: DependencyDiscoverer,
+        workspaceFolderManager: WorkspaceFolderManager
+    ) {
         this.logging = logging
         this.dependencyDiscoverer = dependencyDiscoverer
+        this.workspaceFolderManager = workspaceFolderManager
     }
 
     /**
@@ -32,11 +43,11 @@ export class DependencyEventBundler {
             }
             this.isBundlerWorking = true
             try {
-                const allEvents = this.eventQueue.splice(0)
+                const allEvents = this.eventSendingQueue.splice(0)
 
                 // Form bundles based on unique combination of language and workspaceFolder
                 const dependencyEventBundles = allEvents.reduce((accumulator, event) => {
-                    const key = this.getBundleKey(event)
+                    const key = DependencyEventBundler.getBundleKey(event.language, event.workspaceFolderUri)
                     if (!accumulator.has(key)) {
                         accumulator.set(key, [])
                     }
@@ -46,10 +57,12 @@ export class DependencyEventBundler {
 
                 // Process bundles one by one, concatenating all the paths within the bundle
                 for (const [bundleKey, bundledEvents] of dependencyEventBundles) {
+                    const { language, workspaceFolderUri } = bundledEvents[0]
+                    const workspaceFolder = this.workspaceFolderManager.getWorkspaceFolder(workspaceFolderUri)
                     await this.dependencyDiscoverer.handleDependencyUpdateFromLSP(
-                        bundledEvents[0].language,
+                        language,
                         bundledEvents.flatMap(event => event.paths),
-                        bundledEvents[0].workspaceFolder
+                        workspaceFolder
                     )
                 }
             } catch (err) {
@@ -60,16 +73,39 @@ export class DependencyEventBundler {
         }, this.BUNDLER_PROCESS_INTERVAL)
     }
 
-    private getBundleKey(event: DependencyEvent) {
-        if (event.workspaceFolder === undefined) {
-            return `${event.language}-undefined`
-        }
-        return `${event.language}-${event.workspaceFolder.uri}`
+    public sendDependencyEvent(event: DependencyEvent) {
+        this.eventSendingQueue.push(event)
     }
 
     public dispose(): void {
         if (this.eventBundlerInterval) {
             clearInterval(this.eventBundlerInterval)
         }
+        this.eventSendingQueue = []
+    }
+
+    private static getBundleKey(language: string, workspaceFolderUri: string) {
+        return `${language}-${workspaceFolderUri}`
+    }
+
+    public static recordDependencyEvent(event: DependencyEvent): void {
+        const key = this.getBundleKey(event.language, event.workspaceFolderUri)
+        if (!this.recordedDependencies.has(key)) {
+            this.recordedDependencies.set(key, new Set())
+        }
+        const receivedPaths = this.recordedDependencies.get(key)
+        if (receivedPaths) {
+            event.paths.forEach(path => {
+                receivedPaths.add(path)
+            })
+        }
+    }
+
+    public static getRecordedDependencyPaths(language: string, workspaceFolderUri: string): string[] | undefined {
+        const key = this.getBundleKey(language, workspaceFolderUri)
+        const receivedPaths = this.recordedDependencies.get(key)
+        if (receivedPaths) {
+            return Array.from(receivedPaths)
+        }
     }
 }
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JSTSDependencyHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JSTSDependencyHandler.ts
@@ -186,10 +186,13 @@ export class JSTSDependencyHandler extends LanguageDependencyHandler<JSTSDepende
                 const callBackDependencyUpdate = async (events: string[]) => {
                     this.logging.log(`Change detected in ${packageJsonPath}`)
                     const updatedDependencyMap = this.generateDependencyMap(jstsDependencyInfo)
-                    await this.compareAndUpdateDependencyMap(
+                    const changedDependencyList = this.compareAndUpdateDependencyMap(
                         jstsDependencyInfo.workspaceFolder,
-                        updatedDependencyMap,
-                        true
+                        updatedDependencyMap
+                    )
+                    await this.zipAndUploadDependenciesByChunk(
+                        changedDependencyList,
+                        jstsDependencyInfo.workspaceFolder
                     )
                 }
                 const watcher = new DependencyWatcher(
@@ -206,7 +209,9 @@ export class JSTSDependencyHandler extends LanguageDependencyHandler<JSTSDepende
     }
 
     // JS and TS are not using LSP to sync dependencies
-    override async updateDependencyMapBasedOnLSP(paths: string[], workspaceFolder?: WorkspaceFolder): Promise<void> {}
+    override updateDependencyMapBasedOnLSP(paths: string[], workspaceFolder?: WorkspaceFolder): Dependency[] {
+        return []
+    }
     override transformPathToDependency(
         dependencyName: string,
         dependencyPath: string,
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JavaDependencyHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/JavaDependencyHandler.ts
@@ -57,11 +57,7 @@ export class JavaDependencyHandler extends LanguageDependencyHandler<JavaDepende
             // TODO, check if try catch is necessary here
             try {
                 let generatedDependencyMap: Map<string, Dependency> = this.generateDependencyMap(javaDependencyInfo)
-                this.compareAndUpdateDependencyMap(javaDependencyInfo.workspaceFolder, generatedDependencyMap).catch(
-                    error => {
-                        this.logging.warn(`Error processing Java dependencies: ${error}`)
-                    }
-                )
+                this.compareAndUpdateDependencyMap(javaDependencyInfo.workspaceFolder, generatedDependencyMap)
                 // Log found dependencies
                 this.logging.log(
                     `Total Java dependencies found:  ${generatedDependencyMap.size} under ${javaDependencyInfo.pkgDir}`
@@ -92,10 +88,13 @@ export class JavaDependencyHandler extends LanguageDependencyHandler<JavaDepende
                 const callBackDependencyUpdate = async (events: string[]) => {
                     this.logging.log(`Change detected in ${dotClasspathPath}`)
                     const updatedDependencyMap = this.generateDependencyMap(javaDependencyInfo)
-                    await this.compareAndUpdateDependencyMap(
+                    const changedDependencyList = this.compareAndUpdateDependencyMap(
                         javaDependencyInfo.workspaceFolder,
-                        updatedDependencyMap,
-                        true
+                        updatedDependencyMap
+                    )
+                    await this.zipAndUploadDependenciesByChunk(
+                        changedDependencyList,
+                        javaDependencyInfo.workspaceFolder
                     )
                 }
                 const watcher = new DependencyWatcher(
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/LanguageDependencyHandler.ts
@@ -117,18 +117,17 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
      * @param paths
      * @param workspaceRoot
      */
-    async updateDependencyMapBasedOnLSP(paths: string[], workspaceFolder?: WorkspaceFolder): Promise<void> {
+    updateDependencyMapBasedOnLSP(paths: string[], workspaceFolder: WorkspaceFolder): Dependency[] {
         const dependencyMap = new Map<string, Dependency>()
         paths.forEach((dependencyPath: string) => {
             // basename of the path should be the dependency name
             const dependencyName = path.basename(dependencyPath)
             this.transformPathToDependency(dependencyName, dependencyPath, dependencyMap)
         })
 
-        if (workspaceFolder) {
-            await this.compareAndUpdateDependencyMap(workspaceFolder, dependencyMap, true)
-        }
+        return this.compareAndUpdateDependencyMap(workspaceFolder, dependencyMap)
     }
+
     async zipDependencyMap(folders: WorkspaceFolder[]): Promise<void> {
         // Process each workspace folder sequentially
         for (const [workspaceFolder, correspondingDependencyMap] of this.dependencyMap) {
@@ -143,7 +142,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
         }
     }
 
-    private async zipAndUploadDependenciesByChunk(
+    async zipAndUploadDependenciesByChunk(
         dependencyList: Dependency[],
         workspaceFolder: WorkspaceFolder
     ): Promise<void> {
@@ -251,11 +250,10 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
      */
     protected abstract generateDependencyMap(dependencyInfo: T, dependencyMap: Map<string, Dependency>): void
 
-    protected async compareAndUpdateDependencyMap(
+    protected compareAndUpdateDependencyMap(
         workspaceFolder: WorkspaceFolder,
-        updatedDependencyMap: Map<string, Dependency>,
-        zipChanges: boolean = false
-    ): Promise<void> {
+        updatedDependencyMap: Map<string, Dependency>
+    ): Dependency[] {
         const changes = {
             added: [] as Dependency[],
             updated: [] as Dependency[],
@@ -290,9 +288,7 @@ export abstract class LanguageDependencyHandler<T extends BaseDependencyInfo> {
             this.dependencyMap.get(workspaceFolder)?.set(name, newDep)
         })
 
-        if (zipChanges) {
-            await this.zipAndUploadDependenciesByChunk([...changes.added, ...changes.updated], workspaceFolder)
-        }
+        return [...changes.added, ...changes.updated]
     }
 
     private validateSingleDependencySize(workspaceFolder: WorkspaceFolder, dependency: Dependency): boolean {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/PythonDependencyHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/dependency/dependencyHandler/PythonDependencyHandler.ts
@@ -78,11 +78,7 @@ export class PythonDependencyHandler extends LanguageDependencyHandler<PythonDep
             // TODO, check if the try catch is necessary here
             try {
                 let generatedDependencyMap: Map<string, Dependency> = this.generateDependencyMap(pythonDependencyInfo)
-                this.compareAndUpdateDependencyMap(pythonDependencyInfo.workspaceFolder, generatedDependencyMap).catch(
-                    error => {
-                        this.logging.warn(`Error processing Python dependencies: ${error}`)
-                    }
-                )
+                this.compareAndUpdateDependencyMap(pythonDependencyInfo.workspaceFolder, generatedDependencyMap)
                 // Log found dependencies
                 this.logging.log(
                     `Total Python dependencies found: ${generatedDependencyMap.size} under ${pythonDependencyInfo.pkgDir}`
@@ -122,10 +118,13 @@ export class PythonDependencyHandler extends LanguageDependencyHandler<PythonDep
                                 this.handlePackageChange(sitePackagesPath, fileName, updatedDependencyMap)
                             }
                         }
-                        await this.compareAndUpdateDependencyMap(
+                        const changedDependencyList = this.compareAndUpdateDependencyMap(
                             pythonDependencyInfo.workspaceFolder,
-                            updatedDependencyMap,
-                            true
+                            updatedDependencyMap
+                        )
+                        await this.zipAndUploadDependenciesByChunk(
+                            changedDependencyList,
+                            pythonDependencyInfo.workspaceFolder
                         )
                     } // end of callback function
 
diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts
@@ -21,7 +21,7 @@ import { makeUserContextObject } from '../../shared/telemetryUtils'
 import { safeGet } from '../../shared/utils'
 import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 import { FileUploadJobManager, FileUploadJobType } from './fileUploadJobManager'
-import { DependencyEventBundler } from './dependency/dependencyEventBundler'
+import { DependencyEvent, DependencyEventBundler } from './dependency/dependencyEventBundler'
 import ignore = require('ignore')
 import { BUILDER_ID_START_URL, INTERNAL_USER_START_URL } from '../../shared/constants'
 
@@ -244,7 +244,7 @@ export const WorkspaceContextServer = (): Server => features => {
                 workspaceIdentifier
             )
             fileUploadJobManager = new FileUploadJobManager(logging, workspaceFolderManager)
-            dependencyEventBundler = new DependencyEventBundler(logging, dependencyDiscoverer)
+            dependencyEventBundler = new DependencyEventBundler(logging, dependencyDiscoverer, workspaceFolderManager)
             await updateConfiguration()
 
             lsp.workspace.onDidChangeWorkspaceFolders(async params => {
@@ -511,17 +511,22 @@ export const WorkspaceContextServer = (): Server => features => {
 
     lsp.extensions.onDidChangeDependencyPaths(async params => {
         try {
+            const dependencyEvent: DependencyEvent = {
+                language: params.runtimeLanguage,
+                paths: params.paths,
+                workspaceFolderUri: params.moduleName,
+            }
+            DependencyEventBundler.recordDependencyEvent(dependencyEvent)
+
             if (!isUserEligibleForWorkspaceContext()) {
                 return
             }
-            logging.log(`Received onDidChangeDependencyPaths event for ${params.moduleName}`)
 
-            const workspaceFolder = workspaceFolderManager.getWorkspaceFolder(params.moduleName)
-            dependencyEventBundler.eventQueue.push({
-                language: params.runtimeLanguage,
-                paths: params.paths,
-                workspaceFolder: workspaceFolder,
-            })
+            // Only send events separately when dependency discovery has finished ingesting previous recorded events
+            if (dependencyDiscoverer.isDependencyEventsIngested(params.moduleName)) {
+                dependencyEventBundler.sendDependencyEvent(dependencyEvent)
+                logging.log(`Processed onDidChangeDependencyPaths event for ${params.moduleName}`)
+            }
         } catch (error) {
             logging.error(`Error handling didChangeDependencyPaths event: ${error}`)
         }
