Modifying Agentic chat UX for view code diff and accept/reject files

laileni-aws · laileni-aws · commit 8d6467ba7eb9 · 2025-03-24T17:39:49.000-07:00
diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
@@ -3,7 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ChatItemButton, ChatItemFormItem, ChatItemType, MynahUIDataModel, QuickActionCommand } from '@aws/mynah-ui'
+import {
+    ChatItem,
+    ChatItemButton,
+    ChatItemFormItem,
+    ChatItemType,
+    MynahUIDataModel,
+    QuickActionCommand,
+} from '@aws/mynah-ui'
 import { TabType } from '../storages/tabsStorage'
 import { CWCChatItem } from '../connector'
 import { BaseConnector, BaseConnectorProps } from './baseConnector'
@@ -18,12 +25,14 @@ export interface ConnectorProps extends BaseConnectorProps {
         title?: string,
         description?: string
     ) => void
+    onChatAnswerUpdated?: (tabID: string, message: ChatItem) => void
 }
 
 export class Connector extends BaseConnector {
     private readonly onCWCContextCommandMessage
     private readonly onContextCommandDataReceived
     private readonly onShowCustomForm
+    private readonly onChatAnswerUpdated
 
     override getTabType(): TabType {
         return 'cwc'
@@ -34,6 +43,7 @@ export class Connector extends BaseConnector {
         this.onCWCContextCommandMessage = props.onCWCContextCommandMessage
         this.onContextCommandDataReceived = props.onContextCommandDataReceived
         this.onShowCustomForm = props.onShowCustomForm
+        this.onChatAnswerUpdated = props.onChatAnswerUpdated
     }
 
     onSourceLinkClick = (tabID: string, messageId: string, link: string): void => {
@@ -91,16 +101,14 @@ export class Connector extends BaseConnector {
                 messageId: messageData.messageID ?? messageData.triggerID,
                 body: messageData.message,
                 followUp: followUps,
-                canBeVoted: true,
+                canBeVoted: messageData.canBeVoted ?? false,
                 codeReference: messageData.codeReference,
                 userIntent: messageData.userIntent,
                 codeBlockLanguage: messageData.codeBlockLanguage,
                 contextList: messageData.contextList,
-            }
-
-            // If it is not there we will not set it
-            if (messageData.messageType === 'answer-part' || messageData.messageType === 'answer') {
-                answer.canBeVoted = true
+                title: messageData.title,
+                buttons: messageData.buttons ?? undefined,
+                fileList: messageData.fileList ?? undefined,
             }
 
             if (messageData.relatedSuggestions !== undefined) {
@@ -137,6 +145,8 @@ export class Connector extends BaseConnector {
                               options: messageData.followUps,
                           }
                         : undefined,
+                buttons: messageData.buttons ?? undefined,
+                canBeVoted: messageData.canBeVoted ?? false,
             }
             this.onChatAnswerReceived(messageData.tabID, answer, messageData)
 
@@ -204,7 +214,7 @@ export class Connector extends BaseConnector {
         }
 
         if (messageData.type === 'customFormActionMessage') {
-            this.onCustomFormAction(messageData.tabID, messageData.action)
+            this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action)
             return
         }
         // For other message types, call the base class handleMessageReceive
@@ -235,6 +245,7 @@ export class Connector extends BaseConnector {
 
     onCustomFormAction(
         tabId: string,
+        messageId: string,
         action: {
             id: string
             text?: string | undefined
@@ -248,14 +259,53 @@ export class Connector extends BaseConnector {
         this.sendMessageToExtension({
             command: 'form-action-click',
             action: action,
+            formSelectedValues: action.formItemValues,
             tabType: this.getTabType(),
             tabID: tabId,
         })
+
+        if (!this.onChatAnswerUpdated || !['accept-code-diff', 'reject-code-diff'].includes(action.id)) {
+            return
+        }
+        const answer: ChatItem = {
+            type: ChatItemType.ANSWER,
+            messageId: messageId,
+            buttons: [],
+        }
+        switch (action.id) {
+            case 'accept-code-diff':
+                answer.buttons = [
+                    {
+                        keepCardAfterClick: true,
+                        text: 'Accepted code',
+                        id: 'accepted-code-diff',
+                        status: 'success',
+                        position: 'outside',
+                        disabled: true,
+                    },
+                ]
+                break
+            case 'reject-code-diff':
+                answer.buttons = [
+                    {
+                        keepCardAfterClick: true,
+                        text: 'Rejected code',
+                        id: 'rejected-code-diff',
+                        status: 'error',
+                        position: 'outside',
+                        disabled: true,
+                    },
+                ]
+                break
+            default:
+                break
+        }
+        this.onChatAnswerUpdated(tabId, answer)
     }
 
     onFileClick = (tabID: string, filePath: string, messageId?: string) => {
         this.sendMessageToExtension({
-            command: 'file-click',
+            command: messageId === '' ? 'file-click' : 'open-diff',
             tabID,
             messageId,
             filePath,
diff --git a/packages/core/src/amazonq/webview/ui/connector.ts b/packages/core/src/amazonq/webview/ui/connector.ts
@@ -63,6 +63,7 @@ export interface CWCChatItem extends ChatItem {
     userIntent?: UserIntent
     codeBlockLanguage?: string
     contextList?: Context[]
+    title?: string
 }
 
 export interface Context {
@@ -711,7 +712,7 @@ export class Connector {
                         tabType: 'cwc',
                     })
                 } else {
-                    this.cwChatConnector.onCustomFormAction(tabId, action)
+                    this.cwChatConnector.onCustomFormAction(tabId, messageId ?? '', action)
                 }
                 break
             case 'agentWalkthrough': {
diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts
@@ -352,6 +352,7 @@ export const createMynahUI = (
                     ...(item.followUp !== undefined ? { followUp: item.followUp } : {}),
                     ...(item.fileList !== undefined ? { fileList: item.fileList } : {}),
                     ...(item.header !== undefined ? { header: item.header } : { header: undefined }),
+                    ...(item.buttons !== undefined ? { buttons: item.buttons } : { buttons: undefined }),
                 })
                 if (
                     item.messageId !== undefined &&
@@ -374,7 +375,7 @@ export const createMynahUI = (
                     fileList: {
                         fileTreeTitle: '',
                         filePaths: item.contextList.map((file) => file.relativeFilePath),
-                        rootFolderTitle: 'Context',
+                        rootFolderTitle: item.title,
                         flatList: true,
                         collapsed: true,
                         hideFileCount: true,
diff --git a/packages/core/src/codewhispererChat/app.ts b/packages/core/src/codewhispererChat/app.ts
@@ -28,6 +28,7 @@ import {
     AcceptDiff,
     QuickCommandGroupActionClick,
     FileClick,
+    OpenDiff,
 } from './controllers/chat/model'
 import { EditorContextCommand, registerCommands } from './commands/registerCommands'
 import { ContextSelectedMessage, CustomFormActionMessage } from './view/connector/connector'
@@ -41,6 +42,7 @@ export function init(appContext: AmazonQAppInitContext) {
         processInsertCodeAtCursorPosition: new EventEmitter<InsertCodeAtCursorPosition>(),
         processAcceptDiff: new EventEmitter<AcceptDiff>(),
         processViewDiff: new EventEmitter<ViewDiff>(),
+        processOpenDiff: new EventEmitter<OpenDiff>(),
         processCopyCodeToClipboard: new EventEmitter<CopyCodeToClipboard>(),
         processContextMenuCommand: new EventEmitter<EditorContextCommand>(),
         processTriggerTabIDReceived: new EventEmitter<TriggerTabIDReceived>(),
@@ -76,6 +78,7 @@ export function init(appContext: AmazonQAppInitContext) {
         ),
         processAcceptDiff: new MessageListener<AcceptDiff>(cwChatControllerEventEmitters.processAcceptDiff),
         processViewDiff: new MessageListener<ViewDiff>(cwChatControllerEventEmitters.processViewDiff),
+        processOpenDiff: new MessageListener<OpenDiff>(cwChatControllerEventEmitters.processOpenDiff),
         processCopyCodeToClipboard: new MessageListener<CopyCodeToClipboard>(
             cwChatControllerEventEmitters.processCopyCodeToClipboard
         ),
@@ -137,6 +140,7 @@ export function init(appContext: AmazonQAppInitContext) {
         ),
         processAcceptDiff: new MessagePublisher<AcceptDiff>(cwChatControllerEventEmitters.processAcceptDiff),
         processViewDiff: new MessagePublisher<ViewDiff>(cwChatControllerEventEmitters.processViewDiff),
+        processOpenDiff: new MessagePublisher<OpenDiff>(cwChatControllerEventEmitters.processOpenDiff),
         processCopyCodeToClipboard: new MessagePublisher<CopyCodeToClipboard>(
             cwChatControllerEventEmitters.processCopyCodeToClipboard
         ),
diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
@@ -13,6 +13,9 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr
 
 export class ChatSession {
     private sessionId?: string
+    private _listOfReadFiles: string[] = []
+    private _filePath: string | undefined
+    private _tempFilePath: string | undefined
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -35,6 +38,28 @@ export class ChatSession {
     public setSessionID(id?: string) {
         this.sessionId = id
     }
+    public get listOfReadFiles(): string[] {
+        return this._listOfReadFiles
+    }
+    public get getFilePath(): string | undefined {
+        return this._filePath
+    }
+    public get getTempFilePath(): string | undefined {
+        return this._tempFilePath
+    }
+    public setFilePath(filePath: string | undefined) {
+        this._filePath = filePath
+    }
+    public setTempFilePath(tempFilePath: string | undefined) {
+        this._tempFilePath = tempFilePath
+    }
+    public pushToListOfReadFiles(filePath?: string) {
+        if (filePath) {
+            this._listOfReadFiles.push(filePath)
+        } else {
+            this._listOfReadFiles = []
+        }
+    }
     async chatIam(chatRequest: SendMessageRequest): Promise<SendMessageCommandOutput> {
         const client = await createQDeveloperStreamingClient()
 
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -32,6 +32,7 @@ import {
     DocumentReference,
     FileClick,
     RelevantTextDocumentAddition,
+    OpenDiff,
 } from './model'
 import {
     AppToWebViewMessageDispatcher,
@@ -81,6 +82,7 @@ import {
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
 import { ChatHistoryManager } from '../../storages/chatHistory'
+import { amazonQTabSuffix } from '../../../shared/constants'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -90,6 +92,7 @@ export interface ChatControllerMessagePublishers {
     readonly processInsertCodeAtCursorPosition: MessagePublisher<InsertCodeAtCursorPosition>
     readonly processAcceptDiff: MessagePublisher<AcceptDiff>
     readonly processViewDiff: MessagePublisher<ViewDiff>
+    readonly processOpenDiff: MessagePublisher<OpenDiff>
     readonly processCopyCodeToClipboard: MessagePublisher<CopyCodeToClipboard>
     readonly processContextMenuCommand: MessagePublisher<EditorContextCommand>
     readonly processTriggerTabIDReceived: MessagePublisher<TriggerTabIDReceived>
@@ -115,6 +118,7 @@ export interface ChatControllerMessageListeners {
     readonly processInsertCodeAtCursorPosition: MessageListener<InsertCodeAtCursorPosition>
     readonly processAcceptDiff: MessageListener<AcceptDiff>
     readonly processViewDiff: MessageListener<ViewDiff>
+    readonly processOpenDiff: MessageListener<OpenDiff>
     readonly processCopyCodeToClipboard: MessageListener<CopyCodeToClipboard>
     readonly processContextMenuCommand: MessageListener<EditorContextCommand>
     readonly processTriggerTabIDReceived: MessageListener<TriggerTabIDReceived>
@@ -213,6 +217,10 @@ export class ChatController {
             return this.processViewDiff(data)
         })
 
+        this.chatControllerMessageListeners.processOpenDiff.onMessage((data) => {
+            return this.processOpenDiff(data)
+        })
+
         this.chatControllerMessageListeners.processCopyCodeToClipboard.onMessage((data) => {
             return this.processCopyCodeToClipboard(data)
         })
@@ -388,6 +396,45 @@ export class ChatController {
             })
     }
 
+    private async processOpenDiff(message: OpenDiff) {
+        const session = this.sessionStorage.getSession(message.tabID)
+        const filePath = session.getFilePath ?? message.filePath
+        const fileExists = await fs.existsFile(filePath)
+        // Check if fileExists=false, If yes, return instead of showing broken diff experience.
+        if (!session.getTempFilePath) {
+            return
+        }
+        const leftUri = fileExists ? vscode.Uri.file(filePath) : vscode.Uri.from({ scheme: 'untitled' })
+        const rightUri = vscode.Uri.file(session.getTempFilePath ?? filePath)
+        const fileName = path.basename(filePath)
+        await vscode.commands.executeCommand('vscode.diff', leftUri, rightUri, `${fileName} ${amazonQTabSuffix}`)
+    }
+
+    private async processAcceptCodeDiff(message: CustomFormActionMessage) {
+        const session = this.sessionStorage.getSession(message.tabID ?? '')
+        const filePath = session.getFilePath ?? ''
+        const fileExists = await fs.existsFile(filePath)
+        const tempFilePath = session.getTempFilePath
+        const tempFileExists = await fs.existsFile(tempFilePath ?? '')
+        if (fileExists && tempFileExists) {
+            const fileContent = await fs.readFileText(filePath)
+            const tempFileContent = await fs.readFileText(tempFilePath ?? '')
+            if (fileContent !== tempFileContent) {
+                await fs.writeFile(filePath, tempFileContent)
+            }
+            await fs.delete(tempFilePath ?? '')
+            await vscode.commands.executeCommand('vscode.open', vscode.Uri.file(filePath))
+        } else if (!fileExists && tempFileExists) {
+            const fileContent = await fs.readFileText(tempFilePath ?? '')
+            await fs.writeFile(filePath, fileContent)
+            await fs.delete(tempFilePath ?? '')
+            await vscode.commands.executeCommand('vscode.open', vscode.Uri.file(filePath))
+        }
+        // Reset the filePaths to undefined
+        session.setFilePath(undefined)
+        session.setTempFilePath(undefined)
+    }
+
     private async processCopyCodeToClipboard(message: CopyCodeToClipboard) {
         this.telemetryHelper.recordInteractWithMessage(message)
     }
@@ -577,6 +624,12 @@ export class ChatController {
             const newFileDoc = await vscode.workspace.openTextDocument(newFilePath)
             await vscode.window.showTextDocument(newFileDoc)
             telemetry.ui_click.emit({ elementId: 'amazonq_createSavedPrompt' })
+        } else if (message.action.id === 'accept-code-diff') {
+            await this.processAcceptCodeDiff(message)
+        } else if (message.action.id === 'reject-code-diff') {
+            // Reset the filePaths to undefined
+            this.sessionStorage.getSession(message.tabID ?? '').setFilePath(undefined)
+            this.sessionStorage.getSession(message.tabID ?? '').setTempFilePath(undefined)
         }
     }
 
@@ -835,6 +888,8 @@ export class ChatController {
     }
 
     private async processPromptMessageAsNewThread(message: PromptMessage) {
+        const session = this.sessionStorage.getSession(message.tabID)
+        session.pushToListOfReadFiles(undefined)
         this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then((context) => {
diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -85,6 +85,10 @@ export class Messenger {
                     userIntent: undefined,
                     codeBlockLanguage: undefined,
                     contextList: mergedRelevantDocuments,
+                    title: 'Context',
+                    buttons: undefined,
+                    fileList: undefined,
+                    canBeVoted: false,
                 },
                 tabID
             )
@@ -436,6 +440,7 @@ export class Messenger {
                     userIntent: undefined,
                     codeBlockLanguage: undefined,
                     contextList: undefined,
+                    title: undefined,
                 },
                 tabID
             )
diff --git a/packages/core/src/codewhispererChat/controllers/chat/model.ts b/packages/core/src/codewhispererChat/controllers/chat/model.ts
@@ -89,6 +89,14 @@ export interface ViewDiff {
     totalCodeBlocks?: number
 }
 
+export interface OpenDiff {
+    command: string | undefined
+    tabID: string
+    messageId: string
+    filePath: string
+    referenceTrackerInformation?: CodeReference[]
+}
+
 export type ChatPromptCommandType =
     | 'help'
     | 'clear'
diff --git a/packages/core/src/codewhispererChat/view/connector/connector.ts b/packages/core/src/codewhispererChat/view/connector/connector.ts
diff --git a/packages/core/src/codewhispererChat/view/messages/messageListener.ts b/packages/core/src/codewhispererChat/view/messages/messageListener.ts
