Add comments

aws · Sep 25, 2024 · 41ea2db · 41ea2db
1 parent 163c841
commit 41ea2db
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 29 deletions.
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java
@@ -12,6 +12,11 @@
 import software.aws.toolkits.eclipse.amazonq.util.JsonHandler;
 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.
+ */
 public final class ChatCommunicationManager {
 
     private final JsonHandler jsonHandler;
@@ -63,11 +68,17 @@ public void sendMessageToChatUI(final Browser browser, final ChatUIInboundComman
         });
     }
 
-    public boolean isProcessingChatMessage(String partialResultToken) {
+    /*
+     * Checks if a partial result is being processed with the provided token.
+     */
+    public boolean isProcessingPartialChatMessage(String partialResultToken) {
         return chatPartialResultManager.hasKey(partialResultToken);
     }
 
-    public ChatMessage getChatMessage(String partialResultToken) {
+    /*
+     * Gets the partial chat message using the provided token.
+     */
+    public ChatMessage getPartialChatMessage(String partialResultToken) {
         return chatPartialResultManager.getValue(partialResultToken);
     }
 }

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java
@@ -35,6 +35,10 @@ public ChatRequestParams getChatRequestParams() {
         return chatRequestParams;
     }
 
+    public String getPartialResultToken() {
+        return chatRequestParams.getPartialResultToken();
+    }
+
     public ChatResult sendChatMessageWithProgress() {
         try {
             String partialResultToken = UUID.randomUUID().toString();

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultManager.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultManager.java
@@ -8,6 +8,18 @@
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
 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.
+ * 
+ * $/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.
+ *
+ * @see AmazonQLspClientImpl#notifyProgress(ProgressParams)
+ */
 public class ChatPartialResultManager {
     private static ChatPartialResultManager instance;
     private final Map<String, ChatMessage> tokenToChatMessageMap;

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java
@@ -27,16 +27,20 @@ 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 chatPartialResultManager
+     * - Other notifications are ignored at this time.
+     */
     @Override
 	public void notifyProgress(final ProgressParams params) {
-		PluginLogger.info("Notify Progress caught...");
-		
+        AmazonQChatViewActionHandler chatActionHandler = new AmazonQChatViewActionHandler();
+        
 		ThreadingUtils.executeAsyncTask(() -> {
             try {
-                AmazonQChatViewActionHandler chatActionHandler = new AmazonQChatViewActionHandler();
-                chatActionHandler.handleProgressNotification(params);;
+                chatActionHandler.handlePartialResultProgressNotification(params);;
             } catch (Exception e) {
-                PluginLogger.error("Error processing message from Browser", e);
+                PluginLogger.error("Error processing partial result progress notification", e);
             }
         });
 	}

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java
@@ -3,6 +3,8 @@
 package software.aws.toolkits.eclipse.amazonq.views;
 
 
+import java.util.Objects;
+
 import org.eclipse.lsp4j.ProgressParams;
 import org.eclipse.swt.browser.Browser;
 
@@ -66,39 +68,39 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
 
 
     /*
-     * Handles chat progress notifications from the Amazon Q LSP server. Sends a chat prompt message to the webview.
+     * Handles chat progress notifications from the Amazon Q LSP server. Sends a partial chat prompt message to the webview.
      */
-    public final void handleProgressNotification(ProgressParams params) {
+    public final void handlePartialResultProgressNotification(ProgressParams params) {
         String token = ProgressNotficationUtils.getToken(params);
 
-        if (!chatCommunicationManager.isProcessingChatMessage(token)) {
-            PluginLogger.info("Not processing - not a partial result token");
+        // Check to ensure progress notification is a partial result for chat 
+        if (!chatCommunicationManager.isProcessingPartialChatMessage(token)) {
             return;
         }
-
-        if (params.getValue().isLeft()) {
-            String e = "Error occurred while handling partial result notification: Expected Object not WorkDoneProgressNotification value for partial result token " + token;
+
+        // Check to ensure Object is sent in params
+        if (params.getValue().isLeft() || Objects.isNull(params.getValue().getRight())) {
+            String e = "Error occurred while handling partial result notification: expected Object value";
             throw new AmazonQPluginException(e);
         }
-
-        PluginLogger.info("Handling progress notification");
-
+
         ChatResult chatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
-        ChatMessage chatMessage = chatCommunicationManager.getChatMessage(token);
+        ChatMessage chatMessage = chatCommunicationManager.getPartialChatMessage(token);
         Browser browser = chatMessage.getBrowser();
 
-        // Sending a response would clear the loading indicator, so we should ensure that there is content first for a good UX
-        // (Note: Response is accumulative)
-        if (chatResult.body() != null && chatResult.body().length() > 0) {
-            ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
-                    ChatUIInboundCommandName.ChatPrompt.toString(),
-                    chatMessage.getChatRequestParams().getTabId(),
-                    chatResult,
-                    true
-                );
-
-            chatCommunicationManager.sendMessageToChatUI(browser, chatUIInboundCommand);
+        // Check to ensure the body has content in order to keep displaying the spinner while loading
+        if (chatResult.body() == null && chatResult.body().length() == 0) {
+            return;
         }
 
+        Boolean isPartialResult = true;
+        ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
+            ChatUIInboundCommandName.ChatPrompt.toString(),
+            chatMessage.getChatRequestParams().getTabId(),
+            chatResult,
+            isPartialResult
+        );
+
+        chatCommunicationManager.sendMessageToChatUI(browser, chatUIInboundCommand);
     }
 }