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}`)
+ }
}
}