From 3289322734ac01f329080ee4d4b2d5a39bfbc355 Mon Sep 17 00:00:00 2001 From: Roomote Bot Date: Tue, 8 Jul 2025 13:49:25 -0700 Subject: [PATCH] fix: enforce file restrictions for all edit tools in architect mode (#5445) (#5447) --- src/shared/__tests__/modes.spec.ts | 128 +++++++++++++++++++++++++++++ src/shared/modes.ts | 56 +++++++++++-- 2 files changed, 176 insertions(+), 8 deletions(-) diff --git a/src/shared/__tests__/modes.spec.ts b/src/shared/__tests__/modes.spec.ts index 6166c2445c..52e4424d8d 100644 --- a/src/shared/__tests__/modes.spec.ts +++ b/src/shared/__tests__/modes.spec.ts @@ -245,6 +245,112 @@ describe("isToolAllowedForMode", () => { expect(isToolAllowedForMode("browser_action", "architect", [])).toBe(true) expect(isToolAllowedForMode("use_mcp_tool", "architect", [])).toBe(true) }) + + it("applies restrictions to all edit tools including search_and_replace and insert_content", () => { + // Test search_and_replace with matching file + expect( + isToolAllowedForMode("search_and_replace", "architect", [], undefined, { + path: "test.md", + search: "old text", + replace: "new text", + }), + ).toBe(true) + + // Test insert_content with matching file + expect( + isToolAllowedForMode("insert_content", "architect", [], undefined, { + path: "test.md", + line: "1", + content: "# New content", + }), + ).toBe(true) + + // Test search_and_replace with non-matching file - should throw error + expect(() => + isToolAllowedForMode("search_and_replace", "architect", [], undefined, { + path: "test.py", + search: "old text", + replace: "new text", + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("search_and_replace", "architect", [], undefined, { + path: "test.py", + search: "old text", + replace: "new text", + }), + ).toThrow(/Markdown files only/) + + // Test insert_content with non-matching file - should throw error + expect(() => + isToolAllowedForMode("insert_content", "architect", [], undefined, { + path: "test.py", + line: "1", + content: "print('hello')", + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("insert_content", "architect", [], undefined, { + path: "test.py", + line: "1", + content: "print('hello')", + }), + ).toThrow(/Markdown files only/) + }) + + it("applies restrictions to apply_diff with concurrent file edits (MULTI_FILE_APPLY_DIFF experiment)", () => { + // Test apply_diff with args parameter (used when MULTI_FILE_APPLY_DIFF experiment is enabled) + // This simulates concurrent/batch file editing + const xmlArgs = + "test.md- old content\\n+ new content" + + // Should allow markdown files in architect mode + expect( + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgs, + }), + ).toBe(true) + + // Test with non-markdown file - should throw error + const xmlArgsNonMd = + "test.py- old content\\n+ new content" + + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgsNonMd, + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgsNonMd, + }), + ).toThrow(/Markdown files only/) + + // Test with multiple files - should allow only markdown files + const xmlArgsMultiple = + "readme.md- old content\\n+ new contentdocs.md- old content\\n+ new content" + + expect( + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgsMultiple, + }), + ).toBe(true) + + // Test with mixed file types - should throw error for non-markdown + const xmlArgsMixed = + "readme.md- old content\\n+ new contentscript.py- old content\\n+ new content" + + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgsMixed, + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, { + args: xmlArgsMixed, + }), + ).toThrow(/Markdown files only/) + }) }) it("handles non-existent modes", () => { @@ -269,6 +375,14 @@ describe("FileRestrictionError", () => { expect(error.name).toBe("FileRestrictionError") }) + it("formats error message with tool name when provided", () => { + const error = new FileRestrictionError("Markdown Editor", "\\.md$", undefined, "test.js", "write_to_file") + expect(error.message).toBe( + "Tool 'write_to_file' in mode 'Markdown Editor' can only edit files matching pattern: \\.md$. Got: test.js", + ) + expect(error.name).toBe("FileRestrictionError") + }) + describe("debug mode", () => { it("is configured correctly", () => { const debugMode = modes.find((mode) => mode.slug === "debug") @@ -368,6 +482,20 @@ describe("FileRestrictionError", () => { ) expect(error.name).toBe("FileRestrictionError") }) + + it("formats error message with both tool name and description when provided", () => { + const error = new FileRestrictionError( + "Markdown Editor", + "\\.md$", + "Markdown files only", + "test.js", + "apply_diff", + ) + expect(error.message).toBe( + "Tool 'apply_diff' in mode 'Markdown Editor' can only edit files matching pattern: \\.md$ (Markdown files only). Got: test.js", + ) + expect(error.name).toBe("FileRestrictionError") + }) }) describe("getModeSelection", () => { diff --git a/src/shared/modes.ts b/src/shared/modes.ts index f876aa56fa..4666a85d6d 100644 --- a/src/shared/modes.ts +++ b/src/shared/modes.ts @@ -209,11 +209,15 @@ export function getModeSelection(mode: string, promptComponent?: PromptComponent } } +// Edit operation parameters that indicate an actual edit operation +const EDIT_OPERATION_PARAMS = ["diff", "content", "operations", "search", "replace", "args", "line"] as const + // Custom error class for file restrictions export class FileRestrictionError extends Error { - constructor(mode: string, pattern: string, description: string | undefined, filePath: string) { + constructor(mode: string, pattern: string, description: string | undefined, filePath: string, tool?: string) { + const toolInfo = tool ? `Tool '${tool}' in mode '${mode}'` : `This mode (${mode})` super( - `This mode (${mode}) can only edit files matching pattern: ${pattern}${description ? ` (${description})` : ""}. Got: ${filePath}`, + `${toolInfo} can only edit files matching pattern: ${pattern}${description ? ` (${description})` : ""}. Got: ${filePath}`, ) this.name = "FileRestrictionError" } @@ -272,12 +276,48 @@ export function isToolAllowedForMode( // For the edit group, check file regex if specified if (groupName === "edit" && options.fileRegex) { const filePath = toolParams?.path - if ( - filePath && - (toolParams.diff || toolParams.content || toolParams.operations) && - !doesFileMatchRegex(filePath, options.fileRegex) - ) { - throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath) + // Check if this is an actual edit operation (not just path-only for streaming) + const isEditOperation = EDIT_OPERATION_PARAMS.some((param) => toolParams?.[param]) + + // Handle single file path validation + if (filePath && isEditOperation && !doesFileMatchRegex(filePath, options.fileRegex)) { + throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath, tool) + } + + // Handle XML args parameter (used by MULTI_FILE_APPLY_DIFF experiment) + if (toolParams?.args && typeof toolParams.args === "string") { + // Extract file paths from XML args with improved validation + try { + const filePathMatches = toolParams.args.match(/([^<]+)<\/path>/g) + if (filePathMatches) { + for (const match of filePathMatches) { + // More robust path extraction with validation + const pathMatch = match.match(/([^<]+)<\/path>/) + if (pathMatch && pathMatch[1]) { + const extractedPath = pathMatch[1].trim() + // Validate that the path is not empty and doesn't contain invalid characters + if (extractedPath && !extractedPath.includes("<") && !extractedPath.includes(">")) { + if (!doesFileMatchRegex(extractedPath, options.fileRegex)) { + throw new FileRestrictionError( + mode.name, + options.fileRegex, + options.description, + extractedPath, + tool, + ) + } + } + } + } + } + } catch (error) { + // Re-throw FileRestrictionError as it's an expected validation error + if (error instanceof FileRestrictionError) { + throw error + } + // If XML parsing fails, log the error but don't block the operation + console.warn(`Failed to parse XML args for file restriction validation: ${error}`) + } } }