Make command handling in Q chat webview async (#35)

aws · Sep 25, 2024 · 383e6cf · 383e6cf
1 parent f17c365
commit 383e6cf
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 56 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
@@ -2,6 +2,8 @@
 
 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;
@@ -10,51 +12,49 @@
 import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
 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.views.model.Command;
 
 public final class ChatCommunicationManager {
 
     private final JsonHandler jsonHandler;
-    private final ChatMessageProvider chatMessageProvider;
+    private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
 
     public ChatCommunicationManager() {
         this.jsonHandler = new JsonHandler();
-        this.chatMessageProvider = new ChatMessageProvider();
+        this.chatMessageProvider = ChatMessageProvider.createAsync();
     }
 
-    /*
-     * Sends a message to the Amazon Q LSP server.
-     */
-    public ChatResult sendMessageToChatServer(final Command command, final Object params) {
-
-           switch (command) {
-               case CHAT_SEND_PROMPT:
-                   ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
-                   ChatResult result = chatMessageProvider.sendChatPrompt(chatRequestParams);
-                   return result;
-               case CHAT_READY:
-                   chatMessageProvider.sendChatReady();
-                   return null;
-               case CHAT_TAB_ADD:
-                   GenericTabParams tabParams = jsonHandler.convertObject(params, GenericTabParams.class);
-                   chatMessageProvider.sendTabAdd(tabParams);
-                   return null;
-               default:
-                   throw new AmazonQPluginException("Unhandled command in ChatCommunicationManager: " + command.toString());
-           }
+    public CompletableFuture<ChatResult> sendMessageToChatServer(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);
+                    case CHAT_READY:
+                        chatMessageProvider.sendChatReady();
+                        return CompletableFuture.completedFuture(null);
+                    case CHAT_TAB_ADD:
+                        GenericTabParams tabParams = jsonHandler.convertObject(params, GenericTabParams.class);
+                        chatMessageProvider.sendTabAdd(tabParams);
+                        return CompletableFuture.completedFuture(null);
+                    default:
+                        throw new AmazonQPluginException("Unhandled command in ChatCommunicationManager: " + command.toString());
+                }
+            } catch (Exception e) {
+                PluginLogger.error("Error occurred in sendMessageToChatServer", e);
+                return CompletableFuture.failedFuture(new AmazonQPluginException(e));
+            }
+        });
     }
 
-    /**
-     * Sends a message to the webview.
-     *
-     * See handlers in Flare chat-client:
-     * https://github.com/aws/language-servers/blob/9226fb4ed10dc54f1719b14a5b1dac1807641f79/chat-client/src/client/chat.ts#L67-L101
-     */
     public void sendMessageToChatUI(final Browser browser, final ChatUIInboundCommand command) {
         String message = this.jsonHandler.serialize(command);
         String script = "window.postMessage(" + message + ");";
         browser.getDisplay().asyncExec(() -> {
             browser.evaluate(script);
         });
     }
-}
+
+}
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java
@@ -2,7 +2,7 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat;
 
-import java.util.concurrent.ExecutionException;
+import java.util.concurrent.CompletableFuture;
 
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
 import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
@@ -15,25 +15,24 @@
 
 public final class ChatMessageProvider {
 
-    private AmazonQLspServer amazonQLspServer;
+    private final AmazonQLspServer amazonQLspServer;
 
-    public ChatMessageProvider() {
-        try {
-            amazonQLspServer = LspProvider.getAmazonQServer().get();
-        } catch (InterruptedException | ExecutionException e) {
-            PluginLogger.error("Error occurred while retrieving Amazon Q LSP server. Failed to instantiate ChatMessageProvider.", e);
-            throw new AmazonQPluginException(e);
-        }
+    public static CompletableFuture<ChatMessageProvider> createAsync() {
+        return LspProvider.getAmazonQServer()
+                .thenApply(ChatMessageProvider::new);
+    }
+
+    private ChatMessageProvider(AmazonQLspServer amazonQLspServer) {
+        this.amazonQLspServer = amazonQLspServer;
     }
 
-    public ChatResult sendChatPrompt(final ChatRequestParams chatRequestParams) {
+    public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
         try {
             PluginLogger.info("Sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server");
-            ChatResult chatResult = amazonQLspServer.sendChatPrompt(chatRequestParams).get();
-            return chatResult;
-        } catch (InterruptedException | ExecutionException e) {
+            return amazonQLspServer.sendChatPrompt(chatRequestParams);
+        } catch (Exception e) {
             PluginLogger.error("Error occurred while sending " + Command.CHAT_SEND_PROMPT + " message to Amazon Q LSP server", e);
-            throw new AmazonQPluginException(e);
+            return CompletableFuture.failedFuture(new AmazonQPluginException(e));
         }
     }
 
@@ -56,4 +55,5 @@ public void sendTabAdd(final GenericTabParams tabParams) {
             throw new AmazonQPluginException(e);
         }
     }
+
 }
diff --git a/...c/software/aws/toolkits/eclipse/amazonq/lsp/manager/fetcher/RemoteManifestLspFetcher.java b/...c/software/aws/toolkits/eclipse/amazonq/lsp/manager/fetcher/RemoteManifestLspFetcher.java
@@ -35,7 +35,6 @@ public final class RemoteManifestLspFetcher implements LspFetcher {
     private static final int TIMEOUT_SECONDS = 10;
     private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance();
 
-
     private final String manifestUrl;
     private final VersionRange versionRange;
     private final boolean integrityChecking;

diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java
@@ -7,7 +7,6 @@
 
 import software.aws.toolkits.eclipse.amazonq.chat.ChatCommunicationManager;
 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;
@@ -37,16 +36,17 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
 
         switch (command) {
             case CHAT_SEND_PROMPT:
-                ChatResult chatResult = chatCommunicationManager.sendMessageToChatServer(command, params);
-
-                ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
-                ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
-                    ChatUIInboundCommandName.ChatPrompt.toString(),
-                    chatRequestParams.tabId(),
-                    chatResult,
-                    false
-                );
-                chatCommunicationManager.sendMessageToChatUI(browser, chatUIInboundCommand);
+                chatCommunicationManager.sendMessageToChatServer(command, params)
+                    .thenAccept(chatResult -> {
+                        ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
+                        ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
+                            ChatUIInboundCommandName.ChatPrompt.toString(),
+                            chatRequestParams.tabId(),
+                            chatResult,
+                            false
+                        );
+                        chatCommunicationManager.sendMessageToChatUI(browser, chatUIInboundCommand);
+                    });
                 break;
             case CHAT_READY:
                 chatCommunicationManager.sendMessageToChatServer(command, params);
@@ -61,4 +61,4 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
         }
     }
 
-}
+}