fix(amazonq): right click options dont send context if chat not loaded

Problem:

When a user starts up VSC, hasn't opened Q chat yet, then highlights text and
does a right click option like "send to prompt" THEN the whole prompt fails to
make it to chat.

This is because there is a race condition between the message being sent to the webview,
and the webview completed loading and ready to recieve messages.

Solution:

In the appToMessage publisher, verify that the chat is ready to recieve messages based
on if the webview sent the chat ui ready message. Otherwise wait until we get it,
and then send the message.

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/core/src/amazonq/apps/initContext.ts

@@ -4,20 +4,25 @@
  */
 
 import { EventEmitter } from 'vscode'
-import { MessagePublisher } from '../messages/messagePublisher'
+import { MessagePublisher, UiReadyMessagePublisher } from '../messages/messagePublisher'
 import { MessageListener } from '../messages/messageListener'
 import { TabType } from '../webview/ui/storages/tabsStorage'
 
 export interface AmazonQAppInitContext {
     registerWebViewToAppMessagePublisher(eventEmitter: MessagePublisher<any>, tabType: TabType): void
     getAppsToWebViewMessagePublisher(): MessagePublisher<any>
     onDidChangeAmazonQVisibility: EventEmitter<boolean>
+    setAmazonQUiReady: () => void
 }
 
 export class DefaultAmazonQAppInitContext implements AmazonQAppInitContext {
+    private _isUiReady: boolean = false
     private readonly appsToWebViewEventEmitter = new EventEmitter<any>()
     private readonly appsToWebViewMessageListener = new MessageListener<any>(this.appsToWebViewEventEmitter)
-    private readonly appsToWebViewMessagePublisher = new MessagePublisher<any>(this.appsToWebViewEventEmitter)
+    private readonly appsToWebViewMessagePublisher = new UiReadyMessagePublisher<any>(
+        this.appsToWebViewEventEmitter,
+        () => this._isUiReady
+    ) // this should know the state of the webview so it can sleep until webview is loaded, then send message
     private readonly webViewToAppsMessagePublishers: Map<TabType, MessagePublisher<any>> = new Map()
     public readonly onDidChangeAmazonQVisibility = new EventEmitter<boolean>()
 
@@ -44,4 +49,11 @@ export class DefaultAmazonQAppInitContext implements AmazonQAppInitContext {
     getAppsToWebViewMessagePublisher(): MessagePublisher<any> {
         return this.appsToWebViewMessagePublisher
     }
+
+    /**
+     * Indicate the Q Chat UI is ready to recieve messages
+     */
+    setAmazonQUiReady(): void {
+        this._isUiReady = true
+    }
 }

packages/core/src/amazonq/messages/messagePublisher.ts

@@ -4,6 +4,7 @@
  */
 
 import { EventEmitter } from 'vscode'
+import { getLogger, waitUntil } from '../../shared'
 
 export class MessagePublisher<T> {
     constructor(private readonly eventEmitter: EventEmitter<T>) {}
@@ -12,3 +13,45 @@ export class MessagePublisher<T> {
         this.eventEmitter.fire(event)
     }
 }
+
+/**
+ * Same as {@link MessagePublisher}, but will wait until the UI indicates it
+ * is ready to recieve messages, before the message is published.
+ *
+ * This solves a problem when running a right click menu option like
+ * "Send To Prompt" BUT chat is not opened yet, it would result in the prompt failing to
+ * be recieved by chat.
+ */
+export class UiReadyMessagePublisher<T> extends MessagePublisher<T> {
+    constructor(
+        eventEmitter: EventEmitter<T>,
+        private readonly isChatUiReady: () => boolean
+    ) {
+        super(eventEmitter)
+    }
+    public override publish(event: T): void {
+        // immediately send if Chat UI is ready
+        if (this.isChatUiReady()) {
+            super.publish(event)
+            return
+        }
+        // Otherwise wait, since it is probably still loading
+        void waitUntil(
+            async () => {
+                if (this.isChatUiReady()) {
+                    super.publish(event)
+                    return true
+                }
+                return false
+            },
+            {
+                timeout: 5000,
+                interval: 100,
+            }
+        ).then((res) => {
+            if (!res) {
+                getLogger('chat').error('Message publisher timed out waiting for UI to be ready')
+            }
+        })
+    }
+}

packages/core/src/amazonq/webview/messages/messageDispatcher.ts

@@ -13,6 +13,7 @@ import { telemetry } from '../../../shared/telemetry'
 import { AmazonQChatMessageDuration } from '../../messages/chatMessageDuration'
 import { globals, openUrl } from '../../../shared'
 import { isClickTelemetry, isOpenAgentTelemetry } from '../ui/telemetry/actions'
+import { DefaultAmazonQAppInitContext } from '../../apps/initContext'
 
 export function dispatchWebViewMessagesToApps(
     webview: Webview,
@@ -21,12 +22,12 @@ export function dispatchWebViewMessagesToApps(
     webview.onDidReceiveMessage((msg) => {
         switch (msg.command) {
             case 'ui-is-ready': {
+                DefaultAmazonQAppInitContext.instance.setAmazonQUiReady()
                 /**
                  * ui-is-ready isn't associated to any tab so just record the telemetry event and continue.
                  * This would be equivalent of the duration between "user clicked open q" and "ui has become available"
                  * NOTE: Amazon Q UI is only loaded ONCE. The state is saved between each hide/show of the webview.
                  */
-
                 telemetry.webview_load.emit({
                     webviewName: 'amazonq',
                     duration: performance.measure(amazonqMark.uiReady, amazonqMark.open).duration,

packages/core/src/shared/logger/logger.ts

@@ -5,7 +5,7 @@
 
 import * as vscode from 'vscode'
 
-export type LogTopic = 'crashMonitoring' | 'dev/beta' | 'notifications' | 'test' | 'childProcess' | 'unknown'
+export type LogTopic = 'crashMonitoring' | 'dev/beta' | 'notifications' | 'test' | 'childProcess' | 'unknown' | 'chat'
 
 class ErrorLog {
     constructor(