From d1aece6d7c3f79fef26e8f14cabe2730fa156141 Mon Sep 17 00:00:00 2001 From: Jordan Ang Date: Wed, 25 Sep 2024 12:13:44 -0700 Subject: [PATCH] Cleaned up code and added comments --- .../chat/ChatCommunicationManager.java | 44 ++++++++++++++----- .../eclipse/amazonq/chat/ChatMessage.java | 19 ++++---- ...Manager.java => ChatPartialResultMap.java} | 24 ++++------ .../views/AmazonQChatViewActionHandler.java | 13 +++--- 4 files changed, 59 insertions(+), 41 deletions(-) rename plugin/src/software/aws/toolkits/eclipse/amazonq/chat/{ChatPartialResultManager.java => ChatPartialResultMap.java} (63%) diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java index 7ceeefa25..e0bd742ea 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java +++ b/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); } } diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java index f8248b66e..63f800e2a 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java +++ b/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); diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultManager.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java similarity index 63% rename from plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultManager.java rename to plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java index dbe46c914..7a71905e0 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultManager.java +++ b/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 tokenToChatMessageMap; - private ChatPartialResultManager() { + public ChatPartialResultMap() { tokenToChatMessageMap = new ConcurrentHashMap(); } - 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); } diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java index 3cab630eb..03c38f212 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java +++ b/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,12 +84,11 @@ 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; } @@ -97,7 +96,7 @@ public final void handlePartialResultProgressNotification(ProgressParams params) ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand( ChatUIInboundCommandName.ChatPrompt.toString(), chatMessage.getChatRequestParams().getTabId(), - chatResult, + partialChatResult, isPartialResult );