fix: enforce file restrictions for all edit tools in architect mode (#5445) (#5447)

Author: roomote

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 =
+				"<args><file><path>test.md</path><diff><content>- old content\\n+ new content</content></diff></file></args>"
+
+			// 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 =
+				"<args><file><path>test.py</path><diff><content>- old content\\n+ new content</content></diff></file></args>"
+
+			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 =
+				"<args><file><path>readme.md</path><diff><content>- old content\\n+ new content</content></diff></file><file><path>docs.md</path><diff><content>- old content\\n+ new content</content></diff></file></args>"
+
+			expect(
+				isToolAllowedForMode("apply_diff", "architect", [], undefined, {
+					args: xmlArgsMultiple,
+				}),
+			).toBe(true)
+
+			// Test with mixed file types - should throw error for non-markdown
+			const xmlArgsMixed =
+				"<args><file><path>readme.md</path><diff><content>- old content\\n+ new content</content></diff></file><file><path>script.py</path><diff><content>- old content\\n+ new content</content></diff></file></args>"
+
+			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", () => {

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>([^<]+)<\/path>/g)
+					if (filePathMatches) {
+						for (const match of filePathMatches) {
+							// More robust path extraction with validation
+							const pathMatch = match.match(/<path>([^<]+)<\/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}`)
+				}
 			}
 		}