Skip to content

Commit f1b0a8d

Browse files
shruti0085shruti0085
authored
Refactor to simplify and isolate chat communication workflow (#51)
This change refactors the chat communication workflow to simplify and isolate responsibilities per class. This helps restrict access to certain components like browser and server to dedicated classes as defined below which in turn would help in improving testability. This design also helps easily add handling for more message types such as quick actions which returns ChatResult for rendering in the ChatUI. Changes include: * AmazonQViewActionHandler is now solely responsible for parsing and handling any actions that originate from within the WebView i.e messages from WebView/Chat UI. * ChatCommunicationManager orchestrates handling communication between chat server and chat UI. * For requests to server, it takes the JSON object and converts it to the corresponding parameter Class and calls ChatMessageProvider. * For response received from server, it takes the response and invokes event to be sent to Chat UI * AmazonQWebView is responsible for handling any interactions with the webview. It registers a listener ChatUiRequestListener for OnSendToChat event and sends message to webview for display. * ChatMessage is the downstream layer responsbile for final communication with the AmazonQLSPServer. This class will ultimately also hold the encryption/decryption of request/response. * ChatMessageProvider currently only wraps communication over to ChatMessage but eventually it will house more logic that would involve handling different types of messages eg. add tab/remove tab. * ChatPartialResultMap has been simplified to hold a map of partial result token to tabIDs. Each chatMessage object is uniquely identified by it's tabId. If needed, this can also be further replaced in favor of direct access to the map in appropriate class.
1 parent f648b5c commit f1b0a8d

File tree

10 files changed

+173
-159
lines changed

10 files changed

+173
-159
lines changed

‎plugin/META-INF/MANIFEST.MF

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,4 +115,4 @@ Bundle-Classpath: .,
115115
target/dependency/utils-2.25.33.jar,
116116
target/dependency/xml-apis-ext-1.3.04.jar,
117117
target/dependency/xmlgraphics-commons-2.9.jar,
118-
target/dependency/xz-1.9.jar
118+
target/dependency/xz-1.9.jar

‎plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatCommunicationManager.java

Lines changed: 93 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,22 @@
33
package software.aws.toolkits.eclipse.amazonq.chat;
44

55
import java.util.concurrent.CompletableFuture;
6+
import java.util.function.Function;
7+
import java.util.Objects;
68
import java.util.UUID;
79

8-
import org.eclipse.swt.browser.Browser;
10+
import org.eclipse.lsp4j.ProgressParams;
911

1012
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
1113
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
1214
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommand;
15+
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatUIInboundCommandName;
1316
import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
1417
import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
1518
import software.aws.toolkits.eclipse.amazonq.util.JsonHandler;
1619
import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
20+
import software.aws.toolkits.eclipse.amazonq.util.ProgressNotficationUtils;
21+
import software.aws.toolkits.eclipse.amazonq.views.ChatUiRequestListener;
1722
import software.aws.toolkits.eclipse.amazonq.views.model.Command;
1823

1924
/**
@@ -29,6 +34,7 @@ public final class ChatCommunicationManager {
2934
private final JsonHandler jsonHandler;
3035
private final CompletableFuture<ChatMessageProvider> chatMessageProvider;
3136
private final ChatPartialResultMap chatPartialResultMap;
37+
private ChatUiRequestListener chatUiRequestListener;
3238

3339
private ChatCommunicationManager() {
3440
this.jsonHandler = new JsonHandler();
@@ -43,13 +49,17 @@ public static synchronized ChatCommunicationManager getInstance() {
4349
return instance;
4450
}
4551

46-
public CompletableFuture<ChatResult> sendMessageToChatServer(final Browser browser, final Command command, final Object params) {
52+
public CompletableFuture<ChatResult> sendMessageToChatServer(final Command command, final Object params) {
4753
return chatMessageProvider.thenCompose(chatMessageProvider -> {
4854
try {
4955
switch (command) {
5056
case CHAT_SEND_PROMPT:
5157
ChatRequestParams chatRequestParams = jsonHandler.convertObject(params, ChatRequestParams.class);
52-
return chatMessageProvider.sendChatPrompt(browser, chatRequestParams);
58+
return sendChatRequest(chatRequestParams.getTabId(), token -> {
59+
chatRequestParams.setPartialResultToken(token);
60+
61+
return chatMessageProvider.sendChatPrompt(chatRequestParams);
62+
});
5363
case CHAT_READY:
5464
chatMessageProvider.sendChatReady();
5565
return CompletableFuture.completedFuture(null);
@@ -67,39 +77,102 @@ public CompletableFuture<ChatResult> sendMessageToChatServer(final Browser brows
6777
});
6878
}
6979

70-
public void sendMessageToChatUI(final Browser browser, final ChatUIInboundCommand command) {
71-
String message = jsonHandler.serialize(command);
72-
73-
String script = "window.postMessage(" + message + ");";
74-
browser.getDisplay().asyncExec(() -> {
75-
browser.evaluate(script);
80+
private CompletableFuture<ChatResult> sendChatRequest(final String tabId,
81+
final Function<String, CompletableFuture<ChatResult>> action) {
82+
// Retrieving the chat result is expected to be a long-running process with
83+
// intermittent progress notifications being sent
84+
// from the LSP server. The progress notifications provide a token and a partial
85+
// result Object - we are utilizing a token to
86+
// ChatMessage mapping to acquire the associated ChatMessage so we can formulate
87+
// a message for the UI.
88+
String partialResultToken = addPartialChatMessage(tabId);
89+
90+
return action.apply(partialResultToken).thenApply(result -> {
91+
// The mapping entry no longer needs to be maintained once the final result is
92+
// retrieved.
93+
removePartialChatMessage(partialResultToken);
94+
// show chat response in Chat UI
95+
ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
96+
ChatUIInboundCommandName.ChatPrompt.toString(), tabId, result, false);
97+
sendMessageToChatUI(chatUIInboundCommand);
98+
return result;
7699
});
77100
}
78101

102+
public void setChatUiRequestListener(final ChatUiRequestListener listener) {
103+
chatUiRequestListener = listener;
104+
}
105+
106+
public void removeListener() {
107+
chatUiRequestListener = null;
108+
}
109+
79110
/*
80-
* Gets the partial chat message using the provided token.
111+
* Sends message to Chat UI to show in webview
81112
*/
82-
public ChatMessage getPartialChatMessage(final String partialResultToken) {
83-
return chatPartialResultMap.getValue(partialResultToken);
113+
public void sendMessageToChatUI(final ChatUIInboundCommand command) {
114+
if (chatUiRequestListener != null) {
115+
String message = jsonHandler.serialize(command);
116+
chatUiRequestListener.onSendToChatUi(message);
117+
}
84118
}
85119

86120
/*
87-
* Adds an entry to the partialResultToken to ChatMessage map.
121+
* Handles chat progress notifications from the Amazon Q LSP server.
122+
* - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
123+
* - Other notifications are ignored at this time.
124+
* - Sends a partial chat prompt message to the webview.
88125
*/
89-
public String addPartialChatMessage(final ChatMessage chatMessage) {
90-
String partialResultToken = UUID.randomUUID().toString();
126+
public void handlePartialResultProgressNotification(final ProgressParams params) {
127+
String token = ProgressNotficationUtils.getToken(params);
128+
String tabId = getPartialChatMessage(token);
129+
130+
if (tabId == null || tabId.isEmpty()) {
131+
return;
132+
}
133+
134+
// Check to ensure Object is sent in params
135+
if (params.getValue().isLeft() || Objects.isNull(params.getValue().getRight())) {
136+
throw new AmazonQPluginException("Error occurred while handling partial result notification: expected Object value");
137+
}
91138

92-
// Indicator for the server to send partial result notifications
93-
chatMessage.getChatRequestParams().setPartialResultToken(partialResultToken);
139+
ChatResult partialChatResult = ProgressNotficationUtils.getObject(params, ChatResult.class);
94140

95-
chatPartialResultMap.setEntry(partialResultToken, chatMessage);
141+
// Check to ensure the body has content in order to keep displaying the spinner while loading
142+
if (partialChatResult.body() == null || partialChatResult.body().length() == 0) {
143+
return;
144+
}
145+
146+
ChatUIInboundCommand chatUIInboundCommand = new ChatUIInboundCommand(
147+
ChatUIInboundCommandName.ChatPrompt.toString(),
148+
tabId,
149+
partialChatResult,
150+
true
151+
);
152+
153+
sendMessageToChatUI(chatUIInboundCommand);
154+
}
155+
156+
/*
157+
* Gets the partial chat message represented by the tabId using the provided token.
158+
*/
159+
private String getPartialChatMessage(final String partialResultToken) {
160+
return chatPartialResultMap.getValue(partialResultToken);
161+
}
162+
163+
/*
164+
* Adds an entry to the partialResultToken to ChatMessage's tabId map.
165+
*/
166+
private String addPartialChatMessage(final String tabId) {
167+
String partialResultToken = UUID.randomUUID().toString();
168+
chatPartialResultMap.setEntry(partialResultToken, tabId);
96169
return partialResultToken;
97170
}
98171

99172
/*
100-
* Removes an entry from the partialResultToken to ChatMessage map.
173+
* Removes an entry from the partialResultToken to ChatMessage's tabId map.
101174
*/
102-
public void removePartialChatMessage(final String partialResultToken) {
175+
private void removePartialChatMessage(final String partialResultToken) {
103176
chatPartialResultMap.removeEntry(partialResultToken);
104177
}
105178
}

‎plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessage.java

Lines changed: 8 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,50 +4,27 @@
44

55
import java.util.concurrent.CompletableFuture;
66

7-
import org.eclipse.swt.browser.Browser;
8-
97
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
108
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
9+
import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
1110
import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
1211

1312
public final class ChatMessage {
14-
private final Browser browser;
15-
private final ChatRequestParams chatRequestParams;
1613
private final AmazonQLspServer amazonQLspServer;
17-
private final ChatCommunicationManager chatCommunicationManager;
1814

19-
public ChatMessage(final AmazonQLspServer amazonQLspServer, final Browser browser, final ChatRequestParams chatRequestParams) {
15+
public ChatMessage(final AmazonQLspServer amazonQLspServer) {
2016
this.amazonQLspServer = amazonQLspServer;
21-
this.browser = browser;
22-
this.chatRequestParams = chatRequestParams;
23-
this.chatCommunicationManager = ChatCommunicationManager.getInstance();
24-
}
25-
26-
public Browser getBrowser() {
27-
return browser;
2817
}
2918

30-
public ChatRequestParams getChatRequestParams() {
31-
return chatRequestParams;
19+
public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
20+
return amazonQLspServer.sendChatPrompt(chatRequestParams);
3221
}
3322

34-
public String getPartialResultToken() {
35-
return chatRequestParams.getPartialResultToken();
23+
public void sendChatReady() {
24+
amazonQLspServer.chatReady();
3625
}
3726

38-
public CompletableFuture<ChatResult> sendChatMessageWithProgress() {
39-
// Retrieving the chat result is expected to be a long-running process with intermittent progress notifications being sent
40-
// from the LSP server. The progress notifications provide a token and a partial result Object - we are utilizing a token to
41-
// ChatMessage mapping to acquire the associated ChatMessage so we can formulate a message for the UI.
42-
String partialResultToken = chatCommunicationManager.addPartialChatMessage(this);
43-
44-
CompletableFuture<ChatResult> chatResult = amazonQLspServer.sendChatPrompt(chatRequestParams)
45-
.thenApply(result -> {
46-
// The mapping entry no longer needs to be maintained once the final result is retrieved.
47-
chatCommunicationManager.removePartialChatMessage(partialResultToken);
48-
return result;
49-
});
50-
51-
return chatResult;
27+
public void sendTabAdd(final GenericTabParams tabParams) {
28+
amazonQLspServer.tabAdd(tabParams);
5229
}
5330
}

‎plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatMessageProvider.java

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,11 @@
33
package software.aws.toolkits.eclipse.amazonq.chat;
44

55
import java.util.concurrent.CompletableFuture;
6-
7-
import org.eclipse.swt.browser.Browser;
8-
96
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatRequestParams;
107
import software.aws.toolkits.eclipse.amazonq.chat.models.ChatResult;
118
import software.aws.toolkits.eclipse.amazonq.chat.models.GenericTabParams;
12-
import software.aws.toolkits.eclipse.amazonq.exception.AmazonQPluginException;
139
import software.aws.toolkits.eclipse.amazonq.lsp.AmazonQLspServer;
1410
import software.aws.toolkits.eclipse.amazonq.providers.LspProvider;
15-
import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
16-
import software.aws.toolkits.eclipse.amazonq.views.model.Command;
1711

1812
public final class ChatMessageProvider {
1913

@@ -28,29 +22,19 @@ private ChatMessageProvider(final AmazonQLspServer amazonQLspServer) {
2822
this.amazonQLspServer = amazonQLspServer;
2923
}
3024

31-
public CompletableFuture<ChatResult> sendChatPrompt(final Browser browser, final ChatRequestParams chatRequestParams) {
32-
ChatMessage chatMessage = new ChatMessage(amazonQLspServer, browser, chatRequestParams);
33-
return chatMessage.sendChatMessageWithProgress();
25+
public CompletableFuture<ChatResult> sendChatPrompt(final ChatRequestParams chatRequestParams) {
26+
ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
27+
return chatMessage.sendChatPrompt(chatRequestParams);
3428
}
3529

3630
public void sendChatReady() {
37-
try {
38-
PluginLogger.info("Sending " + Command.CHAT_READY + " message to Amazon Q LSP server");
39-
amazonQLspServer.chatReady();
40-
} catch (Exception e) {
41-
PluginLogger.error("Error occurred while sending " + Command.CHAT_READY + " message to Amazon Q LSP server", e);
42-
throw new AmazonQPluginException(e);
43-
}
31+
ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
32+
chatMessage.sendChatReady();
4433
}
4534

4635
public void sendTabAdd(final GenericTabParams tabParams) {
47-
try {
48-
PluginLogger.info("Sending " + Command.CHAT_TAB_ADD + " message to Amazon Q LSP server");
49-
amazonQLspServer.tabAdd(tabParams);
50-
} catch (Exception e) {
51-
PluginLogger.error("Error occurred while sending " + Command.CHAT_TAB_ADD + " message to Amazon Q LSP server", e);
52-
throw new AmazonQPluginException(e);
53-
}
36+
ChatMessage chatMessage = new ChatMessage(amazonQLspServer);
37+
chatMessage.sendTabAdd(tabParams);
5438
}
5539

5640
}

‎plugin/src/software/aws/toolkits/eclipse/amazonq/chat/ChatPartialResultMap.java

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,37 +8,37 @@
88

99
/**
1010
* ChatPartialResultMap is a utility class responsible for managing the mapping between
11-
* partial result tokens and their corresponding ChatMessage objects in the Amazon Q plugin for Eclipse.
11+
* partial result tokens and their corresponding ChatMessage objects represented by tabId in the Amazon Q plugin for Eclipse.
1212
*
1313
* The Language Server Protocol (LSP) server sends progress notifications during long-running operations,
1414
* such as processing chat requests. These notifications include a token that identifies the specific operation
1515
* and a partial result object containing the progress information.
1616
*
1717
* This class maintains a concurrent map (tokenToChatMessageMap) that associates each token with
18-
* its respective ChatMessage object. This mapping is crucial for correctly updating the chat UI
18+
* its respective ChatMessage object identified via the tabId. This mapping is crucial for correctly updating the chat UI
1919
* with the latest progress information as it becomes available from the LSP server.
2020
*
2121
* The progress notifications are handled by the {@link AmazonQLspClientImpl#notifyProgress(ProgressParams)}
22-
* method, which retrieves the corresponding ChatMessage object from the tokenToChatMessageMap using
22+
* method, which retrieves the corresponding tabId associated with the ChatMessage object from the tokenToChatMessageMap using
2323
* the token provided in the ProgressParams. The ChatMessage can then be updated with the partial result.
2424
*/
2525
public final class ChatPartialResultMap {
2626

27-
private final Map<String, ChatMessage> tokenToChatMessageMap;
27+
private final Map<String, String> tokenToChatMessageMap;
2828

2929
public ChatPartialResultMap() {
30-
tokenToChatMessageMap = new ConcurrentHashMap<String, ChatMessage>();
30+
tokenToChatMessageMap = new ConcurrentHashMap<String, String>();
3131
}
3232

33-
public void setEntry(final String token, final ChatMessage chatMessage) {
34-
tokenToChatMessageMap.put(token, chatMessage);
33+
public void setEntry(final String token, final String tabId) {
34+
tokenToChatMessageMap.put(token, tabId);
3535
}
3636

3737
public void removeEntry(final String token) {
3838
tokenToChatMessageMap.remove(token);
3939
}
4040

41-
public ChatMessage getValue(final String token) {
41+
public String getValue(final String token) {
4242
return tokenToChatMessageMap.getOrDefault(token, null);
4343
}
4444

‎plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/AmazonQLspClientImpl.java

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,13 @@
1313
import org.eclipse.lsp4j.ConfigurationParams;
1414
import org.eclipse.lsp4j.ProgressParams;
1515

16+
import software.aws.toolkits.eclipse.amazonq.chat.ChatCommunicationManager;
1617
import software.aws.toolkits.eclipse.amazonq.configuration.PluginStore;
1718
import software.aws.toolkits.eclipse.amazonq.lsp.model.ConnectionMetadata;
1819
import software.aws.toolkits.eclipse.amazonq.lsp.model.SsoProfileData;
1920
import software.aws.toolkits.eclipse.amazonq.util.Constants;
2021
import software.aws.toolkits.eclipse.amazonq.util.PluginLogger;
2122
import software.aws.toolkits.eclipse.amazonq.util.ThreadingUtils;
22-
import software.aws.toolkits.eclipse.amazonq.views.AmazonQChatViewActionHandler;
2323

2424
@SuppressWarnings("restriction")
2525
public class AmazonQLspClientImpl extends LanguageClientImpl implements AmazonQLspClient {
@@ -34,11 +34,6 @@ public final CompletableFuture<ConnectionMetadata> getConnectionMetadata() {
3434
return CompletableFuture.completedFuture(metadata);
3535
}
3636

37-
/*
38-
* Handles the progress notifications received from the LSP server.
39-
* - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
40-
* - Other notifications are ignored at this time.
41-
*/
4237
@Override
4338
public final CompletableFuture<List<Object>> configuration(final ConfigurationParams configurationParams) {
4439
if (configurationParams.getItems().size() == 0) {
@@ -58,13 +53,18 @@ public final CompletableFuture<List<Object>> configuration(final ConfigurationPa
5853
return CompletableFuture.completedFuture(output);
5954
}
6055

56+
/*
57+
* Handles the progress notifications received from the LSP server.
58+
* - Process partial results for Chat messages if provided token is maintained by ChatCommunicationManager
59+
* - Other notifications are ignored at this time.
60+
*/
6161
@Override
6262
public final void notifyProgress(final ProgressParams params) {
63-
AmazonQChatViewActionHandler chatActionHandler = new AmazonQChatViewActionHandler();
63+
var chatCommunicationManager = ChatCommunicationManager.getInstance();
6464

6565
ThreadingUtils.executeAsyncTask(() -> {
6666
try {
67-
chatActionHandler.handlePartialResultProgressNotification(params);
67+
chatCommunicationManager.handlePartialResultProgressNotification(params);
6868
} catch (Exception e) {
6969
PluginLogger.error("Error processing partial result progress notification", e);
7070
}

‎plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/manager/LspConstants.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ private LspConstants() {
1010
// Prevent instantiation
1111
}
1212

13-
public static final String CW_MANIFEST_URL = "https://dtqjflaii39q2.cloudfront.net/codewhisperer/0/manifest.json";
13+
public static final String CW_MANIFEST_URL = "https://aws-toolkit-language-servers.amazonaws.com/eclipse/0/manifest.json";
1414

1515
public static final String CW_LSP_FILENAME = "aws-lsp-codewhisperer.js";
1616
public static final String NODE_EXECUTABLE_PREFIX = "node";

0 commit comments

Comments
 (0)