Fix chat streaming bug (#299)

aws · Dec 10, 2024 · cdbbbf5 · cdbbbf5
1 parent be041bc
commit cdbbbf5
Showing 1 changed file with 12 additions and 13 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
@@ -52,7 +52,7 @@ public final class ChatCommunicationManager {
     private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
     private final ChatPartialResultMap chatPartialResultMap;
     private final LspEncryptionManager lspEncryptionManager;
-    private ChatUiRequestListener chatUiRequestListener;
+    private CompletableFuture<ChatUiRequestListener> chatUiRequestListenerFuture;
 
     private ChatCommunicationManager(final Builder builder) {
         this.jsonHandler = builder.jsonHandler != null ? builder.jsonHandler : new JsonHandler();
@@ -62,6 +62,7 @@ private ChatCommunicationManager(final Builder builder) {
                 : new ChatPartialResultMap();
         this.lspEncryptionManager = builder.lspEncryptionManager != null ? builder.lspEncryptionManager
                 : DefaultLspEncryptionManager.getInstance();
+        chatUiRequestListenerFuture = new CompletableFuture<>();
     }
 
     public static Builder builder() {
@@ -148,8 +149,7 @@ private ChatRequestParams addEditorState(final ChatRequestParams chatRequestPara
         // only include files that are accessible via lsp which have absolute paths
         getOpenFileUri().ifPresent(filePathUri -> {
             chatRequestParams.setTextDocument(new TextDocumentIdentifier(filePathUri));
-            getSelectionRangeCursorState()
-                    .ifPresent(cursorState -> chatRequestParams.setCursorState(Arrays.asList(cursorState)));
+            getSelectionRangeCursorState().ifPresent(cursorState -> chatRequestParams.setCursorState(Arrays.asList(cursorState)));
         });
         return chatRequestParams;
     }
@@ -193,8 +193,7 @@ private CompletableFuture<ChatResult> sendEncryptedChatMessage(final String tabI
             removePartialChatMessage(partialResultToken);
 
             if (exception != null) {
-                Activator.getLogger()
-                        .error("An error occurred while processing chat request: " + exception.getMessage());
+                Activator.getLogger().error("An error occurred while processing chat request: " + exception.getMessage());
                 sendErrorToUi(tabId, exception);
                 return null;
             } else {
@@ -213,8 +212,7 @@ private CompletableFuture<ChatResult> sendEncryptedChatMessage(final String tabI
                     sendMessageToChatUI(chatUIInboundCommand);
                     return result;
                 } catch (Exception e) {
-                    Activator.getLogger()
-                            .error("An error occurred while processing chat response received: " + e.getMessage());
+                    Activator.getLogger().error("An error occurred while processing chat response received: " + e.getMessage());
                     sendErrorToUi(tabId, e);
                     return null;
                 }
@@ -233,21 +231,22 @@ private void sendErrorToUi(final String tabId, final Throwable exception) {
     }
 
     public void setChatUiRequestListener(final ChatUiRequestListener listener) {
-        chatUiRequestListener = listener;
+        chatUiRequestListenerFuture.complete(listener);
     }
 
     public void removeListener() {
-        chatUiRequestListener = null;
+        chatUiRequestListenerFuture = new CompletableFuture<>();
     }
 
     /*
      * Sends message to Chat UI to show in webview
      */
     public void sendMessageToChatUI(final ChatUIInboundCommand command) {
-        if (chatUiRequestListener != null) {
-            String message = jsonHandler.serialize(command);
-            chatUiRequestListener.onSendToChatUi(message);
-        }
+        String message = jsonHandler.serialize(command);
+        chatUiRequestListenerFuture.thenApply(listener -> {
+            listener.onSendToChatUi(message);
+            return listener;
+        });
     }
 
     /*