Skip to content

Commit 921d2f7

Browse files
bywang56bywang56
committed
Add validation to fileRead, ListDir and ExecBash tool
1 parent c490994 commit 921d2f7

File tree

12 files changed

+235
-8
lines changed

12 files changed

+235
-8
lines changed

‎packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -314,10 +314,14 @@ export class Connector extends BaseConnector {
314314

315315
if (
316316
!this.onChatAnswerUpdated ||
317-
!(
318-
['accept-code-diff', 'run-shell-command', 'reject-shell-command'].includes(action.id) ||
319-
action.id.startsWith('reject-code-diff')
320-
)
317+
![
318+
'accept-code-diff',
319+
'reject-code-diff',
320+
'run-shell-command',
321+
'reject-shell-command',
322+
'confirm-tool-use',
323+
'reject-tool-use',
324+
].includes(action.id)
321325
) {
322326
return
323327
}
@@ -373,6 +377,30 @@ export class Connector extends BaseConnector {
373377
answer.header.buttons = []
374378
}
375379
break
380+
case 'confirm-tool-use':
381+
answer.buttons = [
382+
{
383+
keepCardAfterClick: true,
384+
text: 'Confirmed',
385+
id: 'confirmed-tool-use',
386+
status: 'success',
387+
position: 'outside',
388+
disabled: true,
389+
},
390+
]
391+
break
392+
case 'reject-tool-use':
393+
answer.buttons = [
394+
{
395+
keepCardAfterClick: true,
396+
text: 'Rejected',
397+
id: 'rejected-tool-use',
398+
status: 'error',
399+
position: 'outside',
400+
disabled: true,
401+
},
402+
]
403+
break
376404
default:
377405
break
378406
}

