@probelabs/probe 0.6.0-rc66 → 0.6.0-rc69

package/build/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,14 +729,27 @@ 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...');
         }
         
         try {
           // Step 1: Make a follow-up call to format according to schema
-          const schemaPrompt = `Now you need to respond according to this schema:\n\n${options.schema}\n\nPlease reformat your previous response to match this schema exactly. Only return the formatted response, no additional text.`;
+          const schemaPrompt = `CRITICAL: You MUST respond with ONLY valid JSON that matches this schema. Ignore any previous response format or instructions.
+
+Schema:
+${options.schema}
+
+REQUIREMENTS:
+- Return ONLY the JSON object/array
+- NO additional text, explanations, or markdown
+- NO code blocks or formatting
+- The JSON must be parseable by JSON.parse()
+- DO NOT include any conversational text like "Hello" or explanations
+
+Convert your previous response content into this JSON format now. Return nothing but valid JSON.`;
           
           // Call answer recursively with _schemaFormatted flag to prevent infinite loop
           finalResult = await this.answer(schemaPrompt, [], { 
@@ -770,18 +790,21 @@ When troubleshooting:
           
           // Step 4: Validate and potentially correct JSON responses
           if (isJsonSchema(options.schema)) {
-            const validation = validateJsonResponse(finalResult);
+            let validation = validateJsonResponse(finalResult);
+            let retryCount = 0;
+            const maxRetries = 3;
             
-            if (!validation.isValid) {
+            while (!validation.isValid && retryCount < maxRetries) {
               if (this.debug) {
-                console.log('[DEBUG] JSON validation failed:', validation.error);
+                console.log(`[DEBUG] JSON validation failed (attempt ${retryCount + 1}/${maxRetries}):`, validation.error);
               }
               
-              // Attempt correction once
+              // Create increasingly stronger correction prompts
               const correctionPrompt = createJsonCorrectionPrompt(
                 finalResult, 
                 options.schema, 
-                validation.error
+                validation.error,
+                retryCount
               );
               
               finalResult = await this.answer(correctionPrompt, [], { 
@@ -790,17 +813,55 @@ When troubleshooting:
               });
               finalResult = cleanSchemaResponse(finalResult);
               
-              // Final validation
-              const finalValidation = validateJsonResponse(finalResult);
-              if (!finalValidation.isValid && this.debug) {
-                console.log('[DEBUG] JSON still invalid after correction:', finalValidation.error);
+              // Validate the corrected response
+              validation = validateJsonResponse(finalResult);
+              retryCount++;
+              
+              if (this.debug && !validation.isValid && retryCount < maxRetries) {
+                console.log(`[DEBUG] JSON still invalid after correction ${retryCount}, retrying...`);
               }
             }
+            
+            if (!validation.isValid && this.debug) {
+              console.log(`[DEBUG] JSON still invalid after ${maxRetries} correction attempts:`, validation.error);
+            }
           }
         } catch (error) {
           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/build/agent/index.js

@@ -1409,44 +1409,48 @@ var init_extract = __esm({
 
 // src/tools/common.js
 import { z } from "zod";
-function parseXmlToolCall(xmlString) {
-  const toolMatch = xmlString.match(/<([a-zA-Z0-9_]+)>([\s\S]*?)<\/\1>/);
-  if (!toolMatch) {
-    return null;
-  }
-  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;
-      }
+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;
     }
-    params[paramName] = paramValue;
-  }
-  if (toolName === "attempt_completion") {
-    const resultMatch = innerContent.match(/<result>([\s\S]*?)<\/result>/);
-    if (resultMatch) {
-      params["result"] = resultMatch[1].trim();
+    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;
     }
-    const commandMatch = innerContent.match(/<command>([\s\S]*?)<\/command>/);
-    if (commandMatch) {
-      params["command"] = commandMatch[1].trim();
+    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 { toolName, params };
+  return null;
 }
-var searchSchema, querySchema, extractSchema, attemptCompletionSchema, searchToolDefinition, queryToolDefinition, extractToolDefinition, attemptCompletionToolDefinition, searchDescription, queryDescription, extractDescription;
+var searchSchema, querySchema, extractSchema, attemptCompletionSchema, searchToolDefinition, queryToolDefinition, extractToolDefinition, attemptCompletionToolDefinition, searchDescription, queryDescription, extractDescription, DEFAULT_VALID_TOOLS;
 var init_common = __esm({
   "src/tools/common.js"() {
     "use strict";
@@ -1604,6 +1608,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"
+    ];
   }
 });
 
@@ -2006,11 +2019,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}`);
@@ -2284,11 +2297,53 @@ 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;
   }
   const trimmed = response.trim();
+  const codeBlockPatterns = [
+    /```json\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
+    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
+    /`([{\[][\s\S]*?[}\]])`/
+  ];
+  for (const pattern of codeBlockPatterns) {
+    const match = trimmed.match(pattern);
+    if (match) {
+      return match[1].trim();
+    }
+  }
+  const codeBlockStartPattern = /```(?:json)?\s*\n?\s*([{\[])/;
+  const codeBlockMatch = trimmed.match(codeBlockStartPattern);
+  if (codeBlockMatch) {
+    const startIndex = codeBlockMatch.index + codeBlockMatch[0].length - 1;
+    const openChar = codeBlockMatch[1];
+    const closeChar = openChar === "{" ? "}" : "]";
+    let bracketCount = 1;
+    let endIndex = startIndex + 1;
+    while (endIndex < trimmed.length && bracketCount > 0) {
+      const char = trimmed[endIndex];
+      if (char === openChar) {
+        bracketCount++;
+      } else if (char === closeChar) {
+        bracketCount--;
+      }
+      endIndex++;
+    }
+    if (bracketCount === 0) {
+      return trimmed.substring(startIndex, endIndex);
+    }
+  }
   const firstBracket = Math.min(
     trimmed.indexOf("{") >= 0 ? trimmed.indexOf("{") : Infinity,
     trimmed.indexOf("[") >= 0 ? trimmed.indexOf("[") : Infinity
@@ -2299,7 +2354,7 @@ function cleanSchemaResponse(response) {
   );
   if (firstBracket < Infinity && lastBracket >= 0 && firstBracket < lastBracket) {
     const beforeFirstBracket = trimmed.substring(0, firstBracket).trim();
-    if (beforeFirstBracket === "" || beforeFirstBracket.match(/^```\w*$/)) {
+    if (beforeFirstBracket === "" || beforeFirstBracket.match(/^```\w*$/) || beforeFirstBracket.split("\n").length <= 2) {
       return trimmed.substring(firstBracket, lastBracket + 1);
     }
   }
@@ -2362,22 +2417,38 @@ function isJsonSchema(schema) {
   ];
   return jsonIndicators.some((indicator) => indicator);
 }
-function createJsonCorrectionPrompt(invalidResponse, schema, error, detailedError = "") {
-  let prompt = `Your previous response is not valid JSON and cannot be parsed. Here's what you returned:
+function createJsonCorrectionPrompt(invalidResponse, schema, error, retryCount = 0) {
+  const strengthLevels = [
+    {
+      prefix: "CRITICAL JSON ERROR:",
+      instruction: "You MUST fix this and return ONLY valid JSON.",
+      emphasis: "Return ONLY the corrected JSON, with no additional text or markdown formatting."
+    },
+    {
+      prefix: "URGENT - JSON PARSING FAILED:",
+      instruction: "This is your second chance. Return ONLY valid JSON that can be parsed by JSON.parse().",
+      emphasis: "ABSOLUTELY NO explanatory text, greetings, or formatting. ONLY JSON."
+    },
+    {
+      prefix: "FINAL ATTEMPT - CRITICAL JSON ERROR:",
+      instruction: "This is the final retry. You MUST return ONLY raw JSON without any other content.",
+      emphasis: 'EXAMPLE: {"key": "value"} NOT: ```json{"key": "value"}``` NOT: Here is the JSON: {"key": "value"}'
+    }
+  ];
+  const level = Math.min(retryCount, strengthLevels.length - 1);
+  const currentLevel = strengthLevels[level];
+  let prompt = `${currentLevel.prefix} Your previous response is not valid JSON and cannot be parsed. Here's what you returned:
 
-${invalidResponse}
+${invalidResponse.substring(0, 500)}${invalidResponse.length > 500 ? "..." : ""}
 
-Error: ${error}`;
-  if (detailedError && detailedError !== error) {
-    prompt += `
-Detailed Error: ${detailedError}`;
-  }
-  prompt += `
+Error: ${error}
+
+${currentLevel.instruction}
 
-Please correct your response to be valid JSON that matches this schema:
+Schema to match:
 ${schema}
 
-Return ONLY the corrected JSON, with no additional text or markdown formatting.`;
+${currentLevel.emphasis}`;
   return prompt;
 }
 function extractMermaidFromMarkdown(response) {
@@ -2546,9 +2617,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,
@@ -2556,11 +2683,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,
@@ -2568,7 +2694,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}
@@ -2629,10 +2755,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;
@@ -2718,6 +2852,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 ? `
 
@@ -2725,11 +2875,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.`;
@@ -3256,7 +3407,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);
@@ -3375,16 +3537,24 @@ 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...");
             }
             try {
-              const schemaPrompt = `Now you need to respond according to this schema:
+              const schemaPrompt = `CRITICAL: You MUST respond with ONLY valid JSON that matches this schema. Ignore any previous response format or instructions.
 
+Schema:
 ${options.schema}
 
-Please reformat your previous response to match this schema exactly. Only return the formatted response, no additional text.`;
+REQUIREMENTS:
+- Return ONLY the JSON object/array
+- NO additional text, explanations, or markdown
+- NO code blocks or formatting
+- The JSON must be parseable by JSON.parse()
+- DO NOT include any conversational text like "Hello" or explanations
+
+Convert your previous response content into this JSON format now. Return nothing but valid JSON.`;
               finalResult = await this.answer(schemaPrompt, [], {
                 ...options,
                 _schemaFormatted: true
@@ -3416,30 +3586,63 @@ Please reformat your previous response to match this schema exactly. Only return
                 }
               }
               if (isJsonSchema(options.schema)) {
-                const validation = validateJsonResponse(finalResult);
-                if (!validation.isValid) {
+                let validation = validateJsonResponse(finalResult);
+                let retryCount = 0;
+                const maxRetries = 3;
+                while (!validation.isValid && retryCount < maxRetries) {
                   if (this.debug) {
-                    console.log("[DEBUG] JSON validation failed:", validation.error);
+                    console.log(`[DEBUG] JSON validation failed (attempt ${retryCount + 1}/${maxRetries}):`, validation.error);
                   }
                   const correctionPrompt = createJsonCorrectionPrompt(
                     finalResult,
                     options.schema,
-                    validation.error
+                    validation.error,
+                    retryCount
                   );
                   finalResult = await this.answer(correctionPrompt, [], {
                     ...options,
                     _schemaFormatted: true
                   });
                   finalResult = cleanSchemaResponse(finalResult);
-                  const finalValidation = validateJsonResponse(finalResult);
-                  if (!finalValidation.isValid && this.debug) {
-                    console.log("[DEBUG] JSON still invalid after correction:", finalValidation.error);
+                  validation = validateJsonResponse(finalResult);
+                  retryCount++;
+                  if (this.debug && !validation.isValid && retryCount < maxRetries) {
+                    console.log(`[DEBUG] JSON still invalid after correction ${retryCount}, retrying...`);
                   }
                 }
+                if (!validation.isValid && this.debug) {
+                  console.log(`[DEBUG] JSON still invalid after ${maxRetries} correction attempts:`, validation.error);
+                }
               }
             } 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) {