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 123dd04b..b2514bc6 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java +++ b/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; + private final ChatPartialResultMap chatPartialResultMap; - public ChatCommunicationManager() { + private ChatCommunicationManager() { this.jsonHandler = new JsonHandler(); this.chatMessageProvider = ChatMessageProvider.createAsync(); + this.chatPartialResultMap = new ChatPartialResultMap(); } - public CompletableFuture sendMessageToChatServer(final Command command, final Object params) { + public static synchronized ChatCommunicationManager getInstance() { + if (instance == null) { + instance = new ChatCommunicationManager(); + } + return instance; + } + + public CompletableFuture 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 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); + } } + diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java new file mode 100644 index 00000000..bdd2204c --- /dev/null +++ b/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 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 = 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; + } +} diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java index 5470cd5f..90f611e3 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java +++ b/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 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 sendChatPrompt(final Browser browser, final ChatRequestParams chatRequestParams) { + ChatMessage chatMessage = new ChatMessage(amazonQLspServer, browser, chatRequestParams); + return chatMessage.sendChatMessageWithProgress(); } public void sendChatReady() { diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java new file mode 100644 index 00000000..f0ad7a5d --- /dev/null +++ b/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 tokenToChatMessageMap; + + public ChatPartialResultMap() { + tokenToChatMessageMap = new ConcurrentHashMap(); + } + + 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); + } +} diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatRequestParams.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatRequestParams.java index 89d50f74..de43d2fb 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatRequestParams.java +++ b/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; + } +} diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatResult.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatResult.java index 632a9ff3..81ba03b9 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatResult.java +++ b/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, diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java index 09f47053..194e98ae 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java +++ b/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 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> configuration(final ConfigurationParams configurationParams) { if (configurationParams.getItems().size() == 0) { @@ -49,4 +58,16 @@ public final CompletableFuture> 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); + } + }); + } } diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/util/ProgressNotficationUtils.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/util/ProgressNotficationUtils.java new file mode 100644 index 00000000..c697eece --- /dev/null +++ b/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 getObject(final ProgressParams params, final Class 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; + } +} 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 2e15f79d..46d5a610 100644 --- a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java +++ b/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java @@ -3,15 +3,21 @@ package software.aws.toolkits.eclipse.amazonq.views; +import java.util.Objects; + +import org.eclipse.lsp4j.ProgressParams; import org.eclipse.swt.browser.Browser; 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; import software.aws.toolkits.eclipse.amazonq.util.JsonHandler; import software.aws.toolkits.eclipse.amazonq.util.PluginLogger; +import software.aws.toolkits.eclipse.amazonq.util.ProgressNotficationUtils; import software.aws.toolkits.eclipse.amazonq.views.model.Command; import software.aws.toolkits.eclipse.amazonq.views.model.ParsedCommand; @@ -21,7 +27,7 @@ public class AmazonQChatViewActionHandler implements ViewActionHandler { public AmazonQChatViewActionHandler() { this.jsonHandler = new JsonHandler(); - chatCommunicationManager = new ChatCommunicationManager(); + chatCommunicationManager = ChatCommunicationManager.getInstance(); } /* @@ -36,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 ); @@ -49,10 +55,10 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser }); break; case CHAT_READY: - chatCommunicationManager.sendMessageToChatServer(command, params); + chatCommunicationManager.sendMessageToChatServer(browser, command, params); break; case CHAT_TAB_ADD: - chatCommunicationManager.sendMessageToChatServer(command, params); + chatCommunicationManager.sendMessageToChatServer(browser, command, params); break; case TELEMETRY_EVENT: break; @@ -61,4 +67,38 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser } } + + /* + * Handles chat progress notifications from the Amazon Q LSP server. Sends a partial chat prompt message to the webview. + */ + public final void handlePartialResultProgressNotification(final ProgressParams params) { + String token = ProgressNotficationUtils.getToken(params); + ChatMessage chatMessage = chatCommunicationManager.getPartialChatMessage(token); + + if (chatMessage == null) { + return; + } + + // Check to ensure Object is sent in params + if (params.getValue().isLeft() || Objects.isNull(params.getValue().getRight())) { + throw new AmazonQPluginException("Error occurred while handling partial result notification: expected Object value"); + } + + 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 (partialChatResult.body() == null || partialChatResult.body().length() == 0) { + return; + } + + ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand( + ChatUIInboundCommandName.ChatPrompt.toString(), + chatMessage.getChatRequestParams().getTabId(), + partialChatResult, + true + ); + + chatCommunicationManager.sendMessageToChatUI(browser, chatUIInboundCommand); + } }