Merged with latest changes

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

@@ -13,25 +13,27 @@
 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.views.model.Command;
 
 /**
- * ChatCommunicationManager is responsible for managing communication between
- * the Amazon Q Eclipse Plugin and the LSP server as well as communication
- * between the Amazon Q Eclipse Plugin and the webview. It is implemented 
- * as a singleton to centralize control of all communication in the plugin.
+ * ChatCommunicationManager is a central component of the Amazon Q Eclipse Plugin that
+ * acts as a bridge between the plugin's UI and the LSP server. It is also responsible 
+ * for managing communication between the plugin and the webview used for displaying 
+ * chat conversations. It is implemented as a singleton to centralize control of all 
+ * communication in the plugin.
  */
 public final class ChatCommunicationManager {
     private static ChatCommunicationManager instance;
 
     private final JsonHandler jsonHandler;
     private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
-    private final ChatPartialResultManager chatPartialResultManager;
+    private final ChatPartialResultMap chatPartialResultMap;
 
     private ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
         this.chatMessageProvider = ChatMessageProvider.createAsync();
-        this.chatPartialResultManager = ChatPartialResultManager.getInstance();
+        this.chatPartialResultMap = new ChatPartialResultMap();
     }
 
     public static synchronized ChatCommunicationManager getInstance() {
@@ -41,13 +43,13 @@ public static synchronized ChatCommunicationManager getInstance() {
         return instance;
     }
 
-    public CompletableFuture<ChatResult> sendMessageToChatServer(final Command command, final Object params) {
+    public CompletableFuture<ChatResult> sendMessageToChatServer(final Browser browser, 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(chatRequestParams);
+                        return chatMessageProvider.sendChatPrompt(browser, chatRequestParams);
                     case CHAT_READY:
                         chatMessageProvider.sendChatReady();
                         return CompletableFuture.completedFuture(null);
@@ -68,7 +70,7 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Command comma
     public void sendMessageToChatUI(final Browser browser, final ChatUIInboundCommand command) {
         // Mynah-ui will not render the partial result if null values are included. Must serialize with ignoreNulls set to True.
         Boolean ignoreNull = true;
-        String message = this.jsonHandler.serialize(command, ignoreNull);
+        String message = jsonHandler.serialize(command, ignoreNull);
 
         String script = "window.postMessage(" + message + ");";
         browser.getDisplay().asyncExec(() -> {

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

@@ -2,16 +2,13 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat;
 
-import java.util.concurrent.ExecutionException;
+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.exception.AmazonQPluginException;
 import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
-import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
-import software.aws.toolkits.eclipse.amazonq.views.model.Command;
 
 public class ChatMessage {
     private final Browser browser;
@@ -38,23 +35,19 @@ public String getPartialResultToken() {
         return chatRequestParams.getPartialResultToken();
     }
 
-    public ChatResult sendChatMessageWithProgress() {
-        try {
-            // 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 result - we are utilizing this token to
-            // ChatMessage mapping to acquire the associated ChatMessage.
-            String partialResultToken = chatCommunicationManager.addPartialChatMessage(this);
-
-            PluginLogger.info("Sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server");
-            ChatResult chatResult = amazonQLspServer.sendChatPrompt(chatRequestParams).get();
-
-            // The mapping entry no longer needs to be maintained once the final result is retrieved.
-            chatCommunicationManager.removePartialChatMessage(partialResultToken);
-
-            return chatResult;
-        } catch (InterruptedException | ExecutionException e) {
-            PluginLogger.error("Error occurred while sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server", e);
-            throw new AmazonQPluginException(e);
-        }
+    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;
     }
 }

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

@@ -31,7 +31,7 @@ private ChatMessageProvider(AmazonQLspServer amazonQLspServer) {
     }
 
     public CompletableFuture<ChatResult>  sendChatPrompt(Browser browser, final ChatRequestParams chatRequestParams) {
-        ChatMessage chatMessage = new ChatMessage(amazonQLspServer, chatRequestParams);
+        ChatMessage chatMessage = new ChatMessage(amazonQLspServer, browser, chatRequestParams);
         return chatMessage.sendChatMessageWithProgress();
     }
 

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

@@ -6,18 +6,24 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
+import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspClientImpl;
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
 
 /**
- * ChatPartialResultMap is responsible for maintaining a mapping between
- * partial result tokens and the associated ChatMessage objects.
- * 
- * $/progress notifications are caught and handled in the AmazonQLspClientImpl 
- * notifyProgress method. Within a progress notification, we are provided ProgressParams
- * containing a token and a partial result object. The tokenToChatMessage map in
- * this class allows us to find the associated ChatMessage associated with the token.
+ * 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.
  *
- * @see AmazonQLspClientImpl#notifyProgress(ProgressParams)
+ * 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
+ * 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
+ * the token provided in the ProgressParams. The ChatMessage can then be updated with the partial result.
  */
 public class ChatPartialResultMap {
 

plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java

@@ -11,6 +11,7 @@
 import software.aws.toolkits.eclipse.amazonq.chat.ChatCommunicationManager;
 import software.aws.toolkits.eclipse.amazonq.chat.ChatMessage;
 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.exception.AmazonQPluginException;
@@ -41,12 +42,12 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
 
         switch (command) {
             case CHAT_SEND_PROMPT:
-                chatCommunicationManager.sendMessageToChatServer(command, params)
+                chatCommunicationManager.sendMessageToChatServer(browser, command, params)
                     .thenAccept(chatResult -> {
                         ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
                         ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
                             ChatUIInboundCommandName.ChatPrompt.toString(),
-                            chatRequestParams.tabId(),
+                            chatRequestParams.getTabId(),
                             chatResult,
                             false
                         );