Merge pull request aws#6848 from tsmithsz/feature/agentic-chat

mr-lee · web-flow · commit 6269bce09339 · 2025-03-25T19:52:30.000-04:00
feat(chat): Add executeBash tool for Amazon Q Agentic Chat
diff --git a/packages/core/src/codewhispererChat/tools/executeBash.ts 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<void> {
+        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<InvokeOutput> {
+        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<string> {
+        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
@@ -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,19 +44,19 @@ 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}`)
     }
 
     public async invoke(): Promise<InvokeOutput> {
         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
@@ -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
@@ -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)
+    })
+})
