refactor(amazonq): Pass mynah references to ui handlers (aws#7046)

jpinkney-aws · web-flow · commit 0288f4dffac8 · 2025-04-15T16:38:19.000-04:00
## Problem - Flare's hybrid chat injects MynahUI reference after chat connector creation ## Solution - Implement delayed resolution so that onces MynahUI is injected everything dynamically resolves --- - 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.
diff --git a/packages/core/src/amazonq/commons/types.ts b/packages/core/src/amazonq/commons/types.ts
@@ -11,6 +11,7 @@ import { DiffTreeFileInfo } from '../webview/ui/diffTree/types'
 import { Messenger } from './connector/baseMessenger'
 import { FeatureClient } from '../client/client'
 import { TelemetryHelper } from '../util/telemetryHelper'
+import { MynahUI } from '@aws/mynah-ui'
 
 export enum FollowUpTypes {
     // UnitTestGeneration
@@ -164,3 +165,10 @@ export enum MetricDataResult {
     Error = 'Error',
     LlmFailure = 'LLMFailure',
 }
+
+/**
+ * Note: Passing a reference around allows us to lazily inject mynah UI into
+ * connectors and handlers. This is done to supported "hybrid chat", which
+ * injects mynah UI _after_ the connector has already been created
+ */
+export type MynahUIRef = { mynahUI: MynahUI | undefined }
diff --git a/packages/core/src/amazonq/webview/ui/followUps/handler.ts b/packages/core/src/amazonq/webview/ui/followUps/handler.ts
@@ -8,21 +8,21 @@ import { Connector } from '../connector'
 import { TabsStorage } from '../storages/tabsStorage'
 import { WelcomeFollowupType } from '../apps/amazonqCommonsConnector'
 import { AuthFollowUpType } from './generator'
-import { FollowUpTypes } from '../../../commons/types'
+import { FollowUpTypes, MynahUIRef } from '../../../commons/types'
 
 export interface FollowUpInteractionHandlerProps {
-    mynahUI: MynahUI
+    mynahUIRef: MynahUIRef
     connector: Connector
     tabsStorage: TabsStorage
 }
 
 export class FollowUpInteractionHandler {
-    private mynahUI: MynahUI
+    private mynahUIRef: MynahUIRef
     private connector: Connector
     private tabsStorage: TabsStorage
 
     constructor(props: FollowUpInteractionHandlerProps) {
-        this.mynahUI = props.mynahUI
+        this.mynahUIRef = props.mynahUIRef
         this.connector = props.connector
         this.tabsStorage = props.tabsStorage
     }
@@ -41,6 +41,11 @@ export class FollowUpInteractionHandler {
             this.connector.help(tabID)
             return
         }
+
+        if (!this.mynahUI) {
+            return
+        }
+
         // we need to check if there is a prompt
         // which will cause an api call
         // then we can set the loading state to true
@@ -70,7 +75,7 @@ export class FollowUpInteractionHandler {
         }
 
         const addChatItem = (tabID: string, messageId: string, options: any[]) => {
-            this.mynahUI.addChatItem(tabID, {
+            this.mynahUI?.addChatItem(tabID, {
                 type: ChatItemType.ANSWER_PART,
                 messageId,
                 followUp: {
@@ -150,7 +155,7 @@ export class FollowUpInteractionHandler {
 
     public onWelcomeFollowUpClicked(tabID: string, welcomeFollowUpType: WelcomeFollowupType) {
         if (welcomeFollowUpType === 'continue-to-chat') {
-            this.mynahUI.addChatItem(tabID, {
+            this.mynahUI?.addChatItem(tabID, {
                 type: ChatItemType.ANSWER,
                 body: 'Ok, please write your question below.',
             })
@@ -159,4 +164,8 @@ export class FollowUpInteractionHandler {
             return
         }
     }
+
+    private get mynahUI(): MynahUI | undefined {
+        return this.mynahUIRef.mynahUI
+    }
 }
diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts
@@ -57,7 +57,7 @@ export const createMynahUI = (
 ) => {
     const handler = new WebviewUIHandler({
         postMessage: ideApi.postMessage,
-        mynahUI: undefined,
+        mynahUIRef: { mynahUI: undefined },
         enableAgents: amazonQEnabled,
         featureConfigsSerialized,
         welcomeCount,
@@ -85,8 +85,7 @@ export class WebviewUIHandler {
 
     tabsStorage: TabsStorage
 
-    _mynahUI?: MynahUI
-    mynahUIProps?: MynahUIProps
+    mynahUIProps: MynahUIProps
     connector?: Connector
     tabDataGenerator?: TabDataGenerator
 
@@ -108,9 +107,11 @@ export class WebviewUIHandler {
 
     savedContextCommands: MynahUIDataModel['contextCommands']
 
+    mynahUIRef: { mynahUI: MynahUI | undefined }
+
     constructor({
         postMessage,
-        mynahUI,
+        mynahUIRef,
         enableAgents,
         featureConfigsSerialized,
         welcomeCount,
@@ -121,7 +122,7 @@ export class WebviewUIHandler {
         isSM,
     }: {
         postMessage: any
-        mynahUI: MynahUI | undefined
+        mynahUIRef: { mynahUI: MynahUI | undefined }
         enableAgents: boolean
         featureConfigsSerialized: [string, FeatureContext][]
         welcomeCount: number
@@ -136,6 +137,7 @@ export class WebviewUIHandler {
         this.disclaimerCardActive = !disclaimerAcknowledged
         this.isSMUS = isSMUS ?? false
         this.isSM = isSM ?? false
+        this.mynahUIRef = mynahUIRef
 
         this.responseMetadata = new Map<string, string[]>()
 
@@ -197,7 +199,7 @@ export class WebviewUIHandler {
                 this.isDocEnabled = isAmazonQEnabled
 
                 this.quickActionHandler = new QuickActionHandler({
-                    mynahUI: this.mynahUI!,
+                    mynahUIRef: this.mynahUIRef,
                     connector: this.connector!,
                     tabsStorage: this.tabsStorage,
                     isFeatureDevEnabled: this.isFeatureDevEnabled,
@@ -639,7 +641,7 @@ export class WebviewUIHandler {
             },
         })
 
-        this.mynahUI = new MynahUI({
+        this.mynahUIProps = {
             onReady: this.connector.uiReady,
             onTabAdd: (tabID: string) => {
                 /**
@@ -648,7 +650,7 @@ export class WebviewUIHandler {
                  */
                 if (welcomeCount + 1 >= welcomeCountThreshold) {
                     this.tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
-                    mynahUI?.updateTabDefaults({
+                    this.mynahUI?.updateTabDefaults({
                         store: {
                             ...this.tabDataGenerator?.getTabData('cwc', true, undefined, this.isSMUS),
                             tabHeaderDetails: void 0,
@@ -985,7 +987,9 @@ export class WebviewUIHandler {
                     },
                 ],
             },
-        })
+        }
+
+        this.mynahUIRef = { mynahUI: new MynahUI(this.mynahUIProps) }
 
         /**
          * Update the welcome count if we've initially shown
@@ -996,12 +1000,12 @@ export class WebviewUIHandler {
         }
 
         this.followUpsInteractionHandler = new FollowUpInteractionHandler({
-            mynahUI: this.mynahUI,
+            mynahUIRef: this.mynahUIRef,
             connector: this.connector,
             tabsStorage: this.tabsStorage,
         })
         this.quickActionHandler = new QuickActionHandler({
-            mynahUI: this.mynahUI,
+            mynahUIRef: this.mynahUIRef,
             connector: this.connector,
             tabsStorage: this.tabsStorage,
             isFeatureDevEnabled: this.isFeatureDevEnabled,
@@ -1011,12 +1015,12 @@ export class WebviewUIHandler {
             isDocEnabled: this.isDocEnabled,
         })
         this.textMessageHandler = new TextMessageHandler({
-            mynahUI: this.mynahUI,
+            mynahUIRef: this.mynahUIRef,
             connector: this.connector,
             tabsStorage: this.tabsStorage,
         })
         this.messageController = new MessageController({
-            mynahUI: this.mynahUI,
+            mynahUIRef: this.mynahUIRef,
             connector: this.connector,
             tabsStorage: this.tabsStorage,
             isFeatureDevEnabled: this.isFeatureDevEnabled,
@@ -1078,45 +1082,7 @@ export class WebviewUIHandler {
         return {}
     }
 
-    set mynahUI(mynahUI: MynahUI | undefined) {
-        this._mynahUI = mynahUI
-
-        this.followUpsInteractionHandler = new FollowUpInteractionHandler({
-            mynahUI: this.mynahUI!,
-            connector: this.connector!,
-            tabsStorage: this.tabsStorage,
-        })
-
-        this.quickActionHandler = new QuickActionHandler({
-            mynahUI: this.mynahUI!,
-            connector: this.connector!,
-            tabsStorage: this.tabsStorage,
-            isFeatureDevEnabled: this.isFeatureDevEnabled,
-            isGumbyEnabled: this.isGumbyEnabled,
-            isScanEnabled: this.isScanEnabled,
-            isTestEnabled: this.isTestEnabled,
-            isDocEnabled: this.isDocEnabled,
-        })
-
-        this.textMessageHandler = new TextMessageHandler({
-            mynahUI: this.mynahUI!,
-            connector: this.connector!,
-            tabsStorage: this.tabsStorage,
-        })
-
-        this.messageController = new MessageController({
-            mynahUI: this.mynahUI!,
-            connector: this.connector!,
-            tabsStorage: this.tabsStorage,
-            isFeatureDevEnabled: this.isFeatureDevEnabled,
-            isGumbyEnabled: this.isGumbyEnabled,
-            isScanEnabled: this.isScanEnabled,
-            isTestEnabled: this.isTestEnabled,
-            isDocEnabled: this.isDocEnabled,
-        })
-    }
-
     get mynahUI(): MynahUI | undefined {
-        return this._mynahUI
+        return this.mynahUIRef.mynahUI
     }
 }
diff --git a/packages/core/src/amazonq/webview/ui/messages/controller.ts b/packages/core/src/amazonq/webview/ui/messages/controller.ts
@@ -8,9 +8,10 @@ import { Connector } from '../connector'
 import { TabType, TabsStorage } from '../storages/tabsStorage'
 import { TabDataGenerator } from '../tabs/generator'
 import { uiComponentsTexts } from '../texts/constants'
+import { MynahUIRef } from '../../../commons/types'
 
 export interface MessageControllerProps {
-    mynahUI: MynahUI
+    mynahUIRef: MynahUIRef
     connector: Connector
     tabsStorage: TabsStorage
     isFeatureDevEnabled: boolean
@@ -22,13 +23,13 @@ export interface MessageControllerProps {
 }
 
 export class MessageController {
-    private mynahUI: MynahUI
+    private mynahUIRef: MynahUIRef
     private connector: Connector
     private tabsStorage: TabsStorage
     private tabDataGenerator: TabDataGenerator
 
     constructor(props: MessageControllerProps) {
-        this.mynahUI = props.mynahUI
+        this.mynahUIRef = props.mynahUIRef
         this.connector = props.connector
         this.tabsStorage = props.tabsStorage
         this.tabDataGenerator = new TabDataGenerator({
@@ -43,6 +44,10 @@ export class MessageController {
 
     public sendSelectedCodeToTab(message: ChatItem, command: string = ''): string | undefined {
         const selectedTab = { ...this.tabsStorage.getSelectedTab() }
+        if (!this.mynahUI) {
+            return
+        }
+
         if (
             selectedTab?.id === undefined ||
             selectedTab?.type === undefined ||
@@ -76,6 +81,9 @@ export class MessageController {
 
     public sendMessageToTab(message: ChatItem, tabType: TabType, command: string = ''): string | undefined {
         const selectedTab = this.tabsStorage.getSelectedTab()
+        if (!this.mynahUI) {
+            return
+        }
 
         if (
             selectedTab !== undefined &&
@@ -141,4 +149,8 @@ export class MessageController {
             return newTabID
         }
     }
+
+    private get mynahUI(): MynahUI | undefined {
+        return this.mynahUIRef.mynahUI
+    }
 }
diff --git a/packages/core/src/amazonq/webview/ui/messages/handler.ts b/packages/core/src/amazonq/webview/ui/messages/handler.ts
@@ -6,20 +6,21 @@
 import { ChatItemType, ChatPrompt, MynahUI } from '@aws/mynah-ui'
 import { Connector } from '../connector'
 import { TabsStorage } from '../storages/tabsStorage'
+import { MynahUIRef } from '../../../commons/types'
 
 export interface TextMessageHandlerProps {
-    mynahUI: MynahUI
+    mynahUIRef: MynahUIRef
     connector: Connector
     tabsStorage: TabsStorage
 }
 
 export class TextMessageHandler {
-    private mynahUI: MynahUI
+    private mynahUIRef: MynahUIRef
     private connector: Connector
     private tabsStorage: TabsStorage
 
     constructor(props: TextMessageHandlerProps) {
-        this.mynahUI = props.mynahUI
+        this.mynahUIRef = props.mynahUIRef
         this.connector = props.connector
         this.tabsStorage = props.tabsStorage
     }
@@ -29,6 +30,10 @@ export class TextMessageHandler {
         this.tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
         this.tabsStorage.resetTabTimer(tabID)
         this.connector.onUpdateTabType(tabID)
+        if (!this.mynahUI) {
+            return
+        }
+
         this.mynahUI.addChatItem(tabID, {
             type: ChatItemType.PROMPT,
             body: chatPrompt.escapedPrompt,
@@ -50,4 +55,8 @@ export class TextMessageHandler {
             })
             .then(() => {})
     }
+
+    private get mynahUI(): MynahUI | undefined {
+        return this.mynahUIRef.mynahUI
+    }
 }
diff --git a/packages/core/src/amazonq/webview/ui/quickActions/handler.ts b/packages/core/src/amazonq/webview/ui/quickActions/handler.ts
