Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 106 additions & 0 deletions src/shared/__tests__/modes.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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", () => {
Expand Down
36 changes: 31 additions & 5 deletions src/shared/modes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -266,13 +266,39 @@ 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)
) {
// Check if this is an actual edit operation (not just path-only for streaming)
const isEditOperation = !!(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For better maintainability, would it be helpful to extract these edit operation indicators into a constant? Something like:

const EDIT_OPERATION_PARAMS = ['diff', 'content', 'operations', 'search', 'replace', 'args', 'line'] as const;

const isEditOperation = EDIT_OPERATION_PARAMS.some(param => toolParams?.[param]);

This would make it easier to add new edit-related parameters in the future and keep the logic centralized.

toolParams?.diff ||
toolParams?.content ||
toolParams?.operations ||
toolParams?.search ||
toolParams?.replace ||
toolParams?.args
)

// Handle single file path validation
if (filePath && isEditOperation && !doesFileMatchRegex(filePath, options.fileRegex)) {
throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error message could be more specific by including which tool triggered the restriction. Currently it shows the mode name but not the tool. Would this be more helpful?

throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath, tool)

Then update the error message to include something like: Tool '${tool}' in mode '${mode}' can only edit files...

}

// Handle XML args parameter (used by MULTI_FILE_APPLY_DIFF experiment)
if (toolParams?.args && typeof toolParams.args === "string") {
// Extract file paths from XML args
const filePathMatches = toolParams.args.match(/<path>([^<]+)<\/path>/g)
if (filePathMatches) {
for (const match of filePathMatches) {
const extractedPath = match.replace(/<\/?path>/g, "")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The XML parsing logic for extracting file paths from the args parameter looks functional, but could we make it more robust? Currently it uses a simple regex that might miss edge cases like:

  • Paths with special XML characters that need escaping
  • Malformed XML with nested or unclosed tags

Would it make sense to use a proper XML parser here, or at least add some validation?

if (!doesFileMatchRegex(extractedPath, options.fileRegex)) {
throw new FileRestrictionError(
mode.name,
options.fileRegex,
options.description,
extractedPath,
)
}
}
}
}
}

return true
Expand Down