diff --git a/chat-client/src/client/mynahUi.ts b/chat-client/src/client/mynahUi.ts index 11c656d431..69646c0a8e 100644 --- a/chat-client/src/client/mynahUi.ts +++ b/chat-client/src/client/mynahUi.ts @@ -522,8 +522,8 @@ export const createMynahUi = ( buttons: toMynahButtons(am.buttons), // file diffs in the header need space - fullWidth: am.type === 'tool' && am.header?.fileList ? true : undefined, - padding: am.type === 'tool' && am.header?.fileList ? false : undefined, + fullWidth: am.type === 'tool' ? true : undefined, + padding: am.type === 'tool' ? false : undefined, } if (!chatItems.find(ci => ci.messageId === am.messageId)) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts index 8e98afc945..41622e50f6 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts @@ -16,6 +16,7 @@ import { ToolUse, } from '@amzn/codewhisperer-streaming' import { + Button, ButtonClickParams, ButtonClickResult, ChatMessage, @@ -97,7 +98,9 @@ import { LocalProjectContextController } from '../../shared/localProjectContextC import { workspaceUtils } from '@aws/lsp-core' import { FsReadParams } from './tools/fsRead' import { ListDirectoryParams } from './tools/listDirectory' -import { FsWriteParams, getDiffChanges } from './tools/fsWrite' +import { FsWrite, FsWriteParams, getDiffChanges } from './tools/fsWrite' +import { ExecuteBash, ExecuteBashOutput, ExecuteBashParams } from './tools/executeBash' +import { InvokeOutput } from './tools/toolShared' type ChatHandlers = Omit< LspHandlers, @@ -147,10 +150,28 @@ export class AgenticChatController implements ChatHandlers { } async onButtonClick(params: ButtonClickParams): Promise { - this.#log(`onButtonClick event with params: ${JSON.stringify(params)}`) - return { - success: false, - failureReason: 'not implemented', + if (params.buttonId === 'run-shell-command' || params.buttonId === 'reject-shell-command') { + const session = this.#chatSessionManagementService.getSession(params.tabId) + if (!session.data) { + return { success: false, failureReason: `could not find chat session for tab: ${params.tabId} ` } + } + const handler = session.data.getDeferredToolExecution(params.messageId) + if (!handler?.reject || !handler.resolve) { + return { + success: false, + failureReason: `could not find deferred tool execution for message: ${params.messageId} `, + } + } + params.buttonId === 'reject-shell-command' ? handler.reject() : handler.resolve() + return { + success: true, + } + } else { + this.#log(`onButtonClick event with params: ${JSON.stringify(params)}`) + return { + success: false, + failureReason: 'not implemented', + } } } @@ -361,7 +382,7 @@ export class AgenticChatController implements ChatHandlers { const currentMessage = currentRequestInput.conversationState?.currentMessage // Process tool uses and update the request input for the next iteration - const toolResults = await this.#processToolUses(pendingToolUses, chatResultStream) + const toolResults = await this.#processToolUses(pendingToolUses, chatResultStream, session) currentRequestInput = this.#updateRequestInputWithToolResults(currentRequestInput, toolResults) if (!currentRequestInput.conversationState!.history) { @@ -414,13 +435,15 @@ export class AgenticChatController implements ChatHandlers { */ async #processToolUses( toolUses: Array, - chatResultStream: AgenticChatResultStream + chatResultStream: AgenticChatResultStream, + session: ChatSessionService ): Promise { const results: ToolResult[] = [] for (const toolUse of toolUses) { if (!toolUse.name || !toolUse.toolUseId) continue this.#triggerContext.getToolUseLookup().set(toolUse.toolUseId, toolUse) + let needsConfirmation try { switch (toolUse.name) { @@ -439,6 +462,15 @@ export class AgenticChatController implements ChatHandlers { .set(toolUse.toolUseId, { ...toolUse, oldContent: document?.getText() }) break case 'executeBash': + const bashTool = new ExecuteBash(this.#features) + const { requiresAcceptance, warning } = await bashTool.requiresAcceptance( + toolUse.input as unknown as ExecuteBashParams + ) + if (requiresAcceptance) { + needsConfirmation = true + const confirmationResult = this.#processExecuteBashConfirmation(toolUse, warning) + await chatResultStream.writeResultBlock(confirmationResult) + } break default: await chatResultStream.writeResultBlock({ @@ -449,6 +481,18 @@ export class AgenticChatController implements ChatHandlers { break } + if (needsConfirmation) { + const deferred = this.#createDeferred() + session.setDeferredToolExecution(toolUse.toolUseId, deferred.resolve, deferred.reject) + + // the below line was commented out for now because + // the partial result block from above is not streamed to chat window yet at this point + // so the buttons are not in the window for the promise to be rejected/resolved + // this can to be brought back once intermediate messages are shown + + // await deferred.promise + } + const result = await this.#features.agent.runTool(toolUse.name, toolUse.input) let toolResultContent: ToolResultContentBlock @@ -475,6 +519,9 @@ export class AgenticChatController implements ChatHandlers { const chatResult = await this.#getFsWriteChatResult(toolUse) await chatResultStream.writeResultBlock(chatResult) break + case 'executeBash': + const bashToolResult = this.#getBashExecutionChatResult(toolUse, result) + await chatResultStream.writeResultBlock(bashToolResult) default: await chatResultStream.writeResultBlock({ type: 'tool', @@ -500,6 +547,45 @@ export class AgenticChatController implements ChatHandlers { return results } + #getBashExecutionChatResult(toolUse: ToolUse, result: InvokeOutput): ChatResult { + const outputString = result.output.success + ? (result.output.content as ExecuteBashOutput).stdout + : (result.output.content as ExecuteBashOutput).stderr + return { + type: 'tool', + messageId: toolUse.toolUseId, + body: outputString, + } + } + + #processExecuteBashConfirmation(toolUse: ToolUse, warning?: string): ChatResult { + const buttons: Button[] = [ + { + id: 'reject-shell-command', + text: 'Reject', + icon: 'cancel', + }, + { + id: 'run-shell-command', + text: 'Run', + icon: 'play', + }, + ] + const header = { + body: 'shell', + buttons, + } + + const commandString = (toolUse.input as unknown as ExecuteBashParams).command + const body = '```shell\n' + commandString + '\n```' + return { + type: 'tool', + messageId: toolUse.toolUseId, + header, + body: warning ? warning + body : body, + } + } + async #getFsWriteChatResult(toolUse: ToolUse): Promise { const input = toolUse.input as unknown as FsWriteParams const oldContent = this.#triggerContext.getToolUseLookup().get(toolUse.toolUseId!)?.oldContent ?? '' @@ -1135,6 +1221,16 @@ export class AgenticChatController implements ChatHandlers { return tools } + #createDeferred() { + let resolve + let reject + const promise = new Promise((res, rej) => { + resolve = res + reject = rej + }) + return { promise, resolve, reject } + } + #log(...messages: string[]) { this.#features.logging.log(messages.join(' ')) } diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.ts index 7b5409dc65..3a0001cccf 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.ts @@ -121,6 +121,12 @@ interface TimestampedChunk { isFirst: boolean } +export interface ExecuteBashOutput { + exitStatus: string + stdout: string + stderr: string +} + export class ExecuteBash { private childProcess?: ChildProcess private readonly logging: Features['logging'] @@ -356,7 +362,7 @@ export class ExecuteBash { maxBashToolResponseSize / 3 ) - const outputJson = { + const outputJson: ExecuteBashOutput = { exitStatus: exitStatus.toString(), stdout: stdoutTrunc + (stdoutSuffix ? ' ... truncated' : ''), stderr: stderrTrunc + (stderrSuffix ? ' ... truncated' : ''), diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts index 092e970819..0407bdbe0f 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts @@ -12,12 +12,18 @@ import { } from '../../shared/streamingClientService' export type ChatSessionServiceConfig = CodeWhispererStreamingClientConfig + +type DeferredHandler = { + resolve: () => void + reject: () => void +} export class ChatSessionService { public shareCodeWhispererContentWithAWS = false public pairProgrammingMode: boolean = true #abortController?: AbortController #conversationId?: string #amazonQServiceManager?: AmazonQBaseServiceManager + #deferredToolExecution: Record = {} public get conversationId(): string | undefined { return this.#conversationId @@ -27,6 +33,13 @@ export class ChatSessionService { this.#conversationId = value } + public getDeferredToolExecution(messageId: string): DeferredHandler | undefined { + return this.#deferredToolExecution[messageId] + } + public setDeferredToolExecution(messageId: string, resolve: any, reject: any) { + this.#deferredToolExecution[messageId] = { resolve, reject } + } + constructor(amazonQServiceManager?: AmazonQBaseServiceManager) { this.#amazonQServiceManager = amazonQServiceManager }