@probelabs/probe 0.6.0-rc148 → 0.6.0-rc150

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc148",
+  "version": "0.6.0-rc150",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",
@@ -75,15 +75,16 @@
     "@ai-sdk/google": "^2.0.14",
     "@ai-sdk/openai": "^2.0.10",
     "@modelcontextprotocol/sdk": "^1.0.0",
-    "@probelabs/maid": "^0.0.16",
+    "@probelabs/maid": "^0.0.19",
+    "adm-zip": "^0.5.16",
     "ai": "^5.0.0",
+    "ajv": "^8.17.1",
     "axios": "^1.8.3",
     "dotenv": "^16.4.7",
     "fs-extra": "^11.1.1",
     "glob": "^10.3.10",
     "gpt-tokenizer": "^3.0.1",
     "tar": "^6.2.0",
-    "adm-zip": "^0.5.16",
     "zod": "^3.24.2"
   },
   "devDependencies": {

package/src/agent/ProbeAgent.js

@@ -920,9 +920,7 @@ Examples:
 </search>
 
 <extract>
-<path>src/config.js</path>
-<start_line>15</start_line>
-<end_line>25</end_line>
+<targets>src/config.js:15-25</targets>
 </extract>
 
 <attempt_completion>
@@ -1833,7 +1831,7 @@ Convert your previous response content into actual JSON data that follows this s
               });
             }
 
-            let validation = validateJsonResponse(finalResult, { debug: this.debug });
+            let validation = validateJsonResponse(finalResult, { debug: this.debug, schema: options.schema });
             let retryCount = 0;
             const maxRetries = 3;
 
@@ -1884,7 +1882,7 @@ Convert your previous response content into actual JSON data that follows this s
                   );
 
                   // Validate the corrected response
