Chat Partial Result (#36)

* Log partial result

* Reaching partial result handler

* Rendering partial responses successfully

* Add comments

* Cleaned up code and added comments

* Update comment

* Merged with latest changes

* Move ignore null serialization to ChatResult record

* Fix checkstyle errors

* Fix checkstyle after merge

Author: angjordn

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

@@ -3,6 +3,7 @@
 package software.aws.toolkits.eclipse.amazonq.chat;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.UUID;
 
 import org.eclipse.swt.browser.Browser;
 
@@ -15,23 +16,40 @@
 import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
 import software.aws.toolkits.eclipse.amazonq.views.model.Command;
 
+/**
+ * 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 ChatPartialResultMap chatPartialResultMap;
 
-    public ChatCommunicationManager() {
+    private ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
         this.chatMessageProvider = ChatMessageProvider.createAsync();
+        this.chatPartialResultMap = new ChatPartialResultMap();
     }
 
-    public CompletableFuture<ChatResult> sendMessageToChatServer(final Command command, final Object params) {
+    public static synchronized ChatCommunicationManager getInstance() {
+        if (instance == null) {
+            instance = new ChatCommunicationManager();
+        }
+        return instance;
+    }
+
+    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);
@@ -50,11 +68,39 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Command comma
     }
 
     public void sendMessageToChatUI(final Browser browser, final ChatUIInboundCommand command) {
-        String message = this.jsonHandler.serialize(command);
+        String message = jsonHandler.serialize(command);
+
         String script = "window.postMessage(" + message + ");";
         browser.getDisplay().asyncExec(() -> {
             browser.evaluate(script);
         });
     }
 
+    /*
+     * Gets the partial chat message using the provided token.
+     */
+    public ChatMessage getPartialChatMessage(final String partialResultToken) {
+        return chatPartialResultMap.getValue(partialResultToken);
+    }
+
+    /*
+     * Adds an entry to the partialResultToken to ChatMessage map.
+     */
+    public String addPartialChatMessage(final 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(final String partialResultToken) {
+        chatPartialResultMap.removeEntry(partialResultToken);
+    }
 }
+

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

@@ -0,0 +1,53 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+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.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) {
+        this.amazonQLspServer = amazonQLspServer;
+        this.browser = browser;
+        this.chatRequestParams = chatRequestParams;
+        this.chatCommunicationManager = ChatCommunicationManager.getInstance();
+    }
+
+    public Browser getBrowser() {
+        return browser;
+    }
+
+    public ChatRequestParams getChatRequestParams() {
+        return chatRequestParams;
+    }
+
+    public String getPartialResultToken() {
+        return chatRequestParams.getPartialResultToken();
+    }
+
+    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

@@ -4,6 +4,8 @@
 
 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;
@@ -26,14 +28,9 @@ private ChatMessageProvider(final AmazonQLspServer amazonQLspServer) {
         this.amazonQLspServer = amazonQLspServer;
     }
 
