keystone-cli 1.0.3 → 1.1.1

package/src/runner/standard-tools.ts

@@ -1,8 +1,9 @@
 import * as fs from 'node:fs';
 import * as path from 'node:path';
+import { Lang, parse } from '@ast-grep/napi';
 import type { AgentTool } from '../parser/schema';
 import { LIMITS, TIMEOUTS } from '../utils/constants';
-import { detectShellInjectionRisk } from './shell-executor';
+import { detectShellInjectionRisk } from './executors/shell-executor';
 
 export const STANDARD_TOOLS: AgentTool[] = [
   {
@@ -76,6 +77,25 @@ export const STANDARD_TOOLS: AgentTool[] = [
       content: '${{ args.content }}',
     },
   },
+  {
+    name: 'append_file',
+    description: 'Append content to the end of a file (creates if not exists)',
+    parameters: {
+      type: 'object',
+      properties: {
+        path: { type: 'string', description: 'Path to the file to append to' },
+        content: { type: 'string', description: 'Content to append to the file' },
+      },
+      required: ['path', 'content'],
+    },
+    execution: {
+      id: 'std_append_file',
+      type: 'file',
+      op: 'append',
+      path: '${{ args.path }}',
+      content: '${{ args.content }}',
+    },
+  },
   {
     name: 'list_files',
     description: 'List files in a directory',
@@ -123,22 +143,6 @@ export const STANDARD_TOOLS: AgentTool[] = [
       required: ['pattern'],
     },
     execution: {
-      id: 'std_search_files',
-      type: 'script',
-      run: `
-        (function() {
-          const fs = require('node:fs');
-          const path = require('node:path');
-          const { globSync } = require('glob');
-          const dir = args.dir || '.';
-          const pattern = args.pattern;
-          try {
-            return globSync(pattern, { cwd: dir, nodir: true });
-          } catch (e) {
-            throw new Error('Search failed: ' + e.message);
-          }
-        })();
-      `,
       allowInsecure: true,
     },
   },
@@ -162,106 +166,147 @@ export const STANDARD_TOOLS: AgentTool[] = [
       id: 'std_search_content',
       type: 'script',
       run: `
-        (function() {
+        (async function() {
           const fs = require('node:fs');
           const path = require('node:path');
           const { globSync } = require('glob');
+          // Try to request worker_threads. If not available, we might fall back or fail safely.
+          // Standard Node environment carries it.
+          const { Worker } = require('node:worker_threads');
+          
           const dir = args.dir || '.';
           const pattern = args.pattern || '**/*';
           const query = args.query;
-          const Re2 = (() => {
-            try {
-              return require('re2');
-            } catch {
-              return null;
-            }
-          })();
-          const maxQueryLength = ${LIMITS.MAX_SEARCH_QUERY_LENGTH};
-          const maxPatternLength = ${LIMITS.MAX_REGEX_PATTERN_LENGTH};
-          const maxResults = ${LIMITS.MAX_SEARCH_RESULTS};
-          const maxFileBytes = ${LIMITS.MAX_SEARCH_FILE_BYTES};
-          const maxLineLength = ${LIMITS.MAX_SEARCH_LINE_LENGTH};
-          const regexTimeoutMs = ${TIMEOUTS.REGEX_TIMEOUT_MS};
+          
+          const LIMITS = ${JSON.stringify(LIMITS)};
+          const TIMEOUTS = ${JSON.stringify(TIMEOUTS)};
 
-          if (query.length > maxQueryLength) {
-            throw new Error('Search query exceeds maximum length of ' + maxQueryLength + ' characters');
+          if (query.length > LIMITS.MAX_SEARCH_QUERY_LENGTH) {
+            throw new Error('Search query exceeds maximum length of ' + LIMITS.MAX_SEARCH_QUERY_LENGTH + ' characters');
           }
 
           const isRegex = query.startsWith('/') && query.endsWith('/') && query.length > 1;
-          let regex;
-          let normalizedQuery = '';
-          let regexDeadline;
-          
-          // ReDoS protection: detect dangerous regex patterns
-          if (isRegex) {
-            const pattern = query.slice(1, -1);
-            if (pattern.length > maxPatternLength) {
-              throw new Error('Regex pattern exceeds maximum length of ' + maxPatternLength + ' characters');
-            }
-            // Detect common ReDoS patterns: nested quantifiers, overlapping alternations
-            const dangerousPatterns = [
-              /\\([^)]*(?:[+*]|\\{\\d+(?:,\\d*)?\\})[^)]*\\)(?:[+*]|\\{\\d+(?:,\\d*)?\\})/, // (x+)+, (x{1,3})*
-              /\\([^)]*\\|[^)]*\\)(?:[+*]|\\{\\d+(?:,\\d*)?\\})/, // (a|aa)+
-              /(?:[+*]|\\{\\d+(?:,\\d*)?\\}){2,}/, // consecutive quantifiers
-            ];
-            if (dangerousPatterns.some(p => p.test(pattern))) {
-              throw new Error('Regex pattern contains potentially dangerous constructs (possible ReDoS). Simplify the pattern.');
-            }
 
-            try {
-              regex = Re2 ? new Re2(pattern) : new RegExp(pattern);
-            } catch (e) {
-              throw new Error('Invalid regular expression: ' + e.message);
-            }
-
-          } else {
-            normalizedQuery = query.toLowerCase();
+          // If NOT regex, use simple main-thread search (safe)
+          if (!isRegex) {
+             const normalizedQuery = query.toLowerCase();
+             const files = globSync(pattern, { cwd: dir, nodir: true });
+             const results = [];
+             for (const file of files) {
+                const fullPath = path.join(dir, file);
+                try {
+                  const stat = fs.statSync(fullPath);
+                  if (stat.size > LIMITS.MAX_SEARCH_FILE_BYTES) continue;
+                  const content = fs.readFileSync(fullPath, 'utf8');
+                  const lines = content.split('\\n');
+                  for (let i = 0; i < lines.length; i++) {
+                    const line = lines[i];
+                    if (line.length > LIMITS.MAX_SEARCH_LINE_LENGTH) continue;
+                    if (line.toLowerCase().includes(normalizedQuery)) {
+                      results.push({
+                        file,
+                        line: i + 1,
+                        content: line.trim()
+                      });
+                    }
+                    if (results.length >= LIMITS.MAX_SEARCH_RESULTS) break;
+                  }
+                } catch {}
+                if (results.length >= LIMITS.MAX_SEARCH_RESULTS) break;
+             }
+             return results;
           }
-  
+
+          // For REGEX, use a Worker to handle ReDoS/Timeouts
+          // We pass the files list to the worker so it does the heavy lifting
           const files = globSync(pattern, { cwd: dir, nodir: true });
-          const results = [];
-          for (const file of files) {
-            if (isRegex) {
-              regexDeadline = Date.now() + regexTimeoutMs;
-            }
-            const fullPath = path.join(dir, file);
-            let stat;
-            try {
-              stat = fs.statSync(fullPath);
-            } catch {
-              continue;
-            }
-            if (stat.size > maxFileBytes) {
-              continue;
-            }
-            let content;
+          
+          const workerScript = \`
+            const { parentPort, workerData } = require('node:worker_threads');
+            const fs = require('node:fs');
+            const path = require('node:path');
+            
             try {
-              content = fs.readFileSync(fullPath, 'utf8');
-            } catch {
-              continue;
-            }
-            const lines = content.split('\\n');
-            for (let i = 0; i < lines.length; i++) {
-              if (regexDeadline && Date.now() > regexDeadline) {
-                throw new Error('Regex search exceeded time limit');
-              }
-              const line = lines[i];
-              if (line.length > maxLineLength) {
-                continue;
+              const { files, dir, query, LIMITS } = workerData;
+              let Re2;
+              try { Re2 = require('re2'); } catch {}
+              
+              const patternStr = query.slice(1, -1);
+              let regex;
+              try {
+                regex = Re2 ? new Re2(patternStr) : new RegExp(patternStr);
+              } catch (e) {
+                parentPort.postMessage({ error: 'Invalid regex: ' + e.message });
+                process.exit(0);
               }
-              const matched = isRegex ? regex.test(line) : line.toLowerCase().includes(normalizedQuery);
-              if (matched) {
-                results.push({
-                  file,
-                  line: i + 1,
-                  content: line.trim()
-                });
+
+              const results = [];
+              
+              for (const file of files) {
+                const fullPath = path.join(dir, file);
+                try {
+                  const stat = fs.statSync(fullPath);
+                  if (stat.size > LIMITS.MAX_SEARCH_FILE_BYTES) continue;
+                  const content = fs.readFileSync(fullPath, 'utf8');
+                  const lines = content.split('\\n');
+                  for (let i = 0; i < lines.length; i++) {
+                    const line = lines[i];
+                    if (line.length > LIMITS.MAX_SEARCH_LINE_LENGTH) continue;
+                    
+                    // This matching is interruptible by thread termination
+                    if (regex.test(line)) {
+                      results.push({
+                         file,
+                         line: i + 1,
+                         content: line.trim()
+                      });
+                    }
+                    if (results.length >= LIMITS.MAX_SEARCH_RESULTS) break;
+                  }
+                } catch {}
+                if (results.length >= LIMITS.MAX_SEARCH_RESULTS) break;
               }
-              if (results.length >= maxResults) break;
+              
+              parentPort.postMessage({ results });
+            } catch (err) {
+              parentPort.postMessage({ error: err.message });
             }
-            if (results.length >= maxResults) break;
-          }
-          return results;
+          \`;
+
+          return new Promise((resolve, reject) => {
+            const worker = new Worker(workerScript, {
+              eval: true,
+              workerData: { files, dir, query, LIMITS }
+            });
+
+            // Hard timeout
+            const timeout = setTimeout(() => {
+              worker.terminate();
+              reject(new Error('Regex search timed out (possible ReDoS or large working set).'));
+            }, TIMEOUTS.REGEX_TIMEOUT_MS);
+
+            worker.on('message', (msg) => {
+              clearTimeout(timeout);
+              if (msg.error) {
+                reject(new Error(msg.error));
+              } else {
+                resolve(msg.results);
+              }
+              worker.terminate();
+            });
+
+            worker.on('error', (err) => {
+               clearTimeout(timeout);
+               reject(err);
+            });
+
+            worker.on('exit', (code) => {
+              if (code !== 0) {
+                 clearTimeout(timeout);
+                 reject(new Error(\`Worker stopped with exit code \${code}\`));
+              }
+            });
+          });
         })();
       `,
       allowInsecure: true,
@@ -285,6 +330,189 @@ export const STANDARD_TOOLS: AgentTool[] = [
       dir: '${{ args.dir }}',
     },
   },
+  {
+    name: 'ast_grep_search',
+    description:
+      'Search for structural code patterns using AST pattern matching. More precise than regex for code refactoring.',
+    parameters: {
+      type: 'object',
+      properties: {
+        pattern: {
+          type: 'string',
+          description: 'AST-grep pattern to search for, e.g. "console.log($A)"',
+        },
+        language: {
+          type: 'string',
+          description: 'Programming language (javascript, typescript, python, rust, go, etc.)',
+          default: 'typescript',
+        },
+        paths: {
+          type: 'array',
+          items: { type: 'string' },
+          description: 'File paths to search in',
+        },
+      },
+      required: ['pattern', 'paths'],
+    },
+    execution: {
+      id: 'std_ast_grep_search',
+      type: 'script',
+      run: `
+        (function() {
+          const fs = require('node:fs');
+          const path = require('node:path');
+          const { Lang, parse } = require('@ast-grep/napi');
+
+          const pattern = args.pattern;
+          const language = args.language || 'typescript';
+          const paths = args.paths || [];
+
+          const langMap = {
+            javascript: Lang.JavaScript,
+            typescript: Lang.TypeScript,
+            tsx: Lang.Tsx,
+            python: Lang.Python,
+            rust: Lang.Rust,
+            go: Lang.Go,
+            c: Lang.C,
+            cpp: Lang.Cpp,
+            java: Lang.Java,
+            kotlin: Lang.Kotlin,
+            swift: Lang.Swift,
+            html: Lang.Html,
+            css: Lang.Css,
+            json: Lang.Json,
+          };
+
+          const lang = langMap[language.toLowerCase()];
+          if (!lang) {
+            throw new Error('Unsupported language: ' + language);
+          }
+
+          const results = [];
+          for (const filePath of paths) {
+            if (!fs.existsSync(filePath)) continue;
+            const content = fs.readFileSync(filePath, 'utf8');
+            const tree = parse(lang, content);
+            const root = tree.root();
+            const matches = root.findAll(pattern);
+
+            for (const match of matches) {
+              const range = match.range();
+              results.push({
+                file: filePath,
+                line: range.start.line + 1,
+                column: range.start.column + 1,
+                content: match.text(),
+              });
+            }
+          }
+          return results;
+        })();
+      `,
+      allowInsecure: true,
+    },
+  },
+  {
+    name: 'ast_grep_replace',
+    description:
+      'Replace structural code patterns using AST-aware rewriting. Safer than regex for code refactoring.',
+    parameters: {
+      type: 'object',
+      properties: {
+        pattern: {
+          type: 'string',
+          description: 'AST-grep pattern to match, e.g. "console.log($A)"',
+        },
+        rewrite: { type: 'string', description: 'Replacement pattern, e.g. "logger.info($A)"' },
+        language: {
+          type: 'string',
+          description: 'Programming language (javascript, typescript, python, rust, go, etc.)',
+          default: 'typescript',
+        },
+        paths: {
+          type: 'array',
+          items: { type: 'string' },
+          description: 'File paths to apply replacements to',
+        },
+      },
+      required: ['pattern', 'rewrite', 'paths'],
+    },
+    execution: {
+      id: 'std_ast_grep_replace',
+      type: 'script',
+      run: `
+        (function() {
+          const fs = require('node:fs');
+          const path = require('node:path');
+          const { Lang, parse } = require('@ast-grep/napi');
+
+          const pattern = args.pattern;
+          const rewrite = args.rewrite;
+          const language = args.language || 'typescript';
+          const paths = args.paths || [];
+
+          const langMap = {
+            javascript: Lang.JavaScript,
+            typescript: Lang.TypeScript,
+            tsx: Lang.Tsx,
+            python: Lang.Python,
+            rust: Lang.Rust,
+            go: Lang.Go,
+            c: Lang.C,
+            cpp: Lang.Cpp,
+            java: Lang.Java,
+            kotlin: Lang.Kotlin,
+            swift: Lang.Swift,
+            html: Lang.Html,
+            css: Lang.Css,
+            json: Lang.Json,
+          };
+
+          const lang = langMap[language.toLowerCase()];
+          if (!lang) {
+            throw new Error('Unsupported language: ' + language);
+          }
+
+          const results = [];
+          for (const filePath of paths) {
+            if (!fs.existsSync(filePath)) continue;
+            const content = fs.readFileSync(filePath, 'utf8');
+            const tree = parse(lang, content);
+            const root = tree.root();
+            const edit = root.replace(pattern, rewrite);
+
+            if (edit !== content) {
+              fs.writeFileSync(filePath, edit);
+              results.push({
+                file: filePath,
+                modified: true,
+              });
+            }
+          }
+          return results;
+        })();
+      `,
+      allowInsecure: true,
+    },
+  },
+  {
+    name: 'fetch',
+    description: 'Fetch content from a URL (GET request)',
+    parameters: {
+      type: 'object',
+      properties: {
+        url: { type: 'string', description: 'URL to fetch' },
+      },
+      required: ['url'],
+    },
+    execution: {
+      id: 'std_fetch',
+      type: 'request',
+      url: '${{ args.url }}',
+      method: 'GET',
+    },
+  },
 ];
 
 /**
@@ -292,8 +520,7 @@ export const STANDARD_TOOLS: AgentTool[] = [
  */
 export function validateStandardToolSecurity(
   toolName: string,
-  // biome-ignore lint/suspicious/noExplicitAny: arguments can be any shape
-  args: any,
+  args: unknown,
   options: { allowOutsideCwd?: boolean; allowInsecure?: boolean }
 ): void {
   const cwd = process.cwd();
@@ -330,13 +557,23 @@ export function validateStandardToolSecurity(
       'read_file',
       'read_file_lines',
       'write_file',
+      'append_file',
       'list_files',
       'search_files',
       'search_content',
+      'ast_grep_search',
+      'ast_grep_replace',
     ].includes(toolName)
   ) {
     const rawPath = args.path || args.dir || '.';
     assertWithinCwd(rawPath);
+
+    // For AST tools, validate all paths in the array
+    if (['ast_grep_search', 'ast_grep_replace'].includes(toolName) && Array.isArray(args.paths)) {
+      for (const p of args.paths) {
+        assertWithinCwd(p);
+      }
+    }
   }
 
   // 2. Check shell risk for run_command and guard working directory