Render Chat response in the webview (#28)

aws · Sep 20, 2024 · e90dcfb · e90dcfb
1 parent 0345d26
commit e90dcfb
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 9 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,8 +2,11 @@
 
 package software.aws.toolkits.eclipse.amazonq.chat;
 
+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.ChatUIInboundCommand;
 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;
@@ -19,24 +22,39 @@ public ChatCommunicationManager() {
         this.chatMessageProvider = new ChatMessageProvider();
     }
 
-    public ChatResult sendMessageToChatServerAsync(final Command command, final Object params) {
-
-           String jsonParams = jsonHandler.serialize(params);
+    /*
+     * 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.deserialize(jsonParams, ChatRequestParams.class);
+                   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.deserialize(jsonParams, GenericTabParams.class);
+                   GenericTabParams tabParams = jsonHandler.convertObject(params, GenericTabParams.class);
                    chatMessageProvider.sendTabAdd(tabParams);
                    return null;
                default:
                    throw new AmazonQPluginException("Unhandled command in ChatCommunicationManager: " + command.toString());
            }
     }
+
+    /**
+     * 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/models/ChatUIInboundCommand.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatUIInboundCommand.java
@@ -0,0 +1,14 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+package software.aws.toolkits.eclipse.amazonq.chat.models;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * Represents a command that is being sent to Q Chat UI.
+ */
+public record ChatUIInboundCommand(
+    @JsonProperty("command") String command,
+    @JsonProperty("tabId") String tabId,
+    @JsonProperty("params") Object params,
+    @JsonProperty("isPartialResult") Boolean isPartialResult
+) { };
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatUIInboundCommandName.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/chat/models/ChatUIInboundCommandName.java
@@ -0,0 +1,23 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+package software.aws.toolkits.eclipse.amazonq.chat.models;
+
+public enum ChatUIInboundCommandName {
+    ChatPrompt("aws/chat/sendChatPrompt"), // This is the odd one out, it follows the same message name as the request.
+
+    SendToPrompt("sendToPrompt"),
+    ErrorMessage("errorMessage"),
+    InsertToCursorPosition("insertToCursorPosition"),
+    AuthFollowUpClicked("authFollowUpClicked"),
+    GenericCommand("genericCommand");
+
+    private final String commandString;
+
+    ChatUIInboundCommandName(final String commandString) {
+        this.commandString = commandString;
+    }
+
+    public String toString() {
+        return commandString;
+    }
+}
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/util/JsonHandler.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/util/JsonHandler.java
@@ -32,4 +32,9 @@ public <T> T deserialize(final String jsonString, final Class<T> cls) {
         }
         return null;
     }
+
+    public <T> T convertObject(final Object obj, final Class<T> cls) {
+        T castedObj = deserialize(serialize(obj), cls);
+        return castedObj;
+    }
 }
diff --git a/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java b/plugin/src/software/aws/toolkits/eclipse/amazonq/views/AmazonQChatViewActionHandler.java
@@ -6,19 +6,28 @@
 import org.eclipse.swt.browser.Browser;
 
 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;
+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;
 import software.aws.toolkits.eclipse.amazonq.views.model.ParsedCommand;
 
 public class AmazonQChatViewActionHandler implements ViewActionHandler {
     private ChatCommunicationManager chatCommunicationManager;
+    private final JsonHandler jsonHandler;
 
     public AmazonQChatViewActionHandler() {
+        this.jsonHandler = new JsonHandler();
         chatCommunicationManager = new ChatCommunicationManager();
     }
 
+    /*
+     * Handles the command message received from the webview
+     */
     @Override
     public final void handleCommand(final ParsedCommand parsedCommand, final Browser browser) {
         Command command = parsedCommand.getCommand();
@@ -28,19 +37,28 @@ public final void handleCommand(final ParsedCommand parsedCommand, final Browser
 
         switch (command) {
             case CHAT_SEND_PROMPT:
-                ChatResult chatResult = chatCommunicationManager.sendMessageToChatServerAsync(command, params);
-                PluginLogger.info("Chat result: " + chatResult);
+                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);
                 break;
             case CHAT_READY:
-                chatCommunicationManager.sendMessageToChatServerAsync(command, params);
+                chatCommunicationManager.sendMessageToChatServer(command, params);
                 break;
             case CHAT_TAB_ADD:
-                chatCommunicationManager.sendMessageToChatServerAsync(command, params);
+                chatCommunicationManager.sendMessageToChatServer(command, params);
                 break;
             case TELEMETRY_EVENT:
                 break;
             default:
                 throw new AmazonQPluginException("Unhandled command in AmazonQChatViewActionHandler: " + command.toString());
         }
     }
+
 }