@probelabs/probe 0.6.0-rc163 → 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');
         }
         
@@ -2009,21 +2018,8 @@ When troubleshooting:
           // Add assistant response and ask for tool usage
           currentMessages.push({ role: 'assistant', content: assistantResponseContent });
 
-          // Build appropriate reminder message based on whether schema is provided
-          let reminderContent;
-          if (options.schema) {  // Apply for ANY schema, not just JSON schemas
-            // When schema is provided, use the same instructions as initial message
-            reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
-
-Remember: Use proper XML format with BOTH opening and closing tags:
-
-<tool_name>
-<parameter>value</parameter>
-</tool_name>
-` + generateSchemaInstructions(options.schema, { debug: this.debug }).trim();
-          } else {
-            // Standard reminder without schema
-            reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
+          // Standard reminder - schema was already provided in initial message
+          const reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
 
 Remember: Use proper XML format with BOTH opening and closing tags:
 
@@ -2035,7 +2031,6 @@ Or for quick completion if your previous response was already correct and comple
 <attempt_complete>
 
 IMPORTANT: When using <attempt_complete>, this must be the ONLY content in your response. No additional text, explanations, or other content should be included. This tag signals to reuse your previous response as the final answer.`;
-          }
 
           currentMessages.push({
             role: 'user',

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,
@@ -54315,15 +54317,6 @@ function generateSchemaInstructions(schema, options = {}) {
     instructions += `${JSON.stringify(parsedSchema, null, 2)}
 
 `;
-    const exampleObj = generateExampleFromSchema(parsedSchema, { debug });
-    if (exampleObj) {
-      instructions += `Example:
-<attempt_completion>
-${JSON.stringify(exampleObj, null, 2)}
-</attempt_completion>
-
-`;
-    }
   } catch (e) {
     if (debug) {
       console.error("[DEBUG] generateSchemaInstructions: Failed to parse schema:", e.message);
@@ -54332,7 +54325,7 @@ ${JSON.stringify(exampleObj, null, 2)}
 
 `;
   }
-  instructions += "Your response inside attempt_completion must be ONLY valid JSON - no plain text, no explanations, no markdown.";
+  instructions += "Your response inside attempt_completion must be ONLY valid JSON - no plain text, no explanations, no markdown.\n\nIMPORTANT: First complete the requested analysis/task thoroughly, then provide your final answer in the JSON format above.";
   return instructions;
 }
 function enforceNoAdditionalProperties(schema) {
@@ -54382,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;
@@ -54389,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) {
@@ -54424,7 +54464,7 @@ function cleanSchemaResponse(response) {
       endIndex++;
     }
     if (bracketCount === 0) {
-      return trimmed.substring(startIndex, endIndex);
+      return normalizeJsonQuotes(trimmed.substring(startIndex, endIndex));
     }
   }
   let cleaned = trimmed;
@@ -54436,7 +54476,7 @@ function cleanSchemaResponse(response) {
   const isJsonObject = firstChar === "{" && lastChar === "}";
   const isJsonArray = firstChar === "[" && lastChar === "]";
   if (isJsonObject || isJsonArray) {
-    return cleaned;
+    return normalizeJsonQuotes(cleaned);
   }
   return response;
 }
@@ -54846,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;
@@ -54861,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;
@@ -54873,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) {
@@ -58538,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;
@@ -58791,18 +58945,7 @@ Error: Unknown tool '${toolName}'. Available tools: ${allAvailableTools.join(",
                 break;
               }
               currentMessages.push({ role: "assistant", content: assistantResponseContent });
-              let reminderContent;
-              if (options.schema) {
-                reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
-
-Remember: Use proper XML format with BOTH opening and closing tags:
-
-<tool_name>
-<parameter>value</parameter>
-</tool_name>
-` + generateSchemaInstructions(options.schema, { debug: this.debug }).trim();
-              } else {
-                reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
+              const reminderContent = `Please use one of the available tools to help answer the question, or use attempt_completion if you have enough information to provide a final answer.
 
 Remember: Use proper XML format with BOTH opening and closing tags:
 
@@ -58814,7 +58957,6 @@ Or for quick completion if your previous response was already correct and comple
 <attempt_complete>
 
 IMPORTANT: When using <attempt_complete>, this must be the ONLY content in your response. No additional text, explanations, or other content should be included. This tag signals to reuse your previous response as the final answer.`;
-              }
               currentMessages.push({
                 role: "user",
                 content: reminderContent