feat(amazonq): Add code symbol support in context selection of Q Chat

packages/amazonq/.changes/next-release/Feature-5f829ea1-e33b-4687-bc6d-76ed1558cd7e.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Add support for Code search in Q chat"
+}

packages/core/package.json

@@ -522,7 +522,7 @@
         "@aws-sdk/s3-request-presigner": "<3.696.0",
         "@aws-sdk/smithy-client": "<3.696.0",
         "@aws-sdk/util-arn-parser": "<3.696.0",
-        "@aws/mynah-ui": "^4.25.1",
+        "@aws/mynah-ui": "^4.26.0-beta.5",
         "@gerhobbelt/gitignore-parser": "^0.2.0-9",
         "@iarna/toml": "^2.2.5",
         "@smithy/fetch-http-handler": "^3.0.0",

packages/core/package.nls.json

@@ -329,6 +329,8 @@
     "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.context.code.title": "Code",
+    "AWS.amazonq.context.code.description": "Add code to context",
     "AWS.amazonq.savedPrompts.title": "Prompt name",
     "AWS.amazonq.savedPrompts.create": "Create",
     "AWS.amazonq.savedPrompts.action": "Create a new prompt",

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

@@ -122,7 +122,7 @@ export class LspClient {
         }
     }
 
-    async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add') {
+    async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add' | 'context_command_symbol_update') {
         const payload: UpdateIndexV2RequestPayload = {
             filePaths: filePath,
             updateMode: mode,

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

@@ -15,6 +15,7 @@ import { isAmazonInternalOs } from '../../shared/vscode/env'
 import { WorkspaceLspInstaller } from './workspaceInstaller'
 import { lspSetupStage } from '../../shared/lsp/utils/setupStage'
 import { RelevantTextDocumentAddition } from '../../codewhispererChat/controllers/chat/model'
+import { waitUntil } from '../../shared/utilities/timeoutUtils'
 
 export interface Chunk {
     readonly filePath: string
@@ -45,6 +46,7 @@ export interface BuildIndexConfig {
 export class LspController {
     static #instance: LspController
     private _isIndexingInProgress = false
+    private _contextCommandSymbolsUpdated = false
     private logger = getLogger('amazonqWorkspaceLsp')
 
     public static get instance() {
@@ -192,6 +194,38 @@ export class LspController {
             }
         })
     }
+    /**
+     * Updates context command symbols once per session by synchronizing with the LSP client index.
+     * Context menu will contain file and folders to begin with,
+     * then this asynchronous function should be invoked after the files and folders are found
+     * the LSP then further starts to parse workspace and find symbols, which takes
+     * anywhere from 5 seconds to about 40 seconds, depending on project size.
+     * @returns {Promise<void>}
+     */
+    async updateContextCommandSymbolsOnce() {
+        if (this._contextCommandSymbolsUpdated) {
+            return
+        }
+        this._contextCommandSymbolsUpdated = true
+        getLogger().debug(`LspController: Start adding symbols to context picker menu`)
+        try {
+            const indexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+            await LspClient.instance.updateIndex([], 'context_command_symbol_update')
+            await waitUntil(
+                async () => {
+                    const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber()
+                    if (newIndexSeqNum > indexSeqNum) {
+                        await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`)
+                        return true
+                    }
+                    return false
+                },
+                { interval: 1000, timeout: 60_000, truthy: true }
+            )
+        } finally {
+            this._contextCommandSymbolsUpdated = false
+        }
+    }
 
     private async setupLsp(context: vscode.ExtensionContext) {
         await lspSetupStage('all', async () => {

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

@@ -88,12 +88,44 @@ export const GetIndexSequenceNumberRequestType: RequestType<GetRepomapIndexJSONR
     'lsp/getIndexSequenceNumber'
 )
 
-export type ContextCommandItemType = 'file' | 'folder'
+export type ContextCommandItemType = 'file' | 'folder' | 'code'
+
+export type SymbolType =
+    | 'Class'
+    | 'Function'
+    | 'Interface'
+    | 'Type'
+    | 'Enum'
+    | 'Struct'
+    | 'Delegate'
+    | 'Namespace'
+    | 'Object'
+    | 'Module'
+    | 'Method'
+
+export interface Position {
+    line: number
+    column: number
+}
+export interface Span {
+    start: Position
+    end: Position
+}
+
+// LSP definition of DocumentSymbol
+
+export interface DocumentSymbol {
+    name: string
+    kind: SymbolType
+    range: Span
+}
 
 export interface ContextCommandItem {
     workspaceFolder: string
     type: ContextCommandItemType
     relativePath: string
+    symbol?: DocumentSymbol
+    id?: string
 }
 
 export type GetContextCommandPromptRequestPayload = {

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -446,6 +446,17 @@ export class ChatController {
                         description: i18n('AWS.amazonq.context.files.description'),
                         icon: 'file' as MynahIconsType,
                     },
+                    {
+                        command: i18n('AWS.amazonq.context.code.title'),
+                        children: [
+                            {
+                                groupName: i18n('AWS.amazonq.context.code.title'),
+                                commands: [],
+                            },
+                        ],
+                        description: i18n('AWS.amazonq.context.code.description'),
+                        icon: 'code-block' as MynahIconsType,
+                    },
                     {
                         command: i18n('AWS.amazonq.context.prompts.title'),
                         children: [
@@ -470,7 +481,8 @@ export class ChatController {
                 commands: [{ command: commandName, description: commandDescription }],
             })
         }
-        const promptsCmd: QuickActionCommand = contextCommand[0].commands?.[3]
+        const symbolsCmd: QuickActionCommand = contextCommand[0].commands?.[3]
+        const promptsCmd: QuickActionCommand = contextCommand[0].commands?.[4]
 
         // Check for user prompts
         try {
@@ -513,22 +525,34 @@ export class ChatController {
                         command: path.basename(contextCommandItem.relativePath),
                         description: path.join(wsFolderName, contextCommandItem.relativePath),
                         route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
-                        id: 'file',
+                        label: 'file' as ContextCommandItemType,
+                        id: contextCommandItem.id,
                         icon: 'file' as MynahIconsType,
                     })
-                } else {
+                } else if (contextCommandItem.type === 'folder') {
                     folderCmd.children?.[0].commands.push({
                         command: path.basename(contextCommandItem.relativePath),
                         description: path.join(wsFolderName, contextCommandItem.relativePath),
                         route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
-                        id: 'folder',
+                        label: 'folder' as ContextCommandItemType,
+                        id: contextCommandItem.id,
                         icon: 'folder' as MynahIconsType,
                     })
+                } else if (contextCommandItem.symbol) {
+                    symbolsCmd.children?.[0].commands.push({
+                        command: contextCommandItem.symbol.name,
+                        description: `${contextCommandItem.symbol.kind}, ${contextCommandItem.relativePath}, L${contextCommandItem.symbol.range.start.line}-${contextCommandItem.symbol.range.end.line}`,
+                        route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
+                        label: 'code' as ContextCommandItemType,
+                        id: contextCommandItem.id,
+                        icon: 'code-block' as MynahIconsType,
+                    })
                 }
             }
         }
 
         this.messenger.sendContextCommandData(contextCommand)
+        void LspController.instance.updateContextCommandSymbolsOnce()
     }
 
     private handlePromptCreate(tabID: string) {
@@ -910,9 +934,9 @@ export class ChatController {
     private async resolveContextCommandPayload(
         triggerPayload: TriggerPayload,
         session: ChatSession
-    ): Promise<string[]> {
+    ): Promise<DocumentReference[]> {
         const contextCommands: ContextCommandItem[] = []
-        const relativePaths: string[] = []
+        const result: DocumentReference[] = []
 
         // Check for workspace rules to add to context
         const workspaceRules = await this.collectWorkspaceRules()
@@ -935,10 +959,12 @@ export class ChatController {
         if (triggerPayload.context !== undefined && triggerPayload.context.length > 0) {
             for (const context of triggerPayload.context) {
                 if (typeof context !== 'string' && context.route && context.route.length === 2) {
+                    const itemType = (context.label || '') as ContextCommandItemType
                     contextCommands.push({
                         workspaceFolder: context.route[0] || '',
-                        type: context.icon === 'folder' ? 'folder' : 'file',
+                        type: itemType,
                         relativePath: context.route[1] || '',
+                        id: context.id,
                     })
                 }
             }
@@ -990,6 +1016,8 @@ export class ChatController {
                     name: prompt.name.substring(0, aditionalContentNameLimit),
                     description: description.substring(0, aditionalContentNameLimit),
                     innerContext: prompt.content.substring(0, additionalContentInnerContextLimit),
+                    start: prompt.startLine,
+                    end: prompt.endLine,
                 }
                 // make sure the relevantDocument + additionalContext
                 // combined does not exceed 40k characters before generating the request payload.
@@ -1015,12 +1043,15 @@ export class ChatController {
                 if (prompt.filePath.startsWith(getUserPromptsDirectory())) {
                     relativePath = path.basename(prompt.filePath)
                 }
-                relativePaths.push(relativePath)
+                result.push({
+                    relativeFilePath: relativePath,
+                    lineRanges: [{ first: entry.start, second: entry.end }],
+                })
             }
         }
         getLogger().info(`Retrieved chunks of additional context count: ${triggerPayload.additionalContents.length} `)
 
-        return relativePaths
+        return result
     }
 
     private async generateResponse(
@@ -1057,7 +1088,7 @@ export class ChatController {
         }
 
         const session = this.sessionStorage.getSession(tabID)
-        const relativePathsOfContextCommandFiles = await this.resolveContextCommandPayload(triggerPayload, session)
+        const contextCommandDocumentReferences = await this.resolveContextCommandPayload(triggerPayload, session)
         triggerPayload.useRelevantDocuments =
             triggerPayload.context?.some(
                 (context) => typeof context !== 'string' && context.command === '@workspace'
@@ -1105,24 +1136,22 @@ export class ChatController {
 
         const request = triggerPayloadToChatRequest(triggerPayload)
 
-        if (triggerPayload.documentReferences !== undefined) {
-            const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map(
-                (doc) => doc.relativeFilePath
-            )
-            const seen: string[] = []
-            for (const relativePath of relativePathsOfContextCommandFiles) {
-                if (!relativePathsOfMergedRelevantDocuments.includes(relativePath) && !seen.includes(relativePath)) {
-                    triggerPayload.documentReferences.push({
-                        relativeFilePath: relativePath,
-                        lineRanges: [{ first: -1, second: -1 }],
-                    })
-                    seen.push(relativePath)
-                }
+        const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map(
+            (doc) => doc.relativeFilePath
+        )
+        const seen: string[] = []
+        for (const documentReference of contextCommandDocumentReferences) {
+            if (
+                !relativePathsOfMergedRelevantDocuments.includes(documentReference.relativeFilePath) &&
+                !seen.includes(documentReference.relativeFilePath)
+            ) {
+                triggerPayload.documentReferences.push(documentReference)
+                seen.push(documentReference.relativeFilePath)
             }
-            if (triggerPayload.documentReferences) {
-                for (const doc of triggerPayload.documentReferences) {
-                    session.contexts.set(doc.relativeFilePath, doc.lineRanges)
-                }
+        }
+        if (triggerPayload.documentReferences) {
+            for (const doc of triggerPayload.documentReferences) {
+                session.contexts.set(doc.relativeFilePath, doc.lineRanges)
             }
         }