feat(amazonq): Use @ to add folders, files, and saved prompts as context (aws#6712)

## Problem
Users are unable to specify files or folders to add as context in a
conversation. Additionally, when using `@workspace`, users don't know
exactly which of their local files are being sent to Amazon Q.

## Solution
Users can enter `@` to add folders, files, and saved prompts as context.
The response from Amazon Q now displays a "Context" list with all the
items sent as context.

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: aws-toolkit-automation <[email protected]>
Co-authored-by: Lei Gao <[email protected]>
Co-authored-by: Lei Gao <[email protected]>
Co-authored-by: andrewyuq <[email protected]>
Co-authored-by: Josh Pinkney <[email protected]>

Author: 5 people

packages/amazonq/.changes/next-release/Feature-1707f610-2dde-4af1-b704-0e541efb7c4a.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Amazon Q chat: Use `@` to add folders, files, and saved prompts as context"
+}

packages/amazonq/.changes/next-release/Feature-dc148db8-5936-478c-a8cf-909d0885105b.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Amazon Q chat: Show list of files sent as context in chat response"
+}

packages/amazonq/test/e2e/amazonq/chat.test.ts

@@ -11,7 +11,7 @@ import { MynahUIDataModel } from '@aws/mynah-ui'
 import { assertContextCommands, assertQuickActions } from './assert'
 import { registerAuthHook, using } from 'aws-core-vscode/test'
 import { loginToIdC } from './utils/setup'
