Skip to content

Commit 77761e9

Browse files
laileni-awslaileni-aws
committed
Adding conditional check for Read and List tools for outside of workspace
1 parent 04f3bf9 commit 77761e9

File tree

7 files changed

+59
-23
lines changed

7 files changed

+59
-23
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
@@ -725,13 +725,17 @@ export class Messenger {
725725
}
726726

727727
// Handle read tool and list directory messages
728-
if (toolUse?.name === ToolType.FsRead || toolUse?.name === ToolType.ListDirectory) {
728+
if (
729+
(toolUse?.name === ToolType.FsRead || toolUse?.name === ToolType.ListDirectory) &&
730+
!validation.requiresAcceptance
731+
) {
729732
return this.sendReadAndListDirToolMessage(toolUse, session, tabID, triggerID, messageIdToUpdate)
730733
}
731734

732735
// Handle file write tool, execute bash tool and bash command output log messages
733736
const buttons: ChatItemButton[] = []
734737
let header: ChatItemHeader | undefined = undefined
738+
let messageID: string = toolUse?.toolUseId ?? ''
735739
if (toolUse?.name === ToolType.ExecuteBash && message.startsWith('```shell')) {
736740
if (validation.requiresAcceptance) {
737741
const buttons: ChatItemButton[] = [
@@ -802,6 +806,7 @@ export class Messenger {
802806
}
803807
} else if (toolUse?.name === ToolType.ListDirectory || toolUse?.name === ToolType.FsRead) {
804808
if (validation.requiresAcceptance) {
809+
messageID = 'toolUse'
805810
const buttons: ChatItemButton[] = [
806811
{
807812
id: 'confirm-tool-use',
@@ -817,7 +822,6 @@ export class Messenger {
817822
},
818823
]
819824
header = {
820-
body: 'shell',
821825
buttons,
822826
}
823827
}
@@ -836,7 +840,7 @@ export class Messenger {
836840
followUpsHeader: undefined,
837841
relatedSuggestions: undefined,
838842
triggerID,
839-
messageID: toolUse?.toolUseId ?? '',
843+
messageID,
840844
userIntent: undefined,
841845
codeBlockLanguage: undefined,
842846
contextList: undefined,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ export class ChatStream extends Writable {
2929
private readonly messageIdToUpdate: string | undefined,
3030
// emitEvent decides to show the streaming message or read/list directory tool message to the user.
3131
private readonly emitEvent: boolean,
32-
private readonly validation: CommandValidation,
32+
readonly validation: CommandValidation,
3333
private readonly isReadorList: boolean,
3434
private readonly changeList?: Change[],
3535
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: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,15 @@
22
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
5+
import * as vscode from 'vscode'
56
import { Writable } from 'stream'
67
import { FsRead, FsReadParams } from './fsRead'
78
import { FsWrite, FsWriteParams } from './fsWrite'
89
import { CommandValidation, ExecuteBash, ExecuteBashParams } from './executeBash'
910
import { ToolResult, ToolResultContentBlock, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
1011
import { InvokeOutput, maxToolResponseSize } from './toolShared'
1112
import { ListDirectory, ListDirectoryParams } from './listDirectory'
12-
import * as vscode from 'vscode'
13+
import { ChatStream } from './chatStream'
1314

1415
export enum ToolType {
1516
FsRead = 'fsRead',
@@ -79,7 +80,7 @@ export class ToolUtils {
7980
}
8081
}
8182

82-
static async queueDescription(tool: Tool, updates: Writable): Promise<void> {
83+
static async queueDescription(tool: Tool, updates: ChatStream): Promise<void> {
8384
switch (tool.type) {
8485
case ToolType.FsRead:
8586
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)