@probelabs/probe 0.6.0-rc67 → 0.6.0-rc69

package/build/tools/common.js

@@ -173,52 +173,74 @@ export const searchDescription = 'Search code in the repository using Elasticsea
 export const queryDescription = 'Search code using ast-grep structural pattern matching. Use this tool to find specific code structures like functions, classes, or methods.';
 export const extractDescription = 'Extract code blocks from files based on file paths and optional line numbers. Use this tool to see complete context after finding relevant files.';
 
-// Simple XML parser helper
-export function parseXmlToolCall(xmlString) {
-	const toolMatch = xmlString.match(/<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/);
-	if (!toolMatch) {
-		return null;
-	}
+// Valid tool names that should be parsed as tool calls
+const DEFAULT_VALID_TOOLS = [
+	'search',
+	'query', 
+	'extract',
+	'listFiles',
+	'searchFiles',
+	'implement',
+	'attempt_completion'
+];
 
-	const toolName = toolMatch[1];
-	const innerContent = toolMatch[2];
-	const params = {};
-
-	const paramRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
-	let paramMatch;
-	while ((paramMatch = paramRegex.exec(innerContent)) !== null) {
-		const paramName = paramMatch[1];
-		let paramValue = paramMatch[2].trim();
-
-		// Basic type inference (can be improved)
-		if (paramValue.toLowerCase() === 'true') {
-			paramValue = true;
-		} else if (paramValue.toLowerCase() === 'false') {
-			paramValue = false;
-		} else if (!isNaN(paramValue) && paramValue.trim() !== '') {
-			// Check if it's potentially a number (handle integers and floats)
-			const num = Number(paramValue);
-			if (Number.isFinite(num)) { // Use Number.isFinite to avoid Infinity/NaN
-				paramValue = num;
-			}
-			// Keep as string if not a valid finite number
+// Simple XML parser helper
+export function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
+	// Find all potential XML tag matches
+	const globalRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
+	let match;
+	
+	// Look through all matches to find the first valid tool
+	while ((match = globalRegex.exec(xmlString)) !== null) {
+		const toolName = match[1];
+		
+		// Only parse XML tags that correspond to valid tools
+		if (!validTools.includes(toolName)) {
+			continue; // Skip non-tool tags and look for the next match
 		}
 
-		params[paramName] = paramValue;
-	}
+		const innerContent = match[2];
+		const params = {};
+
+		const paramRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
+		let paramMatch;
+		while ((paramMatch = paramRegex.exec(innerContent)) !== null) {
+			const paramName = paramMatch[1];
+			let paramValue = paramMatch[2].trim();
+
+			// Basic type inference (can be improved)
+			if (paramValue.toLowerCase() === 'true') {
+				paramValue = true;
+			} else if (paramValue.toLowerCase() === 'false') {
+				paramValue = false;
+			} else if (!isNaN(paramValue) && paramValue.trim() !== '') {
+				// Check if it's potentially a number (handle integers and floats)
+				const num = Number(paramValue);
+				if (Number.isFinite(num)) { // Use Number.isFinite to avoid Infinity/NaN
+					paramValue = num;
+				}
+				// Keep as string if not a valid finite number
+			}
 
-	// Special handling for attempt_completion where result might contain nested XML/code
-	if (toolName === 'attempt_completion') {
-		const resultMatch = innerContent.match(/<result>([\s\S]*?)<\/result>/);
-		if (resultMatch) {
-			params['result'] = resultMatch[1].trim(); // Keep result content as is
+			params[paramName] = paramValue;
 		}
-		const commandMatch = innerContent.match(/<command>([\s\S]*?)<\/command>/);
-		if (commandMatch) {
-			params['command'] = commandMatch[1].trim();
+
+		// Special handling for attempt_completion where result might contain nested XML/code
+		if (toolName === 'attempt_completion') {
+			const resultMatch = innerContent.match(/<result>([\s\S]*?)<\/result>/);
+			if (resultMatch) {
+				params['result'] = resultMatch[1].trim(); // Keep result content as is
+			}
+			const commandMatch = innerContent.match(/<command>([\s\S]*?)<\/command>/);
+			if (commandMatch) {
+				params['command'] = commandMatch[1].trim();
+			}
 		}
-	}
 
+		// Return the first valid tool found
+		return { toolName, params };
+	}
 
-	return { toolName, params };
+	// No valid tool found
+	return null;
 }

package/cjs/agent/ProbeAgent.cjs

@@ -1429,44 +1429,48 @@ var init_extract = __esm({
 });
 
 // src/tools/common.js
-function parseXmlToolCall(xmlString) {
-  const toolMatch = xmlString.match(/<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/);
-  if (!toolMatch) {
-    return null;
+function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
+  const globalRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
+  let match;
+  while ((match = globalRegex.exec(xmlString)) !== null) {
+    const toolName = match[1];
+    if (!validTools.includes(toolName)) {
+      continue;
+    }
+    const innerContent = match[2];
+    const params = {};
+    const paramRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
+    let paramMatch;
+    while ((paramMatch = paramRegex.exec(innerContent)) !== null) {
+      const paramName = paramMatch[1];
+      let paramValue = paramMatch[2].trim();
+      if (paramValue.toLowerCase() === "true") {
+        paramValue = true;
+      } else if (paramValue.toLowerCase() === "false") {
+        paramValue = false;
+      } else if (!isNaN(paramValue) && paramValue.trim() !== "") {
+        const num = Number(paramValue);
+        if (Number.isFinite(num)) {
+          paramValue = num;
+        }
+      }
+      params[paramName] = paramValue;
+    }
+    if (toolName === "attempt_completion") {
+      const resultMatch = innerContent.match(/<result>([\s\S]*?)<\/result>/);
+      if (resultMatch) {
+        params["result"] = resultMatch[1].trim();
+      }
+      const commandMatch = innerContent.match(/<command>([\s\S]*?)<\/command>/);
+      if (commandMatch) {
+        params["command"] = commandMatch[1].trim();
+      }
+    }
+    return { toolName, params };
   }
-  const toolName = toolMatch[1];
-  const innerContent = toolMatch[2];
-  const params = {};
-  const paramRegex = /<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/g;
-  let paramMatch;
-  while ((paramMatch = paramRegex.exec(innerContent)) !== null) {
-    const paramName = paramMatch[1];
-    let paramValue = paramMatch[2].trim();
-    if (paramValue.toLowerCase() === "true") {
-      paramValue = true;
-    } else if (paramValue.toLowerCase() === "false") {
-      paramValue = false;
-    } else if (!isNaN(paramValue) && paramValue.trim() !== "") {
-      const num = Number(paramValue);
-      if (Number.isFinite(num)) {
-        paramValue = num;
-      }
-    }
-    params[paramName] = paramValue;
-  }
-  if (toolName === "attempt_completion") {
-    const resultMatch = innerContent.match(/<result>([\s\S]*?)<\/result>/);
-    if (resultMatch) {
-      params["result"] = resultMatch[1].trim();
-    }
-    const commandMatch = innerContent.match(/<command>([\s\S]*?)<\/command>/);
-    if (commandMatch) {
-      params["command"] = commandMatch[1].trim();
-    }
-  }
-  return { toolName, params };
+  return null;
 }
-var import_zod, searchSchema, querySchema, extractSchema, attemptCompletionSchema, searchToolDefinition, queryToolDefinition, extractToolDefinition, attemptCompletionToolDefinition, searchDescription, queryDescription, extractDescription;
+var import_zod, searchSchema, querySchema, extractSchema, attemptCompletionSchema, searchToolDefinition, queryToolDefinition, extractToolDefinition, attemptCompletionToolDefinition, searchDescription, queryDescription, extractDescription, DEFAULT_VALID_TOOLS;
 var init_common = __esm({
   "src/tools/common.js"() {
     "use strict";
@@ -1625,6 +1629,15 @@ Usage Example:
     searchDescription = "Search code in the repository using Elasticsearch-like query syntax. Use this tool first for any code-related questions.";
     queryDescription = "Search code using ast-grep structural pattern matching. Use this tool to find specific code structures like functions, classes, or methods.";
     extractDescription = "Extract code blocks from files based on file paths and optional line numbers. Use this tool to see complete context after finding relevant files.";
+    DEFAULT_VALID_TOOLS = [
+      "search",
+      "query",
+      "extract",
+      "listFiles",
+      "searchFiles",
+      "implement",
+      "attempt_completion"
+    ];
   }
 });
 
@@ -2026,11 +2039,11 @@ function createTools(configOptions) {
     extractTool: extractTool(configOptions)
   };
 }
-function parseXmlToolCallWithThinking(xmlString) {
+function parseXmlToolCallWithThinking(xmlString, validTools) {
   const thinkingMatch = xmlString.match(/<thinking>([\s\S]*?)<\/thinking>/);
   const thinkingContent = thinkingMatch ? thinkingMatch[1].trim() : null;
   const cleanedXmlString = xmlString.replace(/<thinking>[\s\S]*?<\/thinking>/g, "").trim();
-  const parsedTool = parseXmlToolCall(cleanedXmlString);
+  const parsedTool = parseXmlToolCall(cleanedXmlString, validTools);
   if (process.env.DEBUG === "1" && thinkingContent) {
     console.log(`[DEBUG] AI Thinking Process:
 ${thinkingContent}`);
@@ -2305,6 +2318,16 @@ var init_probeTool = __esm({
 });
 
 // src/agent/schemaUtils.js
+function decodeHtmlEntities(text) {
+  if (!text || typeof text !== "string") {
+    return text;
+  }
+  let decoded = text;
+  for (const [entity, character] of Object.entries(HTML_ENTITY_MAP)) {
+    decoded = decoded.replace(new RegExp(entity, "g"), character);
+  }
+  return decoded;
+}
 function cleanSchemaResponse(response) {
   if (!response || typeof response !== "string") {
     return response;
@@ -2581,9 +2604,65 @@ async function validateAndFixMermaidResponse(response, options = {}) {
     };
   }
   if (debug) {
-    console.error("[DEBUG] Invalid Mermaid diagrams detected, starting specialized fixing agent...");
+    console.error("[DEBUG] Invalid Mermaid diagrams detected, trying HTML entity auto-fix first...");
   }
   try {
+    let fixedResponse = response;
+    const fixingResults = [];
+    let htmlEntityFixesApplied = false;
+    const { diagrams } = extractMermaidFromMarkdown(response);
+    const invalidDiagrams = validation.diagrams.map((result, index) => ({ ...result, originalIndex: index })).filter((result) => !result.isValid).reverse();
+    for (const invalidDiagram of invalidDiagrams) {
+      const originalContent = invalidDiagram.content;
+      const decodedContent = decodeHtmlEntities(originalContent);
+      if (decodedContent !== originalContent) {
+        try {
+          const quickValidation = await validateMermaidDiagram(decodedContent);
+          if (quickValidation.isValid) {
+            const originalDiagram = diagrams[invalidDiagram.originalIndex];
+            const attributesStr = originalDiagram.attributes ? ` ${originalDiagram.attributes}` : "";
+            const newCodeBlock = `\`\`\`mermaid${attributesStr}
+${decodedContent}
+\`\`\``;
+            fixedResponse = fixedResponse.slice(0, originalDiagram.startIndex) + newCodeBlock + fixedResponse.slice(originalDiagram.endIndex);
+            fixingResults.push({
+              diagramIndex: invalidDiagram.originalIndex,
+              wasFixed: true,
+              originalContent,
+              fixedContent: decodedContent,
+              originalError: invalidDiagram.error,
+              fixedWithHtmlDecoding: true
+            });
+            htmlEntityFixesApplied = true;
+            if (debug) {
+              console.error(`[DEBUG] Fixed diagram ${invalidDiagram.originalIndex + 1} with HTML entity decoding: ${invalidDiagram.error}`);
+            }
+          }
+        } catch (error) {
+          if (debug) {
+            console.error(`[DEBUG] HTML entity decoding didn't fix diagram ${invalidDiagram.originalIndex + 1}: ${error.message}`);
+          }
+        }
+      }
+    }
+    if (htmlEntityFixesApplied) {
+      const revalidation = await validateMermaidResponse(fixedResponse);
+      if (revalidation.isValid) {
+        if (debug) {
+          console.error("[DEBUG] All diagrams fixed with HTML entity decoding, no AI needed");
+        }
+        return {
+          ...revalidation,
+          wasFixed: true,
+          originalResponse: response,
+          fixedResponse,
+          fixingResults
+        };
+      }
+    }
+    if (debug) {
+      console.error("[DEBUG] Some diagrams still invalid after HTML entity decoding, starting AI fixing...");
+    }
     const mermaidFixer = new MermaidFixingAgent({
       path: path6,
       provider,
@@ -2591,11 +2670,10 @@ async function validateAndFixMermaidResponse(response, options = {}) {
       debug,
       tracer
     });
-    let fixedResponse = response;
-    const fixingResults = [];
-    const { diagrams } = extractMermaidFromMarkdown(response);
-    const invalidDiagrams = validation.diagrams.map((result, index) => ({ ...result, originalIndex: index })).filter((result) => !result.isValid).reverse();
-    for (const invalidDiagram of invalidDiagrams) {
+    const { diagrams: updatedDiagrams } = extractMermaidFromMarkdown(fixedResponse);
+    const updatedValidation = await validateMermaidResponse(fixedResponse);
+    const stillInvalidDiagrams = updatedValidation.diagrams.map((result, index) => ({ ...result, originalIndex: index })).filter((result) => !result.isValid).reverse();
+    for (const invalidDiagram of stillInvalidDiagrams) {
       try {
         const fixedContent = await mermaidFixer.fixMermaidDiagram(
           invalidDiagram.content,
@@ -2603,7 +2681,7 @@ async function validateAndFixMermaidResponse(response, options = {}) {
           { diagramType: invalidDiagram.diagramType }
         );
         if (fixedContent && fixedContent !== invalidDiagram.content) {
-          const originalDiagram = diagrams[invalidDiagram.originalIndex];
+          const originalDiagram = updatedDiagrams[invalidDiagram.originalIndex];
           const attributesStr = originalDiagram.attributes ? ` ${originalDiagram.attributes}` : "";
           const newCodeBlock = `\`\`\`mermaid${attributesStr}
 ${fixedContent}
@@ -2664,10 +2742,18 @@ ${fixedContent}
     };
   }
 }
-var MermaidFixingAgent;
+var HTML_ENTITY_MAP, MermaidFixingAgent;
 var init_schemaUtils = __esm({
   "src/agent/schemaUtils.js"() {
     "use strict";
+    HTML_ENTITY_MAP = {
+      "&lt;": "<",
+      "&gt;": ">",
+      "&amp;": "&",
+      "&quot;": '"',
+      "&#39;": "'",
+      "&nbsp;": " "
+    };
     MermaidFixingAgent = class {
       constructor(options = {}) {
         this.ProbeAgent = null;
@@ -2753,6 +2839,22 @@ When presented with a broken Mermaid diagram, analyze it thoroughly and provide
        * @returns {Promise<string>} - The corrected Mermaid diagram
        */
       async fixMermaidDiagram(diagramContent, originalErrors = [], diagramInfo = {}) {
+        const decodedContent = decodeHtmlEntities(diagramContent);
+        if (decodedContent !== diagramContent) {
+          try {
+            const quickValidation = await validateMermaidDiagram(decodedContent);
+            if (quickValidation.isValid) {
+              if (this.options.debug) {
+                console.error("[DEBUG] Fixed Mermaid diagram with HTML entity decoding only");
+              }
+              return decodedContent;
+            }
+          } catch (error) {
+            if (this.options.debug) {
+              console.error("[DEBUG] HTML entity decoding didn't fully fix diagram, continuing with AI fixing");
+            }
+          }
+        }
         await this.initializeAgent();
         const errorContext = originalErrors.length > 0 ? `
 
@@ -2760,11 +2862,12 @@ Detected errors: ${originalErrors.join(", ")}` : "";
         const diagramTypeHint = diagramInfo.diagramType ? `
 
 Expected diagram type: ${diagramInfo.diagramType}` : "";
+        const contentToFix = decodedContent !== diagramContent ? decodedContent : diagramContent;
         const prompt = `Analyze and fix the following Mermaid diagram.${errorContext}${diagramTypeHint}
 
 Broken Mermaid diagram:
 \`\`\`mermaid
-${diagramContent}
+${contentToFix}
 \`\`\`
 
 Provide only the corrected Mermaid diagram within a mermaid code block. Do not add any explanations or additional text.`;
@@ -3291,7 +3394,18 @@ You are working with a repository located at: ${searchDirectory}
               finalResult = `Error: Failed to get response from AI model during iteration ${currentIteration}. ${error.message}`;
               throw new Error(finalResult);
             }
-            const parsedTool = parseXmlToolCallWithThinking(assistantResponseContent);
+            const validTools = [
+              "search",
+              "query",
+              "extract",
+              "listFiles",
+              "searchFiles",
+              "attempt_completion"
+            ];
+            if (this.allowEdit) {
+              validTools.push("implement");
+            }
+            const parsedTool = parseXmlToolCallWithThinking(assistantResponseContent, validTools);
             if (parsedTool) {
               const { toolName, params } = parsedTool;
               if (this.debug) console.log(`[DEBUG] Parsed tool call: ${toolName} with params:`, params);
@@ -3410,7 +3524,7 @@ Error: Unknown tool '${toolName}'. Available tools: ${Object.keys(this.toolImple
             }
           }
           this.tokenCounter.updateHistory(this.history);
-          if (options.schema && !options._schemaFormatted) {
+          if (options.schema && !options._schemaFormatted && !completionAttempted) {
             if (this.debug) {
               console.log("[DEBUG] Schema provided, applying automatic formatting...");
             }
@@ -3490,6 +3604,32 @@ Convert your previous response content into this JSON format now. Return nothing
             } catch (error) {
               console.error("[ERROR] Schema formatting failed:", error);
             }
+          } else if (completionAttempted && options.schema) {
+            try {
+              finalResult = cleanSchemaResponse(finalResult);
+              const mermaidValidation = await validateAndFixMermaidResponse(finalResult, {
+                debug: this.debug,
+                path: this.allowedFolders[0],
+                provider: this.clientApiProvider,
+                model: this.model
+              });
+              if (mermaidValidation.wasFixed) {
+                finalResult = mermaidValidation.fixedResponse;
+                if (this.debug) {
+                  console.log(`[DEBUG] Mermaid diagrams fixed in attempt_completion result`);
+                }
+              }
+              if (isJsonSchema(options.schema)) {
+                const validation = validateJsonResponse(finalResult);
+                if (!validation.isValid && this.debug) {
+                  console.log(`[DEBUG] attempt_completion result JSON validation failed: ${validation.error}`);
+                }
+              }
+            } catch (error) {
+              if (this.debug) {
+                console.log(`[DEBUG] attempt_completion result cleanup failed: ${error.message}`);
+              }
+            }
           }
           return finalResult;
         } catch (error) {