feat: handle short cut to execute shell command

chat-client/src/client/chat.ts

@@ -155,6 +155,9 @@ export const createChat = (
         }
 
         switch (message?.command) {
+            case 'executeShellCommandShortCut':
+                mynahApi.executeShellCommandShortCut(message.params)
+                break
             case CHAT_REQUEST_METHOD:
                 mynahApi.addChatResponse(message.params, message.tabId, message.isPartialResult)
                 break

chat-client/src/client/mynahUi.ts

@@ -67,6 +67,7 @@ export interface InboundChatApi {
     openTab(requestId: string, params: OpenTabParams): void
     sendContextCommands(params: ContextCommandParams): void
     listConversations(params: ListConversationsResult): void
+    executeShellCommandShortCut(params: any): void
     conversationClicked(params: ConversationClickResult): void
     getSerializedChat(requestId: string, params: GetSerializedChatParams): void
     createTabId(openTab?: boolean): string | undefined
@@ -964,6 +965,38 @@ ${params.message}`,
         messager.onError(params)
     }
 
+    const executeShellCommandShortCut = (params: any) => {
+        const tabId = mynahUi.getSelectedTabId()
+        if (!tabId) return
+
+        const chatItems = mynahUi.getTabData(tabId)?.getStore()?.chatItems || []
+        const buttonId = params.buttonId
+
+        let messageId
+        for (const item of chatItems) {
+            if (buttonId === 'stop-shell-command' && item.buttons && item.buttons.some(b => b.id === buttonId)) {
+                messageId = item.messageId
+                break
+            }
+            if (item.header?.buttons && item.header.buttons.some(b => b.id === buttonId)) {
+                messageId = item.messageId
+                break
+            }
+        }
+
+        if (messageId) {
+            const payload: ButtonClickParams = {
+                tabId,
+                messageId,
+                buttonId,
+            }
+            messager.onButtonClick(payload)
+            if (buttonId === 'stop-shell-command') {
+                messager.onStopChatResponse(tabId)
+            }
+        }
+    }
+
     const openTab = (requestId: string, params: OpenTabParams) => {
         if (params.tabId) {
             if (params.tabId !== mynahUi.getSelectedTabId()) {
@@ -1085,6 +1118,7 @@ ${params.message}`,
         showError: showError,
         openTab: openTab,
         sendContextCommands: sendContextCommands,
+        executeShellCommandShortCut: executeShellCommandShortCut,
         listConversations: listConversations,
         conversationClicked: conversationClicked,
         getSerializedChat: getSerializedChat,

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -3,7 +3,7 @@
  * Will be deleted or merged.
  */
 
 import * as path from 'path'
 import {
     ChatTriggerType,
     CodeWhispererStreamingServiceException,
@@ -1022,7 +1022,6 @@
 
         const toolMsgId = toolUse.toolUseId!
         const chatMsgId = chatResultStream.getResult().messageId
-        let headerEmitted = false
 
         const initialHeader: ChatMessage['header'] = {
             body: 'shell',
@@ -1043,10 +1042,8 @@
                     type: 'tool',
                     messageId: toolMsgId,
                     body: chunk,
-                    header: headerEmitted ? undefined : initialHeader,
+                    header: initialHeader,
                 })
-
-                headerEmitted = true
             },
 
             close: async () => {