Skip to content

Commit 9757bec

Browse files
laileni-awslaileni-aws
committed
Adding comments and adding list directories change
1 parent 234e0c5 commit 9757bec

File tree

5 files changed

+124
-84
lines changed

5 files changed

+124
-84
lines changed

‎packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -392,7 +392,6 @@ export class Connector extends BaseConnector {
392392
break
393393
case 'run-shell-command':
394394
answer.header = {
395-
icon: 'shell' as MynahIconsType,
396395
body: 'shell',
397396
status: {
398397
icon: 'ok' as MynahIconsType,
@@ -403,7 +402,6 @@ export class Connector extends BaseConnector {
403402
break
404403
case 'reject-shell-command':
405404
answer.header = {
406-
icon: 'shell' as MynahIconsType,
407405
body: 'shell',
408406
status: {
409407
icon: 'cancel' as MynahIconsType,

‎packages/core/src/codewhispererChat/clients/chat/v0/chat.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ export class ChatSession {
3333
* _messageIdToUpdate = messageId of a chat message to be updated, used for reducing consecutive tool messages
3434
*/
3535
private _readFiles: DocumentReference[] = []
36+
private _readFolders: DocumentReference[] = []
3637
private _toolUseWithError: ToolUseWithError | undefined
3738
private _showDiffOnFileWrite: boolean = false
3839
private _context: PromptMessage['context']
@@ -43,6 +44,7 @@ export class ChatSession {
4344
*/
4445
localHistoryHydrated: boolean = false
4546
private _messageIdToUpdate: string | undefined
47+
private _messageIdToUpdateListDirectory: string | undefined
4648

4749
contexts: Map<string, { first: number; second: number }[]> = new Map()
4850
// TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -59,6 +61,14 @@ export class ChatSession {
5961
this._messageIdToUpdate = messageId
6062
}
6163

64+
public get messageIdToUpdateListDirectory(): string | undefined {
65+
return this._messageIdToUpdateListDirectory
66+
}
67+
68+
public setMessageIdToUpdateListDirectory(messageId: string | undefined) {
69+
this._messageIdToUpdateListDirectory = messageId
70+
}
71+
6272
public get pairProgrammingModeOn(): boolean {
6373
return this._pairProgrammingModeOn
6474
}
@@ -107,6 +117,9 @@ export class ChatSession {
107117
public get readFiles(): DocumentReference[] {
108118
return this._readFiles
109119
}
120+
public get readFolders(): DocumentReference[] {
121+
return this._readFolders
122+
}
110123
public get showDiffOnFileWrite(): boolean {
111124
return this._showDiffOnFileWrite
112125
}
@@ -119,6 +132,12 @@ export class ChatSession {
119132
public clearListOfReadFiles() {
120133
this._readFiles = []
121134
}
135+
public setReadFolders(folder: DocumentReference) {
136+
this._readFolders.push(folder)
137+
}
138+
public clearListOfReadFolders() {
139+
this._readFolders = []
140+
}
122141
async chatIam(chatRequest: SendMessageRequest): Promise<SendMessageCommandOutput> {
123142
const client = await createQDeveloperStreamingClient()
124143

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1178,6 +1178,7 @@ export class ChatController {
11781178
private async processPromptMessageAsNewThread(message: PromptMessage) {
11791179
const session = this.sessionStorage.getSession(message.tabID)
11801180
session.clearListOfReadFiles()
1181+
session.clearListOfReadFolders()
11811182
session.setShowDiffOnFileWrite(false)
11821183
this.editorContextExtractor
11831184
.extractContextForTrigger('ChatMessage')

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

Lines changed: 97 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector
7171
import { localize } from '../../../../shared/utilities/vsCodeUtils'
7272
import { getDiffLinesFromChanges } from '../../../../shared/utilities/diffUtils'
7373
import { FsReadParams } from '../../../tools/fsRead'
74+
import { ListDirectoryParams } from '../../../tools/listDirectory'
7475

7576
export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
7677

@@ -299,6 +300,12 @@ export class Messenger {
299300
relativeFilePath: input?.path,
300301
lineRanges: [{ first: start, second: end }],
301302
})
303+
} else if (tool.type === ToolType.ListDirectory) {
304+
const input = toolUse.input as unknown as ListDirectoryParams
305+
session.setReadFolders({
306+
relativeFilePath: input?.path,
307+
lineRanges: [{ first: -1, second: -1 }],
308+
})
302309
}
303310
const validation = ToolUtils.requiresAcceptance(tool)
304311
const chatStream = new ChatStream(
@@ -307,7 +314,9 @@ export class Messenger {
307314
triggerID,
308315
toolUse,
309316
session,
310-
session.messageIdToUpdate,
317+
tool.type === ToolType.FsRead
318+
? session.messageIdToUpdate
319+
: session.messageIdToUpdateListDirectory,
311320
true,
312321
validation,
313322
changeList
@@ -318,6 +327,13 @@ export class Messenger {
318327
session.setMessageIdToUpdate(toolUse.toolUseId)
319328
}
320329

330+
if (
331+
session.messageIdToUpdateListDirectory === undefined &&
332+
tool.type === ToolType.ListDirectory
333+
) {
334+
session.setMessageIdToUpdateListDirectory(toolUse.toolUseId)
335+
}
336+
321337
if (!validation.requiresAcceptance) {
322338
// Need separate id for read tool and safe bash command execution as 'run-shell-command' id is required to state in cwChatConnector.ts which will impact generic tool execution.
323339
if (tool.type === ToolType.ExecuteBash) {
@@ -568,6 +584,50 @@ export class Messenger {
568584
)
569585
}
570586

587+
private sendReadAndListDirToolMessage(
588+
toolUse: ToolUse,
589+
session: ChatSession,
590+
tabID: string,
591+
triggerID: string,
592+
messageIdToUpdate?: string
593+
) {
594+
const contextList = toolUse.name === ToolType.ListDirectory ? session.readFolders : session.readFiles
595+
const isFileRead = toolUse.name === ToolType.FsRead
596+
const items = isFileRead ? session.readFiles : session.readFolders
597+
const itemCount = items.length
598+
599+
const title =
600+
itemCount < 1
601+
? 'Gathering context'
602+
: isFileRead
603+
? `${itemCount} file${itemCount > 1 ? 's' : ''} read`
604+
: `${itemCount} ${itemCount === 1 ? 'directory' : 'directories'} listed`
605+
606+
this.dispatcher.sendToolMessage(
607+
new ToolMessage(
608+
{
609+
message: '',
610+
messageType: 'answer-part',
611+
followUps: undefined,
612+
followUpsHeader: undefined,
613+
relatedSuggestions: undefined,
614+
triggerID,
615+
messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '',
616+
userIntent: undefined,
617+
codeBlockLanguage: undefined,
618+
contextList,
619+
canBeVoted: false,
620+
buttons: undefined,
621+
fullWidth: false,
622+
padding: false,
623+
codeBlockActions: undefined,
624+
title,
625+
},
626+
tabID
627+
)
628+
)
629+
}
630+
571631
public sendPartialToolLog(
572632
message: string,
573633
tabID: string,
@@ -578,40 +638,38 @@ export class Messenger {
578638
validation: CommandValidation,
579639
changeList?: Change[]
580640
) {
641+
// Handle read tool and list directory messages
642+
if (toolUse?.name === ToolType.FsRead || toolUse?.name === ToolType.ListDirectory) {
643+
return this.sendReadAndListDirToolMessage(toolUse, session, tabID, triggerID, messageIdToUpdate)
644+
}
645+
646+
// Handle file write tool, execute bash tool and bash command output log messages
581647
const buttons: ChatItemButton[] = []
582648
let header: ChatItemHeader | undefined = undefined
583-
let fullWidth: boolean | undefined = undefined
584-
let padding: boolean | undefined = undefined
585-
let codeBlockActions: ChatItemContent['codeBlockActions'] = undefined
586649
if (toolUse?.name === ToolType.ExecuteBash && message.startsWith('```shell')) {
587650
if (validation.requiresAcceptance) {
588651
const buttons: ChatItemButton[] = [
589-
{
590-
id: 'run-shell-command',
591-
text: localize('AWS.amazonq.executeBash.run', 'Run'),
592-
status: 'main',
593-
icon: 'play' as MynahIconsType,
594-
},
595652
{
596653
id: 'reject-shell-command',
597654
text: localize('AWS.amazonq.executeBash.reject', 'Reject'),
598655
status: 'clear',
599656
icon: 'cancel' as MynahIconsType,
600657
},
658+
{
659+
id: 'run-shell-command',
660+
text: localize('AWS.amazonq.executeBash.run', 'Run'),
661+
status: 'clear',
662+
icon: 'play' as MynahIconsType,
663+
},
601664
]
602665
header = {
603-
icon: 'shell' as MynahIconsType,
604666
body: 'shell',
605667
buttons,
606668
}
607669
}
608670
if (validation.warning) {
609671
message = validation.warning + message
610672
}
611-
fullWidth = true
612-
padding = false
613-
// eslint-disable-next-line unicorn/no-null
614-
codeBlockActions = { 'insert-to-cursor': null, copy: null }
615673
} else if (toolUse?.name === ToolType.FsWrite) {
616674
const input = toolUse.input as unknown as FsWriteParams
617675
const fileName = path.basename(input.path)
@@ -641,69 +699,36 @@ export class Messenger {
641699
},
642700
]
643701
header = {
644-
icon: 'code-block' as MynahIconsType,
645702
buttons,
646703
fileList,
647704
}
648-
fullWidth = true
649-
padding = false
650-
// eslint-disable-next-line unicorn/no-null
651-
codeBlockActions = { 'insert-to-cursor': null, copy: null }
652705
}
653706

654-
if (toolUse?.name === ToolType.FsRead) {
655-
this.dispatcher.sendToolMessage(
656-
new ToolMessage(
657-
{
658-
message: '',
659-
messageType: 'answer-part',
660-
followUps: undefined,
661-
followUpsHeader: undefined,
662-
relatedSuggestions: undefined,
663-
triggerID,
664-
messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '',
665-
userIntent: undefined,
666-
codeBlockLanguage: undefined,
667-
contextList: session.readFiles,
668-
canBeVoted: false,
669-
buttons,
670-
fullWidth: false,
671-
padding: true,
672-
codeBlockActions: undefined,
673-
title:
674-
session.readFiles.length > 1
675-
? `Files Read ${session.readFiles.length} files`
676-
: `File Read 1 file`,
677-
},
678-
tabID
679-
)
680-
)
681-
} else {
682-
this.dispatcher.sendChatMessage(
683-
new ChatMessage(
684-
{
685-
message: message,
686-
messageType: 'answer-part',
687-
followUps: undefined,
688-
followUpsHeader: undefined,
689-
relatedSuggestions: undefined,
690-
triggerID,
691-
messageID: toolUse?.toolUseId ?? '',
692-
userIntent: undefined,
693-
codeBlockLanguage: undefined,
694-
contextList: undefined,
695-
title: undefined,
696-
canBeVoted: false,
697-
buttons,
698-
fullWidth,
699-
padding,
700-
header,
701-
codeBlockActions,
702-
},
703-
tabID
704-
)
707+
this.dispatcher.sendChatMessage(
708+
new ChatMessage(
709+
{
710+
message: message,
711+
messageType: 'answer-part',
712+
followUps: undefined,
713+
followUpsHeader: undefined,
714+
relatedSuggestions: undefined,
715+
triggerID,
716+
messageID: toolUse?.toolUseId ?? '',
717+
userIntent: undefined,
718+
codeBlockLanguage: undefined,
719+
contextList: undefined,
720+
title: undefined,
721+
canBeVoted: false,
722+
buttons,
723+
fullWidth: true,
724+
padding: false,
725+
header,
726+
// eslint-disable-next-line unicorn/no-null
727+
codeBlockActions: { 'insert-to-cursor': null, copy: null },
728+
},
729+
tabID
705730
)
706-
}
731+
)
707732
}
708733

709734
private editorContextMenuCommandVerbs: Map<EditorContextCommandType, string> = new Map([

‎packages/core/src/codewhispererChat/tools/chatStream.ts‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import { ToolUse } from '@amzn/codewhisperer-streaming'
1010
import { CommandValidation } from './executeBash'
1111
import { Change } from 'diff'
1212
import { ChatSession } from '../clients/chat/v0/chat'
13-
import { ToolType } from './toolUtils'
1413

1514
/**
1615
* A writable stream that feeds each chunk/line to the chat UI.
@@ -26,7 +25,7 @@ export class ChatStream extends Writable {
2625
private readonly toolUse: ToolUse | undefined,
2726
private readonly session: ChatSession,
2827
private readonly messageIdToUpdate: string | undefined,
29-
// emitEvent decides to show the read toolMessage to the user.
28+
// emitEvent decides to show the streaming message or read/list directory tool message to the user.
3029
private readonly emitEvent: boolean,
3130
private readonly validation: CommandValidation,
3231
private readonly changeList?: Change[],
@@ -36,15 +35,13 @@ export class ChatStream extends Writable {
3635
this.logger.debug(
3736
`ChatStream created for tabID: ${tabID}, triggerID: ${triggerID}, session: ${session.readFiles}, emitEvent to mynahUI: ${emitEvent}`
3837
)
39-
if (toolUse?.name === ToolType.FsRead && emitEvent) {
40-
if (!messageIdToUpdate) {
41-
// If messageIdToUpdate is undefined, we need to first create an empty message
42-
// with messageId so it can be updated later
43-
this.messenger.sendInitialToolMessage(tabID, triggerID, toolUse?.toolUseId)
44-
}
45-
} else if (!(toolUse?.name === ToolType.FsRead)) {
46-
this.messenger.sendInitalStream(tabID, triggerID)
38+
if (!emitEvent) {
39+
return
4740
}
41+
// If messageIdToUpdate is undefined, we need to first create an empty message with messageId so it can be updated later
42+
messageIdToUpdate
43+
? this.messenger.sendInitalStream(tabID, triggerID)
44+
: this.messenger.sendInitialToolMessage(tabID, triggerID, toolUse?.toolUseId)
4845
}
4946

5047
override _write(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null) => void): void {

0 commit comments

Comments
 (0)