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

This commit is contained in:
Roomote Bot 2025-07-08 13:49:25 -07:00 committed by GitHub
parent 26c2fbbc07
commit 3289322734
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 176 additions and 8 deletions

View file

@ -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", () => {

View file

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