-                  currentValidation = validateJsonResponse(currentResult, { debug: this.debug });
+                  currentValidation = validateJsonResponse(currentResult, { debug: this.debug, schema: options.schema });
                   retryCount++;
 
                   if (this.debug) {

package/src/agent/schemaUtils.js

@@ -5,6 +5,65 @@
 
 import { createMessagePreview } from '../tools/common.js';
 import { validate, fixText, extractMermaidBlocks } from '@probelabs/maid';
+import Ajv from 'ajv';
+
+/**
+ * Recursively apply additionalProperties: false to all object schemas
+ * This ensures strict validation at all nesting levels
+ * @param {Object} schema - JSON schema object
+ * @returns {Object} - Modified schema with additionalProperties enforced
+ */
+function enforceNoAdditionalProperties(schema) {
+  if (!schema || typeof schema !== 'object') {
+    return schema;
+  }
+
+  // Create a deep clone to avoid modifying the original
+  const cloned = JSON.parse(JSON.stringify(schema));
+
+  function applyRecursively(obj) {
+    if (!obj || typeof obj !== 'object') {
+      return;
+    }
+
+    // If this is an object type schema and doesn't have additionalProperties set
+    if (obj.type === 'object' && obj.additionalProperties === undefined) {
+      obj.additionalProperties = false;
+    }
+
+    // Recursively process properties
+    if (obj.properties && typeof obj.properties === 'object') {
+      Object.values(obj.properties).forEach(applyRecursively);
+    }
+
+    // Process items for arrays
+    if (obj.items) {
+      if (Array.isArray(obj.items)) {
+        obj.items.forEach(applyRecursively);
+      } else {
+        applyRecursively(obj.items);
+      }
+    }
+
+    // Process nested schemas in oneOf, anyOf, allOf
+    ['oneOf', 'anyOf', 'allOf'].forEach(key => {
+      if (Array.isArray(obj[key])) {
+        obj[key].forEach(applyRecursively);
+      }
+    });
+
+    // Process definitions/defs
+    if (obj.definitions && typeof obj.definitions === 'object') {
+      Object.values(obj.definitions).forEach(applyRecursively);
+    }
+    if (obj.$defs && typeof obj.$defs === 'object') {
+      Object.values(obj.$defs).forEach(applyRecursively);
+    }
+  }
+
+  applyRecursively(cloned);
+  return cloned;
+}
 
 /**
  * HTML entity decoder map for common entities that might appear in mermaid diagrams
@@ -141,12 +200,15 @@ export function cleanSchemaResponse(response) {
  * @returns {Object} - {isValid: boolean, parsed?: Object, error?: string}
  */
 export function validateJsonResponse(response, options = {}) {
-  const { debug = false } = options;
+  const { debug = false, schema = null, strictSchema = true } = options;
 
   if (debug) {
     console.log(`[DEBUG] JSON validation: Starting validation for response (${response.length} chars)`);
     const preview = createMessagePreview(response);
     console.log(`[DEBUG] JSON validation: Preview: ${preview}`);
+    if (schema) {
+      console.log(`[DEBUG] JSON validation: Schema validation enabled`);
+    }
   }
 
   try {
@@ -159,6 +221,140 @@ export function validateJsonResponse(response, options = {}) {
       console.log(`[DEBUG] JSON validation: Object type: ${typeof parsed}, keys: ${Object.keys(parsed || {}).length}`);
     }
 
+    // If schema provided, validate against it
+    if (schema) {
+      const schemaStart = Date.now();
+      let schemaObj;
+
+      try {
+        // Parse schema if it's a string
+        schemaObj = typeof schema === 'string' ? JSON.parse(schema) : schema;
+
+        // Apply strict mode: enforce additionalProperties: false on all nested objects
+        // unless explicitly disabled via strictSchema: false option
+        if (strictSchema) {
+          schemaObj = enforceNoAdditionalProperties(schemaObj);
+          if (debug) {
+            console.log(`[DEBUG] JSON validation: Applied strict mode - additionalProperties: false enforced on all levels`);
+          }
+        }
+      } catch (schemaParseError) {
+        if (debug) {
+          console.log(`[DEBUG] JSON validation: Failed to parse schema: ${schemaParseError.message}`);
+        }
+        return {
+          isValid: false,
+          error: 'Invalid schema provided',
+          schemaError: schemaParseError.message
+        };
+      }
+
+      // Create AJV validator with strict mode
+      const ajv = new Ajv({
+        strict: false,  // Don't fail on unknown schema keywords
+        allErrors: true,  // Return all errors, not just the first
+        verbose: true  // Include schema and data in errors
+      });
+
+      let validate;
+      try {
+        validate = ajv.compile(schemaObj);
+      } catch (compileError) {
+        if (debug) {
+          console.log(`[DEBUG] JSON validation: Schema compilation failed: ${compileError.message}`);
+        }
+        return {
+          isValid: false,
+          error: 'Schema compilation failed',
+          schemaError: compileError.message
+        };
+      }
+
+      const valid = validate(parsed);
+      const schemaTime = Date.now() - schemaStart;
+
+      if (debug) {
+        console.log(`[DEBUG] JSON validation: Schema validation completed in ${schemaTime}ms, valid: ${valid}`);
+      }
+
+      if (!valid) {
+        // Format schema errors for better readability and actionability
+        const formattedErrors = validate.errors.map(err => {
+          // Convert JSON Pointer path to dot notation for readability
+          // /user/profile/age -> user.profile.age
+          const path = err.instancePath
+            ? err.instancePath.substring(1).replace(/\//g, '.')
+            : '<root>';
+
+          let message = '';
+          let suggestion = '';
+
+          // Create crisp, actionable error messages based on error type
+          if (err.keyword === 'additionalProperties') {
+            const extraField = err.params.additionalProperty;
+            message = `Extra field '${extraField}' is not allowed`;
+            suggestion = `Remove '${extraField}' or add it to the schema`;
+          } else if (err.keyword === 'required') {
+            const missingField = err.params.missingProperty;
+            message = `Missing required field '${missingField}'`;
+            suggestion = `Add '${missingField}' to this object`;
+          } else if (err.keyword === 'type') {
+            const expected = err.params.type;
+            const actual = typeof err.data;
+            const value = JSON.stringify(err.data);
+            message = `Wrong type: expected ${expected}, got ${actual} (value: ${value})`;
+            suggestion = `Change value to ${expected} type`;
+          } else if (err.keyword === 'enum') {
+            const allowed = err.params.allowedValues;
+            const value = JSON.stringify(err.data);
+            message = `Invalid value ${value}. Allowed: ${allowed.map(v => JSON.stringify(v)).join(', ')}`;
+            suggestion = `Use one of the allowed values`;
+          } else if (err.keyword === 'minimum' || err.keyword === 'maximum') {
+            const limit = err.params.limit;
+            const comparison = err.params.comparison;
+            message = `Value ${err.data} ${err.message} (${comparison} ${limit})`;
+            suggestion = `Adjust value to meet constraint`;
+          } else if (err.keyword === 'minLength' || err.keyword === 'maxLength') {
+            const limit = err.params.limit;
+            message = `String length ${err.message} (current: ${err.data.length}, ${err.keyword}: ${limit})`;
+            suggestion = `Adjust string length`;
+          } else if (err.keyword === 'pattern') {
+            message = `Value doesn't match required pattern: ${err.params.pattern}`;
+            suggestion = `Update value to match the pattern`;
+          } else {
+            // Fallback for other error types
+            message = err.message;
+            suggestion = '';
+          }
+
+          // Format: "path: message | suggestion"
+          const location = path ? `at '${path}'` : 'at root';
+          return suggestion
+            ? `${location}: ${message} → ${suggestion}`
+            : `${location}: ${message}`;
+        });
+
+        const errorSummary = formattedErrors.join('\n  ');
+
+        if (debug) {
+          console.log(`[DEBUG] JSON validation: Schema validation errors:\n  ${errorSummary}`);
+        }
+
+        return {
+          isValid: false,
+          error: 'Schema validation failed',
+          schemaErrors: validate.errors,
+          formattedErrors: formattedErrors,
+          errorSummary: `Schema validation failed:\n  ${errorSummary}`,
+          parsed: parsed  // Still return parsed data for debugging
+        };
+      }
+
+      if (debug) {
+        console.log(`[DEBUG] JSON validation: Schema validation passed`);
+      }
+    }
+
     return { isValid: true, parsed };
   } catch (error) {
     // Extract error position from error message if available
@@ -886,9 +1082,21 @@ When presented with broken JSON, analyze it thoroughly and provide the corrected
       errorContext = validationResult.enhancedError;
     }
 
+    // Add schema validation errors if present
+    let schemaErrorDetails = '';
+    if (validationResult.errorSummary) {
+      schemaErrorDetails = `\n\nSchema Validation Errors:\n${validationResult.errorSummary}`;
+    } else if (validationResult.schemaErrors && validationResult.schemaErrors.length > 0) {
+      const errors = validationResult.schemaErrors.map(err => {
+        const path = err.instancePath || '(root)';
+        return `  ${path}: ${err.message}`;
+      }).join('\n');
+      schemaErrorDetails = `\n\nSchema Validation Errors:\n${errors}`;
+    }
+
     const prompt = `Fix the following invalid JSON.
 
-Error: ${errorContext}
+Error: ${errorContext}${schemaErrorDetails}
 
 Invalid JSON:
 ${invalidJson}
@@ -896,7 +1104,7 @@ ${invalidJson}
 Expected schema structure:
 ${schema}
 
-Provide only the corrected JSON without any markdown formatting or explanations.`;
+${schemaErrorDetails ? 'CRITICAL: Pay special attention to the schema validation errors above. The JSON may be syntactically valid but does not conform to the required schema. Make sure to:\n- Include all required fields\n- Use correct data types\n- Remove any additional properties not defined in the schema (if additionalProperties is false)\n- Ensure all values match their schema constraints\n\n' : ''}Provide only the corrected JSON without any markdown formatting or explanations.`;
 
     try {
       if (this.options.debug) {

package/src/extract.js

@@ -48,7 +48,7 @@ export async function extract(options) {
 	const hasContent = options.content !== undefined && options.content !== null;
 
 	if (!hasFiles && !hasInputFile && !hasContent) {
-		throw new Error('Either files array, inputFile, or content must be provided');
+		throw new Error('Extract requires one of: "files" (array of file paths), "inputFile" (path to input file), or "content" (string/buffer for stdin)');
 	}
 
 	// Get the binary path

package/src/mcp/index.ts

@@ -199,18 +199,18 @@ class ProbeServer {
         },
         {
           name: 'extract_code',
-          description: "Extract code from files. Formats: file.js (whole file), file.js:42 (from line), file.js#functionName (symbol).",
+          description: "Extract code blocks from files using tree-sitter AST parsing. Each file path can include optional line numbers or symbol names to extract specific code blocks.",
           inputSchema: {
             type: 'object',
             properties: {
               path: {
                 type: 'string',
-                description: 'Absolute path to the project directory',
+                description: 'Absolute path to the project root directory (used as working directory for relative file paths)',
               },
               files: {
                 type: 'array',
                 items: { type: 'string' },
-                description: 'Array of file paths with optional line/symbol: ["file.rs:10", "file.rs#func_name"]',
+                description: 'Array of file paths to extract from. Formats: "file.js" (entire file), "file.js:42" (code block at line 42), "file.js:10-20" (lines 10-20), "file.js#funcName" (specific symbol). Line numbers and symbols are part of the path string, not separate parameters. Paths can be absolute or relative to the project directory.',
               }
             },
             required: ['path', 'files'],

package/src/tools/common.js

@@ -7,13 +7,8 @@ import { z } from 'zod';
 
 // Common schemas for tool parameters (used for internal execution after XML parsing)
 export const searchSchema = z.object({
-	query: z.string().describe('Search query with Elasticsearch syntax. Use + for important terms.'),
-	path: z.string().optional().default('.').describe('Path to search in. For dependencies use "go:github.com/owner/repo", "js:package_name", or "rust:cargo_name" etc.'),
-	allow_tests: z.boolean().optional().default(false).describe('Allow test files in search results'),
-	exact: z.boolean().optional().default(false).describe('Perform exact search without tokenization (case-insensitive)'),
-	maxResults: z.number().optional().describe('Maximum number of results to return'),
-	maxTokens: z.number().optional().default(10000).describe('Maximum number of tokens to return'),
-	language: z.string().optional().describe('Limit search to files of a specific programming language')
+	query: z.string().describe('Search query with Elasticsearch syntax. Use quotes for exact matches, AND/OR for boolean logic, - for negation.'),
+	path: z.string().optional().default('.').describe('Path to search in. For dependencies use "go:github.com/owner/repo", "js:package_name", or "rust:cargo_name" etc.')
 });
 
 export const querySchema = z.object({
@@ -24,13 +19,8 @@ export const querySchema = z.object({
 });
 
 export const extractSchema = z.object({
-	targets: z.string().optional().describe('File paths or symbols to extract from. Can include line numbers, symbol names, or multiple space-separated targets'),
-	input_content: z.string().optional().describe('Text content to extract file paths from'),
-	line: z.number().optional().describe('Start line number to extract a specific code block'),
-	end_line: z.number().optional().describe('End line number for extracting a range of lines'),
-	allow_tests: z.boolean().optional().default(false).describe('Allow test files and test code blocks'),
-	context_lines: z.number().optional().default(10).describe('Number of context lines to include'),
-	format: z.string().optional().default('plain').describe('Output format (plain, markdown, json, xml, color, outline-xml, outline-diff)')
+	targets: z.string().optional().describe('File paths or symbols to extract from. Formats: "file.js" (whole file), "file.js:42" (line 42), "file.js:10-20" (lines 10-20), "file.js#funcName" (symbol). Multiple targets separated by spaces.'),
+	input_content: z.string().optional().describe('Text content to extract file paths from (alternative to targets)')
 });
 
 export const delegateSchema = z.object({
@@ -120,13 +110,8 @@ You need to focus on main keywords when constructing the query, and always use e
 - Once data is returned, it's cached and won't return on next runs (this is expected behavior)
 
 Parameters:
-- query: (required) Search query with Elasticsearch syntax. You can use + for important terms, and - for negation.
-- path: (required) Path to search in. All dependencies located in /dep folder, under language sub folders, like this: "/dep/go/github.com/owner/repo", "/dep/js/package_name", or "/dep/rust/cargo_name" etc. YOU SHOULD ALWAYS provide FULL PATH when searching dependencies, including depency name.
-- allow_tests: (optional, default: false) Allow test files in search results (true/false).
-- exact: (optional, default: false) Perform exact pricise search. Use it when you already know function or struct name, or some other code block, and want exact match.
-- maxResults: (optional) Maximum number of results to return (number).
-- maxTokens: (optional, default: 10000) Maximum number of tokens to return (number).
-- language: (optional) Limit search to files of a specific programming language (e.g., 'rust', 'js', 'python', 'go' etc.).
+- query: (required) Search query with Elasticsearch syntax. Use quotes for exact matches ("functionName"), AND/OR for boolean logic, - for negation, + for important terms.
+- path: (optional, default: '.') Path to search in. All dependencies located in /dep folder, under language sub folders, like this: "/dep/go/github.com/owner/repo", "/dep/js/package_name", or "/dep/rust/cargo_name" etc.
 
 **Workflow:** Always start with search, then use extract for detailed context when needed.
 
@@ -148,30 +133,24 @@ User: How to calculate the total amount in the payments module?
 <search>
 <query>calculate AND payment</query>
 <path>src/utils</path>
-<allow_tests>false</allow_tests>
 </search>
 
 User: How do the user authentication and authorization work?
 <search>
-<query>+user and (authentification OR authroization OR authz)</query>
+<query>+user AND (authentication OR authorization OR authz)</query>
 <path>.</path>
-<allow_tests>true</allow_tests>
-<language>go</language>
 </search>
 
 User: Find all react imports in the project.
 <search>
-<query>import { react }</query>
+<query>"import" AND "react"</query>
 <path>.</path>
-<exact>true</exact>
-<language>js</language>
 </search>
 
-User: Find how decompoud library works?
+User: Find how decompound library works?
 <search>
-<query>import { react }</query>
+<query>decompound</query>
 <path>/dep/rust/decompound</path>
-<language>rust</language>
 </search>
 
 </examples>
@@ -208,11 +187,9 @@ Full file extraction should be the LAST RESORT! Always prefer search.
 **Session Awareness:** Reuse context from previous tool calls. Don't re-extract the same symbols you already have.
 
 Parameters:
-- targets: (required) File paths or symbols to extract from. Can include line numbers, symbol names, or multiple space-separated targets (e.g., 'src/main.rs:10-20', 'src/utils.js#myFunction').
-  For multiple extractions: 'session.rs#AuthService.login auth.rs:2-100 config.rs#DatabaseConfig'
-- line: (optional) Start line number to extract a specific code block. Use with end_line for ranges.
-- end_line: (optional) End line number for extracting a range of lines.
-- allow_tests: (optional, default: false) Allow test files and test code blocks (true/false).
+- targets: (required) File paths or symbols to extract from. Formats: "file.js" (whole file), "file.js:42" (code block at line 42), "file.js:10-20" (lines 10-20), "file.js#funcName" (specific symbol). Multiple targets separated by spaces.
+- input_content: (optional) Text content to extract file paths from (alternative to targets for processing diffs/logs).
+
 Usage Example:
 
 <examples>
@@ -239,9 +216,7 @@ User: Lets read the whole file
 
 User: Read the first 10 lines of the file
 <extract>
-<targets>src/search/ranking.rs</targets>
-<line>1</line>
-<end_line>10</end_line>
+<targets>src/search/ranking.rs:1-10</targets>
 </extract>
 
 User: Read file inside the dependency
@@ -349,7 +324,7 @@ export const bashDescription = 'Execute bash commands for system exploration and
 // Valid tool names that should be parsed as tool calls
 const DEFAULT_VALID_TOOLS = [
 	'search',
-	'query', 
+	'query',
 	'extract',
 	'delegate',
 	'listFiles',
@@ -358,6 +333,43 @@ const DEFAULT_VALID_TOOLS = [
 	'attempt_completion'
 ];
 
+/**
+ * Get valid parameter names for a specific tool from its schema
+ * @param {string} toolName - Name of the tool
+ * @returns {string[]} - Array of valid parameter names for this tool
+ */
+function getValidParamsForTool(toolName) {
+	// Map tool names to their schemas
+	const schemaMap = {
+		search: searchSchema,
+		query: querySchema,
+		extract: extractSchema,
+		delegate: delegateSchema,
+		bash: bashSchema,
+		attempt_completion: attemptCompletionSchema
+	};
+
+	const schema = schemaMap[toolName];
+	if (!schema) {
+		// For tools without schema (listFiles, searchFiles, implement), return common params
+		// These are the shared params that appear across multiple tools
+		return ['path', 'directory', 'pattern', 'recursive', 'includeHidden', 'task', 'files', 'autoCommits', 'result'];
+	}
+
+	// For attempt_completion, it has custom validation, just return 'result'
+	if (toolName === 'attempt_completion') {
+		return ['result'];
+	}
+
+	// Extract keys from Zod schema
+	if (schema && schema._def && schema._def.shape) {
+		return Object.keys(schema._def.shape());
+	}
+
+	// Fallback: return empty array if we can't extract schema keys
+	return [];
+}
+
 // Simple XML parser helper - safer string-based approach
 export function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
 	// Look for each valid tool name specifically using string search
@@ -387,15 +399,12 @@ export function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
 
 		const params = {};
 
+		// Get valid parameters for this specific tool from its schema
+		const validParams = getValidParamsForTool(toolName);
+
 		// Parse parameters using string-based approach for better safety
-		// Common parameter names to look for (can be extended as needed)
-		// Note: includes both camelCase and underscore_case variants to handle inconsistencies
-		const commonParams = ['query', 'file_path', 'line', 'end_line', 'path', 'recursive', 'includeHidden',
-		                      'max_results', 'maxResults', 'result', 'command', 'description', 'task', 'param', 'pattern',
-		                      'allow_tests', 'exact', 'maxTokens', 'language', 'input_content',
-		                      'context_lines', 'format', 'directory', 'autoCommits', 'files', 'targets'];
-
-		for (const paramName of commonParams) {
+		// Only look for parameters that are valid for this specific tool
+		for (const paramName of validParams) {
 			const paramOpenTag = `<${paramName}>`;
 			const paramCloseTag = `</${paramName}>`;
 
@@ -410,7 +419,7 @@ export function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
 			if (paramCloseIndex === -1) {
 				// Find the next opening tag after this parameter
 				let nextTagIndex = innerContent.length;
-				for (const nextParam of commonParams) {
+				for (const nextParam of validParams) {
 					const nextOpenTag = `<${nextParam}>`;
 					const nextIndex = innerContent.indexOf(nextOpenTag, paramOpenIndex + paramOpenTag.length);
 					if (nextIndex !== -1 && nextIndex < nextTagIndex) {