diff --git a/core/aws-lsp-core/src/index.ts b/core/aws-lsp-core/src/index.ts index 8964a2bab5..1e739785a2 100644 --- a/core/aws-lsp-core/src/index.ts +++ b/core/aws-lsp-core/src/index.ts @@ -14,3 +14,4 @@ export * as textUtils from './util/text' export * as timeoutUtils from './util/timeoutUtils' export * from './util/awsError' export * from './base/index' +export * as testFolder from './test/testFolder' diff --git a/core/aws-lsp-core/src/test/testFolder.ts b/core/aws-lsp-core/src/test/testFolder.ts new file mode 100644 index 0000000000..29e7a0a459 --- /dev/null +++ b/core/aws-lsp-core/src/test/testFolder.ts @@ -0,0 +1,60 @@ +import * as path from 'path' +import * as fs from 'fs' +import * as os from 'os' +import * as crypto from 'crypto' + +/** + * Interface for working with temporary files. + * Simplified port of https://github.com/aws/aws-toolkit-vscode/blob/16477869525fb79f8dc82cb22e301aaea9c5e0c6/packages/core/src/test/testUtil.ts#L77 + * + * Proper usage requires adding the proper logic into the hooks. See example below: + * + * before(async () => { + * ... + * testFolder = await TestFolder.create() + * ... + * } + * // Only necessary if test state should not bleed through. + * afterEach(async () => { + * ... + * await testFolder.clear() + * ... + * }) + * + * after(async () => { + * ... + * await testFolder.delete() + * ... + * }) + */ +export class TestFolder { + private constructor(public readonly path: string) {} + + async write(fileName: string, content: string): Promise { + const filePath = path.join(this.path, fileName) + fs.writeFileSync(filePath, content) + return filePath + } + + static async create() { + const tempDir = path.join( + os.type() === 'Darwin' ? '/tmp' : os.tmpdir(), + 'aws-language-servers', + 'test', + crypto.randomBytes(4).toString('hex') + ) + await fs.promises.mkdir(tempDir, { recursive: true }) + return new TestFolder(tempDir) + } + + async delete() { + fs.rmSync(this.path, { recursive: true, force: true }) + } + + async clear() { + const files = await fs.readdirSync(this.path) + for (const f of files) { + await fs.rmSync(path.join(this.path, f), { recursive: true, force: true }) + } + } +} diff --git a/core/aws-lsp-core/src/util/path.test.ts b/core/aws-lsp-core/src/util/path.test.ts index 741b8d1381..2b6114f32f 100644 --- a/core/aws-lsp-core/src/util/path.test.ts +++ b/core/aws-lsp-core/src/util/path.test.ts @@ -4,7 +4,8 @@ import * as assert from 'assert' import * as path from 'path' import * as os from 'os' -import { normalizeSeparator, normalize, isInDirectory } from './path' +import { normalizeSeparator, normalize, isInDirectory, sanitize } from './path' +import sinon from 'ts-sinon' describe('pathUtils', async function () { it('normalizeSeparator()', function () { @@ -76,4 +77,40 @@ describe('pathUtils', async function () { assert.ok(!isInDirectory('/foo/bar/baz/', '/FOO/BAR/BAZ/A.TXT')) } }) + + describe('sanitizePath', function () { + let sandbox: sinon.SinonSandbox + + beforeEach(function () { + sandbox = sinon.createSandbox() + }) + + afterEach(function () { + sandbox.restore() + }) + + it('trims whitespace from input path', function () { + const result = sanitize(' /test/path ') + assert.strictEqual(result, '/test/path') + }) + + it('expands tilde to user home directory', function () { + const homeDir = '/Users/testuser' + sandbox.stub(os, 'homedir').returns(homeDir) + + const result = sanitize('~/documents/file.txt') + assert.strictEqual(result, path.join(homeDir, 'documents/file.txt')) + }) + + it('converts relative paths to absolute paths', function () { + const result = sanitize('relative/path') + assert.strictEqual(result, path.resolve('relative/path')) + }) + + it('leaves absolute paths unchanged', function () { + const absolutePath = path.resolve('/absolute/path') + const result = sanitize(absolutePath) + assert.strictEqual(result, absolutePath) + }) + }) }) diff --git a/core/aws-lsp-core/src/util/path.ts b/core/aws-lsp-core/src/util/path.ts index a65d1d5703..484f3e5716 100644 --- a/core/aws-lsp-core/src/util/path.ts +++ b/core/aws-lsp-core/src/util/path.ts @@ -78,3 +78,16 @@ export function normalize(p: string): string { } return normalizeSeparator(path.normalize(p)) } + +export function sanitize(inputPath: string): string { + let sanitized = inputPath.trim() + + if (sanitized.startsWith('~')) { + sanitized = path.join(os.homedir(), sanitized.slice(1)) + } + + if (!path.isAbsolute(sanitized)) { + sanitized = path.resolve(sanitized) + } + return sanitized +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.test.ts new file mode 100644 index 0000000000..1842f1dde8 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.test.ts @@ -0,0 +1,103 @@ +import * as assert from 'assert' +import { FsRead } from './fsRead' +import * as path from 'path' +import * as fs from 'fs/promises' +import { TestFeatures } from '@aws/language-server-runtimes/testing' +import { Workspace } from '@aws/language-server-runtimes/server-interface' +import { testFolder } from '@aws/lsp-core' +import { StubbedInstance } from 'ts-sinon' + +describe('FsRead Tool', () => { + let features: TestFeatures + let tempFolder: testFolder.TestFolder + + const stdout = new WritableStream({ + write(chunk) { + process.stdout.write(chunk) + }, + }) + + before(async () => { + features = new TestFeatures() + features.workspace = { + // @ts-ignore reading a file does not require all of fs to be implemented. + fs: { + readFile: (path, options?) => + fs.readFile(path, { encoding: (options?.encoding || 'utf-8') as BufferEncoding }), + readdir: path => fs.readdir(path, { withFileTypes: true }), + exists: path => + fs + .access(path) + .then(() => true) + .catch(() => false), + } as Workspace['fs'], + } as StubbedInstance + tempFolder = await testFolder.TestFolder.create() + }) + + after(async () => { + tempFolder.delete() + }) + + afterEach(async () => { + tempFolder.clear() + }) + + it('throws if path is empty', async () => { + const fsRead = new FsRead(features) + await assert.rejects(() => fsRead.invoke({ path: '' })) + }) + + it('reads entire file', async () => { + const fileContent = 'Line 1\nLine 2\nLine 3' + const filePath = await tempFolder.write('fullFile.txt', fileContent) + + const fsRead = new FsRead(features) + const result = await fsRead.invoke({ path: filePath }) + + assert.strictEqual(result.output.kind, 'text', 'Output kind should be "text"') + assert.strictEqual(result.output.content, fileContent, 'File content should match exactly') + }) + + it('reads partial lines of a file', async () => { + const fileContent = 'A\nB\nC\nD\nE\nF' + const filePath = await tempFolder.write('partialFile.txt', fileContent) + + const fsRead = new FsRead(features) + const result = await fsRead.invoke({ path: filePath, readRange: [2, 4] }) + + assert.strictEqual(result.output.kind, 'text') + assert.strictEqual(result.output.content, 'B\nC\nD') + }) + + it('throws error if path does not exist', async () => { + const filePath = path.join(tempFolder.path, 'no_such_file.txt') + const fsRead = new FsRead(features) + + await assert.rejects(fsRead.invoke({ path: filePath })) + }) + + it('throws error if content exceeds 30KB', async function () { + const bigContent = 'x'.repeat(35_000) + + const filePath = await tempFolder.write('bigFile.txt', bigContent) + + const fsRead = new FsRead(features) + + await assert.rejects( + fsRead.invoke({ path: filePath }), + /This tool only supports reading \d+ bytes at a time/i, + 'Expected a size-limit error' + ) + }) + + it('invalid line range', async () => { + const filePath = await tempFolder.write('rangeTest.txt', '1\n2\n3') + const fsRead = new FsRead(features) + + await fsRead.invoke({ path: filePath, readRange: [3, 2] }) + const result = await fsRead.invoke({ path: filePath, readRange: [3, 2] }) + assert.strictEqual(result.output.kind, 'text') + assert.strictEqual(result.output.content, '') + }) +}) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.ts new file mode 100644 index 0000000000..3ab82df3c8 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsRead.ts @@ -0,0 +1,117 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import { sanitize } from '@aws/lsp-core/out/util/path' +import { InvokeOutput, maxToolResponseSize } from './toolShared' +import { Features } from '@aws/language-server-runtimes/server-interface/server' + +// Port of https://github.com/aws/aws-toolkit-vscode/blob/10bb1c7dc45f128df14d749d95905c0e9b808096/packages/core/src/codewhispererChat/tools/fsRead.ts#L17 + +export interface FsReadParams { + path: string + readRange?: number[] +} + +export class FsRead { + private readonly logging: Features['logging'] + private readonly workspace: Features['workspace'] + + constructor(features: Pick & Partial) { + this.logging = features.logging + this.workspace = features.workspace + } + + public async invoke(params: FsReadParams): Promise { + const path = sanitize(params.path) + const fileContents = await this.readFile(path) + this.logging.info(`Read file: ${path}, size: ${fileContents.length}`) + return this.handleFileRange(params, fileContents) + } + + private async readFile(filePath: string): Promise { + this.logging.info(`Reading file: ${filePath}`) + return await this.workspace.fs.readFile(filePath) + } + + private handleFileRange(params: FsReadParams, fullText: string): InvokeOutput { + if (!params.readRange || params.readRange.length === 0) { + this.logging.log('No range provided. returning entire file.') + return this.createOutput(this.enforceMaxSize(fullText)) + } + + const lines = fullText.split('\n') + const [start, end] = this.parseLineRange(lines.length, params.readRange) + if (start > end) { + this.logging.error(`Invalid range: ${params.readRange.join('-')}`) + return this.createOutput('') + } + + this.logging.log(`Reading file: ${params.path}, lines ${start + 1}-${end + 1}`) + const slice = lines.slice(start, end + 1).join('\n') + return this.createOutput(this.enforceMaxSize(slice)) + } + + private parseLineRange(lineCount: number, range: number[]): [number, number] { + const startIdx = range[0] + let endIdx = range.length >= 2 ? range[1] : undefined + + if (endIdx === undefined) { + endIdx = -1 + } + + const convert = (i: number): number => { + return i < 0 ? lineCount + i : i - 1 + } + + const finalStart = Math.max(0, Math.min(lineCount - 1, convert(startIdx))) + const finalEnd = Math.max(0, Math.min(lineCount - 1, convert(endIdx))) + return [finalStart, finalEnd] + } + + private enforceMaxSize(content: string): string { + const byteCount = Buffer.byteLength(content, 'utf8') + if (byteCount > maxToolResponseSize) { + throw new Error( + `This tool only supports reading ${maxToolResponseSize} bytes at a time. + You tried to read ${byteCount} bytes. Try executing with fewer lines specified.` + ) + } + return content + } + + private createOutput(content: string): InvokeOutput { + return { + output: { + kind: 'text', + content: content, + }, + } + } + + public getSpec() { + return { + name: 'fsRead', + description: + 'A tool for reading a file. \n* This tool returns the contents of a file, and the optional `readRange` determines what range of lines will be read from the specified file.', + inputSchema: { + type: 'object', + parameters: { + path: { + description: 'Absolute path to a file, e.g. `/repo/file.py`.', + type: 'string', + }, + readRange: { + description: + 'Optional parameter when reading files.\n* If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[startLine, -1]` shows all lines from `startLine` to the end of the file.', + items: { + type: 'integer', + }, + type: 'array', + }, + }, + required: ['path'], + }, + } as const + } +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.test.ts new file mode 100644 index 0000000000..4a75ad3e3d --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.test.ts @@ -0,0 +1,418 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +import { AppendParams, CreateParams, FsWrite, InsertParams, StrReplaceParams } from './fsWrite' +import { testFolder } from '@aws/lsp-core' +import * as path from 'path' +import * as assert from 'assert' +import * as fs from 'fs/promises' +import { InvokeOutput } from './toolShared' +import { TestFeatures } from '@aws/language-server-runtimes/testing' +import { Workspace } from '@aws/language-server-runtimes/server-interface' +import { StubbedInstance } from 'ts-sinon' + +describe('FsWrite Tool', function () { + let tempFolder: testFolder.TestFolder + let features: TestFeatures + const expectedOutput: InvokeOutput = { + output: { + kind: 'text', + content: '', + }, + } + + const stdout = new WritableStream({ + write(chunk) { + process.stdout.write(chunk) + }, + }) + + before(async function () { + features = new TestFeatures() + features.workspace = { + // @ts-ignore writing a file does not require all of fs to be implemented + fs: { + writeFile: fs.writeFile, + readFile: (path, options?) => + fs.readFile(path, { encoding: (options?.encoding || 'utf-8') as BufferEncoding }), + exists: path => + fs + .access(path) + .then(() => true) + .catch(() => false), + } as Workspace['fs'], + } as StubbedInstance + tempFolder = await testFolder.TestFolder.create() + }) + + after(async function () { + await tempFolder.delete() + }) + + describe('handleCreate', function () { + it('creates a new file with fileText content', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + const fileExists = await features.workspace.fs.exists(filePath) + assert.ok(!fileExists) + + const params: CreateParams = { + command: 'create', + fileText: 'Hello World', + path: filePath, + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Hello World') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('replaces existing file with fileText content', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + const fileExists = await features.workspace.fs.exists(filePath) + assert.ok(fileExists) + + const params: CreateParams = { + command: 'create', + fileText: 'Goodbye', + path: filePath, + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Goodbye') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('uses newStr when fileText is not provided', async function () { + const filePath = path.join(tempFolder.path, 'file2.txt') + + const params: CreateParams = { + command: 'create', + newStr: 'Hello World', + path: filePath, + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Hello World') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('creates an empty file when no content is provided', async function () { + const filePath = path.join(tempFolder.path, 'file3.txt') + + const params: CreateParams = { + command: 'create', + path: filePath, + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, '') + + assert.deepStrictEqual(output, expectedOutput) + }) + }) + + describe('handleStrReplace', async function () { + before(async function () { + tempFolder = await testFolder.TestFolder.create() + }) + + it('replaces a single occurrence of a string', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + await fs.writeFile(filePath, 'Hello World') + + const params: StrReplaceParams = { + command: 'strReplace', + path: filePath, + oldStr: 'Hello', + newStr: 'Goodbye', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Goodbye World') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('throws error when no matches are found', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + + const params: StrReplaceParams = { + command: 'strReplace', + path: filePath, + oldStr: 'Invalid', + newStr: 'Goodbye', + } + + const fsWrite = new FsWrite(features) + await assert.rejects(() => fsWrite.invoke(params), /No occurrences of "Invalid" were found/) + }) + + it('throws error when multiple matches are found', async function () { + const filePath = path.join(tempFolder.path, 'file2.txt') + await fs.writeFile(filePath, 'Hello Hello World') + + const params: StrReplaceParams = { + command: 'strReplace', + path: filePath, + oldStr: 'Hello', + newStr: 'Goodbye', + } + + const fsWrite = new FsWrite(features) + await assert.rejects( + () => fsWrite.invoke(params), + /2 occurrences of oldStr were found when only 1 is expected/ + ) + }) + + it('handles regular expression special characters correctly', async function () { + const filePath = path.join(tempFolder.path, 'file3.txt') + await fs.writeFile(filePath, 'Text with special chars: .*+?^${}()|[]\\') + + const params: StrReplaceParams = { + command: 'strReplace', + path: filePath, + oldStr: '.*+?^${}()|[]\\', + newStr: 'REPLACED', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Text with special chars: REPLACED') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('preserves whitespace and newlines during replacement', async function () { + const filePath = path.join(tempFolder.path, 'file4.txt') + await fs.writeFile(filePath, 'Line 1\n Indented line\nLine 3') + + const params: StrReplaceParams = { + command: 'strReplace', + path: filePath, + oldStr: ' Indented line\n', + newStr: ' Double indented\n', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const content = await features.workspace.fs.readFile(filePath) + assert.strictEqual(content, 'Line 1\n Double indented\nLine 3') + + assert.deepStrictEqual(output, expectedOutput) + }) + }) + + describe('handleInsert', function () { + before(async function () { + tempFolder = await testFolder.TestFolder.create() + }) + + it('inserts text after the specified line number', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + await fs.writeFile(filePath, 'Line 1\nLine 2\nLine 3\nLine 4') + + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: 2, + newStr: 'New Line', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'Line 1\nLine 2\nNew Line\nLine 3\nLine 4') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('inserts text at the beginning when line number is 0', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: 0, + newStr: 'New First Line', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'New First Line\nLine 1\nLine 2\nNew Line\nLine 3\nLine 4') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('inserts text at the end when line number exceeds file length', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: 10, + newStr: 'New Last Line', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'New First Line\nLine 1\nLine 2\nNew Line\nLine 3\nLine 4\nNew Last Line') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('handles insertion into an empty file', async function () { + const filePath = path.join(tempFolder.path, 'file2.txt') + await fs.writeFile(filePath, '') + + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: 0, + newStr: 'First Line', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'First Line\n') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('handles negative line numbers by inserting at the beginning', async function () { + const filePath = path.join(tempFolder.path, 'file2.txt') + + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: -1, + newStr: 'New First Line', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'New First Line\nFirst Line\n') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('throws error when file does not exist', async function () { + const filePath = path.join(tempFolder.path, 'nonexistent.txt') + + const params: InsertParams = { + command: 'insert', + path: filePath, + insertLine: 1, + newStr: 'New Line', + } + + const fsWrite = new FsWrite(features) + await assert.rejects(() => fsWrite.invoke(params), /no such file or directory/) + }) + }) + + describe('handleAppend', function () { + it('appends text to the end of a file', async function () { + const filePath = path.join(tempFolder.path, 'file1.txt') + await fs.writeFile(filePath, 'Line 1\nLine 2\nLine 3\n') + + const params: AppendParams = { + command: 'append', + path: filePath, + newStr: 'Line 4', + } + + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'Line 1\nLine 2\nLine 3\nLine 4') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('adds a newline before appending if file does not end with one', async function () { + const filePath = path.join(tempFolder.path, 'file2.txt') + await fs.writeFile(filePath, 'Line 1\nLine 2\nLine 3') + + const params: AppendParams = { + command: 'append', + path: filePath, + newStr: 'Line 4', + } + + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'Line 1\nLine 2\nLine 3\nLine 4') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('appends to an empty file', async function () { + const filePath = path.join(tempFolder.path, 'file3.txt') + await fs.writeFile(filePath, '') + + const params: AppendParams = { + command: 'append', + path: filePath, + newStr: 'Line 1', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'Line 1') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('appends multiple lines correctly', async function () { + const filePath = path.join(tempFolder.path, 'file3.txt') + + const params: AppendParams = { + command: 'append', + path: filePath, + newStr: 'Line 2\nLine 3', + } + const fsWrite = new FsWrite(features) + const output = await fsWrite.invoke(params) + + const newContent = await features.workspace.fs.readFile(filePath) + assert.strictEqual(newContent, 'Line 1\nLine 2\nLine 3') + + assert.deepStrictEqual(output, expectedOutput) + }) + + it('throws error when file does not exist', async function () { + const filePath = path.join(tempFolder.path, 'nonexistent.txt') + + const params: AppendParams = { + command: 'append', + path: filePath, + newStr: 'New Line', + } + + const fsWrite = new FsWrite(features) + await assert.rejects(() => fsWrite.invoke(params), /no such file or directory/) + }) + }) +}) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.ts new file mode 100644 index 0000000000..cade189e53 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/fsWrite.ts @@ -0,0 +1,182 @@ +import { InvokeOutput } from './toolShared' +import { Features } from '@aws/language-server-runtimes/server-interface/server' +import { sanitize } from '@aws/lsp-core/out/util/path' + +// Port of https://github.com/aws/aws-toolkit-vscode/blob/10bb1c7dc45f128df14d749d95905c0e9b808096/packages/core/src/codewhispererChat/tools/fsWrite.ts#L42 + +interface BaseParams { + path: string +} + +export interface CreateParams extends BaseParams { + command: 'create' + fileText?: string + newStr?: string +} + +export interface StrReplaceParams extends BaseParams { + command: 'strReplace' + oldStr: string + newStr: string +} + +export interface InsertParams extends BaseParams { + command: 'insert' + insertLine: number + newStr: string +} + +export interface AppendParams extends BaseParams { + command: 'append' + newStr: string +} + +export type FsWriteParams = CreateParams | StrReplaceParams | InsertParams | AppendParams + +export class FsWrite { + private readonly logging: Features['logging'] + private readonly workspace: Features['workspace'] + + constructor(features: Pick & Partial) { + this.logging = features.logging + this.workspace = features.workspace + } + + public async invoke(params: FsWriteParams): Promise { + const sanitizedPath = sanitize(params.path) + + switch (params.command) { + case 'create': + await this.handleCreate(params, sanitizedPath) + break + case 'strReplace': + await this.handleStrReplace(params, sanitizedPath) + break + case 'insert': + await this.handleInsert(params, sanitizedPath) + break + case 'append': + await this.handleAppend(params, sanitizedPath) + break + } + + return { + output: { + kind: 'text', + content: '', + }, + } + } + + private async handleCreate(params: CreateParams, sanitizedPath: string): Promise { + const content = this.getCreateCommandText(params) + + await this.workspace.fs.writeFile(sanitizedPath, content) + } + + private async handleStrReplace(params: StrReplaceParams, sanitizedPath: string): Promise { + const fileContent = await this.workspace.fs.readFile(sanitizedPath) + + const matches = [...fileContent.matchAll(new RegExp(this.escapeRegExp(params.oldStr), 'g'))] + + if (matches.length === 0) { + throw new Error(`No occurrences of "${params.oldStr}" were found`) + } + if (matches.length > 1) { + throw new Error(`${matches.length} occurrences of oldStr were found when only 1 is expected`) + } + + const newContent = fileContent.replace(params.oldStr, params.newStr) + await this.workspace.fs.writeFile(sanitizedPath, newContent) + } + + private async handleInsert(params: InsertParams, sanitizedPath: string): Promise { + const fileContent = await this.workspace.fs.readFile(sanitizedPath) + const lines = fileContent.split('\n') + + const numLines = lines.length + const insertLine = Math.max(0, Math.min(params.insertLine, numLines)) + + let newContent: string + if (insertLine === 0) { + newContent = params.newStr + '\n' + fileContent + } else { + newContent = [...lines.slice(0, insertLine), params.newStr, ...lines.slice(insertLine)].join('\n') + } + + await this.workspace.fs.writeFile(sanitizedPath, newContent) + } + + private async handleAppend(params: AppendParams, sanitizedPath: string): Promise { + const fileContent = await this.workspace.fs.readFile(sanitizedPath) + const needsNewline = fileContent.length !== 0 && !fileContent.endsWith('\n') + + let contentToAppend = params.newStr + if (needsNewline) { + contentToAppend = '\n' + contentToAppend + } + + const newContent = fileContent + contentToAppend + await this.workspace.fs.writeFile(sanitizedPath, newContent) + } + + private getCreateCommandText(params: CreateParams): string { + if (params.fileText) { + return params.fileText + } + if (params.newStr) { + this.logging.warn('Required field `fileText` is missing, use the provided `newStr` instead') + return params.newStr + } + this.logging.warn('No content provided for the create command') + return '' + } + + private escapeRegExp(string: string): string { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') + } + + public getSpec() { + return { + name: 'fsWrite', + description: + 'A tool for creating and editing a file.\n * The `create` command will override the file at `path` if it already exists as a file, and otherwise create a new file\n * The `append` command will add content to the end of an existing file, automatically adding a newline if the file does not end with one. The file must exist.\n Notes for using the `strReplace` command:\n * The `oldStr` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces!\n * If the `oldStr` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `oldStr` to make it unique\n * The `newStr` parameter should contain the edited lines that should replace the `oldStr`. The `insert` command will insert `newStr` after `insertLine` and place it on its own line.', + inputSchema: { + type: 'object', + parameters: { + command: { + type: 'string', + enum: ['create', 'strReplace', 'insert', 'append'], + description: + 'The commands to run. Allowed options are: `create`, `strReplace`, `insert`, `append`.', + }, + fileText: { + description: + 'Required parameter of `create` command, with the content of the file to be created.', + type: 'string', + }, + insertLine: { + description: + 'Required parameter of `insert` command. The `newStr` will be inserted AFTER the line `insertLine` of `path`.', + type: 'integer', + }, + newStr: { + description: + 'Required parameter of `strReplace` command containing the new string. Required parameter of `insert` command containing the string to insert. Required parameter of `append` command containing the content to append to the file.', + type: 'string', + }, + oldStr: { + description: + 'Required parameter of `strReplace` command containing the string in `path` to replace.', + type: 'string', + }, + path: { + description: 'Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`.', + type: 'string', + }, + }, + required: ['command', 'path'], + }, + } as const + } +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts new file mode 100644 index 0000000000..610e4b8123 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts @@ -0,0 +1,14 @@ +import { Server } from '@aws/language-server-runtimes/server-interface' +import { FsRead, FsReadParams } from './fsRead' +import { FsWrite, FsWriteParams } from './fsWrite' + +export const FsToolsServer: Server = ({ workspace, logging, agent }) => { + const fsReadTool = new FsRead({ workspace, logging }) + const fsWriteTool = new FsWrite({ workspace, logging }) + + agent.addTool(fsReadTool.getSpec(), (input: FsReadParams) => fsReadTool.invoke(input)) + + agent.addTool(fsWriteTool.getSpec(), (input: FsWriteParams) => fsWriteTool.invoke(input)) + + return () => {} +} diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolShared.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolShared.ts new file mode 100644 index 0000000000..aa5733a912 --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolShared.ts @@ -0,0 +1,8 @@ +export const maxToolResponseSize = 30720 // 30KB + +export interface InvokeOutput { + output: { + kind: 'text' | 'json' + content: string + } +}