Skip to content

Commit 011654d

Browse files
leigaolleigaol
authored
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.
1 parent c238710 commit 011654d

File tree

11 files changed

+144
-22
lines changed

11 files changed

+144
-22
lines changed

‎package-lock.json

Lines changed: 4 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "Add support for Code search in Q chat"
4+
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@ describe('triggerPayloadToChatRequest', () => {
5858
relativePath: 'path-prompt',
5959
type: 'prompt',
6060
innerContext: createLargeString(size, 'prompt-'),
61+
startLine: 0,
62+
endLine: 100,
6163
}
6264
}
6365

@@ -68,6 +70,8 @@ describe('triggerPayloadToChatRequest', () => {
6870
relativePath: 'path-rule',
6971
type: 'rule',
7072
innerContext: createLargeString(size, 'rule-'),
73+
startLine: 0,
74+
endLine: 100,
7175
}
7276
}
7377

@@ -78,6 +82,8 @@ describe('triggerPayloadToChatRequest', () => {
7882
relativePath: 'path-file',
7983
type: 'file',
8084
innerContext: createLargeString(size, 'file-'),
85+
startLine: 0,
86+
endLine: 100,
8187
}
8288
}
8389

@@ -116,13 +122,23 @@ describe('triggerPayloadToChatRequest', () => {
116122
const payloadWithEmptyContents: TriggerPayload = {
117123
...mockBasicPayload,
118124
additionalContents: [
119-
{ name: 'prompt1', description: 'prompt1', relativePath: 'path1', type: 'prompt', innerContext: '' },
125+
{
126+
name: 'prompt1',
127+
description: 'prompt1',
128+
relativePath: 'path1',
129+
type: 'prompt',
130+
innerContext: '',
131+
startLine: 0,
132+
endLine: 100,
133+
},
120134
{
121135
name: 'prompt2',
122136
description: 'prompt2',
123137
relativePath: 'path2',
124138
type: 'prompt',
125139
innerContext: 'valid content',
140+
startLine: 0,
141+
endLine: 100,
126142
},
127143
],
128144
}

‎packages/core/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -522,7 +522,7 @@
522522
"@aws-sdk/s3-request-presigner": "<3.696.0",
523523
"@aws-sdk/smithy-client": "<3.696.0",
524524
"@aws-sdk/util-arn-parser": "<3.696.0",
525-
"@aws/mynah-ui": "^4.25.1",
525+
"@aws/mynah-ui": "^4.26.0",
526526
"@gerhobbelt/gitignore-parser": "^0.2.0-9",
527527
"@iarna/toml": "^2.2.5",
528528
"@smithy/fetch-http-handler": "^3.0.0",

‎packages/core/package.nls.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,8 @@
330330
"AWS.amazonq.context.files.description": "Add a file to context",
331331
"AWS.amazonq.context.prompts.title": "Prompts",
332332
"AWS.amazonq.context.prompts.description": "Add a saved prompt to context",
333+
"AWS.amazonq.context.code.title": "Code",
334+
"AWS.amazonq.context.code.description": "Add code to context",
333335
"AWS.amazonq.savedPrompts.title": "Prompt name",
334336
"AWS.amazonq.savedPrompts.create": "Create",
335337
"AWS.amazonq.savedPrompts.action": "Create a new prompt",

‎packages/core/src/amazonq/lsp/config.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ export interface LspConfig {
1515

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

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ export class LspClient {
122122
}
123123
}
124124

125-
async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add') {
125+
async updateIndex(filePath: string[], mode: 'update' | 'remove' | 'add' | 'context_command_symbol_update') {
126126
const payload: UpdateIndexV2RequestPayload = {
127127
filePaths: filePath,
128128
updateMode: mode,

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import { isAmazonInternalOs } from '../../shared/vscode/env'
1515
import { WorkspaceLspInstaller } from './workspaceInstaller'
1616
import { lspSetupStage } from '../../shared/lsp/utils/setupStage'
1717
import { RelevantTextDocumentAddition } from '../../codewhispererChat/controllers/chat/model'
18+
import { waitUntil } from '../../shared/utilities/timeoutUtils'
1819

1920
export interface Chunk {
2021
readonly filePath: string
@@ -45,6 +46,7 @@ export interface BuildIndexConfig {
4546
export class LspController {
4647
static #instance: LspController
4748
private _isIndexingInProgress = false
49+
private _contextCommandSymbolsUpdated = false
4850
private logger = getLogger('amazonqWorkspaceLsp')
4951

5052
public static get instance() {
@@ -192,6 +194,38 @@ export class LspController {
192194
}
193195
})
194196
}
197+
/**
198+
* Updates context command symbols once per session by synchronizing with the LSP client index.
199+
* Context menu will contain file and folders to begin with,
200+
* then this asynchronous function should be invoked after the files and folders are found
201+
* the LSP then further starts to parse workspace and find symbols, which takes
202+
* anywhere from 5 seconds to about 40 seconds, depending on project size.
203+
* @returns {Promise<void>}
204+
*/
205+
async updateContextCommandSymbolsOnce() {
206+
if (this._contextCommandSymbolsUpdated) {
207+
return
208+
}
209+
this._contextCommandSymbolsUpdated = true
210+
getLogger().debug(`LspController: Start adding symbols to context picker menu`)
211+
try {
212+
const indexSeqNum = await LspClient.instance.getIndexSequenceNumber()
213+
await LspClient.instance.updateIndex([], 'context_command_symbol_update')
214+
await waitUntil(
215+
async () => {
216+
const newIndexSeqNum = await LspClient.instance.getIndexSequenceNumber()
217+
if (newIndexSeqNum > indexSeqNum) {
218+
await vscode.commands.executeCommand(`aws.amazonq.updateContextCommandItems`)
219+
return true
220+
}
221+
return false
222+
},
223+
{ interval: 1000, timeout: 60_000, truthy: true }
224+
)
225+
} catch (err) {
226+
getLogger().error(`LspController: Failed to find symbols`)
227+
}
228+
}
195229

196230
private async setupLsp(context: vscode.ExtensionContext) {
197231
await lspSetupStage('all', async () => {

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,44 @@ export const GetIndexSequenceNumberRequestType: RequestType<GetRepomapIndexJSONR
8888
'lsp/getIndexSequenceNumber'
8989
)
9090

91-
export type ContextCommandItemType = 'file' | 'folder'
91+
export type ContextCommandItemType = 'file' | 'folder' | 'code'
92+
93+
export type SymbolType =
94+
| 'Class'
95+
| 'Function'
96+
| 'Interface'
97+
| 'Type'
98+
| 'Enum'
99+
| 'Struct'
100+
| 'Delegate'
101+
| 'Namespace'
102+
| 'Object'
103+
| 'Module'
104+
| 'Method'
105+
106+
export interface Position {
107+
line: number
108+
column: number
109+
}
110+
export interface Span {
111+
start: Position
112+
end: Position
113+
}
114+
115+
// LSP definition of DocumentSymbol
116+
117+
export interface DocumentSymbol {
118+
name: string
119+
kind: SymbolType
120+
range: Span
121+
}
92122

93123
export interface ContextCommandItem {
94124
workspaceFolder: string
95125
type: ContextCommandItemType
96126
relativePath: string
127+
symbol?: DocumentSymbol
128+
id?: string
97129
}
98130

99131
export type GetContextCommandPromptRequestPayload = {

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

Lines changed: 43 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -447,6 +447,17 @@ export class ChatController {
447447
description: i18n('AWS.amazonq.context.files.description'),
448448
icon: 'file' as MynahIconsType,
449449
},
450+
{
451+
command: i18n('AWS.amazonq.context.code.title'),
452+
children: [
453+
{
454+
groupName: i18n('AWS.amazonq.context.code.title'),
455+
commands: [],
456+
},
457+
],
458+
description: i18n('AWS.amazonq.context.code.description'),
459+
icon: 'code-block' as MynahIconsType,
460+
},
450461
{
451462
command: i18n('AWS.amazonq.context.prompts.title'),
452463
children: [
@@ -471,7 +482,8 @@ export class ChatController {
471482
commands: [{ command: commandName, description: commandDescription }],
472483
})
473484
}
474-
const promptsCmd: QuickActionCommand = contextCommand[0].commands?.[3]
485+
const symbolsCmd: QuickActionCommand = contextCommand[0].commands?.[3]
486+
const promptsCmd: QuickActionCommand = contextCommand[0].commands?.[4]
475487

476488
// Check for user prompts
477489
try {
@@ -514,22 +526,40 @@ export class ChatController {
514526
command: path.basename(contextCommandItem.relativePath),
515527
description: path.join(wsFolderName, contextCommandItem.relativePath),
516528
route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
517-
id: 'file',
529+
label: 'file' as ContextCommandItemType,
530+
id: contextCommandItem.id,
518531
icon: 'file' as MynahIconsType,
519532
})
520-
} else {
533+
} else if (contextCommandItem.type === 'folder') {
521534
folderCmd.children?.[0].commands.push({
522535
command: path.basename(contextCommandItem.relativePath),
523536
description: path.join(wsFolderName, contextCommandItem.relativePath),
524537
route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
525-
id: 'folder',
538+
label: 'folder' as ContextCommandItemType,
539+
id: contextCommandItem.id,
526540
icon: 'folder' as MynahIconsType,
527541
})
528542
}
543+
// TODO: Remove the limit of 25k once the performance issue of mynahUI in webview is fixed.
544+
else if (
545+
contextCommandItem.symbol &&
546+
symbolsCmd.children &&
547+
symbolsCmd.children[0].commands.length < 25_000
548+
) {
549+
symbolsCmd.children?.[0].commands.push({
550+
command: contextCommandItem.symbol.name,
551+
description: `${contextCommandItem.symbol.kind}, ${path.join(wsFolderName, contextCommandItem.relativePath)}, L${contextCommandItem.symbol.range.start.line}-${contextCommandItem.symbol.range.end.line}`,
552+
route: [contextCommandItem.workspaceFolder, contextCommandItem.relativePath],
553+
label: 'code' as ContextCommandItemType,
554+
id: contextCommandItem.id,
555+
icon: 'code-block' as MynahIconsType,
556+
})
557+
}
529558
}
530559
}
531560

532561
this.messenger.sendContextCommandData(contextCommand)
562+
void LspController.instance.updateContextCommandSymbolsOnce()
533563
}
534564

535565
private handlePromptCreate(tabID: string) {
@@ -951,13 +981,13 @@ export class ChatController {
951981
}
952982
triggerPayload.workspaceRulesCount = workspaceRules.length
953983

954-
// Add context commands added by user to context
955984
for (const context of triggerPayload.context) {
956985
if (typeof context !== 'string' && context.route && context.route.length === 2) {
957986
contextCommands.push({
958987
workspaceFolder: context.route[0] || '',
959-
type: context.icon === 'folder' ? 'folder' : 'file',
988+
type: (context.label || '') as ContextCommandItemType,
960989
relativePath: context.route[1] || '',
990+
id: context.id,
961991
})
962992
}
963993
}
@@ -972,11 +1002,7 @@ export class ChatController {
9721002
workspaceFolders.sort()
9731003
const workspaceFolder = workspaceFolders[0]
9741004
for (const contextCommand of contextCommands) {
975-
const relativePath = path.relative(
976-
workspaceFolder,
977-
path.join(contextCommand.workspaceFolder, contextCommand.relativePath)
978-
)
979-
session.relativePathToWorkspaceRoot.set(relativePath, contextCommand.workspaceFolder)
1005+
session.relativePathToWorkspaceRoot.set(contextCommand.workspaceFolder, contextCommand.workspaceFolder)
9801006
}
9811007
let prompts: AdditionalContextPrompt[] = []
9821008
try {
@@ -1006,6 +1032,8 @@ export class ChatController {
10061032
innerContext: prompt.content.substring(0, additionalContentInnerContextLimit),
10071033
type: contextType,
10081034
relativePath: relativePath,
1035+
startLine: prompt.startLine,
1036+
endLine: prompt.endLine,
10091037
}
10101038

10111039
triggerPayload.additionalContents.push(entry)
@@ -1094,7 +1122,10 @@ export class ChatController {
10941122
if (!relativePathsOfMergedRelevantDocuments.includes(relativePath) && !seen.includes(relativePath)) {
10951123
triggerPayload.documentReferences.push({
10961124
relativeFilePath: relativePath,
1097-
lineRanges: [{ first: -1, second: -1 }],
1125+
lineRanges:
1126+
additionalContent.name === 'symbol'
1127+
? [{ first: additionalContent.startLine, second: additionalContent.endLine }]
1128+
: [{ first: -1, second: -1 }],
10981129
})
10991130
seen.push(relativePath)
11001131
}

0 commit comments

Comments
 (0)