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 2f17d6a8bfa..39294bbfa27 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, @@ -42,6 +44,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' @@ -97,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`) @@ -181,10 +185,21 @@ export function registerMessageListeners( void openUrl(vscode.Uri.parse(linkParams.link)) break } + case STOP_CHAT_RESPONSE: { + const tabId = (message as StopChatResponseMessage).params.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(chatParams.tabId, cancellationToken) + const chatDisposable = languageClient.onProgress( chatRequestType, partialResultToken, @@ -212,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, @@ -240,6 +259,8 @@ export function registerMessageListeners( chatParams.tabId, chatDisposable ) + } finally { + chatStreamTokens.delete(chatParams.tabId) } break } 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",