Skip to content

fix: add fsReplace tool to batch edits #1533

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
* Will be deleted or merged.
*/

import * as crypto from 'crypto'

Check warning on line 6 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test

Do not import Node.js builtin module "crypto"

Check warning on line 6 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test (Windows)

Do not import Node.js builtin module "crypto"
import * as path from 'path'

Check warning on line 7 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test

Do not import Node.js builtin module "path"

Check warning on line 7 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test (Windows)

Do not import Node.js builtin module "path"
import {
ChatTriggerType,
GenerateAssistantResponseCommandInput,
Expand Down Expand Up @@ -124,6 +124,7 @@
import { AgenticChatError, customerFacingErrorCodes, isRequestAbortedError, unactionableErrorCodes } from './errors'
import { CommandCategory } from './tools/executeBash'
import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
import { FsReplace, FsReplaceParams } from './tools/fsReplace'

type ChatHandlers = Omit<
LspHandlers<Chat>,
Expand Down Expand Up @@ -263,7 +264,7 @@
this.#log(`Reverting file change for tooluseId: ${toolUseId}`)
const toolUse = session?.toolUseLookup.get(toolUseId)

const input = toolUse?.input as unknown as FsWriteParams
const input = toolUse?.input as unknown as FsWriteParams | FsReplaceParams
if (toolUse?.fileChange?.before) {
await this.#features.workspace.fs.writeFile(input.path, toolUse.fileChange.before)
} else {
Expand Down Expand Up @@ -914,11 +915,11 @@
await chatResultStream.removeResultBlockAndUpdateUI(progressPrefix + toolUse.toolUseId)

// fsRead and listDirectory write to an existing card and could show nothing in the current position
if (!['fsWrite', 'fsRead', 'listDirectory'].includes(toolUse.name)) {
if (!['fsWrite', 'fsReplace', 'fsRead', 'listDirectory'].includes(toolUse.name)) {
await this.#showUndoAllIfRequired(chatResultStream, session)
}
// fsWrite can take a long time, so we render fsWrite Explanatory upon partial streaming responses.
if (toolUse.name !== 'fsWrite') {
if (toolUse.name !== 'fsWrite' && toolUse.name !== 'fsReplace') {
const { explanation } = toolUse.input as unknown as ExplanatoryParams
if (explanation) {
await chatResultStream.writeResultBlock({
Expand All @@ -934,11 +935,13 @@
case 'grepSearch':
case 'fileSearch':
case 'fsWrite':
case 'fsReplace':
case 'executeBash': {
const toolMap = {
fsRead: { Tool: FsRead },
listDirectory: { Tool: ListDirectory },
fsWrite: { Tool: FsWrite },
fsReplace: { Tool: FsReplace },
executeBash: { Tool: ExecuteBash },
grepSearch: { Tool: GrepSearch },
fileSearch: { Tool: FileSearch },
Expand Down Expand Up @@ -988,9 +991,6 @@
}
break
}
case 'codeSearch':
// no need to write tool message for code search.
break
default:
this.#features.logging.warn(`Recieved unrecognized tool: ${toolUse.name}`)
await chatResultStream.writeResultBlock({
Expand All @@ -1001,8 +1001,8 @@
break
}

if (toolUse.name === 'fsWrite') {
const input = toolUse.input as unknown as FsWriteParams
if (toolUse.name === 'fsWrite' || toolUse.name === 'fsReplace') {
const input = toolUse.input as unknown as FsWriteParams | FsReplaceParams
const document = await this.#triggerContext.getTextDocument(input.path)
session.toolUseLookup.set(toolUse.toolUseId, {
...toolUse,
Expand Down Expand Up @@ -1056,8 +1056,9 @@
await chatResultStream.writeResultBlock(grepSearchResult)
}
break
case 'fsReplace':
case 'fsWrite':
const input = toolUse.input as unknown as FsWriteParams
const input = toolUse.input as unknown as FsWriteParams | FsReplaceParams
const doc = await this.#triggerContext.getTextDocument(input.path)
const chatResult = await this.#getFsWriteChatResult(toolUse, doc, session)
const cachedToolUse = session.toolUseLookup.get(toolUse.toolUseId)
Expand Down Expand Up @@ -1144,9 +1145,9 @@
}
}
// display fs write failure status in the UX of that file card
if (toolUse.name === 'fsWrite' && toolUse.toolUseId) {
if ((toolUse.name === 'fsWrite' || toolUse.name === 'fsReplace') && toolUse.toolUseId) {
const existingCard = chatResultStream.getMessageBlockId(toolUse.toolUseId)
const fsParam = toolUse.input as unknown as FsWriteParams
const fsParam = toolUse.input as unknown as FsWriteParams | FsReplaceParams
const fileName = path.basename(fsParam.path)
const errorResult = {
type: 'tool',
Expand Down Expand Up @@ -1236,7 +1237,7 @@
if (toolUse.name === 'fsRead' || toolUse.name === 'listDirectory') {
return
}
if (toolUse.name === 'fsWrite') {
if (toolUse.name === 'fsWrite' || toolUse.name === 'fsReplace') {
if (session.currentUndoAllId === undefined) {
session.currentUndoAllId = toolUse.toolUseId
}
Expand Down Expand Up @@ -1440,6 +1441,7 @@
let body: string | undefined

switch (toolName) {
case 'fsReplace':
case 'fsWrite':
header = {
body: undefined,
Expand Down Expand Up @@ -1594,7 +1596,7 @@
const commandString = (toolUse.input as unknown as ExecuteBashParams).command
body = '```shell\n' + commandString
break

case 'fsReplace':
case 'fsWrite':
buttons = [
{
Expand All @@ -1610,10 +1612,9 @@
body: '#### Allow file modification outside of your workspace',
buttons,
}
const writeFilePath = (toolUse.input as unknown as FsWriteParams).path
const writeFilePath = (toolUse.input as unknown as FsWriteParams | FsReplaceParams).path
body = `I need permission to modify files in your workspace.\n\`${writeFilePath}\``
break

case 'fsRead':
case 'listDirectory':
default:
Expand Down Expand Up @@ -1657,7 +1658,7 @@
doc: TextDocument | undefined,
session: ChatSessionService
): Promise<ChatMessage> {
const input = toolUse.input as unknown as FsWriteParams
const input = toolUse.input as unknown as FsWriteParams | FsReplaceParams
const oldContent = session.toolUseLookup.get(toolUse.toolUseId!)?.fileChange?.before ?? ''
// Get just the filename instead of the full path
const fileName = path.basename(input.path)
Expand Down Expand Up @@ -2175,8 +2176,8 @@
const toolUseId = params.messageId
const toolUse = toolUseId ? session.data?.toolUseLookup.get(toolUseId) : undefined

if (toolUse?.name === 'fsWrite') {
const input = toolUse.input as unknown as FsWriteParams
if (toolUse?.name === 'fsWrite' || toolUse?.name === 'fsReplace') {
const input = toolUse.input as unknown as FsWriteParams | FsReplaceParams
this.#features.lsp.workspace.openFileDiff({
originalFileUri: input.path,
originalFileContent: toolUse.fileChange?.before,
Expand Down Expand Up @@ -2453,7 +2454,7 @@
}
const toolUses = Object.values(data.toolUses)
for (const toolUse of toolUses) {
if (toolUse.name === 'fsWrite' && typeof toolUse.input === 'string') {
if ((toolUse.name === 'fsWrite' || toolUse.name === 'fsReplace') && typeof toolUse.input === 'string') {
const filepath = extractKey(toolUse.input, 'path')
const msgId = progressPrefix + toolUse.toolUseId
// render fs write UI as soon as fs write starts
Expand All @@ -2480,7 +2481,7 @@
},
})
}
// render the tool use explanatory as soon as this is received for fsWrite
// render the tool use explanatory as soon as this is received for fsWrite/fsReplace
const explanation = extractKey(toolUse.input, 'explanation')
const messageId = progressPrefix + toolUse.toolUseId + '_explanation'
if (explanation && !chatResultStream.hasMessage(messageId)) {
Expand Down Expand Up @@ -2658,7 +2659,9 @@

// it's disabled so filter out the write tools
if (!session.pairProgrammingMode) {
return tools.filter(tool => !['fsWrite', 'executeBash'].includes(tool.toolSpecification?.name || ''))
return tools.filter(
tool => !['fsWrite', 'fsReplace', 'executeBash'].includes(tool.toolSpecification?.name || '')
)
}
return tools
}
Expand Down
Loading
Loading