feat: added quickAction notification for async flows

Author: volodkevych

runtimes/protocol/chat.ts

@@ -39,6 +39,7 @@ import {
     FILE_CLICK_NOTIFICATION_METHOD,
     ChatUpdateParams,
     FileClickParams,
+    QUICK_ACTION_NOTIFICATION_METHOD,
 } from './lsp'
 
 export const chatRequestType = new AutoParameterStructuresProtocolRequestType<
@@ -53,11 +54,15 @@ export const endChatRequestType = new ProtocolRequestType<EndChatParams, EndChat
 )
 export const quickActionRequestType = new AutoParameterStructuresProtocolRequestType<
     QuickActionParams | EncryptedQuickActionParams,
-    QuickActionResult | string | null,
+    QuickActionResult | string,
     QuickActionResult | string,
     void,
     void
 >(QUICK_ACTION_REQUEST_METHOD)
+export const quickActionNotificationType = new ProtocolNotificationType<
+    QuickActionParams | EncryptedQuickActionParams,
+    void
+>(QUICK_ACTION_NOTIFICATION_METHOD)
 export const readyNotificationType = new ProtocolNotificationType<void, void>(READY_NOTIFICATION_METHOD)
 export const feedbackNotificationType = new ProtocolNotificationType<FeedbackParams, void>(FEEDBACK_NOTIFICATION_METHOD)
 export const tabAddNotificationType = new ProtocolNotificationType<TabAddParams, void>(TAB_ADD_NOTIFICATION_METHOD)

runtimes/runtimes/base-runtime.ts

@@ -38,6 +38,7 @@ import {
     selectWorkspaceItemRequestType,
     chatUpdateNotificationType,
     fileClickNotificationType,
+    quickActionNotificationType,
 } from '../protocol'
 import { createConnection } from 'vscode-languageserver/browser'
 import {
@@ -126,6 +127,7 @@ export const baseRuntime = (connections: { reader: MessageReader; writer: Messag
         onChatPrompt: handler => lspConnection.onRequest(chatRequestType.method, handler),
         onEndChat: handler => lspConnection.onRequest(endChatRequestType.method, handler),
         onQuickAction: handler => lspConnection.onRequest(quickActionRequestType.method, handler),
+        onQuickActionTrigger: handler => lspConnection.onNotification(quickActionNotificationType.method, handler),
         onSendFeedback: handler => lspConnection.onNotification(feedbackNotificationType.method, handler),
         onReady: handler => lspConnection.onNotification(readyNotificationType.method, handler),
         onTabAdd: handler => lspConnection.onNotification(tabAddNotificationType.method, handler),

runtimes/runtimes/chat/baseChat.ts

@@ -36,6 +36,7 @@ import {
     chatUpdateNotificationType,
     FileClickParams,
     fileClickNotificationType,
+    quickActionNotificationType,
 } from '../../protocol'
 import { Chat } from '../../server-interface'
 
@@ -54,6 +55,10 @@ export class BaseChat implements Chat {
         this.connection.onRequest(quickActionRequestType.method, handler)
     }
 
+    public onQuickActionTrigger(handler: NotificationHandler<QuickActionParams>) {
+        this.connection.onRequest(quickActionNotificationType.method, handler)
+    }
+
     public onSendFeedback(handler: NotificationHandler<FeedbackParams>) {
         this.connection.onNotification(feedbackNotificationType.method, handler)
     }

runtimes/runtimes/chat/encryptedChat.test.ts

@@ -88,4 +88,6 @@ describe('EncryptedChat', () => {
         assert(!(result instanceof Error))
         sinonAssert.calledOnce(handler)
     })
+
+    // TODO: Add tests
 })

runtimes/runtimes/chat/encryptedChat.ts

@@ -13,6 +13,8 @@ import {
     quickActionRequestType,
     ResponseError,
     LSPErrorCodes,
+    NotificationHandler,
+    quickActionNotificationType,
 } from '../../protocol'
 import { CredentialsEncoding, encryptObjectWithKey, isMessageJWEEncrypted } from '../auth/standalone/encryption'
 import { BaseChat } from './baseChat'
@@ -52,6 +54,42 @@ export class EncryptedChat extends BaseChat {
         )
     }
 
