Skip to content

Commit 2fa78eb

Browse files
shruti0085shruti0085
committed
Refactor to simplify and isolate chat communication workflow
1 parent f648b5c commit 2fa78eb

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)