Skip to content

Commit d381150

Browse files
laileni-awslaileni-aws
committed
Adding conditional check for Read and List tools for outside of workspace
1 parent b222e0b commit d381150

File tree

7 files changed

+58
-22
lines changed

7 files changed

+58
-22
lines changed

‎packages/core/package.nls.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -459,7 +459,7 @@
459459
"AWS.amazonq.executeBash.reject": "Reject",
460460
"AWS.amazonq.chat.directive.pairProgrammingModeOn": "You are using **pair programming mode**: Q can now list files, preview code diffs and allow you to run shell commands.",
461461
"AWS.amazonq.chat.directive.pairProgrammingModeOff": "You turned off **pair programming mode**. Q will not include code diffs or run commands in the chat.",
462-
"AWS.amazonq.chat.directive.runCommandToProceed": "Run the command to proceed.",
462+
"AWS.amazonq.chat.directive.runCommandToProceed": "I need permission to read files and list directories outside the workspace.",
463463
"AWS.toolkit.lambda.walkthrough.quickpickTitle": "Application Builder Walkthrough",
464464
"AWS.toolkit.lambda.walkthrough.title": "Get started building your application",
465465
"AWS.toolkit.lambda.walkthrough.description": "Your quick guide to build an application visually, iterate locally, and deploy to the cloud!",

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -683,13 +683,17 @@ export class Messenger {
683683
changeList?: Change[]
684684
) {
685685
// Handle read tool and list directory messages
686-
if (toolUse?.name === ToolType.FsRead || toolUse?.name === ToolType.ListDirectory) {
686+
if (
687+
(toolUse?.name === ToolType.FsRead || toolUse?.name === ToolType.ListDirectory) &&
688+
!validation.requiresAcceptance
689+
) {
687690
return this.sendReadAndListDirToolMessage(toolUse, session, tabID, triggerID, messageIdToUpdate)
688691
}
689692

690693
// Handle file write tool, execute bash tool and bash command output log messages
691694
const buttons: ChatItemButton[] = []
692695
let header: ChatItemHeader | undefined = undefined
696+
let messageID: string = toolUse?.toolUseId ?? ''
693697
if (toolUse?.name === ToolType.ExecuteBash && message.startsWith('```shell')) {
694698
if (validation.requiresAcceptance) {
695699
const buttons: ChatItemButton[] = [
@@ -757,6 +761,7 @@ export class Messenger {
757761
}
758762
} else if (toolUse?.name === ToolType.ListDirectory || toolUse?.name === ToolType.FsRead) {
759763
if (validation.requiresAcceptance) {
764+
messageID = 'toolUse'
760765
const buttons: ChatItemButton[] = [
761766
{
762767
id: 'confirm-tool-use',
@@ -772,7 +777,6 @@ export class Messenger {
772777
},
773778
]
774779
header = {
775-
body: 'shell',
776780
buttons,
777781
}
778782
}
@@ -787,7 +791,7 @@ export class Messenger {
787791
followUpsHeader: undefined,
788792
relatedSuggestions: undefined,
789793
triggerID,
790-
messageID: toolUse?.toolUseId ?? '',
794+
messageID,
791795
userIntent: undefined,
792796
codeBlockLanguage: undefined,
793797
contextList: undefined,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export class ChatStream extends Writable {
2828
private readonly messageIdToUpdate: string | undefined,
2929
// emitEvent decides to show the streaming message or read/list directory tool message to the user.
3030
private readonly emitEvent: boolean,
31-
private readonly validation: CommandValidation,
31+
readonly validation: CommandValidation,
3232
private readonly isReadorList: boolean,
3333
private readonly changeList?: Change[],
3434
private readonly logger = getLogger('chatStream')

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

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import fs from '../../shared/fs/fs'
88
import { Writable } from 'stream'
99
import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
1010
import { isInDirectory } from '../../shared/filesystemUtilities'
11+
import path from 'path'
12+
import { ChatStream } from './chatStream'
1113

1214
export interface FsReadParams {
1315
path: string
@@ -47,8 +49,28 @@ export class FsRead {
4749
this.logger.debug(`Validation succeeded for path: ${this.fsPath}`)
4850
}
4951

50-
public queueDescription(updates: Writable): void {
51-
updates.write('')
52+
public queueDescription(updates: ChatStream): void {
53+
if (updates.validation.requiresAcceptance) {
54+
const fileName = path.basename(this.fsPath)
55+
const fileUri = vscode.Uri.file(this.fsPath)
56+
updates.write(`Reading file: [${fileName}](${fileUri}), `)
57+
58+
const [start, end] = this.readRange ?? []
59+
60+
if (start && end) {
61+
updates.write(`from line ${start} to ${end}`)
62+
} else if (start) {
63+
if (start > 0) {
64+
updates.write(`from line ${start} to end of file`)
65+
} else {
66+
updates.write(`${start} line from the end of file to end of file`)
67+
}
68+
} else {
69+
updates.write('all lines')
70+
}
71+
} else {
72+
updates.write('')
73+
}
5274
updates.end()
5375
}
5476

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { Writable } from 'stream'
1010
import path from 'path'
1111
import { InvokeOutput, OutputKind, sanitizePath, CommandValidation } from './toolShared'
1212
import { isInDirectory } from '../../shared/filesystemUtilities'
13+
import { ChatStream } from './chatStream'
1314

1415
export interface ListDirectoryParams {
1516
path: string
@@ -49,16 +50,21 @@ export class ListDirectory {
4950
}
5051
}
5152

52-
public queueDescription(updates: Writable): void {
53-
const fileName = path.basename(this.fsPath)
54-
if (this.maxDepth === undefined) {
55-
updates.write(`Analyzing directories recursively: ${fileName}`)
56-
} else if (this.maxDepth === 0) {
57-
updates.write(`Analyzing directory: ${fileName}`)
53+
public queueDescription(updates: ChatStream): void {
54+
if (updates.validation.requiresAcceptance) {
55+
const fileName = path.basename(this.fsPath)
56+
if (this.maxDepth === undefined) {
57+
updates.write(`Analyzing directories recursively: ${fileName}`)
58+
} else if (this.maxDepth === 0) {
59+
updates.write(`Analyzing directory: ${fileName}`)
60+
} else {
61+
const level = this.maxDepth > 1 ? 'levels' : 'level'
62+
updates.write(`Analyzing directory: ${fileName} limited to ${this.maxDepth} subfolder ${level}`)
63+
}
5864
} else {
59-
const level = this.maxDepth > 1 ? 'levels' : 'level'
60-
updates.write(`Analyzing directory: ${fileName} limited to ${this.maxDepth} subfolder ${level}`)
65+
updates.write('')
6166
}
67+
6268
updates.end()
6369
}
6470

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash
99
import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
1010
import { InvokeOutput, maxToolResponseSize } from './toolShared'
1111
import { ListDirectory, ListDirectoryParams } from './listDirectory'
12+
import { ChatStream } from './chatStream'
1213

1314
export enum ToolType {
1415
FsRead = 'fsRead',
@@ -69,7 +70,7 @@ export class ToolUtils {
6970
}
7071
}
7172

72-
static async queueDescription(tool: Tool, updates: Writable): Promise<void> {
73+
static async queueDescription(tool: Tool, updates: ChatStream): Promise<void> {
7374
switch (tool.type) {
7475
case ToolType.FsRead:
7576
tool.tool.queueDescription(updates)

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import { ToolUse } from '@amzn/codewhisperer-streaming'
1515
import path from 'path'
1616
import fs from '../../../shared/fs/fs'
1717
import { ListDirectory } from '../../../codewhispererChat/tools/listDirectory'
18+
import { ChatStream } from '../../../codewhispererChat/tools/chatStream'
1819

1920
describe('ToolUtils', function () {
2021
let sandbox: sinon.SinonSandbox
@@ -23,6 +24,7 @@ describe('ToolUtils', function () {
2324
let mockExecuteBash: sinon.SinonStubbedInstance<ExecuteBash>
2425
let mockListDirectory: sinon.SinonStubbedInstance<ListDirectory>
2526
let mockWritable: sinon.SinonStubbedInstance<Writable>
27+
let mockChatStream: sinon.SinonStubbedInstance<ChatStream>
2628

2729
beforeEach(function () {
2830
sandbox = sinon.createSandbox()
@@ -33,6 +35,7 @@ describe('ToolUtils', function () {
3335
mockWritable = {
3436
write: sandbox.stub(),
3537
} as unknown as sinon.SinonStubbedInstance<Writable>
38+
mockChatStream = sandbox.createStubInstance(ChatStream)
3639
;(mockFsRead.requiresAcceptance as sinon.SinonStub).returns({ requiresAcceptance: false })
3740
;(mockListDirectory.requiresAcceptance as sinon.SinonStub).returns({ requiresAcceptance: false })
3841
})
@@ -185,30 +188,30 @@ describe('ToolUtils', function () {
185188
// TODO: Adding "void" to the following tests for the current implementation but in the next followup PR I will fix this issue.
186189
it('delegates to FsRead tool queueDescription method', function () {
187190
const tool: Tool = { type: ToolType.FsRead, tool: mockFsRead as unknown as FsRead }
188-
void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable)
191+
void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
189192

190-
assert(mockFsRead.queueDescription.calledOnceWith(mockWritable))
193+
assert(mockFsRead.queueDescription.calledOnceWith(mockChatStream))
191194
})
192195

193196
it('delegates to FsWrite tool queueDescription method', function () {
194197
const tool: Tool = { type: ToolType.FsWrite, tool: mockFsWrite as unknown as FsWrite }
195-
void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable)
198+
void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
196199

197200
assert(mockFsWrite.queueDescription.calledOnceWith(mockWritable))
198201
})
199202

200203
it('delegates to ExecuteBash tool queueDescription method', function () {
201204
const tool: Tool = { type: ToolType.ExecuteBash, tool: mockExecuteBash as unknown as ExecuteBash }
202-
void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable)
205+
void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
203206

204207
assert(mockExecuteBash.queueDescription.calledOnceWith(mockWritable))
205208
})
206209

207210
it('delegates to ListDirectory tool queueDescription method', function () {
208211
const tool: Tool = { type: ToolType.ListDirectory, tool: mockListDirectory as unknown as ListDirectory }
209-
void ToolUtils.queueDescription(tool, mockWritable as unknown as Writable)
212+
void ToolUtils.queueDescription(tool, mockWritable as unknown as ChatStream)
210213

211-
assert(mockListDirectory.queueDescription.calledOnceWith(mockWritable))
214+
assert(mockListDirectory.queueDescription.calledOnceWith(mockChatStream))
212215
})
213216
})
214217

0 commit comments

Comments
 (0)