-    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
-        try {
-            PluginLogger.info("Sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server");
-            return amazonQLspServer.sendChatPrompt(chatRequestParams);
-        } catch (Exception e) {
-            PluginLogger.error("Error occurred while sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server", e);
-            return CompletableFuture.failedFuture(new AmazonQPluginException(e));
-        }
+    public CompletableFuture<ChatResult>  sendChatPrompt(final Browser browser, final ChatRequestParams chatRequestParams) {
+        ChatMessage chatMessage = new ChatMessage(amazonQLspServer, browser, chatRequestParams);
+        return chatMessage.sendChatMessageWithProgress();
     }
 
     public void sendChatReady() {

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

@@ -0,0 +1,48 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.chat;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+
+/**
+ * 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.
+ *
+ * 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 final class ChatPartialResultMap {
+
+    private final Map<String, ChatMessage> tokenToChatMessageMap;
+
+    public ChatPartialResultMap() {
+        tokenToChatMessageMap = new ConcurrentHashMap<String, ChatMessage>();
+    }
+
+    public void setEntry(final String token, final ChatMessage chatMessage) {
+        tokenToChatMessageMap.put(token, chatMessage);
+    }
+
+    public void removeEntry(final String token) {
+        tokenToChatMessageMap.remove(token);
+    }
+
+    public ChatMessage getValue(final String token) {
+        return tokenToChatMessageMap.getOrDefault(token, null);
+    }
+
+    public Boolean hasKey(final String token) {
+        return tokenToChatMessageMap.containsKey(token);
+    }
+}

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatRequestParams.java

@@ -4,8 +4,33 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
-public record ChatRequestParams(
-    @JsonProperty("tabId") String tabId,
-    @JsonProperty("prompt") ChatPrompt prompt
-) { }
+public final class ChatRequestParams {
+    private final String tabId;
+    private final ChatPrompt prompt;
+    private String partialResultToken;
+
+    public ChatRequestParams(
+        @JsonProperty("tabId") final String tabId,
+        @JsonProperty("prompt") final ChatPrompt prompt
+    ) {
+        this.tabId = tabId;
+        this.prompt = prompt;
+    }
+
+    public String getTabId() {
+        return tabId;
+    }
+
+    public ChatPrompt getPrompt() {
+        return prompt;
+    }
+
+    public String getPartialResultToken() {
+        return partialResultToken;
+    }
+
+    public void setPartialResultToken(final String partialResultToken) {
+        this.partialResultToken = partialResultToken;
+    }
+}
 

plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatResult.java

@@ -2,8 +2,11 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat.models;
 
+import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+// Mynah-ui will not render the partial result if null values are included. Must ignore nulls values.
+@JsonInclude(JsonInclude.Include.NON_NULL)
 public record ChatResult(
     @JsonProperty("body") String body,
     @JsonProperty("messageId") String messageId,

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

@@ -11,11 +11,15 @@
 
 import org.eclipse.lsp4e.LanguageClientImpl;
 import org.eclipse.lsp4j.ConfigurationParams;
+import org.eclipse.lsp4j.ProgressParams;
 
 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 {
@@ -30,6 +34,11 @@ 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) {
@@ -49,4 +58,16 @@ public final CompletableFuture<List<Object>> configuration(final ConfigurationPa
         return CompletableFuture.completedFuture(output);
     }
 
+    @Override
+    public final void notifyProgress(final ProgressParams params) {
+        AmazonQChatViewActionHandler chatActionHandler = new AmazonQChatViewActionHandler();
+
+        ThreadingUtils.executeAsyncTask(() -> {
+            try {
+                chatActionHandler.handlePartialResultProgressNotification(params);
+            } catch (Exception e) {
+                PluginLogger.error("Error processing partial result progress notification", e);
+            }
+        });
+    }
 }

plugin/src/software/aws/toolkits/eclipse/amazonq/util/ProgressNotficationUtils.java

@@ -0,0 +1,50 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.util;
+
+import org.eclipse.lsp4j.ProgressParams;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonElement;
+
+
+public final class ProgressNotficationUtils {
+    private ProgressNotficationUtils() {
+        // Prevent instantiation
+    }
+
+    /*
+     * Get the token from the ProgressParams value
+     * @return The token as a String
+     */
+    public static String getToken(final ProgressParams params) {
+        String token;
+
+        if (params.getToken().isLeft()) {
+            token = params.getToken().getLeft();
+        } else {
+            token = params.getToken().getRight().toString();
+        }
+
+        return token;
+    }
+
+    /*
+     * Get the object from the ProgressParams value
+     * @param cls The class of the object to be deserialized
+     * @return The deserialized object, or null if the value is not a JsonElement
+     */
+    public static <T> T getObject(final ProgressParams params, final Class<T> cls) {
+        Object val = params.getValue().getRight();
+
+        if (!(val instanceof JsonElement)) {
+            return null;
+        }
+
+        Gson gson = new Gson();
+        JsonElement element = (JsonElement) val;
+        T obj = gson.fromJson(element, cls);
+
+        return obj;
+    }
+}