Refactor to simplify and isolate chat communication workflow

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";