feat: allow command-specific restrictions for ShellTool (#2605)

1 files changed, 23 insertions, 15 deletions

diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 59c9c1bd..4ee2d23f 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -456,25 +456,33 @@ export class Config {
 export function createToolRegistry(config: Config): Promise<ToolRegistry> {
   const registry = new ToolRegistry(config);
   const targetDir = config.getTargetDir();
-  const tools = config.getCoreTools()
-    ? new Set(config.getCoreTools())
-    : undefined;
-  const excludeTools = config.getExcludeTools()
-    ? new Set(config.getExcludeTools())
-    : undefined;
 
   // helper to create & register core tools that are enabled
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const registerCoreTool = (ToolClass: any, ...args: unknown[]) => {
-    // check both the tool name (.Name) and the class name (.name)
-    if (
-      // coreTools contain tool name
-      (!tools || tools.has(ToolClass.Name) || tools.has(ToolClass.name)) &&
-      // excludeTools don't contain tool name
-      (!excludeTools ||
-        (!excludeTools.has(ToolClass.Name) &&
-          !excludeTools.has(ToolClass.name)))
-    ) {
+    const className = ToolClass.name;
+    const toolName = ToolClass.Name || className;
+    const coreTools = config.getCoreTools();
+    const excludeTools = config.getExcludeTools();
+
+    let isEnabled = false;
+    if (coreTools === undefined) {
+      isEnabled = true;
+    } else {
+      isEnabled = coreTools.some(
+        (tool) =>
+          tool === className ||
+          tool === toolName ||
+          tool.startsWith(`${className}(`) ||
+          tool.startsWith(`${toolName}(`),
+      );
+    }
+
+    if (excludeTools?.includes(className) || excludeTools?.includes(toolName)) {
+      isEnabled = false;
+    }
+
+    if (isEnabled) {
       registry.registerTool(new ToolClass(...args));
     }
   };