feat(chat): Add validation to fileRead, ListDir and ExecBash tool

bywang56 · bywang56 · commit 77c02ac1b570 · 2025-04-09T17:04:05.000-07:00
diff --git a/packages/core/package.nls.json b/packages/core/package.nls.json
@@ -456,6 +456,10 @@
     "AWS.amazonq.opensettings:": "Open settings",
     "AWS.amazonq.executeBash.run": "Run",
     "AWS.amazonq.executeBash.reject": "Reject",
+    "AWS.amazonq.fsRead.run": "Run",
+    "AWS.amazonq.fsRead.reject": "Reject",
+    "AWS.amazonq.listDirectory.run": "Run",
+    "AWS.amazonq.listDirectory.reject": "Reject",
     "AWS.toolkit.lambda.walkthrough.quickpickTitle": "Application Builder Walkthrough",
     "AWS.toolkit.lambda.walkthrough.title": "Get started building your application",
     "AWS.toolkit.lambda.walkthrough.description": "Your quick guide to build an application visually, iterate locally, and deploy to the cloud!",
diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
@@ -321,7 +321,14 @@ export class Connector extends BaseConnector {
 
         if (
             !this.onChatAnswerUpdated ||
-            !['accept-code-diff', 'reject-code-diff', 'run-shell-command', 'reject-shell-command'].includes(action.id)
+            ![
+                'accept-code-diff',
+                'reject-code-diff',
+                'run-shell-command',
+                'reject-shell-command',
+                'confirm-tool-use',
+                'reject-tool-use',
+            ].includes(action.id)
         ) {
             return
         }
@@ -381,6 +388,28 @@ export class Connector extends BaseConnector {
                     },
                 }
                 break
