Skip to content

Commit 36a61be

Browse files
bywang56bywang56
committed
feat(chat): Add validation to fileRead, ListDir and ExecBash tool
1 parent 4ad3e92 commit 36a61be

File tree

12 files changed

+220
-7
lines changed

12 files changed

+220
-7
lines changed

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,7 +321,14 @@ export class Connector extends BaseConnector {
321321

322322
if (
323323
!this.onChatAnswerUpdated ||
324-
!['accept-code-diff', 'reject-code-diff', 'run-shell-command', 'reject-shell-command'].includes(action.id)
324+
![
325+
'accept-code-diff',
326+
'reject-code-diff',
327+
'run-shell-command',
328+
'reject-shell-command',
329+
'confirm-tool-use',
330+
'reject-tool-use',
331+
].includes(action.id)
325332
) {
326333
return
327334
}

‎packages/core/src/codewhisperer/client/codewhisperer.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ export interface CodeWhispererConfig {
3131
}
3232

3333
export const defaultServiceConfig: CodeWhispererConfig = {
34-
region: 'us-east-1',
35-
endpoint: 'https://codewhisperer.us-east-1.amazonaws.com/',
34+
region: 'us-west-2',
35+
endpoint: 'https://rts.alpha-us-west-2.codewhisperer.ai.aws.dev/',
3636
}
3737

3838
export function getCodewhispererConfig(): CodeWhispererConfig {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -822,12 +822,14 @@ export class ChatController {
822822
break
823823
case 'run-shell-command':
824824
case 'generic-tool-execution':
825+
case 'confirm-tool-use':
825826
await this.processToolUseMessage(message)
826827
break
827828
case 'accept-code-diff':
828829
await this.closeDiffView()
829830
break
830831
case 'reject-code-diff':
832+
case 'reject-tool-use':
831833
await this.restoreBackup(message)
832834
await this.closeDiffView()
833835
break

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,17 @@ export class Messenger {
557557
}
558558
if (validation.warning) {
559559
message = validation.warning + message
560+
} else if (validation.requiresAcceptance) {
561+
buttons.push({
562+
id: 'reject-tool-use',
563+
text: 'Reject',
564+
status: 'info',
565+
})
566+
buttons.push({
567+
id: 'confirm-tool-use',
568+
text: 'Confirm',
569+
status: 'info',
570+
})
560571
}
561572
fullWidth = true
562573
padding = false

‎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,
@@ -191,6 +194,27 @@ export class ExecuteBash {
191194
return { requiresAcceptance: true }
192195
}
193196
}
197+
for (const cmdArgs of allCommands) {
198+
for (const arg of cmdArgs) {
199+
if (this.looksLikePath(arg)) {
200+
// If not absolute, resolve using workingDirectory if available.
201+
let fullPath = arg
202+
if (!path.isAbsolute(arg) && this.workingDirectory) {
203+
fullPath = path.join(this.workingDirectory, arg)
204+
}
205+
const workspaceFolders = vscode.workspace.workspaceFolders
206+
if (!workspaceFolders || workspaceFolders.length === 0) {
207+
return { requiresAcceptance: true }
208+
}
209+
const isInWorkspace = workspaceFolders.some((folder) =>
210+
isInDirectory(folder.uri.fsPath, fullPath)
211+
)
212+
if (!isInWorkspace) {
213+
return { requiresAcceptance: true }
214+
}
215+
}
216+
}
217+
}
194218
return { requiresAcceptance: false }
195219
} catch (error) {
196220
this.logger.warn(`Error while checking acceptance: ${(error as Error).message}`)
@@ -306,4 +330,8 @@ export class ExecuteBash {
306330
updates.write('```shell\n' + this.command + '\n```')
307331
updates.end()
308332
}
333+
334+
private looksLikePath(arg: string): boolean {
335+
return arg.startsWith('/') || arg.startsWith('./') || arg.startsWith('../')
336+
}
309337
}

‎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, OutputKind, sanitizePath } from './toolShared'
8+
import { InvokeOutput, 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
@@ -32,3 +32,8 @@ export function sanitizePath(inputPath: string): string {
3232
}
3333
return sanitized
3434
}
35+
36+
export interface CommandValidation {
37+
requiresAcceptance: boolean
38+
warning?: string
39+
}

‎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/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)