feat(amazonq): pass featureConfigs for ViewDiffInChat

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt

@@ -134,7 +134,7 @@ class Browser(parent: Disposable, private val webUri: URI, val project: Project)
                                 $postMessageToJavaJsCode
                             }
                         },
-
+                     ${CodeWhispererFeatureConfigService.getInstance().getFeatureConfigJsonString()},
                      "${activeProfile?.profileName.orEmpty()}")
                     const commands = [hybridChatConnector.initialQuickActions[0], hybridChatConnector.initialQuickActions[1]]
                     amazonQChat.createChat(

plugins/amazonq/mynah-ui/src/mynah-ui/connectorAdapter.ts

@@ -8,6 +8,7 @@ import { WebviewUIHandler } from './ui/main'
 import { TabDataGenerator } from './ui/tabs/generator'
 import { ChatClientAdapter, ChatEventHandler } from '@aws/chat-client'
 import { FqnExtractor } from "./fqn/extractor";
+import {FeatureContext} from "./ui/types";
 
 export * from "./ui/main";
 
@@ -25,8 +26,9 @@ export const initiateAdapter = (showWelcomePage: boolean,
                              isCodeScanEnabled: boolean,
                              isCodeTestEnabled: boolean,
                              ideApiPostMessage: (message: any) => void,
-                             profileName?: string) : HybridChatAdapter => {
-    return new HybridChatAdapter(showWelcomePage, disclaimerAcknowledged, isFeatureDevEnabled, isCodeTransformEnabled, isDocEnabled, isCodeScanEnabled, isCodeTestEnabled, ideApiPostMessage, profileName)
+                             featureConfigsSerialized: [string, FeatureContext][],
+                             profileName?: string,) : HybridChatAdapter => {
+    return new HybridChatAdapter(showWelcomePage, disclaimerAcknowledged, isFeatureDevEnabled, isCodeTransformEnabled, isDocEnabled, isCodeScanEnabled, isCodeTestEnabled, ideApiPostMessage, featureConfigsSerialized, profileName)
 }
 
 
@@ -37,7 +39,6 @@ export class HybridChatAdapter implements ChatClientAdapter {
     private mynahUIRef?: { mynahUI: MynahUI}
 
     constructor(
-
         private showWelcomePage: boolean,
         private disclaimerAcknowledged: boolean,
         private isFeatureDevEnabled: boolean,
@@ -46,8 +47,8 @@ export class HybridChatAdapter implements ChatClientAdapter {
         private isCodeScanEnabled: boolean,
         private isCodeTestEnabled: boolean,
         private ideApiPostMessage: (message: any) => void,
+        private featureConfigsSerialized: [string, FeatureContext][],
         private profileName?: string,
-
     ) {}
 
     /**
@@ -69,6 +70,7 @@ export class HybridChatAdapter implements ChatClientAdapter {
             isCodeTestEnabled: this.isCodeTestEnabled,
             profileName: this.profileName,
             hybridChat: true,
+            featureConfigsSerialized: this.featureConfigsSerialized,
         })
 
         return this.uiHandler.mynahUIProps

plugins/amazonq/mynah-ui/src/mynah-ui/ui/connector.ts

@@ -49,7 +49,7 @@ export interface CWCChatItem extends ChatItem {
 export interface ConnectorProps {
     sendMessageToExtension: (message: ExtensionMessage) => void
     onMessageReceived?: (tabID: string, messageData: any, needToShowAPIDocsTab: boolean) => void
-    onChatAnswerReceived?: (tabID: string, message: ChatItem) => void
+    onChatAnswerReceived?: (tabID: string, message: ChatItem, messageData?: any) => void
     onChatAnswerUpdated?: (tabID: string, message:ChatItem) => void
     onCodeTransformChatDisabled: (tabID: string) => void
     onCodeTransformMessageReceived: (
@@ -62,6 +62,7 @@ export interface ConnectorProps {
     onRunTestMessageReceived?: (tabID: string, showRunTestMessage: boolean) => void
     onWelcomeFollowUpClicked: (tabID: string, welcomeFollowUpType: WelcomeFollowupType) => void
     onAsyncEventProgress: (tabID: string, inProgress: boolean, message: string | undefined, cancelButtonWhenLoading?: boolean) => void
+    onQuickHandlerCommand: (tabID: string, command?: string, eventId?: string) => void
     onCWCContextCommandMessage: (message: ChatItem, command?: string) => string | undefined
     onCWCOnboardingPageInteractionMessage: (message: ChatItem) => string | undefined
     onOpenSettingsMessage: (tabID: string) => void
@@ -281,6 +282,8 @@ export class Connector {
             default:
                 break
         }
+        // Reset lastCommand after message is rendered.
+        this.tabsStorage.updateTabLastCommand(messageData.tabID, '')
     }
 
     onTabAdd = (tabID: string): void => {

plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts

@@ -31,42 +31,11 @@ import {welcomeScreenTabData} from "./walkthrough/welcome";
 import { agentWalkthroughDataModel } from './walkthrough/agent'
 import {createClickTelemetry, createOpenAgentTelemetry} from "./telemetry/actions";
 import {disclaimerAcknowledgeButtonId, disclaimerCard} from "./texts/disclaimer";
+import {FeatureContext, tryNewMap} from "./types";
+
 
 
 // Ref: https://github.com/aws/aws-toolkit-vscode/blob/e9ea8082ffe0b9968a873437407d0b6b31b9e1a5/packages/core/src/amazonq/webview/ui/main.ts
-export const createMynahUI = (
-    ideApi: any,
-    showWelcomePage: boolean,
-    disclaimerAcknowledged: boolean,
-    isFeatureDevEnabled: boolean,
-    isCodeTransformEnabled: boolean,
-    isDocEnabled: boolean,
-    isCodeScanEnabled: boolean,
-    isCodeTestEnabled: boolean,
-    highlightCommand?: QuickActionCommand,
-    profileName?: string,
-
-) => {
-    const handler = new WebviewUIHandler({
-        postMessage: ideApi.postMessage,
-        mynahUIRef: { mynahUI: undefined },
-        showWelcomePage,
-        disclaimerAcknowledged,
-        isFeatureDevEnabled,
-        isCodeTransformEnabled,
-        isDocEnabled,
-        isCodeScanEnabled,
-        isCodeTestEnabled,
-        highlightCommand,
-        profileName,
-        hybridChat: false,
-    })
-
-    return {
-        mynahUI: handler.mynahUI,
-        messageReceiver: handler.connector?.handleMessageReceive,
-    }
-}
 
 export class WebviewUIHandler {
     postMessage: any
@@ -81,6 +50,7 @@ export class WebviewUIHandler {
     profileName?: string
     responseMetadata: Map<string, string[]>
     tabsStorage: TabsStorage
+    featureConfigs?: Map<string, FeatureContext>
 
     mynahUIProps: MynahUIProps
     connector?: Connector
@@ -90,14 +60,14 @@ export class WebviewUIHandler {
     textMessageHandler?: TextMessageHandler
     messageController?: MessageController
 
-    savedContextCommands: MynahUIDataModel['contextCommands']
     disclaimerCardActive : boolean
 
 
     mynahUIRef: { mynahUI: MynahUI | undefined }
     constructor({
                     postMessage,
                     mynahUIRef,
+                    featureConfigsSerialized,
                     showWelcomePage,
                     disclaimerAcknowledged,
                     isFeatureDevEnabled,
@@ -112,6 +82,7 @@ export class WebviewUIHandler {
                 } : {
         postMessage: any
         mynahUIRef: { mynahUI: MynahUI | undefined }
+        featureConfigsSerialized: [string, FeatureContext][]
         showWelcomePage: boolean,
         disclaimerAcknowledged: boolean,
         isFeatureDevEnabled: boolean
@@ -127,6 +98,7 @@ export class WebviewUIHandler {
     }) {
         this.postMessage = postMessage
         this.mynahUIRef = mynahUIRef
+        this.featureConfigs = tryNewMap(featureConfigsSerialized)
         this.showWelcomePage = showWelcomePage;
         this.disclaimerAcknowledged = disclaimerAcknowledged
         this.isFeatureDevEnabled = isFeatureDevEnabled
@@ -206,6 +178,8 @@ export class WebviewUIHandler {
                     profileName
                 })
 
+                this.featureConfigs = tryNewMap(featureConfigsSerialized)
+
                 // Set the new defaults for the quick action commands in all tabs now that isFeatureDevEnabled and isCodeTransformEnabled were enabled/disabled
                 for (const tab of this.tabsStorage.getTabs()) {
                     this.mynahUI?.updateStore(tab.id, {
@@ -238,7 +212,17 @@ export class WebviewUIHandler {
             onCWCOnboardingPageInteractionMessage: (message: ChatItem): string | undefined => {
                 return this.messageController?.sendMessageToTab(message, 'cwc')
             },
+            onQuickHandlerCommand: (tabID: string, command?: string, eventId?: string) => {
+                this.tabsStorage.updateTabLastCommand(tabID, command)
+                if (command === 'aws.awsq.transform') {
+                    this.quickActionHandler?.handleCommand({ command: '/transform' }, tabID, eventId)
+                } else if (command === 'aws.awsq.clearchat') {
+                    this.quickActionHandler?.handleCommand({ command: '/clear' }, tabID)
+                }
+            },
             onCWCContextCommandMessage: (message: ChatItem, command?: string): string | undefined => {
+                const selectedTab = this.tabsStorage.getSelectedTab()
+                this.tabsStorage.updateTabLastCommand(selectedTab?.id || '', command || '')
                 if (command === 'aws.amazonq.sendToPrompt') {
                     return this.messageController?.sendSelectedCodeToTab(message)
                 } else {
@@ -398,7 +382,7 @@ export class WebviewUIHandler {
                     } as ChatItem)
                 }
             },
-            onChatAnswerReceived: (tabID: string, item: CWCChatItem) => {
+            onChatAnswerReceived: (tabID: string, item: CWCChatItem, messageData: any) => {
                 if (item.type === ChatItemType.ANSWER_PART || item.type === ChatItemType.CODE_RESULT) {
                     this.mynahUI?.updateLastChatAnswer(tabID, {
                         ...(item.messageId !== undefined ? { messageId: item.messageId } : {}),
@@ -417,8 +401,12 @@ export class WebviewUIHandler {
                     return
                 }
 
-                if (item.body !== undefined || item.relatedContent !== undefined || item.followUp !== undefined) {
-                    this.mynahUI?.addChatItem(tabID, item)
+                if (item.body !== undefined || item.relatedContent !== undefined || item.followUp !== undefined || item.formItems !== undefined || item.buttons !== undefined) {
+                    this.mynahUI?.addChatItem(tabID, {
+                        ...item,
+                        messageId: item.messageId,
+                        codeBlockActions: this.getCodeBlockActions(messageData),
+                    })
                 }
 
                 if (
@@ -543,7 +531,7 @@ export class WebviewUIHandler {
                 this.tabsStorage.updateTabTypeFromUnknown(newTabID, tabType)
                 this.connector?.onKnownTabOpen(newTabID)
                 this.connector?.onUpdateTabType(newTabID)
-
+                this.featureConfigs = tryNewMap(featureConfigsSerialized)
                 this.mynahUI?.updateStore(newTabID, this.tabDataGenerator!.getTabData(tabType, true))
             },
             onStartNewTransform: (tabID: string) => {
@@ -903,6 +891,43 @@ export class WebviewUIHandler {
         })
 
     }
+
+    private getCodeBlockActions(messageData: any) {
+        // Show ViewDiff and AcceptDiff for allowedCommands in CWC
+        const isEnabled = this.featureConfigs?.get('ViewDiffInChat')?.variation === 'TREATMENT'
+        const tab = this.tabsStorage.getTab(messageData?.tabID || '')
+        const allowedCommands = [
+            'aws.amazonq.refactorCode',
+            'aws.amazonq.fixCode',
+            'aws.amazonq.optimizeCode',
+            'aws.amazonq.sendToPrompt',
+        ]
+        if (isEnabled && tab?.type === 'cwc' && allowedCommands.includes(tab.lastCommand || '')) {
+            return {
+                'insert-to-cursor': undefined,
+                accept_diff: {
+                    id: 'accept_diff',
+                    label: 'Apply Diff',
+                    icon: MynahIcons.OK_CIRCLED,
+                    data: messageData,
+                },
+                view_diff: {
+                    id: 'view_diff',
+                    label: 'View Diff',
+                    icon: MynahIcons.EYE,
+                    data: messageData,
+                },
+            }
+        }
+        // Show only "Copy" option for codeblocks in Q Test Tab
+        if (tab?.type === 'testgen') {
+            return {
+                'insert-to-cursor': undefined,
+            }
+        }
+        // Default will show "Copy" and "Insert at cursor" for codeblocks
+        return {}
+    }
     get mynahUI(): MynahUI | undefined {
         return this.mynahUIRef.mynahUI
     }

plugins/amazonq/mynah-ui/src/mynah-ui/ui/quickActions/generator.ts

@@ -142,6 +142,10 @@ export class QuickActionGenerator {
                 description: "This command isn't available in /doc",
                 unavailableItems: ['/help', '/clear'],
             },
+            testgen: {
+                description: "This command isn't available",
+                unavailableItems: ['/help', '/clear'],
+            },
             welcome: {
                 description: '',
                 unavailableItems: ['/clear'],

plugins/amazonq/mynah-ui/src/mynah-ui/ui/storages/tabsStorage.ts

@@ -14,6 +14,7 @@ const TabTypes = [
     'agentWalkthrough',
     'welcome',
     'unknown',
+    'testgen'
 ] as const
 export type TabType = (typeof TabTypes)[number]
 export function isTabType(value: string): value is TabType {
@@ -45,6 +46,7 @@ export interface Tab {
     type: TabType
     isSelected: boolean
     openInteractionType?: TabOpenType
+    lastCommand?: string
 }
 
 export class TabsStorage {
@@ -93,6 +95,17 @@ export class TabsStorage {
         return this.tabs.get(tabID)?.status === 'dead'
     }
 
+    public updateTabLastCommand(tabID: string, command?: string) {
+        if (command === undefined) {
+            return
+        }
+        const currentTabValue = this.tabs.get(tabID)
+        if (currentTabValue === undefined || currentTabValue.status === 'dead') {
+            return
+        }
+        currentTabValue.lastCommand = command
+        this.tabs.set(tabID, currentTabValue)
+    }
     public updateTabStatus(tabID: string, tabStatus: TabStatus) {
         const currentTabValue = this.tabs.get(tabID)
         if (currentTabValue === undefined || currentTabValue.status === 'dead') {

plugins/amazonq/mynah-ui/src/mynah-ui/ui/types.ts

@@ -0,0 +1,26 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+export interface FeatureValue {
+    boolValue?: boolean;
+    doubleValue?: number;
+    longValue?: number;
+    stringValue?: string;
+}
+
+export class FeatureContext {
+    constructor(
+        public name: string,
+        public variation: string,
+        public value: FeatureValue
+    ) {}
+}
+
+export function tryNewMap(arr: [string, FeatureContext][]) {
+    try {
+        return new Map(arr)
+    } catch (error) {
+        return new Map()
+    }
+}