feat(amazonq): Pair programming mode (#21)

* pair programming mode

* pairProgrammingModeOn boolean

* use session in chatHistory

* revert chatHistory

* client side toolName validation

* remove chatHistory from processUnavailableToolUseMessage

* unused import

---------

Co-authored-by: Jacob Chung <[email protected]>

Author: laileni-aws

packages/core/src/amazonq/webview/ui/apps/baseConnector.ts

@@ -194,6 +194,15 @@ export abstract class BaseConnector {
         })
     }
 
+    onPromptInputOptionChange = (tabId: string, optionsValues: Record<string, string>): void => {
+        this.sendMessageToExtension({
+            command: 'prompt-input-option-change',
+            optionsValues,
+            tabType: this.getTabType(),
+            tabID: tabId,
+        })
+    }
+
     requestGenerativeAIAnswer = (tabID: string, messageId: string, payload: ChatPayload): Promise<any> => {
         /**
          * When a user presses "enter" send an event that indicates

packages/core/src/amazonq/webview/ui/commands.ts

@@ -25,6 +25,7 @@ type MessageCommand =
     | 'help'
     | 'chat-item-voted'
     | 'chat-item-feedback'
+    | 'prompt-input-option-change'
     | 'link-was-clicked'
     | 'onboarding-page-interaction'
     | 'source-link-click'

packages/core/src/amazonq/webview/ui/connector.ts

@@ -97,6 +97,7 @@ export interface ConnectorProps {
     onUpdateAuthentication: (featureDevEnabled: boolean, authenticatingTabIDs: string[]) => void
     onNewTab: (tabType: TabType) => void
     onFileActionClick: (tabID: string, messageId: string, filePath: string, actionName: string) => void
+    onPromptInputOptionChange: (tabId: string, optionsValues: Record<string, string>, eventId?: string) => void
     handleCommand: (chatPrompt: ChatPrompt, tabId: string) => void
     sendStaticMessages: (tabID: string, messages: ChatItem[]) => void
     onContextCommandDataReceived: (message: MynahUIDataModel['contextCommands']) => void
@@ -617,6 +618,15 @@ export class Connector {
         }
     }
 
+    onPromptInputOptionChange = (tabId: string, optionsValues: Record<string, string>): void => {
+        switch (this.tabsStorage.getTab(tabId)?.type) {
+            case 'unknown':
+            case 'cwc':
+                this.cwChatConnector.onPromptInputOptionChange(tabId, optionsValues)
+                break
+        }
+    }
+
     sendFeedback = (tabId: string, feedbackPayload: FeedbackPayload): void | undefined => {
         switch (this.tabsStorage.getTab(tabId)?.type) {
             case 'featuredev':

packages/core/src/amazonq/webview/ui/main.ts

@@ -239,6 +239,7 @@ export const createMynahUI = (
             }
         },
         onFileActionClick: (tabID: string, messageId: string, filePath: string, actionName: string): void => {},
+        onPromptInputOptionChange: (tabId: string, optionsValues: Record<string, string>, eventId?: string): void => {},
         onQuickHandlerCommand: (tabID: string, command?: string, eventId?: string) => {
             tabsStorage.updateTabLastCommand(tabID, command)
             if (command === 'aws.awsq.transform') {
@@ -940,6 +941,9 @@ export const createMynahUI = (
         onFileActionClick: async (tabID: string, messageId: string, filePath: string, actionName: string) => {
             connector.onFileActionClick(tabID, messageId, filePath, actionName)
         },
+        onPromptInputOptionChange: (tabId, optionsValues) => {
+            connector.onPromptInputOptionChange(tabId, optionsValues)
+        },
         onFileClick: connector.onFileClick,
         tabs: {
             'tab-1': {

packages/core/src/amazonq/webview/ui/tabs/generator.ts

@@ -72,6 +72,23 @@ export class TabDataGenerator {
                       },
                   ]
                 : [],
+            promptInputOptions: [
+                {
+                    type: 'toggle',
+                    id: 'prompt-type',
+                    value: 'ask',
+                    options: [
+                        {
+                            value: 'pair-programming-on',
+                            icon: 'code-block', // TODO: correct icons
+                        },
+                        {
+                            value: 'pair-programming-off',
+                            icon: 'chat', // TODO: correct icons
+                        },
+                    ],
+                },
+            ],
         }
         return tabData
     }

packages/core/src/codewhispererChat/app.ts

@@ -28,6 +28,7 @@ import {
     AcceptDiff,
     QuickCommandGroupActionClick,
     FileClick,
+    PromptInputOptionChange,
 } from './controllers/chat/model'
 import { EditorContextCommand, registerCommands } from './commands/registerCommands'
 import { ContextSelectedMessage, CustomFormActionMessage } from './view/connector/connector'
@@ -56,6 +57,7 @@ export function init(appContext: AmazonQAppInitContext) {
         processCustomFormAction: new EventEmitter<CustomFormActionMessage>(),
         processContextSelected: new EventEmitter<ContextSelectedMessage>(),
         processFileClick: new EventEmitter<FileClick>(),
+        processPromptInputOptionChange: new EventEmitter<PromptInputOptionChange>(),
     }
 
     const cwChatControllerMessageListeners = {
@@ -117,6 +119,9 @@ export function init(appContext: AmazonQAppInitContext) {
             cwChatControllerEventEmitters.processContextSelected
         ),
         processFileClick: new MessageListener<FileClick>(cwChatControllerEventEmitters.processFileClick),
+        processPromptInputOptionChange: new MessageListener<PromptInputOptionChange>(
+            cwChatControllerEventEmitters.processPromptInputOptionChange
+        ),
     }
 
     const cwChatControllerMessagePublishers = {
@@ -180,6 +185,9 @@ export function init(appContext: AmazonQAppInitContext) {
             cwChatControllerEventEmitters.processContextSelected
         ),
         processFileClick: new MessagePublisher<FileClick>(cwChatControllerEventEmitters.processFileClick),
+        processPromptInputOptionChange: new MessagePublisher<PromptInputOptionChange>(
+            cwChatControllerEventEmitters.processPromptInputOptionChange
+        ),
     }
 
     new CwChatController(

packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

@@ -28,6 +28,7 @@ export class ChatSession {
     private _toolUse: ToolUse | undefined
     private _showDiffOnFileWrite: boolean = false
     private _context: PromptMessage['context']
+    private _pairProgrammingModeOn: boolean = true
     private _messageIdToUpdate: string | undefined
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
@@ -38,6 +39,14 @@ export class ChatSession {
         return this.sessionId
     }
 
+    public get pairProgrammingModeOn(): boolean {
+        return this._pairProgrammingModeOn
+    }
+
+    public setPairProgrammingModeOn(pairProgrammingModeOn: boolean) {
+        this._pairProgrammingModeOn = pairProgrammingModeOn
+    }
+
     public get toolUse(): ToolUse | undefined {
         return this._toolUse
     }

packages/core/src/codewhispererChat/constants.ts

@@ -30,6 +30,11 @@ export const tools: Tool[] = Object.entries(toolsJson).map(([, toolSpec]) => ({
         inputSchema: { json: toolSpec.inputSchema },
     },
 }))
+
+export const noWriteTools: Tool[] = tools.filter(
+    (tool) => !['fsWrite', 'executeBash'].includes(tool.toolSpecification?.name || '')
+)
+
 export const defaultContextLengths: ContextLengths = {
     additionalContextLengths: {
         fileContextLength: 0,

packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts

@@ -6,9 +6,9 @@
 import { ConversationState, CursorState, DocumentSymbol, SymbolType, TextDocument } from '@amzn/codewhisperer-streaming'
 import { AdditionalContentEntryAddition, ChatTriggerType, RelevantTextDocumentAddition, TriggerPayload } from '../model'
 import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
-import { tools } from '../../../constants'
 import { getLogger } from '../../../../shared/logger/logger'
 import vscode from 'vscode'
+import { noWriteTools, tools } from '../../../constants'
 
 const fqnNameSizeDownLimit = 1
 const fqnNameSizeUpLimit = 256
@@ -164,7 +164,7 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
                             workspaceFolders: vscode.workspace.workspaceFolders?.map(({ uri }) => uri.fsPath) ?? [],
                         },
                         additionalContext: triggerPayload.additionalContents,
-                        tools,
+                        tools: triggerPayload.pairProgrammingModeOn ? tools : noWriteTools,
                         ...(triggerPayload.toolResults !== undefined &&
                             triggerPayload.toolResults !== null && { toolResults: triggerPayload.toolResults }),
                     },

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -32,6 +32,7 @@ import {
     DocumentReference,
     FileClick,
     RelevantTextDocumentAddition,
+    PromptInputOptionChange,
 } from './model'
 import {
     AppToWebViewMessageDispatcher,
@@ -117,6 +118,7 @@ export interface ChatControllerMessagePublishers {
     readonly processCustomFormAction: MessagePublisher<CustomFormActionMessage>
     readonly processContextSelected: MessagePublisher<ContextSelectedMessage>
     readonly processFileClick: MessagePublisher<FileClick>
+    readonly processPromptInputOptionChange: MessagePublisher<PromptInputOptionChange>
 }
 
 export interface ChatControllerMessageListeners {
@@ -142,6 +144,7 @@ export interface ChatControllerMessageListeners {
     readonly processCustomFormAction: MessageListener<CustomFormActionMessage>
     readonly processContextSelected: MessageListener<ContextSelectedMessage>
     readonly processFileClick: MessageListener<FileClick>
+    readonly processPromptInputOptionChange: MessageListener<PromptInputOptionChange>
 }
 
 export class ChatController {
@@ -277,6 +280,9 @@ export class ChatController {
         this.chatControllerMessageListeners.processFileClick.onMessage((data) => {
             return this.processFileClickMessage(data)
         })
+        this.chatControllerMessageListeners.processPromptInputOptionChange.onMessage((data) => {
+            return this.processPromptInputOptionChange(data)
+        })
     }
 
     private registerUserPromptsWatcher() {
@@ -631,6 +637,69 @@ export class ChatController {
         telemetry.ui_click.emit({ elementId: 'amazonq_createSavedPrompt' })
     }
 
+    private async processUnavailableToolUseMessage(message: CustomFormActionMessage) {
+        const tabID = message.tabID
+        if (!tabID) {
+            return
+        }
+        this.editorContextExtractor
+            .extractContextForTrigger('ChatMessage')
+            .then(async (context) => {
+                const triggerID = randomUUID()
+                this.triggerEventsStorage.addTriggerEvent({
+                    id: triggerID,
+                    tabID: message.tabID,
+                    message: undefined,
+                    type: 'chat_message',
+                    context,
+                })
+                const session = this.sessionStorage.getSession(tabID)
+                const toolUse = session.toolUse
+                if (!toolUse || !toolUse.input) {
+                    return
+                }
+                session.setToolUse(undefined)
+
+                const toolResults: ToolResult[] = []
+
+                toolResults.push({
+                    content: [{ text: 'This tool is not an available tool in this mode' }],
+                    toolUseId: toolUse.toolUseId,
+                    status: ToolResultStatus.ERROR,
+                })
+
+                await this.generateResponse(
+                    {
+                        message: '',
+                        trigger: ChatTriggerType.ChatMessage,
+                        query: undefined,
+                        codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
+                        fileText: context?.focusAreaContext?.extendedCodeBlock ?? '',
+                        fileLanguage: context?.activeFileContext?.fileLanguage,
+                        filePath: context?.activeFileContext?.filePath,
+                        matchPolicy: context?.activeFileContext?.matchPolicy,
+                        codeQuery: context?.focusAreaContext?.names,
+                        userIntent: undefined,
+                        customization: getSelectedCustomization(),
+                        toolResults: toolResults,
+                        origin: Origin.IDE,
+                        context: session.context ?? [],
+                        relevantTextDocuments: [],
+                        additionalContents: [],
+                        documentReferences: [],
+                        useRelevantDocuments: false,
+                        contextLengths: {
+                            ...defaultContextLengths,
+                        },
+                    },
+                    triggerID
+                )
+            })
+            .catch((e) => {
+                this.processException(e, tabID)
+            })
+    }
+
     private async processToolUseMessage(message: CustomFormActionMessage) {
         const tabID = message.tabID
         if (!tabID) {
@@ -763,6 +832,9 @@ export class ChatController {
             case 'reject-code-diff':
                 await this.closeDiffView()
                 break
+            case 'tool-unavailable':
+                await this.processUnavailableToolUseMessage(message)
+                break
             default:
                 getLogger().warn(`Unhandled action: ${message.action.id}`)
         }
@@ -773,6 +845,18 @@ export class ChatController {
             this.handlePromptCreate(message.tabID)
         }
     }
+
+    private async processPromptInputOptionChange(message: PromptInputOptionChange) {
+        const session = this.sessionStorage.getSession(message.tabID)
+        const promptTypeValue = message.optionsValues['prompt-type']
+        // TODO: display message: You turned off pair programmer mode. Q will not include code diffs or run commands in the chat.
+        if (promptTypeValue === 'pair-programming-on') {
+            session.setPairProgrammingModeOn(true)
+        } else {
+            session.setPairProgrammingModeOn(false)
+        }
+    }
+
     private async processFileClickMessage(message: FileClick) {
         const session = this.sessionStorage.getSession(message.tabID)
         // Check if user clicked on filePath in the contextList or in the fileListTree and perform the functionality accordingly.
@@ -1353,6 +1437,7 @@ export class ChatController {
 
         triggerPayload.contextLengths.userInputContextLength = triggerPayload.message.length
         triggerPayload.contextLengths.focusFileContextLength = triggerPayload.fileText.length
+        triggerPayload.pairProgrammingModeOn = session.pairProgrammingModeOn
 
         const request = triggerPayloadToChatRequest(triggerPayload)