Cleaned up code and added comments

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

@@ -2,6 +2,8 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat;
 
+import java.util.UUID;
+
 import org.eclipse.swt.browser.Browser;
 
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
@@ -15,18 +17,27 @@
 /**
  * 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.
+ * between the Amazon Q Eclipse Plugin and the webview. 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 ChatMessageProvider chatMessageProvider;
-    private final ChatPartialResultManager chatPartialResultManager;
+    private final ChatPartialResultMap chatPartialResultMap;
 
-    public ChatCommunicationManager() {
+    private ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
         this.chatMessageProvider = new ChatMessageProvider();
-        this.chatPartialResultManager = ChatPartialResultManager.getInstance();
+        this.chatPartialResultMap = new ChatPartialResultMap();
+    }
+
+    public static synchronized ChatCommunicationManager getInstance() {
+        if (instance == null) {
+            instance = new ChatCommunicationManager();
+        }
+        return instance;
     }
 
     /*
@@ -69,17 +80,30 @@ public void sendMessageToChatUI(final Browser browser, final ChatUIInboundComman
     }
 
     /*
-     * Checks if a partial result is being processed with the provided token.
+     * Gets the partial chat message using the provided token.
      */
-    public boolean isProcessingPartialChatMessage(String partialResultToken) {
-        return chatPartialResultManager.hasKey(partialResultToken);
+    public ChatMessage getPartialChatMessage(String partialResultToken) {
+        return chatPartialResultMap.getValue(partialResultToken);
     }
 
     /*
-     * Gets the partial chat message using the provided token.
+     * Adds an entry to the partialResultToken to ChatMessage map.
      */
-    public ChatMessage getPartialChatMessage(String partialResultToken) {
-        return chatPartialResultManager.getValue(partialResultToken);
+    public String addPartialChatMessage(ChatMessage chatMessage) {
+        String partialResultToken = UUID.randomUUID().toString();
+
+        // Indicator for the server to send partial result notifications
+        chatMessage.getChatRequestParams().setPartialResultToken(partialResultToken);
+
+        chatPartialResultMap.setEntry(partialResultToken, chatMessage);
+        return partialResultToken;
+    }
+
+    /*
+     * Removes an entry from the partialResultToken to ChatMessage map.
+     */
+    public void removePartialChatMessage(String partialResultToken) {
+        chatPartialResultMap.removeEntry(partialResultToken);
     }
 }
 

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

@@ -2,7 +2,6 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat;
 
-import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 
 import org.eclipse.swt.browser.Browser;
@@ -17,14 +16,14 @@
 public class ChatMessage {
     private final Browser browser;
     private final ChatRequestParams chatRequestParams;
-    private final ChatPartialResultManager chatPartialResultManager;
     private final AmazonQLspServer amazonQLspServer;
+    private final ChatCommunicationManager chatCommunicationManager;
 
     public ChatMessage(final AmazonQLspServer amazonQLspServer, final Browser browser, final ChatRequestParams chatRequestParams) {
         this.amazonQLspServer = amazonQLspServer;
         this.browser = browser;
         this.chatRequestParams = chatRequestParams;
-        this.chatPartialResultManager = ChatPartialResultManager.getInstance();
+        this.chatCommunicationManager = ChatCommunicationManager.getInstance();
     }
 
     public Browser getBrowser() {
@@ -41,13 +40,17 @@ public String getPartialResultToken() {
 
     public ChatResult sendChatMessageWithProgress() {
         try {
-            String partialResultToken = UUID.randomUUID().toString();
-            chatRequestParams.setPartialResultToken(partialResultToken);
-
-            chatPartialResultManager.setMapEntry(partialResultToken, this);
-
+            // 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);

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

@@ -9,37 +9,29 @@
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
 
 /**
- * ChatPartialResultManager is responsible for maintaining a mapping between
- * partial result tokens and the associated ChatMessage objects. It is implemented 
- * as a singleton to centralize control of all partial results in the plugin.
+ * 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 original ChatMessage associated with the token.
+ * this class allows us to find the associated ChatMessage associated with the token.
  *
  * @see AmazonQLspClientImpl#notifyProgress(ProgressParams)
  */
-public class ChatPartialResultManager {
-    private static ChatPartialResultManager instance;
+public class ChatPartialResultMap {
+
     private final Map<String, ChatMessage> tokenToChatMessageMap;
 
-    private ChatPartialResultManager() {
+    public ChatPartialResultMap() {
         tokenToChatMessageMap = new ConcurrentHashMap<String, ChatMessage>();
     }
 
-    public static synchronized ChatPartialResultManager getInstance() {
-        if (instance == null) {
-            instance = new ChatPartialResultManager();
-        }
-        return instance;
-    }
-
-    public void setMapEntry(String token, ChatMessage chatMessage) {
+    public void setEntry(String token, ChatMessage chatMessage) {
         tokenToChatMessageMap.put(token, chatMessage);
     }
 
-    public void deleteMapEntry(String token) {
+    public void removeEntry(String token) {
         tokenToChatMessageMap.remove(token);
     }
 

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

@@ -27,7 +27,7 @@ public class AmazonQChatViewActionHandler implements ViewActionHandler {
 
     public AmazonQChatViewActionHandler() {
         this.jsonHandler = new JsonHandler();
-        chatCommunicationManager = new ChatCommunicationManager();
+        chatCommunicationManager = ChatCommunicationManager.getInstance();
     }
 
     /*
@@ -72,9 +72,9 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
      */
     public final void handlePartialResultProgressNotification(ProgressParams params) {
         String token = ProgressNotficationUtils.getToken(params);
+        ChatMessage chatMessage = chatCommunicationManager.getPartialChatMessage(token);
 
-        // Check to ensure progress notification is a partial result for chat 
-        if (!chatCommunicationManager.isProcessingPartialChatMessage(token)) {
+        if (chatMessage == null) {
             return;
         }
 
@@ -84,20 +84,19 @@ public final void handlePartialResultProgressNotification(ProgressParams params)
             throw new AmazonQPluginException(e);
         }
 
-        ChatResult chatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
-        ChatMessage chatMessage = chatCommunicationManager.getPartialChatMessage(token);
+        ChatResult partialChatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
         Browser browser = chatMessage.getBrowser();
 
         // Check to ensure the body has content in order to keep displaying the spinner while loading
-        if (chatResult.body() == null && chatResult.body().length() == 0) {
+        if (partialChatResult.body() == null || partialChatResult.body().length() == 0) {
             return;
         }
 
         Boolean isPartialResult = true;
         ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
             ChatUIInboundCommandName.ChatPrompt.toString(),
             chatMessage.getChatRequestParams().getTabId(),
-            chatResult,
+            partialChatResult,
             isPartialResult
         );