@probelabs/probe 0.6.0-rc164 → 0.6.0-rc165

package/build/agent/ProbeAgent.js

@@ -1706,16 +1706,25 @@ When troubleshooting:
         }
 
         // Parse tool call from response with valid tools list
-        const validTools = [
-          'search', 'query', 'extract', 'listFiles', 'searchFiles', 'attempt_completion'
-        ];
-        if (this.allowEdit) {
+        // Build validTools based on allowedTools configuration (same pattern as getSystemMessage)
+        const validTools = [];
+        if (this.allowedTools.isEnabled('search')) validTools.push('search');
+        if (this.allowedTools.isEnabled('query')) validTools.push('query');
+        if (this.allowedTools.isEnabled('extract')) validTools.push('extract');
+        if (this.allowedTools.isEnabled('listFiles')) validTools.push('listFiles');
+        if (this.allowedTools.isEnabled('searchFiles')) validTools.push('searchFiles');
+        if (this.allowedTools.isEnabled('attempt_completion')) validTools.push('attempt_completion');
+
+        // Edit tools (require both allowEdit flag AND allowedTools permission)
+        if (this.allowEdit && this.allowedTools.isEnabled('implement')) {
           validTools.push('implement', 'edit', 'create');
         }
-        if (this.enableBash) {
+        // Bash tool (require both enableBash flag AND allowedTools permission)
+        if (this.enableBash && this.allowedTools.isEnabled('bash')) {
           validTools.push('bash');
         }
-        if (this.enableDelegate) {
+        // Delegate tool (require both enableDelegate flag AND allowedTools permission)
+        if (this.enableDelegate && this.allowedTools.isEnabled('delegate')) {
           validTools.push('delegate');
         }
         

package/build/agent/index.js

@@ -54263,6 +54263,7 @@ __export(schemaUtils_exports, {
   createMermaidCorrectionPrompt: () => createMermaidCorrectionPrompt,
   createSchemaDefinitionCorrectionPrompt: () => createSchemaDefinitionCorrectionPrompt,
   decodeHtmlEntities: () => decodeHtmlEntities,
+  extractMermaidFromJson: () => extractMermaidFromJson,
   extractMermaidFromMarkdown: () => extractMermaidFromMarkdown,
   generateExampleFromSchema: () => generateExampleFromSchema,
   generateSchemaInstructions: () => generateSchemaInstructions,
@@ -54270,6 +54271,7 @@ __export(schemaUtils_exports, {
   isJsonSchemaDefinition: () => isJsonSchemaDefinition,
   isMermaidSchema: () => isMermaidSchema,
   processSchemaResponse: () => processSchemaResponse,
+  replaceMermaidDiagramsInJson: () => replaceMermaidDiagramsInJson,
   replaceMermaidDiagramsInMarkdown: () => replaceMermaidDiagramsInMarkdown,
   tryMaidAutoFix: () => tryMaidAutoFix,
   validateAndFixMermaidResponse: () => validateAndFixMermaidResponse,
@@ -54373,6 +54375,49 @@ function decodeHtmlEntities(text) {
   }
   return decoded;
 }
+function normalizeJsonQuotes(str) {
+  if (!str || typeof str !== "string") {
+    return str;
+  }
+  if (!str.includes("'")) {
+    return str;
+  }
+  let result = "";
+  let inDoubleQuote = false;
+  let inSingleQuote = false;
+  let escaped = false;
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i];
+    const prevChar = i > 0 ? str[i - 1] : "";
+    if (escaped) {
+      result += char;
+      escaped = false;
+      continue;
+    }
+    if (char === "\\") {
+      escaped = true;
+      result += char;
+      continue;
+    }
+    if (char === '"' && !inSingleQuote) {
+      inDoubleQuote = !inDoubleQuote;
+      result += char;
+      continue;
+    }
+    if (char === "'" && !inDoubleQuote) {
+      if (inSingleQuote) {
+        result += '"';
+        inSingleQuote = false;
+      } else {
+        result += '"';
+        inSingleQuote = true;
+      }
+      continue;
+    }
+    result += char;
+  }
+  return result;
+}
 function cleanSchemaResponse(response) {
   if (!response || typeof response !== "string") {
     return response;
@@ -54380,23 +54425,27 @@ function cleanSchemaResponse(response) {
   const trimmed = response.trim();
   const jsonBlockMatch = trimmed.match(/```json\s*\n([\s\S]*?)\n```/);
   if (jsonBlockMatch) {
-    return jsonBlockMatch[1].trim();
+    return normalizeJsonQuotes(jsonBlockMatch[1].trim());
   }
   const anyBlockMatch = trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
   if (anyBlockMatch) {
-    return anyBlockMatch[1].trim();
+    return normalizeJsonQuotes(anyBlockMatch[1].trim());
   }
   const codeBlockPatterns = [
     /```json\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /`([{\[][\s\S]*?[}\]])`/
+    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/
   ];
   for (const pattern of codeBlockPatterns) {
     const match2 = trimmed.match(pattern);
     if (match2) {
-      return match2[1].trim();
+      return normalizeJsonQuotes(match2[1].trim());
     }
   }
+  const singleBacktickPattern = /^`([{\[][\s\S]*?[}\]])`$/;
+  const singleBacktickMatch = trimmed.match(singleBacktickPattern);
+  if (singleBacktickMatch) {
+    return normalizeJsonQuotes(singleBacktickMatch[1].trim());
+  }
   const codeBlockStartPattern = /```(?:json)?\s*\n?\s*([{\[])/;
   const codeBlockMatch = trimmed.match(codeBlockStartPattern);
   if (codeBlockMatch) {
@@ -54415,7 +54464,7 @@ function cleanSchemaResponse(response) {
       endIndex++;
     }
     if (bracketCount === 0) {
-      return trimmed.substring(startIndex, endIndex);
+      return normalizeJsonQuotes(trimmed.substring(startIndex, endIndex));
     }
   }
   let cleaned = trimmed;
@@ -54427,7 +54476,7 @@ function cleanSchemaResponse(response) {
   const isJsonObject = firstChar === "{" && lastChar === "}";
   const isJsonArray = firstChar === "[" && lastChar === "]";
   if (isJsonObject || isJsonArray) {
-    return cleaned;
+    return normalizeJsonQuotes(cleaned);
   }
   return response;
 }
@@ -54837,10 +54886,66 @@ function isMermaidSchema(schema) {
   ];
   return mermaidIndicators.some((indicator) => indicator);
 }
+function extractMermaidFromJson(response) {
+  if (!response || typeof response !== "string") {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+  let jsonContent = response.trim();
+  const jsonBlockMatch = jsonContent.match(/```json\s*\n([\s\S]*?)\n```/);
+  if (jsonBlockMatch) {
+    jsonContent = jsonBlockMatch[1].trim();
+  } else {
+    const anyBlockMatch = jsonContent.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
+    if (anyBlockMatch) {
+      jsonContent = anyBlockMatch[1].trim();
+    }
+  }
+  let parsedJson;
+  try {
+    parsedJson = JSON.parse(jsonContent);
+  } catch (e) {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+  const diagrams = [];
+  const jsonPaths = [];
+  function searchObject(obj, path7 = []) {
+    if (typeof obj === "string") {
+      const mermaidPattern = /```mermaid([^\n`]*?)(?:\n|\\n)([\s\S]*?)```/gi;
+      let match2;
+      while ((match2 = mermaidPattern.exec(obj)) !== null) {
+        const attributes = match2[1] ? match2[1].trim() : "";
+        const content = match2[2].replace(/\\n/g, "\n");
+        diagrams.push({
+          content,
+          fullMatch: match2[0],
+          startIndex: match2.index,
+          endIndex: match2.index + match2[0].length,
+          attributes,
+          isInJson: true,
+          jsonPath: path7.join(".")
+        });
+        jsonPaths.push(path7.join("."));
+      }
+    } else if (Array.isArray(obj)) {
+      obj.forEach((item, index) => searchObject(item, [...path7, `[${index}]`]));
+    } else if (obj && typeof obj === "object") {
+      Object.entries(obj).forEach(([key, value]) => searchObject(value, [...path7, key]));
+    }
+  }
+  searchObject(parsedJson);
+  return { diagrams, jsonPaths, parsedJson };
+}
 function extractMermaidFromMarkdown(response) {
   if (!response || typeof response !== "string") {
     return { diagrams: [], cleanedResponse: response };
   }
+  const trimmed = response.trim();
+  if (trimmed.startsWith("{") || trimmed.startsWith("[") || trimmed.includes("```json")) {
+    const jsonResult = extractMermaidFromJson(response);
+    if (jsonResult.diagrams.length > 0) {
+      return { diagrams: jsonResult.diagrams, cleanedResponse: response };
+    }
+  }
   const mermaidBlockRegex = /```mermaid([^\n]*)\n([\s\S]*?)```/gi;
   const diagrams = [];
   let match2;
@@ -54852,11 +54957,66 @@ function extractMermaidFromMarkdown(response) {
       fullMatch: match2[0],
       startIndex: match2.index,
       endIndex: match2.index + match2[0].length,
-      attributes
+      attributes,
+      isInJson: false
     });
   }
   return { diagrams, cleanedResponse: response };
 }
+function replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams) {
+  if (!originalResponse || typeof originalResponse !== "string") {
+    return originalResponse;
+  }
+  if (!correctedDiagrams || correctedDiagrams.length === 0) {
+    return originalResponse;
+  }
+  const jsonResult = extractMermaidFromJson(originalResponse);
+  if (!jsonResult.parsedJson) {
+    return originalResponse;
+  }
+  let modifiedJson = jsonResult.parsedJson;
+  for (const diagram of correctedDiagrams) {
+    if (!diagram.jsonPath || !diagram.isInJson) {
+      continue;
+    }
+    const pathParts = diagram.jsonPath.split(".").filter((p) => p);
+    let current = modifiedJson;
+    for (let i = 0; i < pathParts.length - 1; i++) {
+      const part = pathParts[i];
+      if (part.startsWith("[") && part.endsWith("]")) {
+        const index = parseInt(part.slice(1, -1), 10);
+        current = current[index];
+      } else {
+        current = current[part];
+      }
+    }
+    const lastPart = pathParts[pathParts.length - 1];
+    const attributesStr = diagram.attributes ? ` ${diagram.attributes}` : "";
+    const newCodeBlock = `\`\`\`mermaid${attributesStr}
+${diagram.content}
+\`\`\``;
+    if (lastPart.startsWith("[") && lastPart.endsWith("]")) {
+      const index = parseInt(lastPart.slice(1, -1), 10);
+      const originalString = current[index];
+      current[index] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    } else {
+      const originalString = current[lastPart];
+      current[lastPart] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    }
+  }
+  const modifiedJsonString = JSON.stringify(modifiedJson, null, 2);
+  const trimmed = originalResponse.trim();
+  if (trimmed.match(/```json\s*\n([\s\S]*?)\n```/)) {
+    return originalResponse.replace(/```json\s*\n([\s\S]*?)\n```/, `\`\`\`json
+${modifiedJsonString}
+\`\`\``);
+  } else if (trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/)) {
+    return originalResponse.replace(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/, `\`\`\`
+${modifiedJsonString}
+\`\`\``);
+  }
+  return modifiedJsonString;
+}
 function replaceMermaidDiagramsInMarkdown(originalResponse, correctedDiagrams) {
   if (!originalResponse || typeof originalResponse !== "string") {
     return originalResponse;
@@ -54864,6 +55024,10 @@ function replaceMermaidDiagramsInMarkdown(originalResponse, correctedDiagrams) {
   if (!correctedDiagrams || correctedDiagrams.length === 0) {
     return originalResponse;
   }
+  const hasJsonDiagrams = correctedDiagrams.some((d) => d.isInJson);
+  if (hasJsonDiagrams) {
+    return replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams);
+  }
   let modifiedResponse = originalResponse;
   const sortedDiagrams = [...correctedDiagrams].sort((a, b) => b.startIndex - a.startIndex);
   for (const diagram of sortedDiagrams) {
@@ -58529,21 +58693,20 @@ You are working with a repository located at: ${searchDirectory}
             if (assistantResponseContent) {
               await this.processImageReferences(assistantResponseContent);
             }
-            const validTools = [
-              "search",
-              "query",
-              "extract",
-              "listFiles",
-              "searchFiles",
-              "attempt_completion"
-            ];
-            if (this.allowEdit) {
+            const validTools = [];
+            if (this.allowedTools.isEnabled("search")) validTools.push("search");
+            if (this.allowedTools.isEnabled("query")) validTools.push("query");
+            if (this.allowedTools.isEnabled("extract")) validTools.push("extract");
+            if (this.allowedTools.isEnabled("listFiles")) validTools.push("listFiles");
+            if (this.allowedTools.isEnabled("searchFiles")) validTools.push("searchFiles");
+            if (this.allowedTools.isEnabled("attempt_completion")) validTools.push("attempt_completion");
+            if (this.allowEdit && this.allowedTools.isEnabled("implement")) {
               validTools.push("implement", "edit", "create");
             }
-            if (this.enableBash) {
+            if (this.enableBash && this.allowedTools.isEnabled("bash")) {
               validTools.push("bash");
             }
-            if (this.enableDelegate) {
+            if (this.enableDelegate && this.allowedTools.isEnabled("delegate")) {
               validTools.push("delegate");
             }
             const nativeTools = validTools;

package/build/agent/schemaUtils.js

@@ -165,6 +165,74 @@ export function decodeHtmlEntities(text) {
   return decoded;
 }
 
+/**
+ * Normalize JavaScript syntax to valid JSON syntax
+ * Converts single quotes to double quotes for strings in JSON-like structures
+ *
+ * @param {string} str - String that might contain JavaScript array/object syntax
+ * @returns {string} - String with single quotes normalized to double quotes
+ */
+function normalizeJsonQuotes(str) {
+  if (!str || typeof str !== 'string') {
+    return str;
+  }
+
+  // Quick check: if there are no single quotes, no need to normalize
+  if (!str.includes("'")) {
+    return str;
+  }
+
+  let result = '';
+  let inDoubleQuote = false;
+  let inSingleQuote = false;
+  let escaped = false;
+
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i];
+    const prevChar = i > 0 ? str[i - 1] : '';
+
+    // Handle escape sequences
+    if (escaped) {
+      result += char;
+      escaped = false;
+      continue;
+    }
+
+    if (char === '\\') {
+      escaped = true;
+      result += char;
+      continue;
+    }
+
+    // Track when we're inside double-quoted strings
+    if (char === '"' && !inSingleQuote) {
+      inDoubleQuote = !inDoubleQuote;
+      result += char;
+      continue;
+    }
+
+    // Convert single quotes to double quotes (when not inside double quotes)
+    if (char === "'" && !inDoubleQuote) {
+      // Check if this is a single quote inside a string value (like "It's")
+      // If we're already in a single-quoted string, toggle the state
+      if (inSingleQuote) {
+        // Closing single quote - convert to double quote
+        result += '"';
+        inSingleQuote = false;
+      } else {
+        // Opening single quote - convert to double quote
+        result += '"';
+        inSingleQuote = true;
+      }
+      continue;
+    }
+
+    result += char;
+  }
+
+  return result;
+}
+
 /**
  * Clean AI response by extracting JSON content when response contains JSON
  * Only processes responses that contain JSON structures { or [
@@ -189,29 +257,36 @@ export function cleanSchemaResponse(response) {
   // Try with json language specifier
   const jsonBlockMatch = trimmed.match(/```json\s*\n([\s\S]*?)\n```/);
   if (jsonBlockMatch) {
-    return jsonBlockMatch[1].trim();
+    return normalizeJsonQuotes(jsonBlockMatch[1].trim());
   }
 
   // Try any code block with JSON content
   const anyBlockMatch = trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
   if (anyBlockMatch) {
-    return anyBlockMatch[1].trim();
+    return normalizeJsonQuotes(anyBlockMatch[1].trim());
   }
 
   // Legacy patterns for more specific matching
   const codeBlockPatterns = [
     /```json\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/,
-    /`([{\[][\s\S]*?[}\]])`/
+    /```\s*\n?([{\[][\s\S]*?[}\]])\s*\n?```/
   ];
 
   for (const pattern of codeBlockPatterns) {
     const match = trimmed.match(pattern);
     if (match) {
-      return match[1].trim();
+      return normalizeJsonQuotes(match[1].trim());
     }
   }
 
+  // Single backtick pattern - ONLY if the entire input is just the code block
+  // This prevents extracting inline code from within JSON objects (e.g., `['*']` from markdown text)
+  const singleBacktickPattern = /^`([{\[][\s\S]*?[}\]])`$/;
+  const singleBacktickMatch = trimmed.match(singleBacktickPattern);
+  if (singleBacktickMatch) {
+    return normalizeJsonQuotes(singleBacktickMatch[1].trim());
+  }
+
   // Look for code block start followed immediately by JSON
   const codeBlockStartPattern = /```(?:json)?\s*\n?\s*([{\[])/;
   const codeBlockMatch = trimmed.match(codeBlockStartPattern);
@@ -236,7 +311,7 @@ export function cleanSchemaResponse(response) {
     }
 
     if (bracketCount === 0) {
-      return trimmed.substring(startIndex, endIndex);
+      return normalizeJsonQuotes(trimmed.substring(startIndex, endIndex));
     }
   }
 
@@ -261,7 +336,8 @@ export function cleanSchemaResponse(response) {
   const isJsonArray = firstChar === '[' && lastChar === ']';
 
   if (isJsonObject || isJsonArray) {
-    return cleaned;
+    // Normalize JavaScript syntax (single quotes) to valid JSON syntax (double quotes)
+    return normalizeJsonQuotes(cleaned);
   }
 
   return response; // Return original if no extractable JSON found
@@ -826,6 +902,76 @@ export function isMermaidSchema(schema) {
   return mermaidIndicators.some(indicator => indicator);
 }
 
+/**
+ * Extract Mermaid diagrams from JSON string values
+ * Handles escaped newlines and backticks within JSON strings
+ * @param {string} response - Response that may contain JSON with mermaid blocks in string values
+ * @returns {Object} - {diagrams: Array, jsonPaths: Array, parsedJson: Object|null}
+ */
+export function extractMermaidFromJson(response) {
+  if (!response || typeof response !== 'string') {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+
+  // Try to extract JSON from code blocks first
+  let jsonContent = response.trim();
+  const jsonBlockMatch = jsonContent.match(/```json\s*\n([\s\S]*?)\n```/);
+  if (jsonBlockMatch) {
+    jsonContent = jsonBlockMatch[1].trim();
+  } else {
+    const anyBlockMatch = jsonContent.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/);
+    if (anyBlockMatch) {
+      jsonContent = anyBlockMatch[1].trim();
+    }
+  }
+
+  // Try to parse as JSON
+  let parsedJson;
+  try {
+    parsedJson = JSON.parse(jsonContent);
+  } catch (e) {
+    return { diagrams: [], jsonPaths: [], parsedJson: null };
+  }
+
+  const diagrams = [];
+  const jsonPaths = [];
+
+  // Recursively search for mermaid diagrams in JSON string values
+  function searchObject(obj, path = []) {
+    if (typeof obj === 'string') {
+      // Look for mermaid code blocks in the string value
+      // Handle both escaped (\n) and literal newlines
+      const mermaidPattern = /```mermaid([^\n`]*?)(?:\n|\\n)([\s\S]*?)```/gi;
+      let match;
+
+      while ((match = mermaidPattern.exec(obj)) !== null) {
+        const attributes = match[1] ? match[1].trim() : '';
+        // Unescape the content (replace \\n with actual newlines)
+        const content = match[2].replace(/\\n/g, '\n');
+
+        diagrams.push({
+          content: content,
+          fullMatch: match[0],
+          startIndex: match.index,
+          endIndex: match.index + match[0].length,
+          attributes: attributes,
+          isInJson: true,
+          jsonPath: path.join('.')
+        });
+        jsonPaths.push(path.join('.'));
+      }
+    } else if (Array.isArray(obj)) {
+      obj.forEach((item, index) => searchObject(item, [...path, `[${index}]`]));
+    } else if (obj && typeof obj === 'object') {
+      Object.entries(obj).forEach(([key, value]) => searchObject(value, [...path, key]));
+    }
+  }
+
+  searchObject(parsedJson);
+
+  return { diagrams, jsonPaths, parsedJson };
+}
+
 /**
  * Extract Mermaid diagrams from markdown code blocks with position tracking
  * @param {string} response - Response that may contain markdown with mermaid blocks
@@ -836,6 +982,16 @@ export function extractMermaidFromMarkdown(response) {
     return { diagrams: [], cleanedResponse: response };
   }
 
+  // First check if this looks like a JSON response - if so, use JSON-aware extraction
+  const trimmed = response.trim();
+  if ((trimmed.startsWith('{') || trimmed.startsWith('[')) ||
+      trimmed.includes('```json')) {
+    const jsonResult = extractMermaidFromJson(response);
+    if (jsonResult.diagrams.length > 0) {
+      return { diagrams: jsonResult.diagrams, cleanedResponse: response };
+    }
+  }
+
   // Find all mermaid code blocks with enhanced regex to capture more variations
   // This regex captures optional attributes on same line as ```mermaid, and all diagram content
   const mermaidBlockRegex = /```mermaid([^\n]*)\n([\s\S]*?)```/gi;
@@ -854,7 +1010,8 @@ export function extractMermaidFromMarkdown(response) {
       fullMatch: match[0],
       startIndex: match.index,
       endIndex: match.index + match[0].length,
-      attributes: attributes
+      attributes: attributes,
+      isInJson: false
     });
   }
 
@@ -862,9 +1019,84 @@ export function extractMermaidFromMarkdown(response) {
   return { diagrams, cleanedResponse: response };
 }
 
+/**
+ * Replace mermaid diagrams in JSON string values with corrected versions
+ * @param {string} originalResponse - Original response with JSON
+ * @param {Array} correctedDiagrams - Array of corrected diagram objects with jsonPath
+ * @returns {string} - Response with corrected diagrams properly escaped in JSON
+ */
+export function replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams) {
+  if (!originalResponse || typeof originalResponse !== 'string') {
+    return originalResponse;
+  }
+
+  if (!correctedDiagrams || correctedDiagrams.length === 0) {
+    return originalResponse;
+  }
+
+  // Extract and parse JSON
+  const jsonResult = extractMermaidFromJson(originalResponse);
+  if (!jsonResult.parsedJson) {
+    return originalResponse;
+  }
+
+  let modifiedJson = jsonResult.parsedJson;
+
+  // Replace diagrams in the JSON object
+  for (const diagram of correctedDiagrams) {
+    if (!diagram.jsonPath || !diagram.isInJson) {
+      continue;
+    }
+
+    // Navigate to the path and replace the content
+    const pathParts = diagram.jsonPath.split('.').filter(p => p);
+    let current = modifiedJson;
+
+    for (let i = 0; i < pathParts.length - 1; i++) {
+      const part = pathParts[i];
+      if (part.startsWith('[') && part.endsWith(']')) {
+        const index = parseInt(part.slice(1, -1), 10);
+        current = current[index];
+      } else {
+        current = current[part];
+      }
+    }
+
+    // Get the last key/index
+    const lastPart = pathParts[pathParts.length - 1];
+    const attributesStr = diagram.attributes ? ` ${diagram.attributes}` : '';
+    const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${diagram.content}\n\`\`\``;
+
+    if (lastPart.startsWith('[') && lastPart.endsWith(']')) {
+      const index = parseInt(lastPart.slice(1, -1), 10);
+      const originalString = current[index];
+      // The fullMatch from extraction has unescaped newlines, so we need to match that
+      current[index] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    } else {
+      const originalString = current[lastPart];
+      // The fullMatch from extraction has unescaped newlines, so we need to match that
+      current[lastPart] = originalString.replace(diagram.fullMatch, newCodeBlock);
+    }
+  }
+
+  // Reconstruct the response with modified JSON
+  const modifiedJsonString = JSON.stringify(modifiedJson, null, 2);
+
+  // Check if original was in a code block
+  const trimmed = originalResponse.trim();
+  if (trimmed.match(/```json\s*\n([\s\S]*?)\n```/)) {
+    return originalResponse.replace(/```json\s*\n([\s\S]*?)\n```/, `\`\`\`json\n${modifiedJsonString}\n\`\`\``);
+  } else if (trimmed.match(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/)) {
+    return originalResponse.replace(/```\s*\n([{\[][\s\S]*?[}\]])\s*```/, `\`\`\`\n${modifiedJsonString}\n\`\`\``);
+  }
+
+  return modifiedJsonString;
+}
+
 /**
  * Replace mermaid diagrams in original markdown with corrected versions
- * @param {string} originalResponse - Original response with markdown
+ * Automatically detects JSON vs markdown format and uses appropriate replacement
+ * @param {string} originalResponse - Original response with markdown or JSON
  * @param {Array} correctedDiagrams - Array of corrected diagram objects
  * @returns {string} - Response with corrected diagrams in original format
  */
@@ -877,22 +1109,28 @@ export function replaceMermaidDiagramsInMarkdown(originalResponse, correctedDiag
     return originalResponse;
   }
 
+  // Check if any diagrams are in JSON format
+  const hasJsonDiagrams = correctedDiagrams.some(d => d.isInJson);
+  if (hasJsonDiagrams) {
+    return replaceMermaidDiagramsInJson(originalResponse, correctedDiagrams);
+  }
+
   let modifiedResponse = originalResponse;
-  
+
   // Sort diagrams by start index in reverse order to preserve indices during replacement
   const sortedDiagrams = [...correctedDiagrams].sort((a, b) => b.startIndex - a.startIndex);
-  
+
   for (const diagram of sortedDiagrams) {
     // Reconstruct the code block with original attributes if they existed
     const attributesStr = diagram.attributes ? ` ${diagram.attributes}` : '';
     const newCodeBlock = `\`\`\`mermaid${attributesStr}\n${diagram.content}\n\`\`\``;
-    
+
     // Replace the original code block
-    modifiedResponse = modifiedResponse.slice(0, diagram.startIndex) + 
-                     newCodeBlock + 
+    modifiedResponse = modifiedResponse.slice(0, diagram.startIndex) +
+                     newCodeBlock +
                      modifiedResponse.slice(diagram.endIndex);
   }
-  
+
   return modifiedResponse;
 }