From ef7da204085b94b3cd35af880c083a0f233c8204 Mon Sep 17 00:00:00 2001 From: Tyrone Smith Date: Tue, 25 Mar 2025 02:04:16 -0700 Subject: [PATCH] feat(chat): Add executeBash tool for Amazon Q Agentic Chat --- .../codewhispererChat/tools/executeBash.ts | 208 ++++++++++++++++++ .../src/codewhispererChat/tools/fsRead.ts | 8 +- packages/core/src/shared/logger/logger.ts | 1 + .../tools/executeBash.test.ts | 112 ++++++++++ 4 files changed, 325 insertions(+), 4 deletions(-) create mode 100644 packages/core/src/codewhispererChat/tools/executeBash.ts create mode 100644 packages/core/src/test/codewhispererChat/tools/executeBash.test.ts diff --git a/packages/core/src/codewhispererChat/tools/executeBash.ts b/packages/core/src/codewhispererChat/tools/executeBash.ts new file mode 100644 index 00000000000..ca164c90e8b --- /dev/null +++ b/packages/core/src/codewhispererChat/tools/executeBash.ts @@ -0,0 +1,208 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Writable } from 'stream' +import { getLogger } from '../../shared/logger/logger' +import { fs } from '../../shared/fs/fs' // e.g. for getUserHomeDir() +import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils' +import { InvokeOutput, OutputKind, sanitizePath } from './toolShared' + +export const readOnlyCommands: string[] = ['ls', 'cat', 'echo', 'pwd', 'which', 'head', 'tail'] +export const maxBashToolResponseSize: number = 1024 * 1024 // 1MB +export const lineCount: number = 1024 +export const dangerousPatterns: string[] = ['|', '<(', '$(', '`', '>', '&&', '||'] + +export interface ExecuteBashParams { + command: string + cwd?: string +} + +export class ExecuteBash { + private readonly command: string + private readonly workingDirectory?: string + private readonly logger = getLogger('executeBash') + + constructor(params: ExecuteBashParams) { + this.command = params.command + this.workingDirectory = params.cwd ? sanitizePath(params.cwd) : fs.getUserHomeDir() + } + + public async validate(): Promise { + if (!this.command.trim()) { + throw new Error('Bash command cannot be empty.') + } + + const args = ExecuteBash.parseCommand(this.command) + if (!args || args.length === 0) { + throw new Error('No command found.') + } + + try { + await ExecuteBash.whichCommand(args[0]) + } catch { + throw new Error(`Command "${args[0]}" not found on PATH.`) + } + } + + public requiresAcceptance(): boolean { + try { + const args = ExecuteBash.parseCommand(this.command) + if (!args || args.length === 0) { + return true + } + + if (args.some((arg) => dangerousPatterns.some((pattern) => arg.includes(pattern)))) { + return true + } + + const command = args[0] + return !readOnlyCommands.includes(command) + } catch (error) { + this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`) + return true + } + } + + public async invoke(updates: Writable): Promise { + this.logger.info(`Invoking bash command: "${this.command}" in cwd: "${this.workingDirectory}"`) + + return new Promise(async (resolve, reject) => { + this.logger.debug(`Spawning process with command: bash -c "${this.command}" (cwd=${this.workingDirectory})`) + + const stdoutBuffer: string[] = [] + const stderrBuffer: string[] = [] + + const childProcessOptions: ChildProcessOptions = { + spawnOptions: { + cwd: this.workingDirectory, + stdio: ['pipe', 'pipe', 'pipe'], + }, + collect: false, + waitForStreams: true, + onStdout: (chunk: string) => { + ExecuteBash.handleChunk(chunk, stdoutBuffer, updates) + }, + onStderr: (chunk: string) => { + ExecuteBash.handleChunk(chunk, stderrBuffer, updates) + }, + } + + const childProcess = new ChildProcess('bash', ['-c', this.command], childProcessOptions) + + try { + const result = await childProcess.run() + const exitStatus = result.exitCode ?? 0 + const stdout = stdoutBuffer.join('\n') + const stderr = stderrBuffer.join('\n') + const [stdoutTrunc, stdoutSuffix] = ExecuteBash.truncateSafelyWithSuffix( + stdout, + maxBashToolResponseSize / 3 + ) + const [stderrTrunc, stderrSuffix] = ExecuteBash.truncateSafelyWithSuffix( + stderr, + maxBashToolResponseSize / 3 + ) + + const outputJson = { + exitStatus: exitStatus.toString(), + stdout: stdoutTrunc + (stdoutSuffix ? ' ... truncated' : ''), + stderr: stderrTrunc + (stderrSuffix ? ' ... truncated' : ''), + } + + return { + output: { + kind: OutputKind.Json, + content: outputJson, + }, + } + } catch (err: any) { + this.logger.error(`Failed to execute bash command '${this.command}': ${err.message}`) + throw new Error(`Failed to execute command: ${err.message}`) + } + }) + } + + private static handleChunk(chunk: string, buffer: string[], updates: Writable) { + const lines = chunk.split(/\r?\n/) + for (const line of lines) { + updates.write(`${line}\n`) + buffer.push(line) + if (buffer.length > lineCount) { + buffer.shift() + } + } + } + + private static truncateSafelyWithSuffix(str: string, maxLength: number): [string, boolean] { + if (str.length > maxLength) { + return [str.substring(0, maxLength), true] + } + return [str, false] + } + + private static async whichCommand(cmd: string): Promise { + const cp = new ChildProcess('which', [cmd], { + collect: true, + waitForStreams: true, + }) + const result = await cp.run() + + if (result.exitCode !== 0) { + throw new Error(`Command "${cmd}" not found on PATH.`) + } + + const output = result.stdout.trim() + if (!output) { + throw new Error(`Command "${cmd}" found but 'which' returned empty output.`) + } + return output + } + + private static parseCommand(command: string): string[] | undefined { + const result: string[] = [] + let current = '' + let inQuote: string | undefined + let escaped = false + + for (const char of command) { + if (escaped) { + current += char + escaped = false + } else if (char === '\\') { + escaped = true + } else if (inQuote) { + if (char === inQuote) { + inQuote = undefined + } else { + current += char + } + } else if (char === '"' || char === "'") { + inQuote = char + } else if (char === ' ' || char === '\t') { + if (current) { + result.push(current) + current = '' + } + } else { + current += char + } + } + + if (current) { + result.push(current) + } + + return result + } + + public queueDescription(updates: Writable): void { + updates.write(`I will run the following shell command: `) + + if (this.command.length > 20) { + updates.write('\n') + } + updates.write(`\x1b[32m${this.command}\x1b[0m\n`) + } +} diff --git a/packages/core/src/codewhispererChat/tools/fsRead.ts b/packages/core/src/codewhispererChat/tools/fsRead.ts index 166145d8e1e..bac2715e40e 100644 --- a/packages/core/src/codewhispererChat/tools/fsRead.ts +++ b/packages/core/src/codewhispererChat/tools/fsRead.ts @@ -16,7 +16,7 @@ export interface FsReadParams { export class FsRead { private fsPath: string private readonly readRange?: number[] - private type?: boolean // true for file, false for directory + private isFile?: boolean // true for file, false for directory private readonly logger = getLogger('fsRead') constructor(params: FsReadParams) { @@ -44,7 +44,7 @@ export class FsRead { throw new Error(`Path: "${this.fsPath}" does not exist or cannot be accessed. (${err})`) } - this.type = await fs.existsFile(fileUri) + this.isFile = await fs.existsFile(fileUri) this.logger.debug(`Validation succeeded for path: ${this.fsPath}`) } @@ -52,11 +52,11 @@ export class FsRead { try { const fileUri = vscode.Uri.file(this.fsPath) - if (this.type) { + if (this.isFile) { const fileContents = await this.readFile(fileUri) this.logger.info(`Read file: ${this.fsPath}, size: ${fileContents.length}`) return this.handleFileRange(fileContents) - } else if (!this.type) { + } else if (!this.isFile) { const maxDepth = this.getDirectoryDepth() ?? 0 const listing = await readDirectoryRecursively(fileUri, maxDepth) return this.createOutput(listing.join('\n')) diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts index a4ed0ecf1e4..e2f16cd4daa 100644 --- a/packages/core/src/shared/logger/logger.ts +++ b/packages/core/src/shared/logger/logger.ts @@ -16,6 +16,7 @@ export type LogTopic = | 'stepfunctions' | 'fsRead' | 'fsWrite' + | 'executeBash' class ErrorLog { constructor( diff --git a/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts b/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts new file mode 100644 index 00000000000..68ec48d8851 --- /dev/null +++ b/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts @@ -0,0 +1,112 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { strict as assert } from 'assert' +import sinon from 'sinon' +import { ExecuteBash } from '../../../codewhispererChat/tools/executeBash' +import { ChildProcess } from '../../../shared/utilities/processUtils' + +describe('ExecuteBash Tool', () => { + let runStub: sinon.SinonStub + let invokeStub: sinon.SinonStub + + beforeEach(() => { + runStub = sinon.stub(ChildProcess.prototype, 'run') + invokeStub = sinon.stub(ExecuteBash.prototype, 'invoke') + }) + + afterEach(() => { + sinon.restore() + }) + + it('pass validation for a safe command (read-only)', async () => { + runStub.resolves({ + exitCode: 0, + stdout: '/bin/ls', + stderr: '', + error: undefined, + signal: undefined, + }) + const execBash = new ExecuteBash({ command: 'ls' }) + await execBash.validate() + }) + + it('fail validation if the command is empty', async () => { + const execBash = new ExecuteBash({ command: ' ' }) + await assert.rejects( + execBash.validate(), + /Bash command cannot be empty/i, + 'Expected an error for empty command' + ) + }) + + it('set requiresAcceptance=true if the command has dangerous patterns', () => { + const execBash = new ExecuteBash({ command: 'ls && rm -rf /' }) + const needsAcceptance = execBash.requiresAcceptance() + assert.equal(needsAcceptance, true, 'Should require acceptance for dangerous pattern') + }) + + it('set requiresAcceptance=false if it is a read-only command', () => { + const execBash = new ExecuteBash({ command: 'cat file.txt' }) + const needsAcceptance = execBash.requiresAcceptance() + assert.equal(needsAcceptance, false, 'Read-only command should not require acceptance') + }) + + it('whichCommand cannot find the first arg', async () => { + runStub.resolves({ + exitCode: 1, + stdout: '', + stderr: '', + error: undefined, + signal: undefined, + }) + + const execBash = new ExecuteBash({ command: 'noSuchCmd' }) + await assert.rejects(execBash.validate(), /not found on PATH/i, 'Expected not found error from whichCommand') + }) + + it('whichCommand sees first arg on PATH', async () => { + runStub.resolves({ + exitCode: 0, + stdout: '/usr/bin/noSuchCmd\n', + stderr: '', + error: undefined, + signal: undefined, + }) + + const execBash = new ExecuteBash({ command: 'noSuchCmd' }) + await execBash.validate() + }) + + it('stub invoke() call', async () => { + invokeStub.resolves({ + output: { + kind: 'json', + content: { + exitStatus: '0', + stdout: 'mocked stdout lines', + stderr: '', + }, + }, + }) + + const execBash = new ExecuteBash({ command: 'ls' }) + + const dummyWritable = { write: () => {} } as any + const result = await execBash.invoke(dummyWritable) + + assert.strictEqual(result.output.kind, 'json') + const out = result.output.content as unknown as { + exitStatus: string + stdout: string + stderr: string + } + assert.strictEqual(out.exitStatus, '0') + assert.strictEqual(out.stdout, 'mocked stdout lines') + assert.strictEqual(out.stderr, '') + + assert.strictEqual(invokeStub.callCount, 1) + }) +})