+    public onQuickActionTrigger(handler: NotificationHandler<QuickActionParams>) {
+        this.registerEncryptedNotificationHandler<EncryptedQuickActionParams, QuickActionParams>(
+            quickActionNotificationType,
+            handler
+        )
+    }
+
+    private registerEncryptedNotificationHandler<
+        EncryptedRequestType extends EncryptedRequestParams,
+        DecryptedRequestType extends ChatParams | QuickActionParams,
+    >(notificationType: any, handler: NotificationHandler<DecryptedRequestType>) {
+        this.connection.onNotification(
+            notificationType,
+            async (request: EncryptedRequestType | DecryptedRequestType, cancellationToken: CancellationToken) => {
+                // Verify the request is encrypted as expected
+                if (this.instanceOfEncryptedParams<EncryptedRequestType>(request)) {
+                    // Decrypt request
+                    let decryptedRequest
+                    try {
+                        decryptedRequest = await this.decodeRequest<DecryptedRequestType>(request)
+                    } catch (err: unknown) {
+                        let errorMessage = 'Request could not be decrypted'
+                        if (err instanceof Error) errorMessage = err.message
+                        return new ResponseError<ResponseType>(LSPErrorCodes.ServerCancelled, errorMessage)
+                    }
+
+                    // Preserve the partial result token
+                    decryptedRequest.partialResultToken = request.partialResultToken
+
+                    // Call the handler with decrypted params
+                    await handler(decryptedRequest)
+                }
+            }
+        )
+    }
+
     private registerEncryptedRequestHandler<
         EncryptedRequestType extends EncryptedRequestParams,
         DecryptedRequestType extends ChatParams | QuickActionParams,

runtimes/server-interface/chat.ts

@@ -32,6 +32,7 @@ export type Chat = {
     onQuickAction: (handler: RequestHandler<QuickActionParams, QuickActionResult, void>) => void
     openTab: (params: OpenTabParams) => Promise<OpenTabResult>
     // Notifications
+    onQuickActionTrigger: (handler: NotificationHandler<QuickActionParams>) => void
     onSendFeedback: (handler: NotificationHandler<FeedbackParams>) => void
     onReady: (handler: NotificationHandler<void>) => void
     onTabAdd: (handler: NotificationHandler<TabAddParams>) => void

types/chat.ts

@@ -4,6 +4,7 @@ import { Position, Range, TextDocumentIdentifier } from './lsp'
 export const CHAT_REQUEST_METHOD = 'aws/chat/sendChatPrompt'
 export const END_CHAT_REQUEST_METHOD = 'aws/chat/endChat'
 export const QUICK_ACTION_REQUEST_METHOD = 'aws/chat/sendChatQuickAction'
+export const QUICK_ACTION_NOTIFICATION_METHOD = 'aws/chat/triggerChatQuickAction'
 export const READY_NOTIFICATION_METHOD = 'aws/chat/ready'
 export const FEEDBACK_NOTIFICATION_METHOD = 'aws/chat/feedback'
 export const TAB_ADD_NOTIFICATION_METHOD = 'aws/chat/tabAdd'
@@ -107,9 +108,10 @@ export type EndChatResult = boolean
  */
 export interface QuickActionCommand {
     command: string
-    disabled?: boolean
     description?: string
-    placeholder?: string
+    disabled?: boolean
+    async?: boolean
+    defaultTabData?: TabData
 }
 
 /**
@@ -142,7 +144,6 @@ export interface TabData {
  */
 export interface ChatOptions {
     quickActions?: QuickActions
-    defaultTabData?: TabData
 }
 
 export interface QuickActionParams extends PartialResultParams {