Adds capability to send new context commands to AB groups (#5206)

Add the ability to send a new command that will appear when a user types @ in the chat window
aws · Dec 14, 2024 · 6370cf2 · 6370cf2
1 parent 2226148
commit 6370cf2
Show file tree

Hide file tree

Showing 8 changed files with 108 additions and 10 deletions.
diff --git a/.changes/next-release/feature-fab622c6-240e-49a2-9750-dee956759a0a.json b/.changes/next-release/feature-fab622c6-240e-49a2-9750-dee956759a0a.json
@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Adds capability to send new context commands to AB groups"
+}
diff --git a/...nity/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt b/...nity/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt
@@ -24,6 +24,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
+import software.aws.toolkits.jetbrains.services.amazonq.util.highlightCommand
 import software.aws.toolkits.jetbrains.services.amazonq.webview.BrowserConnector
 import software.aws.toolkits.jetbrains.services.amazonq.webview.FqnWebviewAdapter
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.EditorThemeAdapter
@@ -125,7 +126,8 @@ class AmazonQToolWindow private constructor(
             isFeatureDevAvailable = isFeatureDevAvailable(project),
             isCodeScanAvailable = isCodeScanAvailable(project),
             isCodeTestAvailable = isCodeTestAvailable(project),
-            isDocAvailable = isDocAvailable(project)
+            isDocAvailable = isDocAvailable(project),
+            highlightCommand = highlightCommand()
         )
 
         scope.launch {

diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/util/HighlightCommand.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/util/HighlightCommand.kt
@@ -0,0 +1,16 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.util
+
+import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
+
+data class HighlightCommand(val command: String, val description: String)
+
+fun highlightCommand(): HighlightCommand? {
+    val feature = CodeWhispererFeatureConfigService.getInstance().getHighlightCommandFeature()
+
+    if (feature == null || feature.value.stringValue().isEmpty()) return null
+
+    return HighlightCommand(feature.value.stringValue(), feature.variation)
+}
diff --git a/...tbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt b/...tbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt
@@ -3,10 +3,12 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.webview
 
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.jcef.JBCefJSQuery
 import org.cef.CefApp
+import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
 import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
 
@@ -25,6 +27,7 @@ class Browser(parent: Disposable) : Disposable {
         isDocAvailable: Boolean,
         isCodeScanAvailable: Boolean,
         isCodeTestAvailable: Boolean,
+        highlightCommand: HighlightCommand?,
     ) {
         // register the scheme handler to route http://mynah/ URIs to the resources/assets directory on classpath
         CefApp.getInstance()
@@ -34,7 +37,7 @@ class Browser(parent: Disposable) : Disposable {
                 AssetResourceHandler.AssetResourceHandlerFactory(),
             )
 
-        loadWebView(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable)
+        loadWebView(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand)
     }
 
     override fun dispose() {
@@ -55,12 +58,15 @@ class Browser(parent: Disposable) : Disposable {
         isDocAvailable: Boolean,
         isCodeScanAvailable: Boolean,
         isCodeTestAvailable: Boolean,
+        highlightCommand: HighlightCommand?,
     ) {
         // setup empty state. The message request handlers use this for storing state
         // that's persistent between page loads.
         jcefBrowser.setProperty("state", "")
         // load the web app
-        jcefBrowser.loadHTML(getWebviewHTML(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable))
+        jcefBrowser.loadHTML(
+            getWebviewHTML(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand)
+        )
     }
 
     /**
@@ -73,6 +79,7 @@ class Browser(parent: Disposable) : Disposable {
         isDocAvailable: Boolean,
         isCodeScanAvailable: Boolean,
         isCodeTestAvailable: Boolean,
+        highlightCommand: HighlightCommand?,
     ): String {
         val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)")
 
@@ -92,7 +99,8 @@ class Browser(parent: Disposable) : Disposable {
                         $isCodeTransformAvailable, // whether /transform is available
                         $isDocAvailable, // whether /doc is available
                         $isCodeScanAvailable, // whether /scan is available
-                        $isCodeTestAvailable // whether /test is available
+                        $isCodeTestAvailable, // whether /test is available
+                        ${OBJECT_MAPPER.writeValueAsString(highlightCommand)}
                     );
                 }
             </script>        
@@ -114,5 +122,6 @@ class Browser(parent: Disposable) : Disposable {
     companion object {
         private const val WEB_SCRIPT_URI = "http://mynah/js/mynah-ui.js"
         private const val MAX_ONBOARDING_PAGE_COUNT = 3
+        private val OBJECT_MAPPER = jacksonObjectMapper()
     }
 }
diff --git a/...re/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt b/...re/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt
@@ -26,6 +26,8 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableC
 import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsResponse
 import software.amazon.awssdk.services.codewhispererruntime.paginators.ListAvailableCustomizationsIterable
+import software.aws.toolkits.core.TokenConnectionSettings
+import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.jetbrains.core.MockClientManagerRule
 import software.aws.toolkits.jetbrains.core.credentials.LegacyManagedBearerSsoConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -58,6 +60,44 @@ class CodeWhispererFeatureConfigServiceTest {
         assertThat(CodeWhispererFeatureConfigService.FEATURE_DEFINITIONS).containsKeys("testFeature")
     }
 
+    @Test
+    fun `test highlightCommand returns non-empty`() {
+        mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
+            on { listFeatureEvaluations(any<ListFeatureEvaluationsRequest>()) } doReturn ListFeatureEvaluationsResponse.builder().featureEvaluations(
+                listOf(
+                    FeatureEvaluation.builder()
+                        .feature("highlightCommand")
+                        .variation("a new command")
+                        .value(FeatureValue.fromStringValue("@highlight"))
+                        .build()
+                )
+            ).build()
+        }
+
+        val mockTokenSettings = mock<TokenConnectionSettings> {
+            on { providerId } doReturn "mock"
+            on { region } doReturn AwsRegion.GLOBAL
+        }
+
+        val mockSsoConnection = mock<LegacyManagedBearerSsoConnection> {
+            on { startUrl } doReturn "fake sso url"
+            on { getConnectionSettings() } doReturn mockTokenSettings
+        }
+
+        projectRule.project.replaceService(
+            ToolkitConnectionManager::class.java,
+            mock { on { activeConnectionForFeature(eq(QConnection.getInstance())) } doReturn mockSsoConnection },
+            disposableRule.disposable
+        )
+
+        runBlocking {
+            CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(projectRule.project)
+        }
+
+        assertThat(CodeWhispererFeatureConfigService.getInstance().getHighlightCommandFeature()?.value?.stringValue()).isEqualTo("@highlight")
+        assertThat(CodeWhispererFeatureConfigService.getInstance().getHighlightCommandFeature()?.variation).isEqualTo("a new command")
+    }
+
     @Test
     fun `test customizationArnOverride returns empty for BID users`() {
         testCustomizationArnOverrideABHelper(isIdc = false, isInListAvailableCustomizations = false)
@@ -80,7 +120,7 @@ class CodeWhispererFeatureConfigServiceTest {
             on { listFeatureEvaluations(any<ListFeatureEvaluationsRequest>()) } doReturn ListFeatureEvaluationsResponse.builder().featureEvaluations(
                 listOf(
                     FeatureEvaluation.builder()
-                        .feature(CodeWhispererFeatureConfigService.CUSTOMIZATION_ARN_OVERRIDE_NAME)
+                        .feature("customizationArnOverride")
                         .variation("customization-name")
                         .value(FeatureValue.fromStringValue("test arn"))
                         .build()

diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts
@@ -10,7 +10,7 @@ import {
     MynahUI,
     MynahUIDataModel,
     NotificationType,
-    ProgressField,
+    ProgressField, QuickActionCommand,
     ReferenceTrackerInformation
 } from '@aws/mynah-ui-chat'
 import './styles/dark.scss'
@@ -40,7 +40,8 @@ export const createMynahUI = (
     codeTransformInitEnabled: boolean,
     docInitEnabled: boolean,
     codeScanEnabled: boolean,
-    codeTestEnabled: boolean
+    codeTestEnabled: boolean,
+    highlightCommand?: QuickActionCommand,
 ) => {
     let disclaimerCardActive = !disclaimerAcknowledged
 
@@ -87,6 +88,7 @@ export const createMynahUI = (
         isDocEnabled,
         isCodeScanEnabled,
         isCodeTestEnabled,
+        highlightCommand
     })
 
     // eslint-disable-next-line prefer-const

diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ChatItemType, MynahUIDataModel, QuickActionCommandGroup } from '@aws/mynah-ui-chat'
+import { ChatItemType, MynahUIDataModel, QuickActionCommandGroup, QuickActionCommand } from '@aws/mynah-ui-chat'
 import { TabType } from '../storages/tabsStorage'
 import { FollowUpGenerator } from '../followUps/generator'
 import { QuickActionGenerator } from '../quickActions/generator'
@@ -15,11 +15,13 @@ export interface TabDataGeneratorProps {
     isDocEnabled: boolean
     isCodeScanEnabled: boolean
     isCodeTestEnabled: boolean
+    highlightCommand?: QuickActionCommand
 }
 
 export class TabDataGenerator {
     private followUpsGenerator: FollowUpGenerator
     public quickActionsGenerator: QuickActionGenerator
+    private highlightCommand?: QuickActionCommand
 
     private tabTitle: Map<TabType, string> = new Map([
         ['unknown', 'Chat'],
@@ -88,6 +90,7 @@ What would you like to work on?`,
             isCodeScanEnabled: props.isCodeScanEnabled,
             isCodeTestEnabled: props.isCodeTestEnabled,
         })
+        this.highlightCommand = props.highlightCommand
     }
 
     public getTabData(tabType: TabType, needWelcomeMessages: boolean, taskName?: string): MynahUIDataModel {
@@ -97,7 +100,7 @@ What would you like to work on?`,
                 'Amazon Q Developer uses generative AI. You may need to verify responses. See the [AWS Responsible AI Policy](https://aws.amazon.com/machine-learning/responsible-ai/policy/).',
             quickActionCommands: this.quickActionsGenerator.generateForTab(tabType),
             promptInputPlaceholder: this.tabInputPlaceholder.get(tabType),
-            contextCommands: this.tabContextCommand.get(tabType),
+            contextCommands: this.getContextCommands(tabType),
             chatItems: needWelcomeMessages
                 ? [
                       {
@@ -112,4 +115,23 @@ What would you like to work on?`,
                 : [],
         }
     }
+
+    private getContextCommands(tabType: TabType): QuickActionCommandGroup[] | undefined {
+        const contextCommands = this.tabContextCommand.get(tabType)
+
+        if (this.highlightCommand) {
+            const commandHighlight: QuickActionCommandGroup = {
+                groupName: 'Additional Commands',
+                commands: [this.highlightCommand],
+            }
+
+            if (contextCommands !== undefined) {
+                return [...contextCommands, commandHighlight]
+            }
+
+            return [commandHighlight]
+        }
+
+        return contextCommands
+    }
 }
diff --git a/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt
@@ -111,6 +111,8 @@ class CodeWhispererFeatureConfigService {
 
     fun getCustomizationFeature(): FeatureContext? = getFeature(CUSTOMIZATION_ARN_OVERRIDE_NAME)
 
+    fun getHighlightCommandFeature(): FeatureContext? = getFeature(HIGHLIGHT_COMMAND_NAME)
+
     fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
 
     fun getInlineCompletion(): Boolean = getFeatureValueForKey(INLINE_COMPLETION).stringValue() == "TREATMENT"
@@ -131,7 +133,8 @@ class CodeWhispererFeatureConfigService {
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"
         private const val INLINE_COMPLETION = "ProjectContextV2"
-        const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
+        private const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
+        private const val HIGHLIGHT_COMMAND_NAME = "highlightCommand"
         private const val NEW_AUTO_TRIGGER_UX = "newAutoTriggerUX"
         private val LOG = getLogger<CodeWhispererFeatureConfigService>()