-import { webviewConstants } from 'aws-core-vscode/amazonq'
+import { webviewConstants, webviewTabConstants } from 'aws-core-vscode/amazonq'
 
 describe('Amazon Q Chat', function () {
     let framework: qTestingFramework
@@ -60,7 +60,7 @@ describe('Amazon Q Chat', function () {
     })
 
     it('Shows placeholder', () => {
-        assert.deepStrictEqual(store.promptInputPlaceholder, 'Ask a question or enter "/" for quick actions')
+        assert.deepStrictEqual(store.promptInputPlaceholder, webviewTabConstants.commonTabData.placeholder)
     })
 
     it('Sends message', async () => {

packages/amazonq/test/e2e/amazonq/framework/jsdomInjector.ts

@@ -40,6 +40,9 @@ export function injectJSDOM() {
         get() {
             return this.textContent
         },
+        set(value) {
+            this.textContent = value
+        },
     })
 
     // jsdom doesn't have support for structuredClone. See https://github.com/jsdom/jsdom/issues/3363

packages/core/package.nls.json

@@ -309,6 +309,17 @@
     "AWS.codewhisperer.customization.notification.new_customizations.learn_more": "Learn More",
     "AWS.amazonq.title": "Amazon Q",
     "AWS.amazonq.chat": "Chat",
+    "AWS.amazonq.context.folders.title": "Folders",
+    "AWS.amazonq.context.folders.description": "Add all files in a folder to context",
+    "AWS.amazonq.context.files.title": "Files",
+    "AWS.amazonq.context.files.description": "Add a file to context",
+    "AWS.amazonq.context.prompts.title": "Prompts",
+    "AWS.amazonq.context.prompts.description": "Add a saved prompt to context",
+    "AWS.amazonq.savedPrompts.title": "Prompt name",
+    "AWS.amazonq.savedPrompts.create": "Create",
+    "AWS.amazonq.savedPrompts.action": "Create a new prompt",
+    "AWS.amazonq.savedPrompts.placeholder": "Enter prompt name",
+    "AWS.amazonq.savedPrompts.description": "Use this prompt by typing '@' followed by the prompt name.",
     "AWS.amazonq.chat.workspacecontext.enable.message": "Amazon Q: Workspace index is now enabled. You can disable it from Amazon Q settings.",
     "AWS.amazonq.security": "Code Issues",
     "AWS.amazonq.login": "Login",

packages/core/src/amazonq/index.ts

@@ -26,6 +26,7 @@ export { init as testChatAppInit } from '../amazonqTest/app'
 export { init as docChatAppInit } from '../amazonqDoc/app'
 export { amazonQHelpUrl } from '../shared/constants'
 export * as webviewConstants from './webview/ui/texts/constants'
+export * as webviewTabConstants from './webview/ui/tabs/constants'
 export { listCodeWhispererCommandsWalkthrough } from '../codewhisperer/ui/statusBarMenu'
 export { focusAmazonQPanel, focusAmazonQPanelKeybinding } from '../codewhispererChat/commands/registerCommands'
 export { TryChatCodeLensProvider, tryChatCodeLensCommand } from '../codewhispererChat/editor/codelens'

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

@@ -29,12 +29,18 @@ import {
     QueryRepomapIndexRequestType,
     GetRepomapIndexJSONRequestType,
     Usage,
+    GetContextCommandItemsRequestType,
+    ContextCommandItem,
+    GetIndexSequenceNumberRequestType,
+    GetContextCommandPromptRequestType,
+    AdditionalContextPrompt,
 } from './types'
 import { Writable } from 'stream'
 import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { fs } from '../../shared/fs/fs'
 import { getLogger } from '../../shared/logger/logger'
 import globals from '../../shared/extensionGlobals'
+import { waitUntil } from '../../shared/utilities/timeoutUtils'
 
 const localize = nls.loadMessageBundle()
 
@@ -168,6 +174,66 @@ export class LspClient {
             throw e
         }
     }
+
+    async getContextCommandItems(): Promise<ContextCommandItem[]> {
+        try {
+            const workspaceFolders = vscode.workspace.workspaceFolders || []
+            const request = JSON.stringify({
+                workspaceFolders: workspaceFolders.map((it) => it.uri.fsPath),
+            })
+            const resp: any = await this.client?.sendRequest(
+                GetContextCommandItemsRequestType,
+                await this.encrypt(request)
+            )
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: getContextCommandItems error: ${e}`)
+            throw e
+        }
+    }
+
+    async getContextCommandPrompt(contextCommandItems: ContextCommandItem[]): Promise<AdditionalContextPrompt[]> {
+        try {
+            const request = JSON.stringify({
+                contextCommands: contextCommandItems,
+            })
+            const resp: any = await this.client?.sendRequest(
+                GetContextCommandPromptRequestType,
+                await this.encrypt(request)
+            )
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: getContextCommandPrompt error: ${e}`)
+            throw e
+        }
+    }
+
+    async getIndexSequenceNumber(): Promise<number> {
+        try {
+            const request = JSON.stringify({})
+            const resp: any = await this.client?.sendRequest(
+                GetIndexSequenceNumberRequestType,
+                await this.encrypt(request)
+            )
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: getIndexSequenceNumber error: ${e}`)
+            throw e
+        }
+    }
+
+    async waitUntilReady() {
+        return waitUntil(
+            async () => {
+                if (this.client === undefined) {
+                    return false
+                }
+                await this.client.onReady()
+                return true
+            },
+            { interval: 500, timeout: 60_000 * 3, truthy: true }
+        )
+    }
 }
 /**
  * Activates the language server, this will start LSP server running over IPC protocol.
@@ -249,6 +315,37 @@ export async function activate(extensionContext: ExtensionContext) {
 
     let savedDocument: vscode.Uri | undefined = undefined
 
+    const onAdd = async (filePaths: string[]) => {
+        const indexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+        await LspClient.instance.updateIndex(filePaths, 'add')
+        await waitUntil(
+            async () => {
+                const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+                if (newIndexSeqNum > indexSeqNum) {
+                    await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`)
+                    return true
+                }
+                return false
+            },
+            { interval: 500, timeout: 5_000, truthy: true }
+        )
+    }
+    const onRemove = async (filePaths: string[]) => {
+        const indexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+        await LspClient.instance.updateIndex(filePaths, 'remove')
+        await waitUntil(
+            async () => {
+                const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+                if (newIndexSeqNum > indexSeqNum) {
+                    await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`)
+                    return true
+                }
+                return false
+            },
+            { interval: 500, timeout: 5_000, truthy: true }
+        )
+    }
+
     toDispose.push(
         vscode.workspace.onDidSaveTextDocument((document) => {
             if (document.uri.scheme !== 'file') {
@@ -260,18 +357,23 @@ export async function activate(extensionContext: ExtensionContext) {
             if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
                 void LspClient.instance.updateIndex([savedDocument.fsPath], 'update')
             }
+            // user created a new empty file using File -> New File
+            // these events will not be captured by vscode.workspace.onDidCreateFiles
+            // because it was created by File Explorer(Win) or Finder(MacOS)
+            // TODO: consider using a high performance fs watcher
+            if (editor?.document.getText().length === 0) {
+                void onAdd([editor.document.uri.fsPath])
+            }
         }),
-        vscode.workspace.onDidCreateFiles((e) => {
-            void LspClient.instance.updateIndex(
-                e.files.map((f) => f.fsPath),
-                'add'
-            )
+        vscode.workspace.onDidCreateFiles(async (e) => {
+            await onAdd(e.files.map((f) => f.fsPath))
         }),
-        vscode.workspace.onDidDeleteFiles((e) => {
-            void LspClient.instance.updateIndex(
-                e.files.map((f) => f.fsPath),
-                'remove'
-            )
+        vscode.workspace.onDidDeleteFiles(async (e) => {
+            await onRemove(e.files.map((f) => f.fsPath))
+        }),
+        vscode.workspace.onDidRenameFiles(async (e) => {
+            await onRemove(e.files.map((f) => f.oldUri.fsPath))
+            await onAdd(e.files.map((f) => f.newUri.fsPath))
         })
     )
 

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

@@ -13,7 +13,6 @@ import fetch from 'node-fetch'
 import request from '../../shared/request'
 import { LspClient } from './lspClient'
 import AdmZip from 'adm-zip'
-import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
 import { makeTemporaryToolkitFolder, tryRemoveFolder } from '../../shared/filesystemUtilities'
 import { activate as activateLsp } from './lspClient'
 import { telemetry } from '../../shared/telemetry/telemetry'
@@ -24,13 +23,16 @@ import { ToolkitError } from '../../shared/errors'
 import { isWeb } from '../../shared/extensionGlobals'
 import { getUserAgent } from '../../shared/telemetry/util'
 import { isAmazonInternalOs } from '../../shared/vscode/env'
+import { RelevantTextDocumentAddition } from '../../codewhispererChat/controllers/chat/model'
 
 export interface Chunk {
     readonly filePath: string
     readonly content: string
     readonly context?: string
     readonly relativePath?: string
     readonly programmingLanguage?: string
+    readonly startLine?: number
+    readonly endLine?: number
 }
 
 export interface Content {
@@ -60,7 +62,7 @@ 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.35']
+const supportedLspServerVersions = ['0.1.42']
 
 const nodeBinName = process.platform === 'win32' ? 'node.exe' : 'node'
 
@@ -279,9 +281,9 @@ export class LspController {
         }
     }
 
-    async query(s: string): Promise<RelevantTextDocument[]> {
+    async query(s: string): Promise<RelevantTextDocumentAddition[]> {
         const chunks: Chunk[] | undefined = await LspClient.instance.queryVectorIndex(s)
-        const resp: RelevantTextDocument[] = []
+        const resp: RelevantTextDocumentAddition[] = []
         if (chunks) {
             for (const chunk of chunks) {
                 const text = chunk.context ? chunk.context : chunk.content
@@ -292,11 +294,15 @@ export class LspController {
                         programmingLanguage: {
                             languageName: chunk.programmingLanguage,
                         },
+                        startLine: chunk.startLine ?? -1,
+                        endLine: chunk.endLine ?? -1,
                     })
                 } else {
                     resp.push({
                         text: text,
                         relativeFilePath: chunk.relativePath ? chunk.relativePath : path.basename(chunk.filePath),
+                        startLine: chunk.startLine ?? -1,
+                        endLine: chunk.endLine ?? -1,
                     })
                 }
             }
@@ -393,6 +399,7 @@ export class LspController {
             try {
                 await activateLsp(context)
                 getLogger().info('LspController: LSP activated')
+                await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`)
                 void LspController.instance.buildIndex(buildIndexConfig)
                 // log the LSP server CPU and Memory usage per 30 minutes.
                 globals.clock.setInterval(

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

@@ -76,3 +76,43 @@ export type GetRepomapIndexJSONRequest = string
 export const GetRepomapIndexJSONRequestType: RequestType<GetRepomapIndexJSONRequest, any, any> = new RequestType(
     'lsp/getRepomapIndexJSON'
 )
+
+export type GetContextCommandItemsRequestPayload = { workspaceFolders: string[] }
+export type GetContextCommandItemsRequest = string
+export const GetContextCommandItemsRequestType: RequestType<GetContextCommandItemsRequest, any, any> = new RequestType(
+    'lsp/getContextCommandItems'
+)
+
+export type GetIndexSequenceNumberRequest = string
+export const GetIndexSequenceNumberRequestType: RequestType<GetRepomapIndexJSONRequest, any, any> = new RequestType(
+    'lsp/getIndexSequenceNumber'
+)
+
+export type ContextCommandItemType = 'file' | 'folder'
+
+export interface ContextCommandItem {
+    workspaceFolder: string
+    type: ContextCommandItemType
+    relativePath: string
+}
+
+export type GetContextCommandPromptRequestPayload = {
+    contextCommands: {
+        workspaceFolder: string
+        type: 'file' | 'folder'
+        relativePath: string
+    }[]
+}
+export type GetContextCommandPromptRequest = string
+export const GetContextCommandPromptRequestType: RequestType<GetContextCommandPromptRequest, any, any> =
+    new RequestType('lsp/getContextCommandPrompt')
+
+export interface AdditionalContextPrompt {
+    content: string
+    name: string
+    description: string
+    startLine: number
+    endLine: number
+    filePath: string
+    relativePath: string
+}

packages/core/src/amazonq/webview/messages/messageDispatcher.ts

@@ -36,6 +36,8 @@ export function dispatchWebViewMessagesToApps(
                 })
                 performance.clearMarks(amazonqMark.uiReady)
                 performance.clearMarks(amazonqMark.open)
+                // let cwcController know the ui is ready
+                webViewToAppsMessagePublishers.get('cwc')?.publish(msg)
                 return
             }
             case 'start-chat-message-telemetry': {