From 62455ade9d609569ffde6015dea0fc501bf960cd Mon Sep 17 00:00:00 2001 From: Taylor Mullen Date: Thu, 15 May 2025 00:53:52 -0700 Subject: [PATCH] Fix(write-file): Ensure correct validation method is called in WriteFileTool - The `WriteFileTool` had a validation method named `validateParams`. - However, its `shouldConfirmExecute` method was attempting to call `this.validateToolParams`, which would have invoked the placeholder implementation from `BaseTool` instead of `WriteFileTool`'s own, more specific validation. - This commit renames `WriteFileTool`'s `validateParams` to `validateToolParams`, correctly overriding the `BaseTool` method. - Internal calls within `WriteFileTool` now correctly use `this.validateToolParams`, ensuring its specific validation logic is used. - Adds tests to verify the validation logic within `WriteFileTool`. Fixes https://b.corp.google.com/issues/417883702 Signed-off and authored by: Gemini "My code may not be perfect, but at least it is not trying to take over the world... yet." --- packages/server/src/tools/write-file.test.ts | 184 +++++++++++++++++++ packages/server/src/tools/write-file.ts | 4 +- 2 files changed, 186 insertions(+), 2 deletions(-) create mode 100644 packages/server/src/tools/write-file.test.ts diff --git a/packages/server/src/tools/write-file.test.ts b/packages/server/src/tools/write-file.test.ts new file mode 100644 index 00000000..2ce5cdf5 --- /dev/null +++ b/packages/server/src/tools/write-file.test.ts @@ -0,0 +1,184 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { WriteFileTool } from './write-file.js'; +import { ToolConfirmationOutcome } from './tools.js'; +import path from 'path'; +import fs from 'fs'; +import os from 'os'; + +describe('WriteFileTool', () => { + let tool: WriteFileTool; + let tempDir: string; + // Using a subdirectory within the OS temp directory for the root to avoid potential permission issues. + const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); + + beforeEach(() => { + // Create a unique temporary directory for files created outside the root (for testing boundary conditions) + tempDir = fs.mkdtempSync( + path.join(os.tmpdir(), 'write-file-test-external-'), + ); + // Ensure the rootDir for the tool exists + if (!fs.existsSync(rootDir)) { + fs.mkdirSync(rootDir, { recursive: true }); + } + tool = new WriteFileTool(rootDir); + }); + + afterEach(() => { + // Clean up the temporary directories + if (fs.existsSync(tempDir)) { + fs.rmSync(tempDir, { recursive: true, force: true }); + } + if (fs.existsSync(rootDir)) { + fs.rmSync(rootDir, { recursive: true, force: true }); + } + }); + + describe('validateToolParams', () => { + it('should return null for valid absolute path within root', () => { + const params = { + file_path: path.join(rootDir, 'test.txt'), + content: 'hello', + }; + expect(tool.validateToolParams(params)).toBeNull(); + }); + + it('should return error for relative path', () => { + const params = { file_path: 'test.txt', content: 'hello' }; + expect(tool.validateToolParams(params)).toMatch( + /File path must be absolute/, + ); + }); + + it('should return error for path outside root', () => { + const outsidePath = path.resolve(tempDir, 'outside-root.txt'); + const params = { + file_path: outsidePath, + content: 'hello', + }; + expect(tool.validateToolParams(params)).toMatch( + /File path must be within the root directory/, + ); + }); + + it('should return null for path that is the root itself', () => { + const params = { + file_path: rootDir, // Attempting to write to the root directory itself (as a file) + content: 'hello', + }; + // This is a tricky case. The validation should allow it if it's treated as a file path. + // The actual write operation might fail if it's a directory, but validation should pass. + expect(tool.validateToolParams(params)).toBeNull(); + }); + + it('should return error for path that is just / and root is not /', () => { + const params = { file_path: path.resolve('/'), content: 'hello' }; + if (rootDir === path.resolve('/')) { + // This case would only occur if the test runner somehow sets rootDir to actual '/', which is highly unlikely and unsafe. + expect(tool.validateToolParams(params)).toBeNull(); + } else { + expect(tool.validateToolParams(params)).toMatch( + /File path must be within the root directory/, + ); + } + }); + }); + + describe('shouldConfirmExecute', () => { + it('should return false if params are invalid (relative path)', async () => { + const params = { file_path: 'relative.txt', content: 'test' }; + const confirmation = await tool.shouldConfirmExecute(params); + expect(confirmation).toBe(false); + }); + + it('should return false if params are invalid (outside root)', async () => { + const outsidePath = path.resolve(tempDir, 'outside-root.txt'); + const params = { file_path: outsidePath, content: 'test' }; + const confirmation = await tool.shouldConfirmExecute(params); + expect(confirmation).toBe(false); + }); + + it('should request confirmation for valid params if file does not exist', async () => { + const filePath = path.join(rootDir, 'new_file.txt'); + const params = { file_path: filePath, content: 'new content' }; + const confirmation = await tool.shouldConfirmExecute(params); + expect(confirmation).toEqual( + expect.objectContaining({ + title: `Confirm Write: ${path.basename(filePath)}`, + fileName: 'new_file.txt', + fileDiff: expect.any(String), + }), + ); + }); + }); + + describe('execute', () => { + it('should return error if params are invalid (relative path)', async () => { + const params = { file_path: 'relative.txt', content: 'test' }; + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch(/Error: Invalid parameters provided/); + expect(result.returnDisplay).toMatch(/Error: File path must be absolute/); + }); + + it('should return error if params are invalid (path outside root)', async () => { + const outsidePath = path.resolve(tempDir, 'outside-root.txt'); + const params = { file_path: outsidePath, content: 'test' }; + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch(/Error: Invalid parameters provided/); + expect(result.returnDisplay).toMatch( + /Error: File path must be within the root directory/, + ); + }); + + it('should write a new file and return diff', async () => { + const filePath = path.join(rootDir, 'execute_new_file.txt'); + const content = 'Hello from execute!'; + const params = { file_path: filePath, content }; + + const confirmDetails = await tool.shouldConfirmExecute(params); + if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); + } + + const result = await tool.execute(params, new AbortController().signal); + + expect(result.llmContent).toMatch( + /Successfully created and wrote to new file/, + ); + expect(fs.existsSync(filePath)).toBe(true); + expect(fs.readFileSync(filePath, 'utf8')).toBe(content); + const display = result.returnDisplay as { fileDiff: string }; // Type assertion + // For new files, the diff will include the filename in the "Original" header + expect(display.fileDiff).toMatch(/--- execute_new_file.txt\tOriginal/); + expect(display.fileDiff).toMatch(/\+\+\+ execute_new_file.txt\tWritten/); + expect(display.fileDiff).toMatch(content); + }); + + it('should overwrite an existing file and return diff', async () => { + const filePath = path.join(rootDir, 'execute_existing_file.txt'); + const initialContent = 'Initial content.'; + const newContent = 'Overwritten content!'; + fs.writeFileSync(filePath, initialContent, 'utf8'); + + const params = { file_path: filePath, content: newContent }; + + const confirmDetails = await tool.shouldConfirmExecute(params); + if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); + } + + const result = await tool.execute(params, new AbortController().signal); + + expect(result.llmContent).toMatch(/Successfully overwrote file/); + expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent); + const display = result.returnDisplay as { fileDiff: string }; // Type assertion + expect(display.fileDiff).toMatch(initialContent); + expect(display.fileDiff).toMatch(newContent); + }); + }); +}); diff --git a/packages/server/src/tools/write-file.ts b/packages/server/src/tools/write-file.ts index 1f4c0d94..44794a2d 100644 --- a/packages/server/src/tools/write-file.ts +++ b/packages/server/src/tools/write-file.ts @@ -76,7 +76,7 @@ export class WriteFileTool extends BaseTool { ); } - validateParams(params: WriteFileToolParams): string | null { + validateToolParams(params: WriteFileToolParams): string | null { if ( this.schema.parameters && !SchemaValidator.validate( @@ -154,7 +154,7 @@ export class WriteFileTool extends BaseTool { params: WriteFileToolParams, _signal: AbortSignal, ): Promise { - const validationError = this.validateParams(params); + const validationError = this.validateToolParams(params); if (validationError) { return { llmContent: `Error: Invalid parameters provided. Reason: ${validationError}`,