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

bywang56 · bywang56 · commit 641655190e2b · 2025-04-08T14:43:25.000-07:00
diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
@@ -314,10 +314,14 @@ export class Connector extends BaseConnector {
 
         if (
             !this.onChatAnswerUpdated ||
-            !(
-                ['accept-code-diff', 'run-shell-command', 'reject-shell-command'].includes(action.id) ||
-                action.id.startsWith('reject-code-diff')
-            )
+            ![
+                'accept-code-diff',
+                'reject-code-diff',
+                'run-shell-command',
+                'reject-shell-command',
+                'confirm-tool-use',
+                'reject-tool-use',
+            ].includes(action.id)
         ) {
             return
         }
@@ -373,6 +377,30 @@ export class Connector extends BaseConnector {
                     answer.header.buttons = []
                 }
                 break
+            case 'confirm-tool-use':
+                answer.buttons = [
+                    {
+                        keepCardAfterClick: true,
+                        text: 'Confirmed',
+                        id: 'confirmed-tool-use',
+                        status: 'success',
+                        position: 'outside',
+                        disabled: true,
+                    },
+                ]
+                break
+            case 'reject-tool-use':
+                answer.buttons = [
+                    {
+                        keepCardAfterClick: true,
+                        text: 'Rejected',
+                        id: 'rejected-tool-use',
+                        status: 'error',
+                        position: 'outside',
+                        disabled: true,
+                    },
+                ]
+                break
             default:
                 break
         }
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -849,9 +849,14 @@ export class ChatController {
                 await this.handleCreatePrompt(message)
                 break
             case 'run-shell-command':
+            case 'confirm-tool-use':
             case 'generic-tool-execution':
                 await this.processToolUseMessage(message)
                 break
+            case 'reject-code-diff':
+            case 'reject-tool-use':
+                await this.closeDiffView()
+                break
             case 'reject-shell-command':
                 await this.rejectShellCommand(message)
                 break
diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -482,6 +482,17 @@ export class Messenger {
             if (validation.warning) {
                 message = validation.warning + message + '\nRun the command to proceed.\n'
             }
+        } else if (validation.requiresAcceptance) {
+            buttons.push({
+                id: 'reject-tool-use',
+                text: 'Reject',
+                status: 'info',
+            })
+            buttons.push({
+                id: 'confirm-tool-use',
+                text: 'Confirm',
+                status: 'info',
+            })
         } else if (toolUse?.name === ToolType.FsWrite) {
             const input = toolUse.input as unknown as FsWriteParams
             const fileName = path.basename(input.path)
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,
@@ -215,6 +218,27 @@ export class ExecuteBash {
                         return { requiresAcceptance: true, warning: highRiskCommandWarningMessage }
                 }
             }
+            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}`)
@@ -330,4 +354,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, maxToolResponseSize, OutputKind, sanitizePath } from './toolShared'
+import { InvokeOutput, maxToolResponseSize, 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
@@ -33,3 +33,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')
@@ -77,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
@@ -6,6 +6,8 @@ import assert from 'assert'
 import { ListDirectory } from '../../../codewhispererChat/tools/listDirectory'
 import { TestFolder } from '../../testUtil'
 import path from 'path'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
 
 describe('ListDirectory Tool', () => {
     let testFolder: TestFolder
@@ -14,6 +16,10 @@ describe('ListDirectory Tool', () => {
         testFolder = await TestFolder.create()
     })
 
+    afterEach(() => {
+        sinon.restore()
+    })
+
     it('throws if path is empty', async () => {
         const listDirectory = new ListDirectory({ path: '', maxDepth: 0 })
         await assert.rejects(listDirectory.validate(), /Path cannot be empty/i, 'Expected an error about empty path')
@@ -86,4 +92,32 @@ describe('ListDirectory Tool', () => {
         assert.strictEqual(result.output.kind, 'text')
         assert.ok(result.output.content.length > 0)
     })
+
+    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 listDir = new ListDirectory({ path: '/not/in/workspace/dir', maxDepth: 0 })
+        const result = listDir.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 listDir = new ListDirectory({ path: '/workspace/folder/mydir', maxDepth: 0 })
+        const result = listDir.requiresAcceptance()
+        assert.equal(
+            result.requiresAcceptance,
+            false,
+            'Expected requiresAcceptance to be false for a path inside the workspace'
+        )
+        workspaceStub.restore()
+    })
 })

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