diff --git a/packages/types/src/mode.ts b/packages/types/src/mode.ts index f981ba7bf9a..0973b52d5d5 100644 --- a/packages/types/src/mode.ts +++ b/packages/types/src/mode.ts @@ -25,6 +25,24 @@ export const groupOptionsSchema = z.object({ }, { message: "Invalid regular expression pattern" }, ), + commandRegex: z + .string() + .optional() + .refine( + (pattern) => { + if (!pattern) { + return true // Optional, so empty is valid. + } + + try { + new RegExp(pattern) + return true + } catch { + return false + } + }, + { message: "Invalid regular expression pattern" }, + ), description: z.string().optional(), }) diff --git a/src/core/config/__tests__/ModeConfig.spec.ts b/src/core/config/__tests__/ModeConfig.spec.ts index 74cbc0c4373..809d117265d 100644 --- a/src/core/config/__tests__/ModeConfig.spec.ts +++ b/src/core/config/__tests__/ModeConfig.spec.ts @@ -172,6 +172,55 @@ describe("CustomModeSchema", () => { }) }) + describe("commandRegex", () => { + it("validates a mode with command restrictions and descriptions", () => { + const modeWithJustRegex = { + slug: "npm-only", + name: "NPM Only", + roleDefinition: "NPM command mode", + groups: ["read", ["command", { commandRegex: "^npm\\s" }]], + } + + const modeWithDescription = { + slug: "package-manager", + name: "Package Manager", + roleDefinition: "Package manager command mode", + groups: [ + "read", + ["command", { commandRegex: "^(npm|yarn|pnpm)\\s", description: "Package manager commands only" }], + ], + } + + expect(() => modeConfigSchema.parse(modeWithJustRegex)).not.toThrow() + expect(() => modeConfigSchema.parse(modeWithDescription)).not.toThrow() + }) + + it("validates command regex patterns", () => { + const validPatterns = ["^npm\\s", "^git\\s", ".*", "npm|yarn|pnpm", "^docker\\s+", "[a-z]+"] + const invalidPatterns = ["[", "(unclosed", "\\", "*", "+*", "?"] + + validPatterns.forEach((pattern) => { + const mode = { + slug: "test", + name: "Test", + roleDefinition: "Test", + groups: ["read", ["command", { commandRegex: pattern }]], + } + expect(() => modeConfigSchema.parse(mode)).not.toThrow() + }) + + invalidPatterns.forEach((pattern) => { + const mode = { + slug: "test", + name: "Test", + roleDefinition: "Test", + groups: ["read", ["command", { commandRegex: pattern }]], + } + expect(() => modeConfigSchema.parse(mode)).toThrow() + }) + }) + }) + const validBaseMode = { slug: "123e4567-e89b-12d3-a456-426614174000", name: "Test Mode", diff --git a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/architect-mode-prompt.snap b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/architect-mode-prompt.snap index 5bed6df09d1..dbe7114a8e2 100644 --- a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/architect-mode-prompt.snap +++ b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/architect-mode-prompt.snap @@ -43,7 +43,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/ask-mode-prompt.snap b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/ask-mode-prompt.snap index 243dfc19b7b..e98512d13e5 100644 --- a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/ask-mode-prompt.snap +++ b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/ask-mode-prompt.snap @@ -43,7 +43,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/mcp-server-creation-disabled.snap b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/mcp-server-creation-disabled.snap index 5bed6df09d1..dbe7114a8e2 100644 --- a/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/mcp-server-creation-disabled.snap +++ b/src/core/prompts/__tests__/__snapshots__/add-custom-instructions/mcp-server-creation-disabled.snap @@ -43,7 +43,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/__tests__/__snapshots__/system-prompt/consistent-system-prompt.snap b/src/core/prompts/__tests__/__snapshots__/system-prompt/consistent-system-prompt.snap index 42e8bba9c68..d6575b18adf 100644 --- a/src/core/prompts/__tests__/__snapshots__/system-prompt/consistent-system-prompt.snap +++ b/src/core/prompts/__tests__/__snapshots__/system-prompt/consistent-system-prompt.snap @@ -43,7 +43,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/__tests__/__snapshots__/system-prompt/with-mcp-hub-provided.snap b/src/core/prompts/__tests__/__snapshots__/system-prompt/with-mcp-hub-provided.snap index 5aa6677ab03..e6e22379f2a 100644 --- a/src/core/prompts/__tests__/__snapshots__/system-prompt/with-mcp-hub-provided.snap +++ b/src/core/prompts/__tests__/__snapshots__/system-prompt/with-mcp-hub-provided.snap @@ -45,7 +45,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/__tests__/__snapshots__/system-prompt/with-undefined-mcp-hub.snap b/src/core/prompts/__tests__/__snapshots__/system-prompt/with-undefined-mcp-hub.snap index 42e8bba9c68..d6575b18adf 100644 --- a/src/core/prompts/__tests__/__snapshots__/system-prompt/with-undefined-mcp-hub.snap +++ b/src/core/prompts/__tests__/__snapshots__/system-prompt/with-undefined-mcp-hub.snap @@ -43,7 +43,7 @@ RULES - You cannot `cd` into a different directory to complete a task. You are stuck operating from '/test/path', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '/test/path', and if so prepend with `cd`'ing into that directory && then executing the command (as one command since you are stuck operating from '/test/path'). For example, if you needed to run `npm install` in a project outside of '/test/path', you would need to prepend with a `cd` i.e. pseudocode for this would be `cd (path to project) && (command, in this case npm install)`. -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/prompts/sections/rules.ts b/src/core/prompts/sections/rules.ts index 4f6e573fa73..4df1a30d781 100644 --- a/src/core/prompts/sections/rules.ts +++ b/src/core/prompts/sections/rules.ts @@ -76,7 +76,7 @@ RULES - You cannot \`cd\` into a different directory to complete a task. You are stuck operating from '${cwd.toPosix()}', so be sure to pass in the correct 'path' parameter when using tools that require a path. - Do not use the ~ character or $HOME to refer to the home directory. - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '${cwd.toPosix()}', and if so prepend with \`cd\`'ing into that directory ${chainOp} then executing the command (as one command since you are stuck operating from '${cwd.toPosix()}'). For example, if you needed to run \`npm install\` in a project outside of '${cwd.toPosix()}', you would need to prepend with a \`cd\` i.e. pseudocode for this would be \`cd (path to project) ${chainOp} (command, in this case npm install)\`.${chainNote ? ` ${chainNote}` : ""} -- Some modes have restrictions on which files they can edit. If you attempt to edit a restricted file, the operation will be rejected with a FileRestrictionError that will specify which file patterns are allowed for the current mode. +- Some modes have restrictions on which files can be edited or commands that can be executed. If you attempt to edit a restricted file or execute a restricted command, the operation will be rejected with a restriction error that will specify which patterns are allowed for this mode. - Be sure to consider the type of project (e.g. Python, JavaScript, web application) when determining the appropriate structure and files to include. Also consider what files may be most relevant to accomplishing the task, for example looking at a project's manifest file would help you understand the project's dependencies, which you could incorporate into any code you write. * For example, in architect mode trying to edit app.js would be rejected because architect mode can only edit files matching "\\.md$" - When making changes to code, always consider the context in which the code is being used. Ensure that your changes are compatible with the existing codebase and that they follow the project's coding standards and best practices. diff --git a/src/core/tools/__tests__/validateToolUse.spec.ts b/src/core/tools/__tests__/validateToolUse.spec.ts index 29455e36883..601bdbb58e8 100644 --- a/src/core/tools/__tests__/validateToolUse.spec.ts +++ b/src/core/tools/__tests__/validateToolUse.spec.ts @@ -2,7 +2,7 @@ import type { ModeConfig } from "@roo-code/types" -import { modes } from "../../../shared/modes" +import { modes, CommandRestrictionError } from "../../../shared/modes" import { TOOL_GROUPS } from "../../../shared/tools" import { validateToolUse, isToolAllowedForMode } from "../validateToolUse" @@ -250,5 +250,165 @@ describe("mode-validator", () => { expect(() => validateToolUse("execute_command", codeMode, [], toolRequirements)).not.toThrow() }) + + describe("commandRegex restrictions", () => { + it("should allow commands matching the regex", () => { + const customModes: ModeConfig[] = [ + { + slug: "npm-only", + name: "NPM Only", + roleDefinition: "You can only run npm commands", + groups: [ + "read", + ["command", { commandRegex: "^npm\\s", description: "Only npm commands" }], + ] as const, + }, + ] + expect( + isToolAllowedForMode("execute_command", "npm-only", customModes, undefined, { + command: "npm install", + }), + ).toBe(true) + }) + + it("should block commands not matching the regex", () => { + const customModes: ModeConfig[] = [ + { + slug: "npm-only", + name: "NPM Only", + roleDefinition: "You can only run npm commands", + groups: [ + "read", + ["command", { commandRegex: "^npm\\s", description: "Only npm commands" }], + ] as const, + }, + ] + expect(() => + isToolAllowedForMode("execute_command", "npm-only", customModes, undefined, { + command: "git push", + }), + ).toThrow() + }) + + it("should include description in error message", () => { + const customModes: ModeConfig[] = [ + { + slug: "npm-only", + name: "NPM Only", + roleDefinition: "You can only run npm commands", + groups: [ + "read", + ["command", { commandRegex: "^npm\\s", description: "Only npm commands" }], + ] as const, + }, + ] + try { + isToolAllowedForMode("execute_command", "npm-only", customModes, undefined, { command: "git push" }) + expect(true).toBe(false) // Should not reach here + } catch (error) { + expect(error).toBeInstanceOf(CommandRestrictionError) + expect(error.message).toContain("Only npm commands") + } + }) + + it("should work without commandRegex (no restriction)", () => { + const customModes: ModeConfig[] = [ + { + slug: "command-mode", + name: "Command Mode", + roleDefinition: "You can run any commands", + groups: ["read", "command"] as const, + }, + ] + // Should allow all commands when no commandRegex is specified + expect( + isToolAllowedForMode("execute_command", "command-mode", customModes, undefined, { + command: "git push", + }), + ).toBe(true) + expect( + isToolAllowedForMode("execute_command", "command-mode", customModes, undefined, { + command: "npm install", + }), + ).toBe(true) + expect( + isToolAllowedForMode("execute_command", "command-mode", customModes, undefined, { + command: "ls -la", + }), + ).toBe(true) + }) + + it("should allow commands when commandRegex is undefined in options", () => { + const customModes: ModeConfig[] = [ + { + slug: "command-mode", + name: "Command Mode", + roleDefinition: "You can run any commands", + groups: ["read", ["command", { description: "All commands allowed" }]] as const, + }, + ] + // Should allow all commands when commandRegex is not specified + expect( + isToolAllowedForMode("execute_command", "command-mode", customModes, undefined, { + command: "any command", + }), + ).toBe(true) + }) + + it("should handle multiple regex patterns", () => { + const customModes: ModeConfig[] = [ + { + slug: "package-manager", + name: "Package Manager", + roleDefinition: "You can run package manager commands", + groups: [ + "read", + [ + "command", + { commandRegex: "^(npm|yarn|pnpm)\\s", description: "Package manager commands" }, + ], + ] as const, + }, + ] + expect( + isToolAllowedForMode("execute_command", "package-manager", customModes, undefined, { + command: "npm install", + }), + ).toBe(true) + expect( + isToolAllowedForMode("execute_command", "package-manager", customModes, undefined, { + command: "yarn add lodash", + }), + ).toBe(true) + expect( + isToolAllowedForMode("execute_command", "package-manager", customModes, undefined, { + command: "pnpm build", + }), + ).toBe(true) + expect(() => + isToolAllowedForMode("execute_command", "package-manager", customModes, undefined, { + command: "git push", + }), + ).toThrow() + }) + + it("should throw CommandRestrictionError with correct tool name", () => { + const customModes: ModeConfig[] = [ + { + slug: "npm-only", + name: "NPM Only", + roleDefinition: "You can only run npm commands", + groups: ["read", ["command", { commandRegex: "^npm\\s" }]] as const, + }, + ] + try { + isToolAllowedForMode("execute_command", "npm-only", customModes, undefined, { command: "git push" }) + expect(true).toBe(false) // Should not reach here + } catch (error) { + expect(error).toBeInstanceOf(CommandRestrictionError) + expect(error.message).toContain("execute_command") + } + }) + }) }) }) diff --git a/src/core/tools/validateToolUse.ts b/src/core/tools/validateToolUse.ts index 243a170ed90..118e2beec87 100644 --- a/src/core/tools/validateToolUse.ts +++ b/src/core/tools/validateToolUse.ts @@ -2,7 +2,13 @@ import type { ToolName, ModeConfig, ExperimentId, GroupOptions, GroupEntry } fro import { toolNames as validToolNames } from "@roo-code/types" import { customToolRegistry } from "@roo-code/core" -import { type Mode, FileRestrictionError, getModeBySlug, getGroupName } from "../../shared/modes" +import { + type Mode, + FileRestrictionError, + CommandRestrictionError, + getModeBySlug, + getGroupName, +} from "../../shared/modes" import { EXPERIMENT_IDS } from "../../shared/experiments" import { TOOL_GROUPS, ALWAYS_AVAILABLE_TOOLS, TOOL_ALIASES } from "../../shared/tools" @@ -117,6 +123,16 @@ function doesFileMatchRegex(filePath: string, pattern: string): boolean { } } +function doesCommandMatchRegex(command: string, pattern: string): boolean { + try { + const regex = new RegExp(pattern) + return regex.test(command) + } catch (error) { + console.error(`Invalid regex pattern: ${pattern}`, error) + return false + } +} + export function isToolAllowedForMode( tool: string, modeSlug: string, @@ -232,6 +248,14 @@ export function isToolAllowedForMode( // Native-only: multi-file edits provide structured params; no legacy XML args parsing. } + // For the command group, check command regex if specified + if (groupName === "command" && options.commandRegex) { + const command = toolParams?.command + if (command && !doesCommandMatchRegex(command, options.commandRegex)) { + throw new CommandRestrictionError(mode.name, options.commandRegex, options.description, command, tool) + } + } + return true } diff --git a/src/shared/__tests__/modes.spec.ts b/src/shared/__tests__/modes.spec.ts index ceb3cacb4d9..8a6d5b21672 100644 --- a/src/shared/__tests__/modes.spec.ts +++ b/src/shared/__tests__/modes.spec.ts @@ -9,7 +9,7 @@ vi.mock("../../core/prompts/sections/custom-instructions", () => ({ addCustomInstructions: vi.fn().mockResolvedValue("Combined instructions"), })) -import { FileRestrictionError, getFullModeDetails, modes, getModeSelection } from "../modes" +import { FileRestrictionError, CommandRestrictionError, getFullModeDetails, modes, getModeSelection } from "../modes" import { isToolAllowedForMode } from "../../core/tools/validateToolUse" import { addCustomInstructions } from "../../core/prompts/sections/custom-instructions" @@ -925,3 +925,40 @@ describe("getModeSelection", () => { expect(selection.baseInstructions).toBe(promptComponentAsk.customInstructions) }) }) + +describe("CommandRestrictionError", () => { + it("formats error message with pattern when no description provided", () => { + const error = new CommandRestrictionError("NPM Mode", "^npm\\s", undefined, "git push") + expect(error.message).toBe( + "This mode (NPM Mode) can only execute commands matching pattern: ^npm\\s. Got: git push", + ) + expect(error.name).toBe("CommandRestrictionError") + }) + + it("formats error message with description when provided", () => { + const error = new CommandRestrictionError("NPM Mode", "^npm\\s", "Only npm commands", "git push") + expect(error.message).toBe( + "This mode (NPM Mode) can only execute commands matching pattern: ^npm\\s (Only npm commands). Got: git push", + ) + }) + + it("formats error message with tool name when provided", () => { + const error = new CommandRestrictionError("NPM Mode", "^npm\\s", undefined, "git push", "execute_command") + expect(error.message).toBe( + "Tool 'execute_command' in mode 'NPM Mode' can only execute commands matching pattern: ^npm\\s. Got: git push", + ) + }) + + it("formats error message with tool name and description when both provided", () => { + const error = new CommandRestrictionError( + "NPM Mode", + "^npm\\s", + "Only npm commands", + "git push", + "execute_command", + ) + expect(error.message).toBe( + "Tool 'execute_command' in mode 'NPM Mode' can only execute commands matching pattern: ^npm\\s (Only npm commands). Got: git push", + ) + }) +}) diff --git a/src/shared/modes.ts b/src/shared/modes.ts index a94aa47ed0b..0fa12067f39 100644 --- a/src/shared/modes.ts +++ b/src/shared/modes.ts @@ -142,6 +142,17 @@ export class FileRestrictionError extends Error { } } +// Custom error class for command restrictions +export class CommandRestrictionError extends Error { + constructor(mode: string, pattern: string, description: string | undefined, command: string, tool?: string) { + const toolInfo = tool ? `Tool '${tool}' in mode '${mode}'` : `This mode (${mode})` + super( + `${toolInfo} can only execute commands matching pattern: ${pattern}${description ? ` (${description})` : ""}. Got: ${command}`, + ) + this.name = "CommandRestrictionError" + } +} + // Create the mode-specific default prompts export const defaultPrompts: Readonly = Object.freeze( Object.fromEntries(