feat: handle short cut to execute shell command

Author: yueny2020

packages/amazonq/package.json

@@ -544,8 +544,18 @@
         },
         "commands": [
             {
-                "command": "aws.amazonq.stopToolExecution",
-                "title": "Stop Amazon Q Tool Execution",
+                "command": "aws.amazonq.stopCmdExecution",
+                "title": "Stop Amazon Q Command Execution",
+                "category": "%AWS.amazonq.title%"
+            },
+            {
+                "command": "aws.amazonq.runCmdExecution",
+                "title": "Run Amazon Q Command Execution",
+                "category": "%AWS.amazonq.title%"
+            },
+            {
+                "command": "aws.amazonq.rejectCmdExecution",
+                "title": "Reject Amazon Q Command Execution",
                 "category": "%AWS.amazonq.title%"
             },
             {
@@ -901,9 +911,22 @@
                 "when": "editorTextFocus && aws.codewhisperer.connected && amazonq.inline.codelensShortcutEnabled"
             },
             {
-                "command": "aws.amazonq.stopToolExecution",
+                "command": "aws.amazonq.stopCmdExecution",
+                "key": "cmd+alt+s",
+                "mac": "cmd+alt+s",
+                "when": "aws.amazonq.amazonqChatLSP.isRunning"
+            },
+            {
+                "command": "aws.amazonq.runCmdExecution",
+                "key": "cmd+alt+r",
+                "mac": "cmd+alt+r",
+                "when": "aws.amazonq.amazonqChatLSP.isRunning"
+            },
+            {
+                "command": "aws.amazonq.rejectCmdExecution",
                 "key": "cmd+alt+h",
-                "mac": "cmd+alt+h"
+                "mac": "cmd+alt+h",
+                "when": "aws.amazonq.amazonqChatLSP.isRunning"
             }
         ],
         "icons": {

packages/amazonq/src/lsp/chat/commands.ts

@@ -16,14 +16,6 @@ import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
  */
 export function registerCommands(provider: AmazonQChatViewProvider) {
     globals.context.subscriptions.push(
-        Commands.register('aws.amazonq.stopToolExecution', async () => {
-            void focusAmazonQPanel().then(() => {
-                void provider.webview?.postMessage({
-                    command: 'stopChatResponse',
-                    params: {}, // how to get the current tabId?
-                })
-            })
-        }),
         registerGenericCommand('aws.amazonq.explainCode', 'Explain', provider),
         registerGenericCommand('aws.amazonq.refactorCode', 'Refactor', provider),
         registerGenericCommand('aws.amazonq.fixCode', 'Fix', provider),
@@ -92,7 +84,10 @@ export function registerCommands(provider: AmazonQChatViewProvider) {
                     params: {},
                 })
             })
-        })
+        }),
+        registerShellCommandShortCut('aws.amazonq.runCmdExecution', 'run-shell-command', provider),
+        registerShellCommandShortCut('aws.amazonq.rejectCmdExecution', 'reject-shell-command', provider),
+        registerShellCommandShortCut('aws.amazonq.stopCmdExecution', 'stop-shell-command', provider)
     )
 }
 
@@ -137,3 +132,14 @@ async function focusAmazonQPanel() {
     await Commands.tryExecute('aws.amazonq.AmazonQChatView.focus')
     await Commands.tryExecute('aws.amazonq.AmazonCommonAuth.focus')
 }
+
+function registerShellCommandShortCut(commandName: string, buttonId: string, provider: AmazonQChatViewProvider) {
+    return Commands.register(commandName, async () => {
+        void focusAmazonQPanel().then(() => {
+            void provider.webview?.postMessage({
+                command: 'executeShellCommandShortCut',
+                params: { buttonId },
+            })
+        })
+    })
+}

packages/amazonq/src/lsp/chat/messages.ts

@@ -133,13 +133,6 @@ export function registerMessageListeners(
         }
 
         const webview = provider.webview
-        if (message.command === 'stopChatResponse') {
-            const tabId = (message as StopChatResponseMessage).params.tabId
-            const token = chatStreamTokens.get(tabId)
-            token?.cancel()
-            token?.dispose()
-            chatStreamTokens.delete(tabId)
-        }
         switch (message.command) {
             case COPY_TO_CLIPBOARD:
                 languageClient.info('[VSCode Client] Copy to clipboard event received')

packages/amazonq/src/lsp/client.ts

@@ -32,6 +32,7 @@ import {
     getLogger,
     undefinedIfEmpty,
     getOptOutPreference,
+    setContext,
 } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
@@ -157,6 +158,17 @@ export async function startLanguageServer(
 
         if (Experiments.instance.get('amazonqChatLSP', true)) {
             await activate(client, encryptionKey, resourcePaths.ui)
+
+            await setContext('aws.amazonq.amazonqChatLSP.isRunning', true)
+            getLogger().info('Amazon Q Chat LSP context flag set on client activated')
+
+            // Add a disposable to reset the context flag when the client stops
+            toDispose.push({
+                dispose: async () => {
+                    await setContext('aws.amazonq.amazonqChatLSP.isRunning', false)
+                    getLogger().info('Amazon Q Chat LSP context flag reset on client disposal')
+                },
+            })
         }
 
         const refreshInterval = auth.startTokenRefreshInterval(10 * oneSecond)

packages/core/src/shared/vscode/setContext.ts

@@ -39,6 +39,7 @@ export type contextKey =
     | 'gumby.wasQCodeTransformationUsed'
     | 'amazonq.inline.codelensShortcutEnabled'
     | 'aws.toolkit.lambda.walkthroughSelected'
+    | 'aws.amazonq.amazonqChatLSP.isRunning'
 
 const contextMap: Partial<Record<contextKey, any>> = {}