From 5f8a8887fb0825060d47a8ef7a2d15822100e2f5 Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Fri, 24 Jan 2025 10:03:35 -0500 Subject: [PATCH] Validation fixes --- src/core/Cline.ts | 30 ++-- src/core/config/CustomModesSchema.ts | 2 +- .../__tests__/CustomModesSchema.test.ts | 2 +- src/core/mode-validator.ts | 5 +- src/shared/__tests__/modes.test.ts | 153 +++++++++--------- src/shared/modes.ts | 14 +- 6 files changed, 108 insertions(+), 98 deletions(-) diff --git a/src/core/Cline.ts b/src/core/Cline.ts index 0e2872b..9e569ba 100644 --- a/src/core/Cline.ts +++ b/src/core/Cline.ts @@ -1142,17 +1142,25 @@ export class Cline { await this.browserSession.closeBrowser() } - // Validate tool use before execution - const { mode } = (await this.providerRef.deref()?.getState()) ?? {} - const { customModes } = (await this.providerRef.deref()?.getState()) ?? {} - try { - validateToolUse(block.name as ToolName, mode ?? defaultModeSlug, customModes ?? [], { - apply_diff: this.diffEnabled, - }) - } catch (error) { - this.consecutiveMistakeCount++ - pushToolResult(formatResponse.toolError(error.message)) - break + // Only validate complete tool uses + if (!block.partial) { + const { mode } = (await this.providerRef.deref()?.getState()) ?? {} + const { customModes } = (await this.providerRef.deref()?.getState()) ?? {} + try { + validateToolUse( + block.name as ToolName, + mode ?? defaultModeSlug, + customModes ?? [], + { + apply_diff: this.diffEnabled, + }, + block.params, + ) + } catch (error) { + this.consecutiveMistakeCount++ + pushToolResult(formatResponse.toolError(error.message)) + break + } } switch (block.name) { diff --git a/src/core/config/CustomModesSchema.ts b/src/core/config/CustomModesSchema.ts index 0d2d14b..29f5f3e 100644 --- a/src/core/config/CustomModesSchema.ts +++ b/src/core/config/CustomModesSchema.ts @@ -22,7 +22,7 @@ const GroupOptionsSchema = z.object({ }, { message: "Invalid regular expression pattern" }, ), - fileRegexDescription: z.string().optional(), + description: z.string().optional(), }) // Schema for a group entry - either a tool group string or a tuple of [group, options] diff --git a/src/core/config/__tests__/CustomModesSchema.test.ts b/src/core/config/__tests__/CustomModesSchema.test.ts index 3a9d7a1..9af3ba7 100644 --- a/src/core/config/__tests__/CustomModesSchema.test.ts +++ b/src/core/config/__tests__/CustomModesSchema.test.ts @@ -135,7 +135,7 @@ describe("CustomModeSchema", () => { roleDefinition: "Documentation editing mode", groups: [ "read", - ["edit", { fileRegex: "\\.(md|txt)$", fileRegexDescription: "Documentation files only" }], + ["edit", { fileRegex: "\\.(md|txt)$", description: "Documentation files only" }], "browser", ], } diff --git a/src/core/mode-validator.ts b/src/core/mode-validator.ts index 4432997..509496f 100644 --- a/src/core/mode-validator.ts +++ b/src/core/mode-validator.ts @@ -1,4 +1,4 @@ -import { Mode, isToolAllowedForMode, getModeConfig, ModeConfig } from "../shared/modes" +import { Mode, isToolAllowedForMode, getModeConfig, ModeConfig, FileRestrictionError } from "../shared/modes" import { ToolName } from "../shared/tool-groups" export { isToolAllowedForMode } @@ -9,8 +9,9 @@ export function validateToolUse( mode: Mode, customModes?: ModeConfig[], toolRequirements?: Record, + toolParams?: Record, ): void { - if (!isToolAllowedForMode(toolName, mode, customModes ?? [], toolRequirements)) { + if (!isToolAllowedForMode(toolName, mode, customModes ?? [], toolRequirements, toolParams)) { throw new Error(`Tool "${toolName}" is not allowed in ${mode} mode.`) } } diff --git a/src/shared/__tests__/modes.test.ts b/src/shared/__tests__/modes.test.ts index 56e0fe8..ac9d1a5 100644 --- a/src/shared/__tests__/modes.test.ts +++ b/src/shared/__tests__/modes.test.ts @@ -29,59 +29,57 @@ describe("isToolAllowedForMode", () => { describe("file restrictions", () => { it("allows editing matching files", () => { // Test markdown editor mode - const mdResult = isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, "test.md") + const mdResult = isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { + path: "test.md", + }) expect(mdResult).toBe(true) // Test CSS editor mode - const cssResult = isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, "styles.css") + const cssResult = isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, { + path: "styles.css", + }) expect(cssResult).toBe(true) }) it("rejects editing non-matching files", () => { // Test markdown editor mode with non-markdown file - const mdError = isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, "test.js") - expect(mdError).toBeInstanceOf(FileRestrictionError) - expect((mdError as FileRestrictionError).message).toContain("\\.md$") + expect(() => + isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(/\\.md\$/) // Test CSS editor mode with non-CSS file - const cssError = isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, "test.js") - expect(cssError).toBeInstanceOf(FileRestrictionError) - expect((cssError as FileRestrictionError).message).toContain("\\.css$") - }) - - it("requires file path for restricted edit operations", () => { - const result = isToolAllowedForMode("write_to_file", "markdown-editor", customModes) - expect(result).toBeInstanceOf(FileRestrictionError) - expect((result as FileRestrictionError).message).toContain("\\.md$") + expect(() => + isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(/\\.css\$/) }) it("applies restrictions to both write_to_file and apply_diff", () => { // Test write_to_file - const writeResult = isToolAllowedForMode( - "write_to_file", - "markdown-editor", - customModes, - undefined, - "test.md", - ) + const writeResult = isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { + path: "test.md", + }) expect(writeResult).toBe(true) // Test apply_diff - const diffResult = isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, "test.md") + const diffResult = isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, { + path: "test.md", + }) expect(diffResult).toBe(true) // Test both with non-matching file - const writeError = isToolAllowedForMode( - "write_to_file", - "markdown-editor", - customModes, - undefined, - "test.js", - ) - expect(writeError).toBeInstanceOf(FileRestrictionError) + expect(() => + isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(FileRestrictionError) - const diffError = isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, "test.js") - expect(diffError).toBeInstanceOf(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, { path: "test.js" }), + ).toThrow(FileRestrictionError) }) it("uses description in file restriction error for custom modes", () => { @@ -99,60 +97,57 @@ describe("isToolAllowedForMode", () => { ] // Test write_to_file with non-matching file - const writeError = isToolAllowedForMode( - "write_to_file", - "docs-editor", - customModesWithDescription, - undefined, - "test.js", - ) - expect(writeError).toBeInstanceOf(FileRestrictionError) - expect((writeError as FileRestrictionError).message).toContain("Documentation files only") + expect(() => + isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, { + path: "test.js", + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, { + path: "test.js", + }), + ).toThrow(/Documentation files only/) // Test apply_diff with non-matching file - const diffError = isToolAllowedForMode( - "apply_diff", - "docs-editor", - customModesWithDescription, - undefined, - "test.js", - ) - expect(diffError).toBeInstanceOf(FileRestrictionError) - expect((diffError as FileRestrictionError).message).toContain("Documentation files only") + expect(() => + isToolAllowedForMode("apply_diff", "docs-editor", customModesWithDescription, undefined, { + path: "test.js", + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "docs-editor", customModesWithDescription, undefined, { + path: "test.js", + }), + ).toThrow(/Documentation files only/) // Test that matching files are allowed - const mdResult = isToolAllowedForMode( - "write_to_file", - "docs-editor", - customModesWithDescription, - undefined, - "test.md", - ) - expect(mdResult).toBe(true) + expect( + isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, { + path: "test.md", + }), + ).toBe(true) - const txtResult = isToolAllowedForMode( - "write_to_file", - "docs-editor", - customModesWithDescription, - undefined, - "test.txt", - ) - expect(txtResult).toBe(true) + expect( + isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, { + path: "test.txt", + }), + ).toBe(true) }) it("allows ask mode to edit markdown files only", () => { // Should allow editing markdown files - const mdResult = isToolAllowedForMode("write_to_file", "ask", [], undefined, "test.md") - expect(mdResult).toBe(true) + expect(isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.md" })).toBe(true) // Should allow applying diffs to markdown files - const diffResult = isToolAllowedForMode("apply_diff", "ask", [], undefined, "readme.md") - expect(diffResult).toBe(true) + expect(isToolAllowedForMode("apply_diff", "ask", [], undefined, { path: "readme.md" })).toBe(true) // Should reject non-markdown files - const jsResult = isToolAllowedForMode("write_to_file", "ask", [], undefined, "test.js") - expect(jsResult).toBeInstanceOf(FileRestrictionError) - expect((jsResult as FileRestrictionError).message).toContain("Markdown files only") + expect(() => isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.js" })).toThrow( + FileRestrictionError, + ) + expect(() => isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.js" })).toThrow( + /Markdown files only/, + ) // Should maintain read capabilities expect(isToolAllowedForMode("read_file", "ask", [])).toBe(true) @@ -176,14 +171,18 @@ describe("isToolAllowedForMode", () => { describe("FileRestrictionError", () => { it("formats error message with pattern when no description provided", () => { - const error = new FileRestrictionError("Markdown Editor", "\\.md$") - expect(error.message).toBe("This mode (Markdown Editor) can only edit files matching the pattern: \\.md$") + const error = new FileRestrictionError("Markdown Editor", "\\.md$", undefined, "test.js") + expect(error.message).toBe( + "This mode (Markdown Editor) can only edit files matching pattern: \\.md$. Got: test.js", + ) expect(error.name).toBe("FileRestrictionError") }) it("formats error message with description when provided", () => { - const error = new FileRestrictionError("Markdown Editor", "\\.md$", "Markdown files only") - expect(error.message).toBe("This mode (Markdown Editor) can only edit files matching Markdown files only") + const error = new FileRestrictionError("Markdown Editor", "\\.md$", "Markdown files only", "test.js") + expect(error.message).toBe( + "This mode (Markdown Editor) can only edit files matching pattern: \\.md$ (Markdown files only). Got: test.js", + ) expect(error.name).toBe("FileRestrictionError") }) }) diff --git a/src/shared/modes.ts b/src/shared/modes.ts index 2064771..c29800c 100644 --- a/src/shared/modes.ts +++ b/src/shared/modes.ts @@ -146,9 +146,9 @@ export function isCustomMode(slug: string, customModes?: ModeConfig[]): boolean // Custom error class for file restrictions export class FileRestrictionError extends Error { - constructor(mode: string, pattern: string, description?: string) { + constructor(mode: string, pattern: string, description: string | undefined, filePath: string) { super( - `This mode (${mode}) can only edit files matching ${description ? description : `the pattern: ${pattern}`}`, + `This mode (${mode}) can only edit files matching pattern: ${pattern}${description ? ` (${description})` : ""}. Got: ${filePath}`, ) this.name = "FileRestrictionError" } @@ -159,8 +159,8 @@ export function isToolAllowedForMode( modeSlug: string, customModes: ModeConfig[], toolRequirements?: Record, - filePath?: string, // Optional file path for checking regex patterns -): boolean | FileRestrictionError { + toolParams?: Record, // All tool parameters +): boolean { // Always allow these tools if (ALWAYS_AVAILABLE_TOOLS.includes(tool as any)) { return true @@ -195,8 +195,10 @@ export function isToolAllowedForMode( // For the edit group, check file regex if specified if (groupName === "edit" && options.fileRegex) { - if (!filePath || !doesFileMatchRegex(filePath, options.fileRegex)) { - return new FileRestrictionError(mode.name, options.fileRegex, options.description) + const filePath = toolParams?.path + // Only validate regex if a path is provided + if (filePath && !doesFileMatchRegex(filePath, options.fileRegex)) { + throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath) } return true }