+            case 'confirm-tool-use':
+                answer.header = {
+                    icon: 'shell' as MynahIconsType,
+                    body: 'shell',
+                    status: {
+                        icon: 'ok' as MynahIconsType,
+                        text: 'Accepted',
+                        status: 'success',
+                    },
+                }
+                break
+            case 'reject-tool-use':
+                answer.header = {
+                    icon: 'shell' as MynahIconsType,
+                    body: 'shell',
+                    status: {
+                        icon: 'cancel' as MynahIconsType,
+                        text: 'Rejected',
+                        status: 'error',
+                    },
+                }
+                break
             default:
                 break
         }
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -813,7 +813,12 @@ export class ChatController {
 
         const session = this.sessionStorage.getSession(message.tabID!)
         const currentToolUse = session.toolUseWithError?.toolUse
-        if (currentToolUse && currentToolUse.name === ToolType.ExecuteBash) {
+        if (
+            currentToolUse &&
+            (currentToolUse.name === ToolType.ExecuteBash ||
+                currentToolUse.name === ToolType.FsRead ||
+                currentToolUse.name === ToolType.ListDirectory)
+        ) {
             session.toolUseWithError.error = new Error('Tool use was rejected by the user.')
         } else {
             getLogger().error(
@@ -829,6 +834,7 @@ export class ChatController {
                 break
             case 'run-shell-command':
             case 'generic-tool-execution':
+            case 'confirm-tool-use':
                 await this.processToolUseMessage(message)
                 break
             case 'accept-code-diff':
@@ -839,6 +845,7 @@ export class ChatController {
                 await this.closeDiffView(message)
                 break
             case 'reject-shell-command':
+            case 'reject-tool-use':
                 await this.rejectShellCommand(message)
                 await this.processToolUseMessage(message)
                 break
diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -596,6 +596,37 @@ export class Messenger {
             }
             fullWidth = true
             padding = false
+            // eslint-disable-next-line unicorn/no-null
+            codeBlockActions = { 'insert-to-cursor': null, copy: null }
+        } else if (toolUse?.name === ToolType.ListDirectory || toolUse?.name === ToolType.FsRead) {
+            if (validation.requiresAcceptance) {
+                const buttons: ChatItemButton[] = [
+                    {
+                        id: 'confirm-tool-use',
+                        text: localize(`AWS.amazonq.${toolUse?.name}.run`, 'Run'),
+                        status: 'main',
+                        icon: 'play' as MynahIconsType,
+                    },
+                    {
+                        id: 'reject-tool-use',
+                        text: localize(`AWS.amazonq.${toolUse?.name}.reject`, 'Reject'),
+                        status: 'clear',
+                        icon: 'cancel' as MynahIconsType,
+                    },
+                ]
+                header = {
+                    icon: toolUse?.name === ToolType.ListDirectory ? 'folder' : ('file' as MynahIconsType),
+                    body: toolUse?.name === ToolType.ListDirectory ? 'folder' : 'file',
+                    buttons,
+                }
+            }
+            if (validation.warning) {
+                message = validation.warning + message
+            }
+            fullWidth = true
+            padding = false
+            // eslint-disable-next-line unicorn/no-null
+            codeBlockActions = { 'insert-to-cursor': null, copy: null }
         }
 
         this.dispatcher.sendChatMessage(
diff --git a/packages/core/src/codewhispererChat/tools/executeBash.ts b/packages/core/src/codewhispererChat/tools/executeBash.ts
@@ -9,6 +9,9 @@ import { fs } from '../../shared/fs/fs'
 import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
 import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
 import { split } from 'shlex'
+import path from 'path'
+import * as vscode from 'vscode'
+import { isInDirectory } from '../../shared/filesystemUtilities'
 
 export enum CommandCategory {
     ReadOnly,
@@ -191,6 +194,27 @@ export class ExecuteBash {
                         return { requiresAcceptance: true }
                 }
             }
+            for (const cmdArgs of allCommands) {
+                for (const arg of cmdArgs) {
+                    if (this.looksLikePath(arg)) {
+                        // If not absolute, resolve using workingDirectory if available.
+                        let fullPath = arg
+                        if (!path.isAbsolute(arg) && this.workingDirectory) {
+                            fullPath = path.join(this.workingDirectory, arg)
+                        }
+                        const workspaceFolders = vscode.workspace.workspaceFolders
+                        if (!workspaceFolders || workspaceFolders.length === 0) {
+                            return { requiresAcceptance: true }
+                        }
+                        const isInWorkspace = workspaceFolders.some((folder) =>
+                            isInDirectory(folder.uri.fsPath, fullPath)
+                        )
+                        if (!isInWorkspace) {
+                            return { requiresAcceptance: true }
+                        }
+                    }
+                }
+            }
             return { requiresAcceptance: false }
         } catch (error) {
             this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
@@ -306,4 +330,8 @@ export class ExecuteBash {
         updates.write('```shell\n' + this.command + '\n```')
         updates.end()
     }
+
+    private looksLikePath(arg: string): boolean {
+        return arg.startsWith('/') || arg.startsWith('./') || arg.startsWith('../')
+    }
 }
diff --git a/packages/core/src/codewhispererChat/tools/fsRead.ts b/packages/core/src/codewhispererChat/tools/fsRead.ts
@@ -5,7 +5,8 @@
 import * as vscode from 'vscode'
 import { getLogger } from '../../shared/logger/logger'
 import fs from '../../shared/fs/fs'
-import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
+import { isInDirectory } from '../../shared/filesystemUtilities'
 import { Writable } from 'stream'
 import path from 'path'
 
@@ -68,6 +69,18 @@ export class FsRead {
         updates.end()
     }
 
+    public requiresAcceptance(): CommandValidation {
+        const workspaceFolders = vscode.workspace.workspaceFolders
+        if (!workspaceFolders || workspaceFolders.length === 0) {
+            return { requiresAcceptance: true }
+        }
+        const isInWorkspace = workspaceFolders.some((folder) => isInDirectory(folder.uri.fsPath, this.fsPath))
+        if (!isInWorkspace) {
+            return { requiresAcceptance: true }
+        }
+        return { requiresAcceptance: false }
+    }
+
     public async invoke(updates?: Writable): Promise<InvokeOutput> {
         try {
             const fileUri = vscode.Uri.file(this.fsPath)
diff --git a/packages/core/src/codewhispererChat/tools/listDirectory.ts b/packages/core/src/codewhispererChat/tools/listDirectory.ts
@@ -6,7 +6,8 @@ import * as vscode from 'vscode'
 import { getLogger } from '../../shared/logger/logger'
 import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
 import fs from '../../shared/fs/fs'
-import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
+import { isInDirectory } from '../../shared/filesystemUtilities'
 import { Writable } from 'stream'
 import path from 'path'
 
@@ -61,6 +62,18 @@ export class ListDirectory {
         updates.end()
     }
 
+    public requiresAcceptance(): CommandValidation {
+        const workspaceFolders = vscode.workspace.workspaceFolders
+        if (!workspaceFolders || workspaceFolders.length === 0) {
+            return { requiresAcceptance: true }
+        }
+        const isInWorkspace = workspaceFolders.some((folder) => isInDirectory(folder.uri.fsPath, this.fsPath))
+        if (!isInWorkspace) {
+            return { requiresAcceptance: true }
+        }
+        return { requiresAcceptance: false }
+    }
+
     public async invoke(updates?: Writable): Promise<InvokeOutput> {
         try {
             const fileUri = vscode.Uri.file(this.fsPath)
diff --git a/packages/core/src/codewhispererChat/tools/toolShared.ts b/packages/core/src/codewhispererChat/tools/toolShared.ts
@@ -32,3 +32,8 @@ export function sanitizePath(inputPath: string): string {
     }
     return sanitized
 }
+
+export interface CommandValidation {
+    requiresAcceptance: boolean
+    warning?: string
+}
diff --git a/packages/core/src/codewhispererChat/tools/toolUtils.ts b/packages/core/src/codewhispererChat/tools/toolUtils.ts
@@ -40,13 +40,13 @@ export class ToolUtils {
     static requiresAcceptance(tool: Tool): CommandValidation {
         switch (tool.type) {
             case ToolType.FsRead:
-                return { requiresAcceptance: false }
+                return tool.tool.requiresAcceptance()
             case ToolType.FsWrite:
                 return { requiresAcceptance: false }
             case ToolType.ExecuteBash:
                 return tool.tool.requiresAcceptance()
             case ToolType.ListDirectory:
-                return { requiresAcceptance: false }
+                return tool.tool.requiresAcceptance()
         }
     }
 
diff --git a/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts b/packages/core/src/test/codewhispererChat/tools/executeBash.test.ts
@@ -7,6 +7,7 @@ import { strict as assert } from 'assert'
 import sinon from 'sinon'
 import { destructiveCommandWarningMessage, ExecuteBash } from '../../../codewhispererChat/tools/executeBash'
 import { ChildProcess } from '../../../shared/utilities/processUtils'
+import * as vscode from 'vscode'
 
 describe('ExecuteBash Tool', () => {
     let runStub: sinon.SinonStub
@@ -114,4 +115,55 @@ describe('ExecuteBash Tool', () => {
 
         assert.strictEqual(invokeStub.callCount, 1)
     })
+
+    it('requires acceptance if the command references an absolute file path outside the workspace', () => {
+        // Stub workspace folders to simulate a workspace at '/workspace/folder'
+        const workspaceStub = sinon
+            .stub(vscode.workspace, 'workspaceFolders')
+            .value([{ uri: { fsPath: '/workspace/folder' } } as any])
+        // Command references an absolute path outside the workspace
+        const execBash = new ExecuteBash({ command: 'cat /not/in/workspace/file.txt', cwd: '/workspace/folder' })
+        const result = execBash.requiresAcceptance()
+
+        assert.equal(
+            result.requiresAcceptance,
+            true,
+            'Should require acceptance for an absolute path outside of workspace'
+        )
+        workspaceStub.restore()
+    })
+
+    it('does NOT require acceptance if the command references a relative file path inside the workspace', () => {
+        // Stub workspace folders to simulate a workspace at '/workspace/folder'
+        const workspaceStub = sinon
+            .stub(vscode.workspace, 'workspaceFolders')
+            .value([{ uri: { fsPath: '/workspace/folder' } } as any])
+
+        // Command references a relative path that resolves within the workspace
+        const execBash = new ExecuteBash({ command: 'cat ./file.txt', cwd: '/workspace/folder' })
+        const result = execBash.requiresAcceptance()
+
+        assert.equal(result.requiresAcceptance, false, 'Relative path inside workspace should not require acceptance')
+
+        workspaceStub.restore()
+    })
+
+    it('does NOT require acceptance if there is no path-like token in the command', () => {
+        // Stub workspace folders (even though they are not used in this case)
+        const workspaceStub = sinon
+            .stub(vscode.workspace, 'workspaceFolders')
+            .value([{ uri: { fsPath: '/workspace/folder' } } as any])
+
+        // Command with tokens that do not look like file paths
+        const execBash = new ExecuteBash({ command: 'echo hello world', cwd: '/workspace/folder' })
+        const result = execBash.requiresAcceptance()
+
+        assert.equal(
+            result.requiresAcceptance,
+            false,
+            'A command without any path-like token should not require acceptance'
+        )
+
+        workspaceStub.restore()
+    })
 })
diff --git a/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts b/packages/core/src/test/codewhispererChat/tools/fsRead.test.ts
@@ -6,6 +6,8 @@ import assert from 'assert'
 import { FsRead } from '../../../codewhispererChat/tools/fsRead'
 import { TestFolder } from '../../testUtil'
 import path from 'path'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
 
 describe('FsRead Tool', () => {
     let testFolder: TestFolder
@@ -14,6 +16,10 @@ describe('FsRead Tool', () => {
         testFolder = await TestFolder.create()
     })
 
+    afterEach(() => {
+        sinon.restore()
+    })
+
     it('throws if path is empty', async () => {
         const fsRead = new FsRead({ path: '' })
         await assert.rejects(fsRead.validate(), /Path cannot be empty/i, 'Expected an error about empty path')
@@ -54,6 +60,20 @@ describe('FsRead Tool', () => {
         )
     })
 
+    it('throws error if content exceeds 30KB', async function () {
+        const bigContent = 'x'.repeat(35_000)
+        const bigFilePath = await testFolder.write('bigFile.txt', bigContent)
+
+        const fsRead = new FsRead({ path: bigFilePath })
+        await fsRead.validate()
+
+        await assert.rejects(
+            fsRead.invoke(process.stdout),
+            /This tool only supports reading \d+ bytes at a time/i,
+            'Expected a size-limit error'
+        )
+    })
+
     it('invalid line range', async () => {
         const filePath = await testFolder.write('rangeTest.txt', '1\n2\n3')
         const fsRead = new FsRead({ path: filePath, readRange: [3, 2] })
@@ -63,4 +83,32 @@ describe('FsRead Tool', () => {
         assert.strictEqual(result.output.kind, 'text')
         assert.strictEqual(result.output.content, '')
     })
+
+    it('should require acceptance if fsPath is outside the workspace', () => {
+        const workspaceStub = sinon
+            .stub(vscode.workspace, 'workspaceFolders')
+            .value([{ uri: { fsPath: '/workspace/folder' } } as any])
+        const fsRead = new FsRead({ path: '/not/in/workspace/file.txt' })
+        const result = fsRead.requiresAcceptance()
+        assert.equal(
+            result.requiresAcceptance,
+            true,
+            'Expected requiresAcceptance to be true for a path outside the workspace'
+        )
+        workspaceStub.restore()
+    })
+
+    it('should not require acceptance if fsPath is inside the workspace', () => {
+        const workspaceStub = sinon
+            .stub(vscode.workspace, 'workspaceFolders')
+            .value([{ uri: { fsPath: '/workspace/folder' } } as any])
+        const fsRead = new FsRead({ path: '/workspace/folder/file.txt' })
+        const result = fsRead.requiresAcceptance()
+        assert.equal(
+            result.requiresAcceptance,
+            false,
+            'Expected requiresAcceptance to be false for a path inside the workspace'
+        )
+        workspaceStub.restore()
+    })
 })
diff --git a/packages/core/src/test/codewhispererChat/tools/listDirectory.test.ts b/packages/core/src/test/codewhispererChat/tools/listDirectory.test.ts

Original file line number	Diff line number	Diff line change
`@@ -32,3 +32,8 @@ export function sanitizePath(inputPath: string): string {`
`32`	`32`	`}`
`33`	`33`	`return sanitized`
`34`	`34`	`}`
	`35`	`+`
	`36`	`+export interface CommandValidation {`
	`37`	`+ requiresAcceptance: boolean`
	`38`	`+ warning?: string`
	`39`	`+}`