Procházet zdrojové kódy

Throw FileRestrictionErrors sooner

Matt Rubens před 11 měsíci
rodič
revize
5d20477785

+ 5 - 0
.changeset/late-moose-watch.md

@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+Throw errors sooner when a mode tries to write a restricted file

+ 16 - 19
src/core/Cline.ts

@@ -1144,25 +1144,22 @@ export class Cline {
 					await this.browserSession.closeBrowser()
 				}
 
-				// Only validate complete tool uses
-				if (!block.partial) {
-					const { mode } = (await this.providerRef.deref()?.getState()) ?? {}
-					const { customModes } = (await this.providerRef.deref()?.getState()) ?? {}
-					try {
-						validateToolUse(
-							block.name as ToolName,
-							mode ?? defaultModeSlug,
-							customModes ?? [],
-							{
-								apply_diff: this.diffEnabled,
-							},
-							block.params,
-						)
-					} catch (error) {
-						this.consecutiveMistakeCount++
-						pushToolResult(formatResponse.toolError(error.message))
-						break
-					}
+				// Validate tool use before execution
+				const { mode, customModes } = (await this.providerRef.deref()?.getState()) ?? {}
+				try {
+					validateToolUse(
+						block.name as ToolName,
+						mode ?? defaultModeSlug,
+						customModes ?? [],
+						{
+							apply_diff: this.diffEnabled,
+						},
+						block.params,
+					)
+				} catch (error) {
+					this.consecutiveMistakeCount++
+					pushToolResult(formatResponse.toolError(error.message))
+					break
 				}
 
 				switch (block.name) {

+ 87 - 14
src/shared/__tests__/modes.test.ts

@@ -31,12 +31,14 @@ describe("isToolAllowedForMode", () => {
 			// Test markdown editor mode
 			const mdResult = isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
 				path: "test.md",
+				content: "# Test",
 			})
 			expect(mdResult).toBe(true)
 
 			// Test CSS editor mode
 			const cssResult = isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, {
 				path: "styles.css",
+				content: ".test { color: red; }",
 			})
 			expect(cssResult).toBe(true)
 		})
@@ -44,41 +46,83 @@ describe("isToolAllowedForMode", () => {
 		it("rejects editing non-matching files", () => {
 			// Test markdown editor mode with non-markdown file
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
 			).toThrow(FileRestrictionError)
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
 			).toThrow(/\\.md\$/)
 
 			// Test CSS editor mode with non-CSS file
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
 			).toThrow(FileRestrictionError)
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("write_to_file", "css-editor", customModes, undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
 			).toThrow(/\\.css\$/)
 		})
 
+		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(
+				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+				}),
+			).toBe(true)
+
+			expect(
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+				}),
+			).toBe(true)
+
+			// Should allow path-only for ask mode too
+			expect(
+				isToolAllowedForMode("write_to_file", "ask", [], undefined, {
+					path: "test.js",
+				}),
+			).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",
 			})
 			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)
 
 			// Test both with non-matching file
 			expect(() =>
-				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("write_to_file", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
 			).toThrow(FileRestrictionError)
 
 			expect(() =>
-				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, { path: "test.js" }),
+				isToolAllowedForMode("apply_diff", "markdown-editor", customModes, undefined, {
+					path: "test.js",
+					diff: "- old\n+ new",
+				}),
 			).toThrow(FileRestrictionError)
 		})
 
@@ -100,11 +144,13 @@ describe("isToolAllowedForMode", () => {
 			expect(() =>
 				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.js",
+					content: "console.log('test')",
 				}),
 			).toThrow(FileRestrictionError)
 			expect(() =>
 				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.js",
+					content: "console.log('test')",
 				}),
 			).toThrow(/Documentation files only/)
 
@@ -112,11 +158,13 @@ describe("isToolAllowedForMode", () => {
 			expect(() =>
 				isToolAllowedForMode("apply_diff", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.js",
+					diff: "- old\n+ new",
 				}),
 			).toThrow(FileRestrictionError)
 			expect(() =>
 				isToolAllowedForMode("apply_diff", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.js",
+					diff: "- old\n+ new",
 				}),
 			).toThrow(/Documentation files only/)
 
@@ -124,30 +172,55 @@ describe("isToolAllowedForMode", () => {
 			expect(
 				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.md",
+					content: "# Test",
 				}),
 			).toBe(true)
 
 			expect(
 				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
 					path: "test.txt",
+					content: "Test content",
+				}),
+			).toBe(true)
+
+			// Test partial streaming cases
+			expect(
+				isToolAllowedForMode("write_to_file", "docs-editor", customModesWithDescription, undefined, {
+					path: "test.js",
 				}),
 			).toBe(true)
 		})
 
 		it("allows ask mode to edit markdown files only", () => {
 			// Should allow editing markdown files
-			expect(isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.md" })).toBe(true)
+			expect(
+				isToolAllowedForMode("write_to_file", "ask", [], undefined, {
+					path: "test.md",
+					content: "# Test",
+				}),
+			).toBe(true)
 
 			// Should allow applying diffs to markdown files
-			expect(isToolAllowedForMode("apply_diff", "ask", [], undefined, { path: "readme.md" })).toBe(true)
+			expect(
+				isToolAllowedForMode("apply_diff", "ask", [], undefined, {
+					path: "readme.md",
+					diff: "- old\n+ new",
+				}),
+			).toBe(true)
 
 			// Should reject non-markdown files
-			expect(() => isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.js" })).toThrow(
-				FileRestrictionError,
-			)
-			expect(() => isToolAllowedForMode("write_to_file", "ask", [], undefined, { path: "test.js" })).toThrow(
-				/Markdown files only/,
-			)
+			expect(() =>
+				isToolAllowedForMode("write_to_file", "ask", [], undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
+			).toThrow(FileRestrictionError)
+			expect(() =>
+				isToolAllowedForMode("write_to_file", "ask", [], undefined, {
+					path: "test.js",
+					content: "console.log('test')",
+				}),
+			).toThrow(/Markdown files only/)
 
 			// Should maintain read capabilities
 			expect(isToolAllowedForMode("read_file", "ask", [])).toBe(true)

+ 5 - 3
src/shared/modes.ts

@@ -196,11 +196,13 @@ export function isToolAllowedForMode(
 		// For the edit group, check file regex if specified
 		if (groupName === "edit" && options.fileRegex) {
 			const filePath = toolParams?.path
-			// Only validate regex if a path is provided
-			if (filePath && !doesFileMatchRegex(filePath, options.fileRegex)) {
+			if (
+				filePath &&
+				(toolParams.diff || toolParams.content) &&
+				!doesFileMatchRegex(filePath, options.fileRegex)
+			) {
 				throw new FileRestrictionError(mode.name, options.fileRegex, options.description, filePath)
 			}
-			return true
 		}
 
 		return true