From 86079da07ce2c274825449cd2a2ca61072c86988 Mon Sep 17 00:00:00 2001 From: Josh Pinkney Date: Mon, 21 Apr 2025 11:46:09 -0400 Subject: [PATCH 1/3] add stop response --- packages/amazonq/src/lsp/chat/messages.ts | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index 2f17d6a8bfa..f217e29f12d 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -42,6 +42,7 @@ import { LINK_CLICK_NOTIFICATION_METHOD, LinkClickParams, INFO_LINK_CLICK_NOTIFICATION_METHOD, + CancellationTokenSource, } from '@aws/language-server-runtimes/protocol' import { v4 as uuidv4 } from 'uuid' import * as vscode from 'vscode' @@ -109,6 +110,8 @@ export function registerMessageListeners( return } + const chatStreamTokens = new Map() // tab id -> token + const webview = provider.webview switch (message.command) { case COPY_TO_CLIPBOARD: @@ -181,10 +184,21 @@ export function registerMessageListeners( void openUrl(vscode.Uri.parse(linkParams.link)) break } + case STOP_RESPONSE: { + const tabId = (message as StopResponseMessage).tabId + const token = chatStreamTokens.get(tabId) + token?.cancel() + token?.dispose() + chatStreamTokens.delete(tabId) + break + } case chatRequestType.method: { const chatParams: ChatParams = { ...message.params } const partialResultToken = uuidv4() let lastPartialResult: ChatResult | undefined + const cancellationToken = new CancellationTokenSource() + chatStreamTokens.set(message.params.tabId, cancellationToken) + const chatDisposable = languageClient.onProgress( chatRequestType, partialResultToken, From 43658ff6f2db44438e433efc6af3facd12d824b8 Mon Sep 17 00:00:00 2001 From: Josh Pinkney Date: Mon, 21 Apr 2025 20:06:00 -0400 Subject: [PATCH 2/3] feat(amazonq): add stop chat repsponse --- package-lock.json | 16 ++++++++-------- packages/amazonq/src/lsp/chat/messages.ts | 21 +++++++++++++-------- packages/core/package.json | 2 +- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/package-lock.json b/package-lock.json index 77bfc670449..1050abd6a36 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10833,12 +10833,12 @@ } }, "node_modules/@aws/chat-client-ui-types": { - "version": "0.1.22", - "resolved": "https://registry.npmjs.org/@aws/chat-client-ui-types/-/chat-client-ui-types-0.1.22.tgz", - "integrity": "sha512-vn+UKnh9hgZN1LCMONgeZE8WWxivWXaHQq+oG9wpbFhaTXn/nNBTQ9ON7S2fvMqo0g0Np/6hirxZy5ROcWnB9Q==", + "version": "0.1.25", + "resolved": "https://registry.npmjs.org/@aws/chat-client-ui-types/-/chat-client-ui-types-0.1.25.tgz", + "integrity": "sha512-sxSookCLlhfsamse3x9AkvCei7SSUYDOklAe1O2jiUOYSN79M5JlVVRZShoqiOCHds7bb9nSaz+DMWIwEK1+2w==", "dev": true, "dependencies": { - "@aws/language-server-runtimes-types": "^0.1.19" + "@aws/language-server-runtimes-types": "^0.1.21" } }, "node_modules/@aws/language-server-runtimes": { @@ -10876,9 +10876,9 @@ } }, "node_modules/@aws/language-server-runtimes-types": { - "version": "0.1.19", - "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes-types/-/language-server-runtimes-types-0.1.19.tgz", - "integrity": "sha512-c81J3G3N6JP5A6g70xTpK/XPS1YWwviQBn307Rk3S5fSiALT8INeHM+IPDg9AuONU6w378RJjzQy3+PE0gJvsw==", + "version": "0.1.21", + "resolved": "https://registry.npmjs.org/@aws/language-server-runtimes-types/-/language-server-runtimes-types-0.1.21.tgz", + "integrity": "sha512-03C3dz4MvMyKg4UAgHMNNw675OQJkDq+7TPXUPaiasqPF946ywTDD9xoNPaVOQI+YTtC7Re4vhPRfBzyad3MOg==", "dev": true, "dependencies": { "vscode-languageserver-textdocument": "^1.0.12", @@ -26804,7 +26804,7 @@ "devDependencies": { "@aws-sdk/types": "^3.13.1", "@aws/chat-client": "^0.1.4", - "@aws/chat-client-ui-types": "^0.1.22", + "@aws/chat-client-ui-types": "^0.1.24", "@aws/language-server-runtimes": "^0.2.58", "@aws/language-server-runtimes-types": "^0.1.13", "@cspotcode/source-map-support": "^0.8.1", diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index f217e29f12d..a59faa14a99 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -14,6 +14,8 @@ import { UiMessageResultParams, CHAT_PROMPT_OPTION_ACKNOWLEDGED, ChatPromptOptionAcknowledgedMessage, + STOP_CHAT_RESPONSE, + StopChatResponseMessage, } from '@aws/chat-client-ui-types' import { ChatResult, @@ -98,6 +100,7 @@ export function registerMessageListeners( provider: AmazonQChatViewProvider, encryptionKey: Buffer ) { + const chatStreamTokens = new Map() // tab id -> token provider.webview?.onDidReceiveMessage(async (message) => { languageClient.info(`[VSCode Client] Received ${JSON.stringify(message)} from chat`) @@ -110,8 +113,6 @@ export function registerMessageListeners( return } - const chatStreamTokens = new Map() // tab id -> token - const webview = provider.webview switch (message.command) { case COPY_TO_CLIPBOARD: @@ -184,8 +185,8 @@ export function registerMessageListeners( void openUrl(vscode.Uri.parse(linkParams.link)) break } - case STOP_RESPONSE: { - const tabId = (message as StopResponseMessage).tabId + case STOP_CHAT_RESPONSE: { + const tabId = (message as StopChatResponseMessage).params.tabId const token = chatStreamTokens.get(tabId) token?.cancel() token?.dispose() @@ -226,10 +227,14 @@ export function registerMessageListeners( const chatRequest = await encryptRequest(chatParams, encryptionKey) try { - const chatResult = await languageClient.sendRequest(chatRequestType.method, { - ...chatRequest, - partialResultToken, - }) + const chatResult = await languageClient.sendRequest( + chatRequestType.method, + { + ...chatRequest, + partialResultToken, + }, + cancellationToken.token + ) await handleCompleteResult( chatResult, encryptionKey, diff --git a/packages/core/package.json b/packages/core/package.json index 98f24feae81..b9261d971dc 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -442,7 +442,7 @@ "devDependencies": { "@aws-sdk/types": "^3.13.1", "@aws/chat-client": "^0.1.4", - "@aws/chat-client-ui-types": "^0.1.22", + "@aws/chat-client-ui-types": "^0.1.24", "@aws/language-server-runtimes": "^0.2.58", "@aws/language-server-runtimes-types": "^0.1.13", "@cspotcode/source-map-support": "^0.8.1", From 55d6c91e3bb012209764eaa5b0de13665b9c0bcc Mon Sep 17 00:00:00 2001 From: Josh Pinkney Date: Mon, 21 Apr 2025 20:11:43 -0400 Subject: [PATCH 3/3] fix(amazonq): delete token when you're done --- packages/amazonq/src/lsp/chat/messages.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts index a59faa14a99..39294bbfa27 100644 --- a/packages/amazonq/src/lsp/chat/messages.ts +++ b/packages/amazonq/src/lsp/chat/messages.ts @@ -198,7 +198,7 @@ export function registerMessageListeners( const partialResultToken = uuidv4() let lastPartialResult: ChatResult | undefined const cancellationToken = new CancellationTokenSource() - chatStreamTokens.set(message.params.tabId, cancellationToken) + chatStreamTokens.set(chatParams.tabId, cancellationToken) const chatDisposable = languageClient.onProgress( chatRequestType, @@ -259,6 +259,8 @@ export function registerMessageListeners( chatParams.tabId, chatDisposable ) + } finally { + chatStreamTokens.delete(chatParams.tabId) } break }