diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index ada841c10d..6618439bc6 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -738,6 +738,7 @@ export async function presentAssistantMessage(cline: Task) { block.params, stateExperiments, includedTools, + block.nativeArgs, ) } catch (error) { cline.consecutiveMistakeCount++ diff --git a/src/core/tools/validateToolUse.ts b/src/core/tools/validateToolUse.ts index d057033741..4a784da790 100644 --- a/src/core/tools/validateToolUse.ts +++ b/src/core/tools/validateToolUse.ts @@ -32,6 +32,7 @@ export function validateToolUse( toolParams?: Record, experiments?: Record, includedTools?: string[], + nativeArgs?: Record, ): void { // First, check if the tool name is actually a valid/known tool // This catches completely invalid tool names like "edit_file" that don't exist @@ -51,6 +52,7 @@ export function validateToolUse( toolParams, experiments, includedTools, + nativeArgs, ) ) { throw new Error(`Tool "${toolName}" is not allowed in ${mode} mode.`) @@ -81,6 +83,7 @@ export function isToolAllowedForMode( toolParams?: Record, // All tool parameters experiments?: Record, includedTools?: string[], // Opt-in tools explicitly included (e.g., from modelInfo) + nativeArgs?: Record, // Native protocol arguments (e.g., nativeArgs.files for multi-file apply_diff) ): boolean { // Always allow these tools if (ALWAYS_AVAILABLE_TOOLS.includes(tool as any)) { @@ -188,6 +191,37 @@ export function isToolAllowedForMode( console.warn(`Failed to parse XML args for file restriction validation: ${error}`) } } + + // Handle native protocol multi-file format (nativeArgs.files from multi_apply_diff schema) + if (nativeArgs?.files && Array.isArray(nativeArgs.files)) { + for (const file of nativeArgs.files) { + const filePath = file?.path + if (filePath && typeof filePath === "string") { + if (!doesFileMatchRegex(filePath, options.fileRegex)) { + throw new FileRestrictionError( + mode.name, + options.fileRegex, + options.description, + filePath, + tool, + ) + } + } + } + } + + // Handle native protocol single-file format (nativeArgs.path from apply_diff schema) + if (nativeArgs?.path && typeof nativeArgs.path === "string" && nativeArgs?.diff) { + if (!doesFileMatchRegex(nativeArgs.path, options.fileRegex)) { + throw new FileRestrictionError( + mode.name, + options.fileRegex, + options.description, + nativeArgs.path, + tool, + ) + } + } } return true diff --git a/src/shared/__tests__/modes.spec.ts b/src/shared/__tests__/modes.spec.ts index a00abde787..db42cc2b57 100644 --- a/src/shared/__tests__/modes.spec.ts +++ b/src/shared/__tests__/modes.spec.ts @@ -300,6 +300,191 @@ describe("isToolAllowedForMode", () => { }), ).toThrow(/Markdown files only/) }) + + it("applies restrictions to apply_diff with native protocol multi-file format (nativeArgs.files)", () => { + // Test apply_diff with nativeArgs.files (native protocol multi-file format) + // This simulates the multi_apply_diff schema used with native protocol + + // Should allow markdown files in architect mode + expect( + isToolAllowedForMode( + "apply_diff", + "architect", + [], + undefined, + {}, // toolParams is empty for native protocol + undefined, + undefined, + { files: [{ path: "test.md", diff: "- old\n+ new" }] }, + ), + ).toBe(true) + + // Test with non-markdown file - should throw error + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + files: [{ path: "test.py", diff: "- old\n+ new" }], + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + files: [{ path: "test.py", diff: "- old\n+ new" }], + }), + ).toThrow(/Markdown files only/) + + // Test with multiple markdown files - should allow + expect( + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + files: [ + { path: "readme.md", diff: "- old\n+ new" }, + { path: "docs.md", diff: "- old\n+ new" }, + ], + }), + ).toBe(true) + + // Test with mixed file types - should throw error for non-markdown + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + files: [ + { path: "readme.md", diff: "- old\n+ new" }, + { path: "script.py", diff: "- old\n+ new" }, + ], + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + files: [ + { path: "readme.md", diff: "- old\n+ new" }, + { path: "script.py", diff: "- old\n+ new" }, + ], + }), + ).toThrow(/Markdown files only/) + }) + + it("applies restrictions to apply_diff with native protocol single-file format (nativeArgs.path)", () => { + // Test apply_diff with nativeArgs.path (native protocol single-file format) + // This simulates the apply_diff schema used with native protocol + + // Should allow markdown files in architect mode + expect( + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + path: "test.md", + diff: "- old\n+ new", + }), + ).toBe(true) + + // Test with non-markdown file - should throw error + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + path: "test.py", + diff: "- old\n+ new", + }), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode("apply_diff", "architect", [], undefined, {}, undefined, undefined, { + path: "test.py", + diff: "- old\n+ new", + }), + ).toThrow(/Markdown files only/) + }) + + it("applies native protocol file restrictions to custom modes with fileRegex", () => { + // Test that custom mode file restrictions work with native protocol formats + const customModesWithRegex: ModeConfig[] = [ + { + slug: "ts-editor", + name: "TypeScript Editor", + roleDefinition: "You are a TypeScript editor", + groups: [ + "read", + ["edit", { fileRegex: "\\.tsx?$", description: "TypeScript files only" }], + "browser", + ], + }, + ] + + // Test native multi-file format with valid TS files + expect( + isToolAllowedForMode( + "apply_diff", + "ts-editor", + customModesWithRegex, + undefined, + {}, + undefined, + undefined, + { + files: [ + { path: "app.ts", diff: "- old\n+ new" }, + { path: "component.tsx", diff: "- old\n+ new" }, + ], + }, + ), + ).toBe(true) + + // Test native multi-file format with invalid file + expect(() => + isToolAllowedForMode( + "apply_diff", + "ts-editor", + customModesWithRegex, + undefined, + {}, + undefined, + undefined, + { + files: [ + { path: "app.ts", diff: "- old\n+ new" }, + { path: "styles.css", diff: "- old\n+ new" }, + ], + }, + ), + ).toThrow(FileRestrictionError) + expect(() => + isToolAllowedForMode( + "apply_diff", + "ts-editor", + customModesWithRegex, + undefined, + {}, + undefined, + undefined, + { + files: [ + { path: "app.ts", diff: "- old\n+ new" }, + { path: "styles.css", diff: "- old\n+ new" }, + ], + }, + ), + ).toThrow(/TypeScript files only/) + + // Test native single-file format with valid TS file + expect( + isToolAllowedForMode( + "apply_diff", + "ts-editor", + customModesWithRegex, + undefined, + {}, + undefined, + undefined, + { path: "app.ts", diff: "- old\n+ new" }, + ), + ).toBe(true) + + // Test native single-file format with invalid file + expect(() => + isToolAllowedForMode( + "apply_diff", + "ts-editor", + customModesWithRegex, + undefined, + {}, + undefined, + undefined, + { path: "styles.css", diff: "- old\n+ new" }, + ), + ).toThrow(FileRestrictionError) + }) }) it("handles non-existent modes", () => {