Ensure that the entire chat request is cancellable

chat-client/src/client/chat.ts

@@ -17,6 +17,8 @@ import {
 import {
     CHAT_REQUEST_METHOD,
     ChatParams,
+    END_CHAT_REQUEST_METHOD,
+    EndChatParams,
     FEEDBACK_NOTIFICATION_METHOD,
     FOLLOW_UP_CLICK_NOTIFICATION_METHOD,
     FeedbackParams,
@@ -106,6 +108,9 @@ export const createChat = (
     }
 
     const chatApi: OutboundChatApi = {
+        endChat: (params: EndChatParams) => {
+            sendMessageToClient({ command: END_CHAT_REQUEST_METHOD, params })
+        },
         sendChatPrompt: (params: ChatParams) => {
             sendMessageToClient({ command: CHAT_REQUEST_METHOD, params })
         },

chat-client/src/client/messager.ts

@@ -12,6 +12,7 @@ import {
 } from '@aws/chat-client-ui-types'
 import {
     ChatParams,
+    EndChatParams,
     FeedbackParams,
     FollowUpClickParams,
     InfoLinkClickParams,
@@ -40,6 +41,7 @@ import {
 } from '../contracts/telemetry'
 
 export interface OutboundChatApi {
+    endChat(params: EndChatParams): void
     sendChatPrompt(params: ChatParams): void
     sendQuickActionCommand(params: QuickActionParams): void
     tabAdded(params: TabAddParams): void
@@ -81,6 +83,10 @@ export class Messager {
         this.chatApi.telemetry({ ...params, tabId, name: SEND_TO_PROMPT_TELEMETRY_EVENT })
     }
 
+    onStopChatResponse = (params: EndChatParams): void => {
+        this.chatApi.endChat(params)
+    }
+
     onChatPrompt = (params: ChatParams, triggerType?: string): void => {
         // Let the server know about the latest trigger interaction on the tabId
         this.chatApi.telemetry({

chat-client/src/client/mynahUi.test.ts

@@ -22,6 +22,7 @@ describe('MynahUI', () => {
 
     beforeEach(() => {
         outboundChatApi = {
+            endChat: sinon.stub(),
             sendChatPrompt: sinon.stub(),
             sendQuickActionCommand: sinon.stub(),
             tabAdded: sinon.stub(),

chat-client/src/client/mynahUi.ts

@@ -221,6 +221,13 @@ export const createMynahUi = (messager: Messager, tabFactory: TabFactory): [Myna
             }
             messager.onSourceLinkClick(payload)
         },
+        onStopChatResponse(tabId, _eventId) {
+            messager.onStopChatResponse({ tabId })
+            mynahUi.updateStore(tabId, {
+                loadingChat: false,
+                promptInputDisabledState: false,
+            })
+        },
         onInfoLinkClick: (tabId, link, mouseEvent, eventId) => {
             mouseEvent?.preventDefault()
             mouseEvent?.stopPropagation()

chat-client/src/contracts/serverContracts.ts

@@ -19,11 +19,13 @@ import {
     SOURCE_LINK_CLICK_NOTIFICATION_METHOD,
     INFO_LINK_CLICK_NOTIFICATION_METHOD,
     QUICK_ACTION_REQUEST_METHOD,
+    END_CHAT_REQUEST_METHOD,
 } from '@aws/language-server-runtimes-types'
 
 export const TELEMETRY = 'telemetry/event'
 
 export type ServerMessageCommand =
+    | typeof END_CHAT_REQUEST_METHOD
     | typeof CHAT_REQUEST_METHOD
     | typeof TAB_ADD_NOTIFICATION_METHOD
     | typeof TAB_REMOVE_NOTIFICATION_METHOD

client/vscode/src/chatActivation.ts

@@ -12,6 +12,7 @@ import {
     quickActionRequestType,
     QuickActionResult,
     QuickActionParams,
+    END_CHAT_REQUEST_METHOD,
 } from '@aws/language-server-runtimes/protocol'
 import { v4 as uuidv4 } from 'uuid'
 import { Uri, ViewColumn, Webview, WebviewPanel, commands, window } from 'vscode'
@@ -54,26 +55,64 @@ export function registerChat(languageClient: LanguageClient, extensionUri: Uri,
         languageClient.info(`vscode client: Received ${JSON.stringify(message)} from chat`)
 
         switch (message.command) {
+            case END_CHAT_REQUEST_METHOD:
+                languageClient.sendRequest(END_CHAT_REQUEST_METHOD, message.params)
+                break
             case INSERT_TO_CURSOR_POSITION:
                 insertTextAtCursorPosition(message.params.code)
                 break
             case AUTH_FOLLOW_UP_CLICKED:
                 languageClient.info('AuthFollowUp clicked')
                 break
-            case chatRequestType.method:
+            case chatRequestType.method: {
+                let chatResult: ChatResult | string = ''
+                const textDocument = window.visibleTextEditors.find(editor => editor.document.uri.scheme === 'file')
+                const documentUri = textDocument?.document.uri.toString()
+
                 const partialResultToken = uuidv4()
-                const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, partialResult =>
+                const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, partialResult => {
+                    chatResult = partialResult
                     handlePartialResult<ChatResult>(partialResult, encryptionKey, panel, message.params.tabId)
+                })
+
+                const chatRequest = await encryptRequest<ChatParams>(
+                    {
+                        ...message.params,
+                        textDocument: { uri: documentUri },
+                    },
+                    encryptionKey
                 )
 
-                const chatRequest = await encryptRequest<ChatParams>(message.params, encryptionKey)
-                const chatResult = await languageClient.sendRequest(chatRequestType, {
-                    ...chatRequest,
-                    partialResultToken,
-                })
-                handleCompleteResult<ChatResult>(chatResult, encryptionKey, panel, message.params.tabId, chatDisposable)
+                try {
+                    chatResult = await languageClient.sendRequest(chatRequestType, {
+                        ...chatRequest,
+                        partialResultToken,
+                    })
+                } catch (e) {
+                    if (e instanceof Error) {
+                        languageClient.info(`Client caught error during chat request: ${e.message}`)
+
+                        if (chatResult === '') {
+                            if (e.message === 'Request cancelled') {
+                                languageClient.info('Request cancelled before receiving any partial result')
+                                chatResult = { body: 'Request Cancelled' }
+                            } else {
+                                chatResult = { body: `Error in chat: ${e.message}` }
+                            }
+                        }
+                    }
+                } finally {
+                    handleCompleteResult<ChatResult>(
+                        chatResult,
+                        encryptionKey,
+                        panel,
+                        message.params.tabId,
+                        chatDisposable
+                    )
+                }
                 break
-            case quickActionRequestType.method:
+            }
+            case quickActionRequestType.method: {
                 const quickActionPartialResultToken = uuidv4()
                 const quickActionDisposable = languageClient.onProgress(
                     quickActionRequestType,
@@ -100,6 +139,7 @@ export function registerChat(languageClient: LanguageClient, extensionUri: Uri,
                     quickActionDisposable
                 )
                 break
+            }
             case followUpClickNotificationType.method:
                 if (!isValidAuthFollowUpType(message.params.followUp.type))
                     languageClient.sendNotification(followUpClickNotificationType, message.params)

core/aws-lsp-fqn/src/browser/fqnWorkerPool.ts

@@ -1,25 +1,28 @@
-import { ExtractorResult, FqnExtractorInput, IFqnWorkerPool } from '../common/types'
+import { ExtractorResult, FqnExtractorInput, IFqnWorkerPool, Cancellable } from '../common/types'
 
 // TODO: implement logic for browser/webworker environment
 export class FqnWorkerPool implements IFqnWorkerPool {
-    public async exec(_input: FqnExtractorInput): Promise<ExtractorResult> {
-        return Promise.resolve({
-            success: true,
-            data: {
-                fullyQualified: {
-                    declaredSymbols: [],
-                    usedSymbols: [],
+    public exec(_input: FqnExtractorInput): Cancellable<Promise<ExtractorResult>> {
+        return [
+            Promise.resolve({
+                success: true,
+                data: {
+                    fullyQualified: {
+                        declaredSymbols: [],
+                        usedSymbols: [],
+                    },
+                    simple: {
+                        declaredSymbols: [],
+                        usedSymbols: [],
+                    },
+                    externalSimple: {
+                        declaredSymbols: [],
+                        usedSymbols: [],
+                    },
                 },
-                simple: {
-                    declaredSymbols: [],
-                    usedSymbols: [],
-                },
-                externalSimple: {
-                    declaredSymbols: [],
-                    usedSymbols: [],
-                },
-            },
-        })
+            }),
+            () => {},
+        ]
     }
 
     public dispose() {}

core/aws-lsp-fqn/src/common/commonFqnWorkerPool.ts

@@ -1,6 +1,6 @@
 import { pool, Pool } from 'workerpool'
 import { DEFAULT_MAX_QUEUE_SIZE, DEFAULT_MAX_WORKERS, DEFAULT_TIMEOUT, FQN_WORKER_ID } from './defaults'
-import { ExtractorResult, FqnExtractorInput, IFqnWorkerPool, Logger, WorkerPoolConfig } from './types'
+import { ExtractorResult, FqnExtractorInput, IFqnWorkerPool, Logger, WorkerPoolConfig, Cancellable } from './types'
 
 export class CommonFqnWorkerPool implements IFqnWorkerPool {
     #workerPool: Pool
@@ -17,25 +17,32 @@ export class CommonFqnWorkerPool implements IFqnWorkerPool {
         })
     }
 
-    public async exec(input: FqnExtractorInput): Promise<ExtractorResult> {
+    public exec(input: FqnExtractorInput): Cancellable<Promise<ExtractorResult>> {
         this.#logger?.log(`Extracting fully qualified names for ${input.languageId}`)
 
-        return this.#workerPool
-            .exec(FQN_WORKER_ID, [input])
-            .timeout(this.#timeout)
-            .then(data => data as ExtractorResult)
-            .catch(error => {
-                const errorMessage = `Encountered error while extracting fully qualified names: ${
-                    error instanceof Error ? error.message : 'Unknown'
-                }`
-
-                this.#logger?.error(errorMessage)
-
-                return {
-                    success: false,
-                    error: errorMessage,
-                }
-            })
+        const execPromise = this.#workerPool.exec(FQN_WORKER_ID, [input]).timeout(this.#timeout)
+
+        return [
+            // have to wrap this in promise since exec promise is not a true promise
+            new Promise<ExtractorResult>(resolve => {
+                execPromise
+                    .then(data => resolve(data as ExtractorResult))
+                    .catch(error => {
+                        const errorMessage = `Encountered error while extracting fully qualified names: ${
+                            error instanceof Error ? error.message : 'Unknown'
+                        }`
+
+                        this.#logger?.error(errorMessage)
+
+                        // using result pattern, so we will resolve with success: false
+                        return resolve({
+                            success: false,
+                            error,
+                        })
+                    })
+            }),
+            () => execPromise.cancel(),
+        ]
     }
 
     public dispose() {

core/aws-lsp-fqn/src/common/types.ts

@@ -20,7 +20,7 @@ export type Result<TData, TError> =
           error: TError
       }
 
-export type ExtractorResult = Result<FqnExtractorOutput, string>
+export type ExtractorResult = Result<FqnExtractorOutput, Error>
 
 export interface FullyQualifiedName {
     source: string[]
@@ -72,6 +72,9 @@ export interface FqnExtractorInput {
 }
 
 export interface IFqnWorkerPool {
-    exec(input: FqnExtractorInput): Promise<ExtractorResult>
+    exec(input: FqnExtractorInput): Cancellable<Promise<ExtractorResult>>
     dispose(): void
 }
+
+export type CancelFn = () => void
+export type Cancellable<T> = [T, CancelFn]

core/aws-lsp-fqn/src/index.ts

@@ -1,8 +1,9 @@
-import { ExtractorResult, FqnExtractorInput, IFqnWorkerPool, WorkerPoolConfig } from './common/types'
+import { Cancellable, ExtractorResult, FqnExtractorInput, IFqnWorkerPool, WorkerPoolConfig } from './common/types'
 
 export * from './common/types'
+
 export declare class FqnWorkerPool implements IFqnWorkerPool {
     constructor(workerPoolConfig?: WorkerPoolConfig)
-    exec(input: FqnExtractorInput): Promise<ExtractorResult>
+    exec(input: FqnExtractorInput): Cancellable<Promise<ExtractorResult>>
     dispose(): void
 }

server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts

@@ -12,7 +12,7 @@ import { ChatController } from './chatController'
 import { ChatSessionManagementService } from './chatSessionManagementService'
 import { ChatSessionService } from './chatSessionService'
 import { ChatTelemetryController } from './telemetry/chatTelemetryController'
-import { DocumentContextExtractor } from './contexts/documentContext'
+import { TriggerContextExtractor } from './contexts/triggerContextExtractor'
 import * as utils from './utils'
 import { DEFAULT_HELP_FOLLOW_UP_PROMPT, HELP_MESSAGE } from './constants'
 
@@ -61,6 +61,8 @@ describe('ChatController', () => {
     }
     let removeConversationSpy: sinon.SinonSpy
     let emitConversationMetricStub: sinon.SinonStub
+    let abortRequestStub: sinon.SinonStub
+    let triggerContextCancelSpy: sinon.SinonSpy
 
     let testFeatures: TestFeatures
     let chatSessionManagementService: ChatSessionManagementService
@@ -88,9 +90,10 @@ describe('ChatController', () => {
         activeTabSpy = sinon.spy(ChatTelemetryController.prototype, 'activeTabId', ['get', 'set'])
         removeConversationSpy = sinon.spy(ChatTelemetryController.prototype, 'removeConversation')
         emitConversationMetricStub = sinon.stub(ChatTelemetryController.prototype, 'emitConversationMetric')
+        triggerContextCancelSpy = sinon.spy(TriggerContextExtractor.prototype, 'cancel')
 
         disposeStub = sinon.stub(ChatSessionService.prototype, 'dispose')
-
+        abortRequestStub = sinon.stub(ChatSessionService.prototype, 'abortRequest')
         chatSessionManagementService = ChatSessionManagementService.getInstance().withCredentialsProvider(
             testFeatures.credentialsProvider
         )
@@ -132,11 +135,8 @@ describe('ChatController', () => {
 
         chatController.onEndChat({ tabId: mockTabId }, mockCancellationToken)
 
-        sinon.assert.calledOnce(disposeStub)
-
-        const hasSession = chatSessionManagementService.hasSession(mockTabId)
-
-        assert.ok(!hasSession)
+        sinon.assert.calledOnce(abortRequestStub)
+        sinon.assert.calledOnce(triggerContextCancelSpy)
     })
 
     it('onTabAdd sets active tab id in telemetryController', () => {
@@ -340,7 +340,7 @@ describe('ChatController', () => {
             }
 
             beforeEach(() => {
-                extractDocumentContextStub = sinon.stub(DocumentContextExtractor.prototype, 'extractDocumentContext')
+                extractDocumentContextStub = sinon.stub(TriggerContextExtractor.prototype, 'extractDocumentContext')
                 testFeatures.openDocument(typescriptDocument)
             })