Refactor to simplify and isolate chat communication workflow (#51)

This change refactors the chat communication workflow to simplify and isolate responsibilities per class. This helps restrict access to certain components like browser and server to dedicated classes as defined below which in turn would help in improving testability. This design also helps easily add handling for more message types such as quick actions which returns ChatResult for rendering in the ChatUI.

Changes include:
* AmazonQViewActionHandler is now solely responsible for parsing and handling any actions that originate from within the WebView i.e messages from WebView/Chat UI.
* ChatCommunicationManager orchestrates handling communication between chat server and chat UI.
  * For requests to server, it takes the JSON object and converts it to the corresponding parameter Class and calls ChatMessageProvider.
  * For response received from server, it takes the response and invokes event to be sent to Chat UI
* AmazonQWebView is responsible for handling any interactions with the webview. It registers a listener ChatUiRequestListener for OnSendToChat event and sends message to webview for display.
* ChatMessage is the downstream layer responsbile for final communication with the AmazonQLSPServer. This class will ultimately also hold the encryption/decryption of request/response.
* ChatMessageProvider currently only wraps communication over to ChatMessage but eventually it will house more logic that would involve handling different types of messages eg. add tab/remove tab.
* ChatPartialResultMap has been simplified to hold a map of partial result token to tabIDs. Each chatMessage object is uniquely identified by it's tabId. If needed, this can also be further replaced in favor of direct access to the map in appropriate class.

plugin/META-INF/MANIFEST.MF

@@ -115,4 +115,4 @@ Bundle-Classpath: .,
  target/dependency/utils-2.25.33.jar,
  target/dependency/xml-apis-ext-1.3.04.jar,
  target/dependency/xmlgraphics-commons-2.9.jar,
- target/dependency/xz-1.9.jar
+ target/dependency/xz-1.9.jar

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java

@@ -3,17 +3,22 @@
 package software.aws.toolkits.eclipse.amazonq.chat;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+import java.util.Objects;
 import java.util.UUID;
 
-import org.eclipse.swt.browser.Browser;
+import org.eclipse.lsp4j.ProgressParams;
 
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommand;
+import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommandName;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
 import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
 import software.aws.toolkits.eclipse.amazonq.util.JsonHandler;
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
+import software.aws.toolkits.eclipse.amazonq.util.ProgressNotficationUtils;
+import software.aws.toolkits.eclipse.amazonq.views.ChatUiRequestListener;
 import software.aws.toolkits.eclipse.amazonq.views.model.Command;
 
 /**
@@ -29,6 +34,7 @@ public final class ChatCommunicationManager {
     private final JsonHandler jsonHandler;
     private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
     private final ChatPartialResultMap chatPartialResultMap;
+    private ChatUiRequestListener chatUiRequestListener;
 
     private ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
@@ -43,13 +49,17 @@ public static synchronized ChatCommunicationManager getInstance() {
         return instance;
     }
 
-    public CompletableFuture<ChatResult> sendMessageToChatServer(final Browser browser, final Command command, final Object params) {
+    public CompletableFuture<ChatResult> sendMessageToChatServer(final Command command, final Object params) {
         return chatMessageProvider.thenCompose(chatMessageProvider -> {
             try {
                 switch (command) {
                     case CHAT_SEND_PROMPT:
                         ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
-                        return chatMessageProvider.sendChatPrompt(browser, chatRequestParams);
+                    return sendChatRequest(chatRequestParams.getTabId(), token -> {
+                        chatRequestParams.setPartialResultToken(token);
+
+                        return chatMessageProvider.sendChatPrompt(chatRequestParams);
+                    });
                     case CHAT_READY:
                         chatMessageProvider.sendChatReady();
                         return CompletableFuture.completedFuture(null);
@@ -67,39 +77,102 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Browser brows
         });
     }
 
-    public void sendMessageToChatUI(final Browser browser, final ChatUIInboundCommand command) {
-        String message = jsonHandler.serialize(command);
-
-        String script = "window.postMessage(" + message + ");";
-        browser.getDisplay().asyncExec(() -> {
-            browser.evaluate(script);
+    private CompletableFuture<ChatResult> sendChatRequest(final String tabId,
+            final Function<String, CompletableFuture<ChatResult>> action) {
+        // Retrieving the chat result is expected to be a long-running process with
+        // intermittent progress notifications being sent
+        // from the LSP server. The progress notifications provide a token and a partial
+        // result Object - we are utilizing a token to
+        // ChatMessage mapping to acquire the associated ChatMessage so we can formulate
+        // a message for the UI.
+        String partialResultToken = addPartialChatMessage(tabId);
+
+        return action.apply(partialResultToken).thenApply(result -> {
+            // The mapping entry no longer needs to be maintained once the final result is
+            // retrieved.
+            removePartialChatMessage(partialResultToken);
+            // show chat response in Chat UI
+            ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
+                    ChatUIInboundCommandName.ChatPrompt.toString(), tabId, result, false);
+            sendMessageToChatUI(chatUIInboundCommand);
+            return result;
         });
     }
 
+    public void setChatUiRequestListener(final ChatUiRequestListener listener) {
+        chatUiRequestListener = listener;
+    }
+
+    public void removeListener() {
+        chatUiRequestListener =  null;
+    }
+
     /*
-     * Gets the partial chat message using the provided token.
+     * Sends message to Chat UI to show in webview
      */
-    public ChatMessage getPartialChatMessage(final String partialResultToken) {
-        return chatPartialResultMap.getValue(partialResultToken);
+    public void sendMessageToChatUI(final ChatUIInboundCommand command) {
+        if (chatUiRequestListener != null) {
+            String message = jsonHandler.serialize(command);
+            chatUiRequestListener.onSendToChatUi(message);
+        }
     }
 
     /*
-     * Adds an entry to the partialResultToken to ChatMessage map.
+     * Handles chat progress notifications from the Amazon Q LSP server.
+     * - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
+     * - Other notifications are ignored at this time.
+     * - Sends a partial chat prompt message to the webview.
      */
-    public String addPartialChatMessage(final ChatMessage chatMessage) {
-        String partialResultToken = UUID.randomUUID().toString();
+    public void handlePartialResultProgressNotification(final ProgressParams params) {
+        String token = ProgressNotficationUtils.getToken(params);
+        String tabId = getPartialChatMessage(token);
+
+        if (tabId == null || tabId.isEmpty()) {
+            return;
+        }
+
+        // Check to ensure Object is sent in params
+        if (params.getValue().isLeft() || Objects.isNull(params.getValue().getRight())) {
+            throw new AmazonQPluginException("Error occurred while handling partial result notification: expected Object value");
+        }
 
-        // Indicator for the server to send partial result notifications
-        chatMessage.getChatRequestParams().setPartialResultToken(partialResultToken);
+        ChatResult partialChatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
 
-        chatPartialResultMap.setEntry(partialResultToken, chatMessage);
+        // Check to ensure the body has content in order to keep displaying the spinner while loading
+        if (partialChatResult.body() == null || partialChatResult.body().length() == 0) {
+            return;
+        }
+
+        ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
+            ChatUIInboundCommandName.ChatPrompt.toString(),
+            tabId,
+            partialChatResult,
+            true
+        );
+
+        sendMessageToChatUI(chatUIInboundCommand);
+    }
+
+    /*
+     * Gets the partial chat message represented by the tabId using the provided token.
+     */
+    private String getPartialChatMessage(final String partialResultToken) {
+        return chatPartialResultMap.getValue(partialResultToken);
+    }
+
+    /*
+     * Adds an entry to the partialResultToken to ChatMessage's tabId map.
+     */
+    private String addPartialChatMessage(final String tabId) {
+        String partialResultToken = UUID.randomUUID().toString();
+        chatPartialResultMap.setEntry(partialResultToken, tabId);
         return partialResultToken;
     }
 
     /*
-     * Removes an entry from the partialResultToken to ChatMessage map.
+     * Removes an entry from the partialResultToken to ChatMessage's tabId map.
      */
-    public void removePartialChatMessage(final String partialResultToken) {
+    private void removePartialChatMessage(final String partialResultToken) {
         chatPartialResultMap.removeEntry(partialResultToken);
     }
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java

@@ -4,50 +4,27 @@
 
 import java.util.concurrent.CompletableFuture;
 
-import org.eclipse.swt.browser.Browser;
-
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
+import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
 
 public final class ChatMessage {
-    private final Browser browser;
-    private final ChatRequestParams chatRequestParams;
     private final AmazonQLspServer amazonQLspServer;
-    private final ChatCommunicationManager chatCommunicationManager;
 
-    public ChatMessage(final AmazonQLspServer amazonQLspServer, final Browser browser, final ChatRequestParams chatRequestParams) {
+    public ChatMessage(final AmazonQLspServer amazonQLspServer) {
         this.amazonQLspServer = amazonQLspServer;
-        this.browser = browser;
-        this.chatRequestParams = chatRequestParams;
-        this.chatCommunicationManager = ChatCommunicationManager.getInstance();
-    }
-
-    public Browser getBrowser() {
-        return browser;
     }
 
-    public ChatRequestParams getChatRequestParams() {
-        return chatRequestParams;
+    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
+        return amazonQLspServer.sendChatPrompt(chatRequestParams);
     }
 
-    public String getPartialResultToken() {
-        return chatRequestParams.getPartialResultToken();
+    public void sendChatReady() {
+        amazonQLspServer.chatReady();
     }
 
-    public CompletableFuture<ChatResult> sendChatMessageWithProgress() {
-        // Retrieving the chat result is expected to be a long-running process with intermittent progress notifications being sent
-        // from the LSP server. The progress notifications provide a token and a partial result Object - we are utilizing a token to
-        // ChatMessage mapping to acquire the associated ChatMessage so we can formulate a message for the UI.
-        String partialResultToken = chatCommunicationManager.addPartialChatMessage(this);
-
-        CompletableFuture<ChatResult> chatResult = amazonQLspServer.sendChatPrompt(chatRequestParams)
-            .thenApply(result -> {
-                // The mapping entry no longer needs to be maintained once the final result is retrieved.
-                chatCommunicationManager.removePartialChatMessage(partialResultToken);
-                return result;
-            });
-
-        return chatResult;
+    public void sendTabAdd(final GenericTabParams tabParams) {
+        amazonQLspServer.tabAdd(tabParams);
     }
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java

@@ -3,17 +3,11 @@
 package software.aws.toolkits.eclipse.amazonq.chat;
 
 import java.util.concurrent.CompletableFuture;
-
-import org.eclipse.swt.browser.Browser;
-
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
-import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
 import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
-import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
-import software.aws.toolkits.eclipse.amazonq.views.model.Command;
 
 public final class ChatMessageProvider {
 
@@ -28,29 +22,19 @@ private ChatMessageProvider(final AmazonQLspServer amazonQLspServer) {
         this.amazonQLspServer = amazonQLspServer;
     }
 
-    public CompletableFuture<ChatResult>  sendChatPrompt(final Browser browser, final ChatRequestParams chatRequestParams) {
-        ChatMessage chatMessage = new ChatMessage(amazonQLspServer, browser, chatRequestParams);
-        return chatMessage.sendChatMessageWithProgress();
+    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
+        ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
+        return chatMessage.sendChatPrompt(chatRequestParams);
     }
 
     public void sendChatReady() {
-        try {
-            PluginLogger.info("Sending " + Command.CHAT_READY + " message to Amazon Q LSP server");
-            amazonQLspServer.chatReady();
-        } catch (Exception e) {
-            PluginLogger.error("Error occurred while sending " + Command.CHAT_READY + " message to Amazon Q LSP server", e);
-            throw new AmazonQPluginException(e);
-        }
+        ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
+        chatMessage.sendChatReady();
     }
 
     public void sendTabAdd(final GenericTabParams tabParams) {
-        try {
-            PluginLogger.info("Sending " + Command.CHAT_TAB_ADD + " message to Amazon Q LSP server");
-            amazonQLspServer.tabAdd(tabParams);
-        } catch (Exception e) {
-            PluginLogger.error("Error occurred while sending " + Command.CHAT_TAB_ADD + " message to Amazon Q LSP server", e);
-            throw new AmazonQPluginException(e);
-        }
+        ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
+        chatMessage.sendTabAdd(tabParams);
     }
 
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java

@@ -8,37 +8,37 @@
 
 /**
  * ChatPartialResultMap is a utility class responsible for managing the mapping between
- * partial result tokens and their corresponding ChatMessage objects in the Amazon Q plugin for Eclipse.
+ * partial result tokens and their corresponding ChatMessage objects represented by tabId in the Amazon Q plugin for Eclipse.
  *
  * The Language Server Protocol (LSP) server sends progress notifications during long-running operations,
  * such as processing chat requests. These notifications include a token that identifies the specific operation
  * and a partial result object containing the progress information.
  *
  * This class maintains a concurrent map (tokenToChatMessageMap) that associates each token with
- * its respective ChatMessage object. This mapping is crucial for correctly updating the chat UI
+ * its respective ChatMessage object identified via the tabId. This mapping is crucial for correctly updating the chat UI
  * with the latest progress information as it becomes available from the LSP server.
  *
  * The progress notifications are handled by the {@link AmazonQLspClientImpl#notifyProgress(ProgressParams)}
- * method, which retrieves the corresponding ChatMessage object from the tokenToChatMessageMap using
+ * method, which retrieves the corresponding tabId associated with the ChatMessage object from the tokenToChatMessageMap using
  * the token provided in the ProgressParams. The ChatMessage can then be updated with the partial result.
  */
 public final class ChatPartialResultMap {
 
-    private final Map<String, ChatMessage> tokenToChatMessageMap;
+    private final Map<String, String> tokenToChatMessageMap;
 
     public ChatPartialResultMap() {
-        tokenToChatMessageMap = new ConcurrentHashMap<String, ChatMessage>();
+        tokenToChatMessageMap = new ConcurrentHashMap<String, String>();
     }
 
-    public void setEntry(final String token, final ChatMessage chatMessage) {
-        tokenToChatMessageMap.put(token, chatMessage);
+    public void setEntry(final String token, final String tabId) {
+        tokenToChatMessageMap.put(token, tabId);
     }
 
     public void removeEntry(final String token) {
         tokenToChatMessageMap.remove(token);
     }
 
-    public ChatMessage getValue(final String token) {
+    public String getValue(final String token) {
         return tokenToChatMessageMap.getOrDefault(token, null);
     }
 

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java

@@ -13,13 +13,13 @@
 import org.eclipse.lsp4j.ConfigurationParams;
 import org.eclipse.lsp4j.ProgressParams;
 
+import software.aws.toolkits.eclipse.amazonq.chat.ChatCommunicationManager;
 import software.aws.toolkits.eclipse.amazonq.configuration.PluginStore;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.ConnectionMetadata;
 import software.aws.toolkits.eclipse.amazonq.lsp.model.SsoProfileData;
 import software.aws.toolkits.eclipse.amazonq.util.Constants;
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
 import software.aws.toolkits.eclipse.amazonq.util.ThreadingUtils;
-import software.aws.toolkits.eclipse.amazonq.views.AmazonQChatViewActionHandler;
 
 @SuppressWarnings("restriction")
 public class AmazonQLspClientImpl extends LanguageClientImpl implements AmazonQLspClient {
@@ -34,11 +34,6 @@ public final CompletableFuture<ConnectionMetadata> getConnectionMetadata() {
         return CompletableFuture.completedFuture(metadata);
     }
 
-    /*
-     * Handles the progress notifications received from the LSP server.
-     * - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
-     * - Other notifications are ignored at this time.
-     */
     @Override
     public final CompletableFuture<List<Object>> configuration(final ConfigurationParams configurationParams) {
         if (configurationParams.getItems().size() == 0) {
@@ -58,13 +53,18 @@ public final CompletableFuture<List<Object>> configuration(final ConfigurationPa
         return CompletableFuture.completedFuture(output);
     }
 
+    /*
+     * Handles the progress notifications received from the LSP server.
+     * - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
+     * - Other notifications are ignored at this time.
+     */
     @Override
     public final void notifyProgress(final ProgressParams params) {
-        AmazonQChatViewActionHandler chatActionHandler = new AmazonQChatViewActionHandler();
+        var chatCommunicationManager = ChatCommunicationManager.getInstance();
 
         ThreadingUtils.executeAsyncTask(() -> {
             try {
-                chatActionHandler.handlePartialResultProgressNotification(params);
+                chatCommunicationManager.handlePartialResultProgressNotification(params);
             } catch (Exception e) {
                 PluginLogger.error("Error processing partial result progress notification", e);
             }

plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/manager/LspConstants.java

@@ -10,7 +10,7 @@ private LspConstants() {
         // Prevent instantiation
     }
 
-    public static final String CW_MANIFEST_URL = "https://dtqjflaii39q2.cloudfront.net/codewhisperer/0/manifest.json";
+    public static final String CW_MANIFEST_URL = "https://aws-toolkit-language-servers.amazonaws.com/eclipse/0/manifest.json";
 
     public static final String CW_LSP_FILENAME = "aws-lsp-codewhisperer.js";
     public static final String NODE_EXECUTABLE_PREFIX = "node";