feat: send pinned context button immediately with pending state (#2353)

* feat: context command pending state passed to ui, triggered when indexing complete

* refactored pending flow

* fix: removed extra spaces

* fix: workspace pending turned back on when building index (when config is changed)

* fix: log errors if inital context commands fail to send

* refactor: changed pending state from boolean to string to send to ui

* refactor: update from pending to disabledText

* fix: send intial pending state before any other action in contextCommandsProvider

* fix: send initial pending context commands onReady, instead of before

* test: added unit test for onReady

* test: added unit test for when indexingInProgress is changed

* chore: fix the mynal ui test

* chore: fix the mynal test for lack of disabled field

---------

Co-authored-by: Conor Stewart <[email protected]>
Co-authored-by: constewart9 <[email protected]>
Co-authored-by: manodnyab <[email protected]>

Author: 3 people

chat-client/src/client/mynahUi.test.ts

@@ -634,6 +634,7 @@ describe('MynahUI', () => {
                         route: ['/workspace', 'src/file1.ts'],
                         icon: 'file',
                         children: undefined,
+                        disabled: false,
                     },
                 ],
                 promptTopBarTitle: '@',
@@ -690,6 +691,7 @@ describe('MynahUI', () => {
                         ...activeEditorCommand,
                         description: 'file:///workspace/src/active.ts',
                         children: undefined,
+                        disabled: false,
                     },
                 ],
                 promptTopBarTitle: '@Pin Context',
@@ -729,7 +731,7 @@ describe('MynahUI', () => {
             // Verify updateStore was called with empty context items
             // Active editor should be removed since no textDocument was provided
             sinon.assert.calledWith(updateStoreSpy, tabId, {
-                promptTopBarContextItems: [{ ...fileCommand, children: undefined }],
+                promptTopBarContextItems: [{ ...fileCommand, children: undefined, disabled: false }],
                 promptTopBarTitle: '@',
                 promptTopBarButton: null,
             })

chat-client/src/client/mynahUi.ts

@@ -1560,6 +1560,7 @@ ${params.message}`,
                 commands: toContextCommands(child.commands),
             })),
             icon: toMynahIcon(command.icon),
+            disabled: command.disabledText != null,
         }))
     }
 

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

@@ -3847,9 +3847,11 @@ export class AgenticChatController implements ChatHandlers {
      */
     async onReady() {
         await this.restorePreviousChats()
+        this.#contextCommandsProvider.onReady()
         try {
             const localProjectContextController = await LocalProjectContextController.getInstance()
             const contextItems = await localProjectContextController.getContextCommandItems()
+            this.#contextCommandsProvider.setFilesAndFoldersPending(false)
             await this.#contextCommandsProvider.processContextCommandUpdate(contextItems)
             void this.#contextCommandsProvider.maybeUpdateCodeSymbols()
         } catch (error) {

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

@@ -3,6 +3,7 @@ import * as sinon from 'sinon'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 import * as chokidar from 'chokidar'
 import { ContextCommandItem } from 'local-indexing'
+import { LocalProjectContextController } from '../../../shared/localProjectContextController'
 
 describe('ContextCommandsProvider', () => {
     let provider: ContextCommandsProvider
@@ -21,6 +22,12 @@ describe('ContextCommandsProvider', () => {
 
         testFeatures.workspace.fs.exists = fsExistsStub
         testFeatures.workspace.fs.readdir = fsReadDirStub
+
+        sinon.stub(LocalProjectContextController, 'getInstance').resolves({
+            onContextItemsUpdated: sinon.stub(),
+            onIndexingInProgressChanged: sinon.stub(),
+        } as any)
+
         provider = new ContextCommandsProvider(
             testFeatures.logging,
             testFeatures.chat,
@@ -58,6 +65,26 @@ describe('ContextCommandsProvider', () => {
         })
     })
 
+    describe('onReady', () => {
+        it('should call processContextCommandUpdate with empty array on first call', async () => {
+            const processUpdateSpy = sinon.spy(provider, 'processContextCommandUpdate')
+
+            provider.onReady()
+
+            sinon.assert.calledOnce(processUpdateSpy)
+            sinon.assert.calledWith(processUpdateSpy, [])
+        })
+
+        it('should not call processContextCommandUpdate on subsequent calls', async () => {
+            const processUpdateSpy = sinon.spy(provider, 'processContextCommandUpdate')
+
+            provider.onReady()
+            provider.onReady()
+
+            sinon.assert.calledOnce(processUpdateSpy)
+        })
+    })
+
     describe('onContextItemsUpdated', () => {
         it('should call processContextCommandUpdate when controller raises event', async () => {
             const mockContextItems: ContextCommandItem[] = [
@@ -78,4 +105,29 @@ describe('ContextCommandsProvider', () => {
             sinon.assert.calledWith(processUpdateSpy, mockContextItems)
         })
     })
+
+    describe('onIndexingInProgressChanged', () => {
+        it('should update workspacePending and call processContextCommandUpdate when indexing status changes', async () => {
+            let capturedCallback: ((indexingInProgress: boolean) => void) | undefined
+
+            const mockController = {
+                onContextItemsUpdated: sinon.stub(),
+                set onIndexingInProgressChanged(callback: (indexingInProgress: boolean) => void) {
+                    capturedCallback = callback
+                },
+            }
+
+            const processUpdateSpy = sinon.spy(provider, 'processContextCommandUpdate')
+            ;(LocalProjectContextController.getInstance as sinon.SinonStub).resolves(mockController as any)
+
+            // Set initial state to false so condition is met
+            ;(provider as any).workspacePending = false
+
+            await (provider as any).registerContextCommandHandler()
+
+            capturedCallback?.(true)
+
+            sinon.assert.calledWith(processUpdateSpy, [])
+        })
+    })
 })

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

@@ -12,6 +12,10 @@ import { activeFileCmd } from './additionalContextProvider'
 export class ContextCommandsProvider implements Disposable {
     private promptFileWatcher?: FSWatcher
     private cachedContextCommands?: ContextCommandItem[]
+    private codeSymbolsPending = true
+    private filesAndFoldersPending = true
+    private workspacePending = true
+    private initialStateSent = false
     constructor(
         private readonly logging: Logging,
         private readonly chat: Chat,
@@ -24,12 +28,27 @@ export class ContextCommandsProvider implements Disposable {
         )
     }
 
+    onReady() {
+        if (!this.initialStateSent) {
+            this.initialStateSent = true
+            void this.processContextCommandUpdate([]).catch(e =>
+                this.logging.error(`Failed to send initial context commands: ${e}`)
+            )
+        }
+    }
+
     private async registerContextCommandHandler() {
         try {
             const controller = await LocalProjectContextController.getInstance()
             controller.onContextItemsUpdated = async contextItems => {
                 await this.processContextCommandUpdate(contextItems)
             }
+            controller.onIndexingInProgressChanged = (indexingInProgress: boolean) => {
+                if (this.workspacePending !== indexingInProgress) {
+                    this.workspacePending = indexingInProgress
+                    void this.processContextCommandUpdate(this.cachedContextCommands ?? [])
+                }
+            }
         } catch (e) {
             this.logging.warn(`Error processing context command update: ${e}`)
         }
@@ -105,6 +124,7 @@ export class ContextCommandsProvider implements Disposable {
             ],
             description: 'Add all files in a folder to context',
             icon: 'folder',
+            disabledText: this.filesAndFoldersPending ? 'pending' : undefined,
         }
 
         const fileCmds: ContextCommand[] = [activeFileCmd]
@@ -118,6 +138,7 @@ export class ContextCommandsProvider implements Disposable {
             ],
             description: 'Add a file to context',
             icon: 'file',
+            disabledText: this.filesAndFoldersPending ? 'pending' : undefined,
         }
 
         const codeCmds: ContextCommand[] = []
@@ -131,6 +152,7 @@ export class ContextCommandsProvider implements Disposable {
             ],
             description: 'Add code to context',
             icon: 'code-block',
+            disabledText: this.codeSymbolsPending ? 'pending' : undefined,
         }
 
         const promptCmds: ContextCommand[] = []
@@ -152,10 +174,12 @@ export class ContextCommandsProvider implements Disposable {
             icon: 'image',
             placeholder: 'Select an image file',
         }
-        const workspaceCmd = {
+
+        const workspaceCmd: ContextCommand = {
             command: '@workspace',
             id: '@workspace',
             description: 'Reference all code in workspace',
+            disabledText: this.workspacePending ? 'pending' : undefined,
         }
         const commands = [workspaceCmd, folderCmdGroup, fileCmdGroup, codeCmdGroup, promptCmdGroup]
 
@@ -209,11 +233,16 @@ export class ContextCommandsProvider implements Disposable {
             await LocalProjectContextController.getInstance()
         ).shouldUpdateContextCommandSymbolsOnce()
         if (needUpdate) {
+            this.codeSymbolsPending = false
             const items = await (await LocalProjectContextController.getInstance()).getContextCommandItems()
             await this.processContextCommandUpdate(items)
         }
     }
 
+    setFilesAndFoldersPending(value: boolean) {
+        this.filesAndFoldersPending = value
+    }
+
     dispose() {
         void this.promptFileWatcher?.close()
     }

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

@@ -50,6 +50,8 @@ export interface LocalProjectContextInitializationOptions {
 export class LocalProjectContextController {
     // Event handler for context items updated
     public onContextItemsUpdated: ((contextItems: ContextCommandItem[]) => Promise<void>) | undefined
+    // Event handler for when index is being built
+    public onIndexingInProgressChanged: ((enabled: boolean) => void) | undefined
     private static instance: LocalProjectContextController | undefined
 
     private workspaceFolders: WorkspaceFolder[]
@@ -214,6 +216,7 @@ export class LocalProjectContextController {
         }
         try {
             this._isIndexingInProgress = true
+            this.onIndexingInProgressChanged?.(this._isIndexingInProgress)
             if (this._vecLib) {
                 if (!this.workspaceFolders.length) {
                     this.log.info('skip building index because no workspace folder found')
@@ -234,6 +237,7 @@ export class LocalProjectContextController {
             this.log.error(`Error building index: ${error}`)
         } finally {
             this._isIndexingInProgress = false
+            this.onIndexingInProgressChanged?.(this._isIndexingInProgress)
         }
     }