fix(amazonq): move context command provider to agentic chat controller (#999)

Author: zixlin7

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts

@@ -4,6 +4,7 @@
  */
 
 import * as path from 'path'
+import * as chokidar from 'chokidar'
 import {
     ChatResponseStream,
     CodeWhispererStreaming,
@@ -39,6 +40,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A
 import { TabBarController } from './tabBarController'
 import { getUserPromptsDirectory } from './context/contextUtils'
 import { AdditionalContextProvider } from './context/addtionalContextProvider'
+import { ContextCommandsProvider } from './context/contextCommandsProvider'
 
 describe('AgenticChatController', () => {
     const mockTabId = 'tab-1'
@@ -120,6 +122,11 @@ describe('AgenticChatController', () => {
     const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures)
 
     beforeEach(() => {
+        sinon.stub(chokidar, 'watch').returns({
+            on: sinon.stub(),
+            close: sinon.stub(),
+        } as unknown as chokidar.FSWatcher)
+
         sendMessageStub = sinon.stub(CodeWhispererStreaming.prototype, 'sendMessage').callsFake(() => {
             return new Promise(resolve =>
                 setTimeout(() => {
@@ -190,6 +197,7 @@ describe('AgenticChatController', () => {
         emitConversationMetricStub = sinon.stub(ChatTelemetryController.prototype, 'emitConversationMetric')
 
         disposeStub = sinon.stub(ChatSessionService.prototype, 'dispose')
+        sinon.stub(ContextCommandsProvider.prototype, 'maybeUpdateCodeSymbols').resolves()
 
         AmazonQTokenServiceManager.resetInstance()
 

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -83,6 +83,8 @@ import {
 } from './context/agenticChatTriggerContext'
 import { AdditionalContextProvider } from './context/addtionalContextProvider'
 import { getNewPromptFilePath } from './context/contextUtils'
+import { ContextCommandsProvider } from './context/contextCommandsProvider'
+import { LocalProjectContextController } from '../../shared/localProjectContextController'
 
 type ChatHandlers = Omit<
     LspHandlers<Chat>,
@@ -107,6 +109,7 @@ export class AgenticChatController implements ChatHandlers {
     #tabBarController: TabBarController
     #chatHistoryDb: ChatDatabase
     #additionalContextProvider: AdditionalContextProvider
+    #contextCommandsProvider: ContextCommandsProvider
 
     constructor(
         chatSessionManagementService: ChatSessionManagementService,
@@ -123,6 +126,11 @@ export class AgenticChatController implements ChatHandlers {
         this.#chatHistoryDb = new ChatDatabase(features)
         this.#tabBarController = new TabBarController(features, this.#chatHistoryDb)
         this.#additionalContextProvider = new AdditionalContextProvider(features.workspace)
+        this.#contextCommandsProvider = new ContextCommandsProvider(
+            this.#features.logging,
+            this.#features.chat,
+            this.#features.workspace
+        )
     }
 
     async onCreatePrompt(params: CreatePromptParams): Promise<void> {
@@ -140,6 +148,7 @@ export class AgenticChatController implements ChatHandlers {
         this.#chatSessionManagementService.dispose()
         this.#telemetryController.dispose()
         this.#chatHistoryDb.close()
+        this.#contextCommandsProvider?.dispose()
     }
 
     async onListConversations(params: ListConversationsParams) {
@@ -727,6 +736,9 @@ export class AgenticChatController implements ChatHandlers {
 
     async onReady() {
         await this.#tabBarController.loadChats()
+        const contextItems = await (await LocalProjectContextController.getInstance()).getContextCommandItems()
+        await this.#contextCommandsProvider.processContextCommandUpdate(contextItems)
+        void this.#contextCommandsProvider.maybeUpdateCodeSymbols()
     }
 
     onSendFeedback({ tabId, feedbackPayload }: FeedbackParams) {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.test.ts

@@ -24,7 +24,7 @@ describe('AdditionalContextProvider', () => {
         testFeatures.workspace.fs.readdir = fsReadDirStub
         getContextCommandPromptStub = sinon.stub()
         provider = new AdditionalContextProvider(testFeatures.workspace)
-        localProjectContextControllerInstanceStub = sinon.stub(LocalProjectContextController, 'getInstance').returns({
+        localProjectContextControllerInstanceStub = sinon.stub(LocalProjectContextController, 'getInstance').resolves({
             getContextCommandPrompt: getContextCommandPromptStub,
         } as unknown as LocalProjectContextController)
     })

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/addtionalContextProvider.ts

@@ -81,8 +81,9 @@ export class AdditionalContextProvider {
             return []
         }
 
-        const prompts =
-            await LocalProjectContextController.getInstance().getContextCommandPrompt(additionalContextCommands)
+        const prompts = await (
+            await LocalProjectContextController.getInstance()
+        ).getContextCommandPrompt(additionalContextCommands)
 
         const contextEntry: AdditionalContentEntryAddition[] = []
         for (const prompt of prompts.slice(0, 20)) {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextCommandsProvider.ts

@@ -5,6 +5,7 @@ import { Disposable } from 'vscode-languageclient/node'
 import { Chat, Logging, Workspace } from '@aws/language-server-runtimes/server-interface'
 import { getUserPromptsDirectory, promptFileExtension } from './contextUtils'
 import { ContextCommandItem } from 'local-indexing'
+import { LocalProjectContextController } from '../../../shared/localProjectContextController'
 
 export class ContextCommandsProvider implements Disposable {
     private promptFileWatcher?: FSWatcher
@@ -164,6 +165,16 @@ export class ContextCommandsProvider implements Disposable {
         return allCommands
     }
 
+    async maybeUpdateCodeSymbols() {
+        const needUpdate = await (
+            await LocalProjectContextController.getInstance()
+        ).shouldUpdateContextCommandSymbolsOnce()
+        if (needUpdate) {
+            const items = await (await LocalProjectContextController.getInstance()).getContextCommandItems()
+            await this.processContextCommandUpdate(items)
+        }
+    }
+
     dispose() {
         void this.promptFileWatcher?.close()
     }

server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts

@@ -21,9 +21,7 @@ export const LocalProjectContextServer = (): Server => features => {
         localProjectContextController = new LocalProjectContextController(
             params.clientInfo?.name ?? 'unknown',
             params.workspaceFolders ?? [],
-            logging,
-            chat,
-            workspace
+            logging
         )
 
         const supportedFilePatterns = Object.keys(languageByExtension).map(ext => `**/*${ext}`)

server/aws-lsp-codewhisperer/src/shared/localProjectContextController.test.ts

@@ -6,7 +6,6 @@ import { Dirent } from 'fs'
 import * as path from 'path'
 import { URI } from 'vscode-uri'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
-import * as chokidar from 'chokidar'
 
 class LoggingMock {
     public error: SinonStub
@@ -41,10 +40,6 @@ describe('LocalProjectContextController', () => {
                 name: 'workspace1',
             },
         ]
-        stub(chokidar, 'watch').returns({
-            on: stub(),
-            close: stub(),
-        } as unknown as chokidar.FSWatcher)
 
         vectorLibMock = {
             start: stub().resolves({
@@ -70,14 +65,7 @@ describe('LocalProjectContextController', () => {
             }
         })
 
-        controller = new LocalProjectContextController(
-            'testClient',
-            mockWorkspaceFolders,
-            logging as any,
-            testFeatures.chat,
-            testFeatures.workspace
-        )
-        stub(controller, 'maybeUpdateCodeSymbols').resolves()
+        controller = new LocalProjectContextController('testClient', mockWorkspaceFolders, logging as any)
     })
 
     afterEach(() => {
@@ -113,9 +101,7 @@ describe('LocalProjectContextController', () => {
             const uninitializedController = new LocalProjectContextController(
                 'testClient',
                 mockWorkspaceFolders,
-                logging as any,
-                testFeatures.chat,
-                testFeatures.workspace
+                logging as any
             )
 
             const result = await uninitializedController.queryVectorIndex({ query: 'test' })
@@ -146,9 +132,7 @@ describe('LocalProjectContextController', () => {
             const uninitializedController = new LocalProjectContextController(
                 'testClient',
                 mockWorkspaceFolders,
-                logging as any,
-                testFeatures.chat,
-                testFeatures.workspace
+                logging as any
             )
 
             const result = await uninitializedController.queryInlineProjectContext({
@@ -191,9 +175,7 @@ describe('LocalProjectContextController', () => {
             const uninitializedController = new LocalProjectContextController(
                 'testClient',
                 mockWorkspaceFolders,
-                logging as any,
-                testFeatures.chat,
-                testFeatures.workspace
+                logging as any
             )
 
             await uninitializedController.updateIndex(['test.java'], 'add')

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

@@ -14,7 +14,6 @@ import type {
 } from 'local-indexing'
 import { URI } from 'vscode-uri'
 import { waitUntil } from '@aws/lsp-core/out/util/timeoutUtils'
-import { ContextCommandsProvider } from '../language-server/agenticChat/context/contextCommandsProvider'
 
 import * as fs from 'fs'
 import * as path from 'path'
@@ -51,7 +50,6 @@ export class LocalProjectContextController {
     private workspaceFolders: WorkspaceFolder[]
     private _vecLib?: VectorLibAPI
     private _contextCommandSymbolsUpdated = false
-    private readonly contextCommandsProvider: ContextCommandsProvider
     private readonly clientName: string
     private readonly log: Logging
 
@@ -67,24 +65,28 @@ export class LocalProjectContextController {
     private readonly DEFAULT_MAX_FILE_SIZE_MB = 10
     private readonly MB_TO_BYTES = 1024 * 1024
 
-    constructor(
-        clientName: string,
-        workspaceFolders: WorkspaceFolder[],
-        logging: Logging,
-        chat: Chat,
-        workspace: Workspace
-    ) {
+    constructor(clientName: string, workspaceFolders: WorkspaceFolder[], logging: Logging) {
         this.workspaceFolders = workspaceFolders
         this.clientName = clientName
         this.log = logging
-        this.contextCommandsProvider = new ContextCommandsProvider(logging, chat, workspace)
     }
 
-    public static getInstance() {
-        if (!this.instance) {
-            throw new Error('LocalProjectContextController not initialized')
+    public static async getInstance(): Promise<LocalProjectContextController> {
+        try {
+            await waitUntil(async () => this.instance, {
+                interval: 1000,
+                timeout: 60_000,
+                truthy: true,
+            })
+
+            if (!this.instance) {
+                throw new Error('LocalProjectContextController initialization timeout after 60 seconds')
+            }
+
+            return this.instance
+        } catch (error) {
+            throw new Error(`Failed to get LocalProjectContextController instance: ${error}`)
         }
-        return this.instance
     }
 
     public async init({
@@ -111,10 +113,6 @@ export class LocalProjectContextController {
                 this._vecLib = await vecLib.start(LIBRARY_DIR, this.clientName, this.indexCacheDirPath)
                 await this.buildIndex()
                 LocalProjectContextController.instance = this
-
-                const contextItems = await this.getContextCommandItems()
-                await this.contextCommandsProvider.processContextCommandUpdate(contextItems)
-                void this.maybeUpdateCodeSymbols()
             } else {
                 this.log.warn(`Vector library could not be imported from: ${libraryPath}`)
             }
@@ -128,7 +126,6 @@ export class LocalProjectContextController {
             await this._vecLib?.clear?.()
             this._vecLib = undefined
         }
-        this.contextCommandsProvider?.dispose()
     }
 
     public async updateIndex(filePaths: string[], operation: UpdateMode): Promise<void> {
@@ -269,14 +266,6 @@ export class LocalProjectContextController {
         }
     }
 
-    async maybeUpdateCodeSymbols() {
-        const needUpdate = await LocalProjectContextController.getInstance().shouldUpdateContextCommandSymbolsOnce()
-        if (needUpdate) {
-            const items = await this.getContextCommandItems()
-            await this.contextCommandsProvider.processContextCommandUpdate(items)
-        }
-    }
-
     private fileMeetsFileSizeConstraints(filePath: string, sizeConstraints: SizeConstraints): boolean {
         let fileSize