Merge pull request aws#6842 from ctlai95/fs-write

feat(chat): add fsWrite tool create command

Author: mr-lee

packages/core/src/codewhispererChat/tools/fsWrite.ts

@@ -0,0 +1,110 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { InvokeOutput, OutputKind, sanitizePath } from './toolShared'
+import { getLogger } from '../../shared/logger/logger'
+import vscode from 'vscode'
+import { fs } from '../../shared/fs/fs'
+
+interface BaseCommand {
+    path: string
+}
+
+export interface CreateCommand extends BaseCommand {
+    command: 'create'
+    fileText?: string
+    newStr?: string
+}
+
+export interface StrReplaceCommand extends BaseCommand {
+    command: 'str_replace'
+    oldStr: string
+    newStr: string
+}
+
+export interface InsertCommand extends BaseCommand {
+    command: 'insert'
+    insertLine: number
+    newStr: string
+}
+
+export interface AppendCommand extends BaseCommand {
+    command: 'append'
+    newStr: string
+}
+
+export type FsWriteCommand = CreateCommand | StrReplaceCommand | InsertCommand | AppendCommand
+
+export class FsWrite {
+    private static readonly logger = getLogger('fsWrite')
+
+    public static async invoke(command: FsWriteCommand): Promise<InvokeOutput> {
+        const sanitizedPath = sanitizePath(command.path)
+
+        switch (command.command) {
+            case 'create':
+                await this.handleCreate(command, sanitizedPath)
+                break
+            case 'str_replace':
+                await this.handleStrReplace(command, sanitizedPath)
+                break
+            case 'insert':
+                await this.handleInsert(command, sanitizedPath)
+                break
+            case 'append':
+                await this.handleAppend(command, sanitizedPath)
+                break
+        }
+
+        return {
+            output: {
+                kind: OutputKind.Text,
+                content: '',
+            },
+        }
+    }
+
+    private static async handleCreate(command: CreateCommand, sanitizedPath: string): Promise<void> {
+        const content = this.getCreateCommandText(command)
+
+        const fileExists = await fs.existsFile(sanitizedPath)
+        const actionType = fileExists ? 'Replacing' : 'Creating'
+
+        await vscode.window.withProgress(
+            {
+                location: vscode.ProgressLocation.Notification,
+                title: `${actionType}: ${sanitizedPath}`,
+                cancellable: false,
+            },
+            async () => {
+                await fs.writeFile(sanitizedPath, content)
+            }
+        )
+    }
+
+    private static async handleStrReplace(command: StrReplaceCommand, sanitizedPath: string): Promise<void> {
+        // TODO
+    }
+
+    private static async handleInsert(command: InsertCommand, sanitizedPath: string): Promise<void> {
+        // TODO
+    }
+
+    private static async handleAppend(command: AppendCommand, sanitizedPath: string): Promise<void> {
+        // TODO
+    }
+
+    private static getCreateCommandText(command: CreateCommand): string {
+        if (command.fileText) {
+            return command.fileText
+        }
+        if (command.newStr) {
+            this.logger.warn('Required field `fileText` is missing, use the provided `newStr` instead')
+            return command.newStr
+        }
+        this.logger.warn('No content provided for the create command')
+        return ''
+    }
+}

packages/core/src/shared/logger/logger.ts

@@ -15,6 +15,7 @@ export type LogTopic =
     | 'chat'
     | 'stepfunctions'
     | 'fsRead'
+    | 'fsWrite'
 
 class ErrorLog {
     constructor(

packages/core/src/test/codewhispererChat/tools/fsWrite.test.ts

@@ -0,0 +1,93 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { CreateCommand, FsWrite } from '../../../codewhispererChat/tools/fsWrite'
+import { TestFolder } from '../../testUtil'
+import path from 'path'
+import assert from 'assert'
+import { fs } from '../../../shared/fs/fs'
+import { InvokeOutput, OutputKind } from '../../../codewhispererChat/tools/toolShared'
+
+describe('FsWrite Tool', function () {
+    let testFolder: TestFolder
+    const expectedOutput: InvokeOutput = {
+        output: {
+            kind: OutputKind.Text,
+            content: '',
+        },
+    }
+
+    before(async function () {
+        testFolder = await TestFolder.create()
+    })
+
+    describe('create', function () {
+        it('creates a new file with fileText content', async function () {
+            const filePath = path.join(testFolder.path, 'file1.txt')
+            const fileExists = await fs.existsFile(filePath)
+            assert.ok(!fileExists)
+
+            const command: CreateCommand = {
+                command: 'create',
+                fileText: 'Hello World',
+                path: filePath,
+            }
+            const output = await FsWrite.invoke(command)
+
+            const content = await fs.readFileText(filePath)
+            assert.strictEqual(content, 'Hello World')
+
+            assert.deepStrictEqual(output, expectedOutput)
+        })
+
+        it('replaces existing file with fileText content', async function () {
+            const filePath = path.join(testFolder.path, 'file1.txt')
+            const fileExists = await fs.existsFile(filePath)
+            assert.ok(fileExists)
+
+            const command: CreateCommand = {
+                command: 'create',
+                fileText: 'Goodbye',
+                path: filePath,
+            }
+            const output = await FsWrite.invoke(command)
+
+            const content = await fs.readFileText(filePath)
+            assert.strictEqual(content, 'Goodbye')
+
+            assert.deepStrictEqual(output, expectedOutput)
+        })
+
+        it('uses newStr when fileText is not provided', async function () {
+            const filePath = path.join(testFolder.path, 'file2.txt')
+
+            const command: CreateCommand = {
+                command: 'create',
+                newStr: 'Hello World',
+                path: filePath,
+            }
+            const output = await FsWrite.invoke(command)
+
+            const content = await fs.readFileText(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(testFolder.path, 'file3.txt')
+
+            const command: CreateCommand = {
+                command: 'create',
+                path: filePath,
+            }
+            const output = await FsWrite.invoke(command)
+
+            const content = await fs.readFileText(filePath)
+            assert.strictEqual(content, '')
+
+            assert.deepStrictEqual(output, expectedOutput)
+        })
+    })
+})