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

## Problem
We can select @file and @folder in the Q chat menu. But we should be
able to also let user select class, function and these symbols.

## Solution
Support @code in the context selection of Q chat. 

<img width="742" alt="Screenshot 2025-03-21 at 2 15 34 PM"
src="https://github.com/user-attachments/assets/85f99420-831f-4683-bdfa-7a4cd441cf6a"
/>


---

- 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.

Author: leigaol

package-lock.json

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

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

@@ -58,6 +58,8 @@ describe('triggerPayloadToChatRequest', () => {
             relativePath: 'path-prompt',
             type: 'prompt',
             innerContext: createLargeString(size, 'prompt-'),
+            startLine: 0,
+            endLine: 100,
         }
     }
 
@@ -68,6 +70,8 @@ describe('triggerPayloadToChatRequest', () => {
             relativePath: 'path-rule',
             type: 'rule',
             innerContext: createLargeString(size, 'rule-'),
+            startLine: 0,
+            endLine: 100,
         }
     }
 
@@ -78,6 +82,8 @@ describe('triggerPayloadToChatRequest', () => {
             relativePath: 'path-file',
             type: 'file',
             innerContext: createLargeString(size, 'file-'),
+            startLine: 0,
+            endLine: 100,
         }
     }
 
@@ -116,13 +122,23 @@ describe('triggerPayloadToChatRequest', () => {
         const payloadWithEmptyContents: TriggerPayload = {
             ...mockBasicPayload,
             additionalContents: [
-                { name: 'prompt1', description: 'prompt1', relativePath: 'path1', type: 'prompt', innerContext: '' },
+                {
+                    name: 'prompt1',
+                    description: 'prompt1',
+                    relativePath: 'path1',
+                    type: 'prompt',
+                    innerContext: '',
+                    startLine: 0,
+                    endLine: 100,
+                },
                 {
                     name: 'prompt2',
                     description: 'prompt2',
                     relativePath: 'path2',
                     type: 'prompt',
                     innerContext: 'valid content',
+                    startLine: 0,
+                    endLine: 100,
                 },
             ],
         }

packages/amazonq/test/unit/codewhispererChat/controllers/chat/chatRequest/converter.test.ts

@@ -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",
         "@gerhobbelt/gitignore-parser": "^0.2.0-9",
         "@iarna/toml": "^2.2.5",
         "@smithy/fetch-http-handler": "^3.0.0",

packages/core/package.json

@@ -330,6 +330,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/package.nls.json

@@ -15,7 +15,7 @@ export interface LspConfig {
 
 export const defaultAmazonQWorkspaceLspConfig: LspConfig = {
     manifestUrl: 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json',
-    supportedVersions: '0.1.42',
+    supportedVersions: '0.1.46',
     id: 'AmazonQ-Workspace', // used for identification in global storage/local disk location. Do not change.
     path: undefined,
 }

packages/core/src/amazonq/lsp/config.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/lspClient.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 }
+            )
+        } catch (err) {
+            getLogger().error(`LspController: Failed to find symbols`)
+        }
+    }
 
     private async setupLsp(context: vscode.ExtensionContext) {
         await lspSetupStage('all', async () => {

packages/core/src/amazonq/lsp/lspController.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/amazonq/lsp/types.ts

@@ -447,6 +447,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: [
@@ -471,7 +482,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 {
@@ -514,22 +526,40 @@ 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,
                     })
                 }
+                // TODO: Remove the limit of 25k once the performance issue of mynahUI in webview is fixed.
+                else if (
+                    contextCommandItem.symbol &&
+                    symbolsCmd.children &&
+                    symbolsCmd.children[0].commands.length < 25_000
+                ) {
+                    symbolsCmd.children?.[0].commands.push({
+                        command: contextCommandItem.symbol.name,
+                        description: `${contextCommandItem.symbol.kind}, ${path.join(wsFolderName, 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) {
@@ -951,13 +981,13 @@ export class ChatController {
         }
         triggerPayload.workspaceRulesCount = workspaceRules.length
 
-        // Add context commands added by user to context
         for (const context of triggerPayload.context) {
             if (typeof context !== 'string' && context.route && context.route.length === 2) {
                 contextCommands.push({
                     workspaceFolder: context.route[0] || '',
-                    type: context.icon === 'folder' ? 'folder' : 'file',
+                    type: (context.label || '') as ContextCommandItemType,
                     relativePath: context.route[1] || '',
+                    id: context.id,
                 })
             }
         }
@@ -972,11 +1002,7 @@ export class ChatController {
         workspaceFolders.sort()
         const workspaceFolder = workspaceFolders[0]
         for (const contextCommand of contextCommands) {
-            const relativePath = path.relative(
-                workspaceFolder,
-                path.join(contextCommand.workspaceFolder, contextCommand.relativePath)
-            )
-            session.relativePathToWorkspaceRoot.set(relativePath, contextCommand.workspaceFolder)
+            session.relativePathToWorkspaceRoot.set(contextCommand.workspaceFolder, contextCommand.workspaceFolder)
         }
         let prompts: AdditionalContextPrompt[] = []
         try {
@@ -1006,6 +1032,8 @@ export class ChatController {
                 innerContext: prompt.content.substring(0, additionalContentInnerContextLimit),
                 type: contextType,
                 relativePath: relativePath,
+                startLine: prompt.startLine,
+                endLine: prompt.endLine,
             }
 
             triggerPayload.additionalContents.push(entry)
@@ -1094,7 +1122,10 @@ export class ChatController {
             if (!relativePathsOfMergedRelevantDocuments.includes(relativePath) && !seen.includes(relativePath)) {
                 triggerPayload.documentReferences.push({
                     relativeFilePath: relativePath,
-                    lineRanges: [{ first: -1, second: -1 }],
+                    lineRanges:
+                        additionalContent.name === 'symbol'
+                            ? [{ first: additionalContent.startLine, second: additionalContent.endLine }]
+                            : [{ first: -1, second: -1 }],
                 })
                 seen.push(relativePath)
             }