‎packages/core/src/codewhispererChat/controllers/chat/controller.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -849,9 +849,14 @@ export class ChatController {
849849
await this.handleCreatePrompt(message)
850850
break
851851
case 'run-shell-command':
852+
case 'confirm-tool-use':
852853
case 'generic-tool-execution':
853854
await this.processToolUseMessage(message)
854855
break
856+
case 'reject-code-diff':
857+
case 'reject-tool-use':
858+
await this.closeDiffView()
859+
break
855860
case 'reject-shell-command':
856861
await this.rejectShellCommand(message)
857862
break

‎packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -482,6 +482,17 @@ export class Messenger {
482482
if (validation.warning) {
483483
message = validation.warning + message + '\nRun the command to proceed.\n'
484484
}
485+
} else if (validation.requiresAcceptance) {
486+
buttons.push({
487+
id: 'reject-tool-use',
488+
text: 'Reject',
489+
status: 'info',
490+
})
491+
buttons.push({
492+
id: 'confirm-tool-use',
493+
text: 'Confirm',
494+
status: 'info',
495+
})
485496
} else if (toolUse?.name === ToolType.FsWrite) {
486497
const input = toolUse.input as unknown as FsWriteParams
487498
const fileName = path.basename(input.path)

‎packages/core/src/codewhispererChat/tools/executeBash.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ import { fs } from '../../shared/fs/fs'
99
import { ChildProcess, ChildProcessOptions } from '../../shared/utilities/processUtils'
1010
import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
1111
import { split } from 'shlex'
12+
import path from 'path'
13+
import * as vscode from 'vscode'
14+
import { isInDirectory } from '../../shared/filesystemUtilities'
1215

1316
export enum CommandCategory {
1417
ReadOnly,
@@ -215,6 +218,27 @@ export class ExecuteBash {
215218
return { requiresAcceptance: true, warning: highRiskCommandWarningMessage }
216219
}
217220
}
221+
for (const cmdArgs of allCommands) {
222+
for (const arg of cmdArgs) {
223+
if (this.looksLikePath(arg)) {
224+
// If not absolute, resolve using workingDirectory if available.
225+
let fullPath = arg
226+
if (!path.isAbsolute(arg) && this.workingDirectory) {
227+
fullPath = path.join(this.workingDirectory, arg)
228+
}
229+
const workspaceFolders = vscode.workspace.workspaceFolders
230+
if (!workspaceFolders || workspaceFolders.length === 0) {
231+
return { requiresAcceptance: true }
232+
}
233+
const isInWorkspace = workspaceFolders.some((folder) =>
234+
isInDirectory(folder.uri.fsPath, fullPath)
235+
)
236+
if (!isInWorkspace) {
237+
return { requiresAcceptance: true }
238+
}
239+
}
240+
}
241+
}
218242
return { requiresAcceptance: false }
219243
} catch (error) {
220244
this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
@@ -330,4 +354,8 @@ export class ExecuteBash {
330354
updates.write('```shell\n' + this.command + '\n```')
331355
updates.end()
332356
}
357+
358+
private looksLikePath(arg: string): boolean {
359+
return arg.startsWith('/') || arg.startsWith('./') || arg.startsWith('../')
360+
}
333361
}

‎packages/core/src/codewhispererChat/tools/fsRead.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@
55
import * as vscode from 'vscode'
66
import { getLogger } from '../../shared/logger/logger'
77
import fs from '../../shared/fs/fs'
8-
import { InvokeOutput, maxToolResponseSize, OutputKind, sanitizePath } from './toolShared'
8+
import { InvokeOutput, maxToolResponseSize, OutputKind, sanitizePath, CommandValidation } from './toolShared'
9+
import { isInDirectory } from '../../shared/filesystemUtilities'
910
import { Writable } from 'stream'
1011
import path from 'path'
1112

@@ -68,6 +69,18 @@ export class FsRead {
6869
updates.end()
6970
}
7071

72+
public requiresAcceptance(): CommandValidation {
73+
const workspaceFolders = vscode.workspace.workspaceFolders
74+
if (!workspaceFolders || workspaceFolders.length === 0) {
75+
return { requiresAcceptance: true }
76+
}
77+
const isInWorkspace = workspaceFolders.some((folder) => isInDirectory(folder.uri.fsPath, this.fsPath))
78+
if (!isInWorkspace) {
79+
return { requiresAcceptance: true }
80+
}
81+
return { requiresAcceptance: false }
82+
}
83+
7184
public async invoke(updates?: Writable): Promise<InvokeOutput> {
7285
try {
7386
const fileUri = vscode.Uri.file(this.fsPath)

‎packages/core/src/codewhispererChat/tools/listDirectory.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ import * as vscode from 'vscode'
66
import { getLogger } from '../../shared/logger/logger'
77
import { readDirectoryRecursively } from '../../shared/utilities/workspaceUtils'
88
import fs from '../../shared/fs/fs'
9-
import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
9+
import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
10+
import { isInDirectory } from '../../shared/filesystemUtilities'
1011
import { Writable } from 'stream'
1112
import path from 'path'
1213

@@ -61,6 +62,18 @@ export class ListDirectory {
6162
updates.end()
6263
}
6364

65+
public requiresAcceptance(): CommandValidation {
66+
const workspaceFolders = vscode.workspace.workspaceFolders
67+
if (!workspaceFolders || workspaceFolders.length === 0) {
68+
return { requiresAcceptance: true }
69+
}
70+
const isInWorkspace = workspaceFolders.some((folder) => isInDirectory(folder.uri.fsPath, this.fsPath))
71+
if (!isInWorkspace) {
72+
return { requiresAcceptance: true }
73+
}
74+
return { requiresAcceptance: false }
75+
}
76+
6477
public async invoke(updates?: Writable): Promise<InvokeOutput> {
6578
try {
6679
const fileUri = vscode.Uri.file(this.fsPath)

‎packages/core/src/codewhispererChat/tools/toolShared.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,8 @@ export function sanitizePath(inputPath: string): string {
3333
}
3434
return sanitized
3535
}
36+
37+
export interface CommandValidation {
38+
requiresAcceptance: boolean
39+
warning?: string
40+
}

‎packages/core/src/codewhispererChat/tools/toolUtils.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,13 +40,13 @@ export class ToolUtils {
4040
static requiresAcceptance(tool: Tool): CommandValidation {
4141
switch (tool.type) {
4242
case ToolType.FsRead:
43-
return { requiresAcceptance: false }
43+
return tool.tool.requiresAcceptance()
4444
case ToolType.FsWrite:
4545
return { requiresAcceptance: false }
4646
case ToolType.ExecuteBash:
4747
return tool.tool.requiresAcceptance()
4848
case ToolType.ListDirectory:
49-
return { requiresAcceptance: false }
49+
return tool.tool.requiresAcceptance()
5050
}
5151
}
5252

‎packages/core/src/codewhispererChat/tools/tool_index.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,10 @@
6868
"cwd": {
6969
"type": "string",
7070
"description": "Parameter to set the current working directory for the bash command."
71+
},
72+
"path": {
73+
"description": "Absolute path to a file or directory, This path should be the target path from the bash command.",
74+
"type": "string"
7175
}
7276
},
7377
"required": ["command", "cwd"]

‎packages/core/src/test/codewhispererChat/tools/executeBash.test.ts‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { strict as assert } from 'assert'
77
import sinon from 'sinon'
88
import { destructiveCommandWarningMessage, ExecuteBash } from '../../../codewhispererChat/tools/executeBash'
99
import { ChildProcess } from '../../../shared/utilities/processUtils'
10+
import * as vscode from 'vscode'
1011

1112
describe('ExecuteBash Tool', () => {
1213
let runStub: sinon.SinonStub
@@ -114,4 +115,55 @@ describe('ExecuteBash Tool', () => {
114115

115116
assert.strictEqual(invokeStub.callCount, 1)
116117
})
118+
119+
it('requires acceptance if the command references an absolute file path outside the workspace', () => {
120+
// Stub workspace folders to simulate a workspace at '/workspace/folder'
121+
const workspaceStub = sinon
122+
.stub(vscode.workspace, 'workspaceFolders')
123+
.value([{ uri: { fsPath: '/workspace/folder' } } as any])
124+
// Command references an absolute path outside the workspace
125+
const execBash = new ExecuteBash({ command: 'cat /not/in/workspace/file.txt', cwd: '/workspace/folder' })
126+
const result = execBash.requiresAcceptance()
127+
128+
assert.equal(
129+
result.requiresAcceptance,
130+
true,
131+
'Should require acceptance for an absolute path outside of workspace'
132+
)
133+
workspaceStub.restore()
134+
})
135+
136+
it('does NOT require acceptance if the command references a relative file path inside the workspace', () => {
137+
// Stub workspace folders to simulate a workspace at '/workspace/folder'
138+
const workspaceStub = sinon
139+
.stub(vscode.workspace, 'workspaceFolders')
140+
.value([{ uri: { fsPath: '/workspace/folder' } } as any])
141+
142+
// Command references a relative path that resolves within the workspace
143+
const execBash = new ExecuteBash({ command: 'cat ./file.txt', cwd: '/workspace/folder' })
144+
const result = execBash.requiresAcceptance()
145+
146+
assert.equal(result.requiresAcceptance, false, 'Relative path inside workspace should not require acceptance')
147+
148+
workspaceStub.restore()
149+
})
150+
151+
it('does NOT require acceptance if there is no path-like token in the command', () => {
152+
// Stub workspace folders (even though they are not used in this case)
153+
const workspaceStub = sinon
154+
.stub(vscode.workspace, 'workspaceFolders')
155+
.value([{ uri: { fsPath: '/workspace/folder' } } as any])
156+
157+
// Command with tokens that do not look like file paths
158+
const execBash = new ExecuteBash({ command: 'echo hello world', cwd: '/workspace/folder' })
159+
const result = execBash.requiresAcceptance()
160+
161+
assert.equal(
162+
result.requiresAcceptance,
163+
false,
164+
'A command without any path-like token should not require acceptance'
165+
)
166+
167+
workspaceStub.restore()
168+
})
117169
})

0 commit comments

Comments
 (0)