@probelabs/probe 0.6.0-rc67 → 0.6.0-rc69

package/cjs/index.cjs

@@ -1129,44 +1129,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";
@@ -1325,6 +1329,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"
+    ];
   }
 });
 
@@ -2100,11 +2113,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}`);
@@ -2379,6 +2392,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;
@@ -2655,9 +2678,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,
@@ -2665,11 +2744,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,
@@ -2677,7 +2755,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}
@@ -2738,10 +2816,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;
@@ -2827,6 +2913,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 ? `
 
@@ -2834,11 +2936,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.`;
@@ -3365,7 +3468,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);
@@ -3484,7 +3598,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...");
             }
@@ -3564,6 +3678,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) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc67",
+  "version": "0.6.0-rc69",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",

package/src/agent/ProbeAgent.js

@@ -580,8 +580,15 @@ When troubleshooting:
           throw new Error(finalResult);
         }
 
-        // Parse tool call from response
-        const parsedTool = parseXmlToolCallWithThinking(assistantResponseContent);
+        // Parse tool call from response with valid tools list
+        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);
@@ -722,7 +729,8 @@ When troubleshooting:
       this.tokenCounter.updateHistory(this.history);
 
       // Schema handling - format response according to provided schema
-      if (options.schema && !options._schemaFormatted) {
+      // Skip schema processing if result came from attempt_completion tool
+      if (options.schema && !options._schemaFormatted && !completionAttempted) {
         if (this.debug) {
           console.log('[DEBUG] Schema provided, applying automatic formatting...');
         }
@@ -822,6 +830,38 @@ Convert your previous response content into this JSON format now. Return nothing
           console.error('[ERROR] Schema formatting failed:', error);
           // Return the original result if schema formatting fails
         }
+      } else if (completionAttempted && options.schema) {
+        // For attempt_completion results with schema, still clean markdown if needed
+        try {
+          finalResult = cleanSchemaResponse(finalResult);
+          
+          // Validate and fix Mermaid diagrams if present
+          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`);
+            }
+          }
+          
+          // Validate JSON if schema expects JSON
+          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;

package/src/agent/schemaUtils.js

@@ -3,6 +3,37 @@
  * Supports JSON and Mermaid diagram validation
  */
 
