fix: enforce mode file restrictions for apply_diff tool (#4732)

src/shared/__tests__/modes.test.ts

@@ -94,56 +94,99 @@ describe("isToolAllowedForMode", () => {
 		})
 
 		it("handles partial streaming cases (path only, no content/diff)", () => {
-			// Should allow path-only for matching files (no validation yet since content/diff not provided)
-			expect(
+			// Should enforce file restrictions even when only path is provided (streaming scenario)
+			// This prevents bypassing restrictions when parameters arrive separately
+
+			// Non-matching files should be rejected immediately when path is provided
+			expect(() =>
 				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
 					path: "test.js",
 				}),
+			).toThrow(FileRestrictionError)
+
+			expect(() =>
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+				}),
+			).toThrow(FileRestrictionError)
+
+			// Matching files should be allowed even with path-only
+			expect(
+				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
+					path: "test.md",
+				}),
 			).toBe(true)
 
 			expect(
 				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
-					path: "test.js",
+					path: "test.md",
 				}),
 			).toBe(true)
 
-			// Should allow path-only for architect mode too
+			// Architect mode (built-in) allows markdown files
 			expect(
 				isToolAllowedForMode("write_to_file", "architect", [], undefined, {
-					path: "test.js",
+					path: "test.md",
 				}),
 			).toBe(true)
 		})
 
-		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, {
-				path: "test.md",
-				content: "# Test",
+		it("applies restrictions to all edit tools", () => {
+			// Test all edit tools with matching files
+			const editTools = ["write_to_file", "apply_diff", "insert_content", "search_and_replace"]
+
+			editTools.forEach((tool) => {
+				// Matching file should be allowed
+				const params: any = { path: "test.md" }
+				if (tool === "write_to_file") params.content = "# Test"
+				if (tool === "apply_diff") params.diff = "- old\n+ new"
+				if (tool === "insert_content") params.content = "new line"
+				if (tool === "search_and_replace") {
+					params.search = "old"
+					params.replace = "new"
+				}
+
+				expect(isToolAllowedForMode(tool as any, "markdown-editor", customModes, undefined, params)).toBe(true)
+
+				// Non-matching file should be rejected
+				params.path = "test.js"
+				expect(() =>
+					isToolAllowedForMode(tool as any, "markdown-editor", customModes, undefined, params),
+				).toThrow(FileRestrictionError)
 			})
-			expect(writeResult).toBe(true)
+		})
 
-			// Test apply_diff
-			const diffResult = isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
-				path: "test.md",
-				diff: "- old\n+ new",
-			})
-			expect(diffResult).toBe(true)
+		it("prevents bypassing restrictions via parameter streaming", () => {
+			// This test specifically addresses the bug where apply_diff could bypass
+			// file restrictions when path arrives before diff content
 
-			// Test both with non-matching file
+			// Step 1: Path arrives first (should be rejected for non-matching files)
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
-					path: "test.js",
-					content: "console.log('test')",
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					path: "test.js", // Non-matching file
 				}),
 			).toThrow(FileRestrictionError)
 
+			// Step 2: Even when diff arrives later, it should still be rejected
 			expect(() =>
 				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
 					path: "test.js",
 					diff: "- old\n+ new",
 				}),
 			).toThrow(FileRestrictionError)
+
+			// Same for other edit tools
+			expect(() =>
+				isToolAllowedForMode("insert_content", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+				}),
+			).toThrow(FileRestrictionError)
+
+			expect(() =>
+				isToolAllowedForMode("search_and_replace", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+				}),
+			).toThrow(FileRestrictionError)
 		})
 
 		it("uses description in file restriction error for custom modes", () => {
@@ -203,11 +246,18 @@ describe("isToolAllowedForMode", () => {
 				}),
 			).toBe(true)
 
-			// Test partial streaming cases
-			expect(
+			// Test partial streaming cases - should now reject non-matching files immediately
+			expect(() =>
 				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.js",
 				}),
+			).toThrow(FileRestrictionError)
+
+			// But should allow matching files even with path-only
+			expect(
+				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
+					path: "test.md",
+				}),
 			).toBe(true)
 		})
 
@@ -247,6 +297,72 @@ describe("isToolAllowedForMode", () => {
 			expect(isToolAllowedForMode("browser_action", "architect", [])).toBe(true)
 			expect(isToolAllowedForMode("use_mcp_tool", "architect", [])).toBe(true)
 		})
+		it("validates paths in multi-file apply_diff args parameter", () => {
+			// Test multi-file apply_diff with XML args containing non-matching files
+			const argsWithNonMatchingFile = `<args>
+  <file>
+    <path>test.js</path>
+    <diff>
+      <content>- old
++ new</content>
+    </diff>
+  </file>
+</args>`
+
+			expect(() =>
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					args: argsWithNonMatchingFile,
+				}),
+			).toThrow(FileRestrictionError)
+
+			// Test with multiple files, one non-matching
+			const argsWithMixedFiles = `<args>
+  <file>
+    <path>doc.md</path>
+    <diff>
+      <content>- old
++ new</content>
+    </diff>
+  </file>
+  <file>
+    <path>script.js</path>
+    <diff>
+      <content>- old
++ new</content>
+    </diff>
+  </file>
+</args>`
+
+			expect(() =>
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					args: argsWithMixedFiles,
+				}),
+			).toThrow(FileRestrictionError)
+
+			// Test with only matching files (should pass)
+			const argsWithMatchingFiles = `<args>
+  <file>
+    <path>doc1.md</path>
+    <diff>
+      <content>- old
++ new</content>
+    </diff>
+  </file>
+  <file>
+    <path>doc2.md</path>
+    <diff>
+      <content>- old
++ new</content>
+    </diff>
+  </file>
+</args>`
+
+			expect(
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					args: argsWithMatchingFiles,
+				}),
+			).toBe(true)
+		})
 	})
 
 	it("handles non-existent modes", () => {

src/shared/modes.ts

@@ -257,14 +257,23 @@ export function isToolAllowedForMode(
 
 		// For the edit group, check file regex if specified
 		if (groupName === "edit" && options.fileRegex) {
+			// Check direct path parameter
 			const filePath = toolParams?.path
-			if (
-				filePath &&
-				(toolParams.diff || toolParams.content || toolParams.operations) &&
-				!doesFileMatchRegex(filePath, options.fileRegex)
-			) {
+			if (filePath && !doesFileMatchRegex(filePath, options.fileRegex)) {
 				throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath)
 			}
+
+			// For apply_diff with multi-file support, also check paths in args parameter
+			if (tool === "apply_diff" && toolParams?.args && typeof toolParams.args === "string") {
+				// Extract all file paths from the XML args
+				const pathMatches = toolParams.args.matchAll(/<path>([^<]+)<\/path>/g)
+				for (const match of pathMatches) {
+					const xmlPath = match[1]
+					if (xmlPath && !doesFileMatchRegex(xmlPath, options.fileRegex)) {
+						throw new FileRestrictionError(mode.name, options.fileRegex, options.description, xmlPath)
+					}
+				}
+			}
 		}
 
 		return true