+/**
+ * HTML entity decoder map for common entities that might appear in mermaid diagrams
+ */
+const HTML_ENTITY_MAP = {
+  '&lt;': '<',
+  '&gt;': '>',
+  '&amp;': '&',
+  '&quot;': '"',
+  '&#39;': "'",
+  '&nbsp;': ' '
+};
+
+/**
+ * Decode HTML entities in text without requiring external dependencies
+ * @param {string} text - Text that may contain HTML entities
+ * @returns {string} - Text with HTML entities decoded
+ */
+export function decodeHtmlEntities(text) {
+  if (!text || typeof text !== 'string') {
+    return text;
+  }
+
+  let decoded = text;
+  for (const [entity, character] of Object.entries(HTML_ENTITY_MAP)) {
+    // Use global replacement to catch all instances
+    decoded = decoded.replace(new RegExp(entity, 'g'), character);
+  }
+
+  return decoded;
+}
+
 /**
  * Clean AI response by extracting JSON content when response contains JSON
  * Only processes responses that contain JSON structures { or [ 
@@ -644,6 +675,28 @@ When presented with a broken Mermaid diagram, analyze it thoroughly and provide
    * @returns {Promise<string>} - The corrected Mermaid diagram
    */
   async fixMermaidDiagram(diagramContent, originalErrors = [], diagramInfo = {}) {
+    // First, try auto-fixing HTML entities without AI
+    const decodedContent = decodeHtmlEntities(diagramContent);
+    
+    // If HTML entity decoding changed the content, validate it first
+    if (decodedContent !== diagramContent) {
+      try {
+        const quickValidation = await validateMermaidDiagram(decodedContent);
+        if (quickValidation.isValid) {
+          // HTML entity decoding fixed the issue, no need for AI
+          if (this.options.debug) {
+            console.error('[DEBUG] Fixed Mermaid diagram with HTML entity decoding only');
+          }
+          return decodedContent;
+        }
+      } catch (error) {
+        // If validation fails, continue with AI fixing using decoded content
+        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 
@@ -654,11 +707,14 @@ When presented with a broken Mermaid diagram, analyze it thoroughly and provide
       ? `\n\nExpected diagram type: ${diagramInfo.diagramType}` 
       : '';
 
+    // Use decoded content for AI fixing to ensure HTML entities are handled
+    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.`;
@@ -751,30 +807,104 @@ export async function validateAndFixMermaidResponse(response, options = {}) {
     };
   }
 
-  // Some diagrams are invalid, try to fix them with specialized agent
+  // Some diagrams are invalid, first try HTML entity decoding auto-fix
   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 {
-    // Create specialized fixing agent
-    const mermaidFixer = new MermaidFixingAgent({
-      path, provider, model, debug, tracer
-    });
-
     let fixedResponse = response;
     const fixingResults = [];
+    let htmlEntityFixesApplied = false;
     
     // Extract diagrams with position information for replacement
     const { diagrams } = extractMermaidFromMarkdown(response);
     
-    // Fix invalid diagrams in reverse order to preserve indices
+    // First pass: Try HTML entity decoding on invalid diagrams
     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) {
+        // HTML entities were found and decoded, validate the result
+        try {
+          const quickValidation = await validateMermaidDiagram(decodedContent);
+          if (quickValidation.isValid) {
+            // HTML entity decoding fixed this diagram!
+            const originalDiagram = diagrams[invalidDiagram.originalIndex];
+            const attributesStr = originalDiagram.attributes ? ` ${originalDiagram.attributes}` : '';
+            const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${decodedContent}\n\`\`\``;
+            
+            fixedResponse = fixedResponse.slice(0, originalDiagram.startIndex) + 
+                           newCodeBlock + 
+                           fixedResponse.slice(originalDiagram.endIndex);
+            
+            fixingResults.push({
+              diagramIndex: invalidDiagram.originalIndex,
+              wasFixed: true,
+              originalContent: 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 HTML entity fixes were applied, re-validate the entire response
+    if (htmlEntityFixesApplied) {
+      const revalidation = await validateMermaidResponse(fixedResponse);
+      if (revalidation.isValid) {
+        // All diagrams are now valid, return without AI fixing
+        if (debug) {
+          console.error('[DEBUG] All diagrams fixed with HTML entity decoding, no AI needed');
+        }
+        return {
+          ...revalidation,
+          wasFixed: true,
+          originalResponse: response,
+          fixedResponse: fixedResponse,
+          fixingResults: fixingResults
+        };
+      }
+    }
+    
+    // Still have invalid diagrams after HTML entity decoding, proceed with AI fixing
+    if (debug) {
+      console.error('[DEBUG] Some diagrams still invalid after HTML entity decoding, starting AI fixing...');
+    }
+    
+    // Create specialized fixing agent for remaining invalid diagrams
+    const mermaidFixer = new MermaidFixingAgent({
+      path, provider, model, debug, tracer
+    });
+    
+    // Re-extract diagrams and re-validate after HTML entity fixes
+    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,
@@ -784,7 +914,7 @@ export async function validateAndFixMermaidResponse(response, options = {}) {
 
         if (fixedContent && fixedContent !== invalidDiagram.content) {
           // Replace the diagram in the response
-          const originalDiagram = diagrams[invalidDiagram.originalIndex];
+          const originalDiagram = updatedDiagrams[invalidDiagram.originalIndex];
           const attributesStr = originalDiagram.attributes ? ` ${originalDiagram.attributes}` : '';
           const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${fixedContent}\n\`\`\``;
           

package/src/agent/tools.js

@@ -124,9 +124,10 @@ User: Find all markdown files in the docs directory, but only at the top level.
  * This function removes any <thinking></thinking> tags from the input string
  * before passing it to the original parseXmlToolCall function
  * @param {string} xmlString - The XML string to parse
+ * @param {string[]} [validTools] - List of valid tool names to parse (optional)
  * @returns {Object|null} - The parsed tool call or null if no valid tool call found
  */
-export function parseXmlToolCallWithThinking(xmlString) {
+export function parseXmlToolCallWithThinking(xmlString, validTools) {
   // Extract thinking content if present (for potential logging or analysis)
   const thinkingMatch = xmlString.match(/<thinking>([\s\S]*?)<\/thinking>/);
   const thinkingContent = thinkingMatch ? thinkingMatch[1].trim() : null;
@@ -135,7 +136,7 @@ export function parseXmlToolCallWithThinking(xmlString) {
   const cleanedXmlString = xmlString.replace(/<thinking>[\s\S]*?<\/thinking>/g, '').trim();
 
   // Use the original parseXmlToolCall function to parse the cleaned XML string
-  const parsedTool = parseXmlToolCall(cleanedXmlString);
+  const parsedTool = parseXmlToolCall(cleanedXmlString, validTools);
 
   // If debugging is enabled, log the thinking content
   if (process.env.DEBUG === '1' && thinkingContent) {