@probelabs/probe 0.6.0-rc241 → 0.6.0-rc245

package/package.json

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

package/src/agent/ProbeAgent.js

@@ -102,6 +102,7 @@ import { formatErrorForAI, ParameterError } from '../utils/error-types.js';
 import { getCommonPrefix, toRelativePath, safeRealpath } from '../utils/path-validation.js';
 import { truncateIfNeeded, getMaxOutputTokens } from './outputTruncator.js';
 import { DelegationManager } from '../delegate.js';
+import { extractRawOutputBlocks } from '../tools/executePlan.js';
 import {
   TaskManager,
   createTaskTool,
@@ -3610,6 +3611,18 @@ Follow these instructions carefully:
 
                 let toolResultContent = typeof executionResult === 'string' ? executionResult : JSON.stringify(executionResult, null, 2);
 
+                // Extract raw output blocks and pass them through to output buffer (before truncation)
+                // This prevents LLM from processing/hallucinating large structured output from execute_plan
+                if (this._outputBuffer) {
+                  const { cleanedContent, extractedBlocks } = extractRawOutputBlocks(toolResultContent, this._outputBuffer);
+                  if (extractedBlocks.length > 0) {
+                    toolResultContent = cleanedContent;
+                    if (this.debug) {
+                      console.log(`[DEBUG] Extracted ${extractedBlocks.length} raw output blocks (${extractedBlocks.reduce((sum, b) => sum + b.length, 0)} chars) to output buffer`);
+                    }
+                  }
+                }
+
                 // Truncate if output exceeds token limit
                 try {
                   const truncateResult = await truncateIfNeeded(toolResultContent, this.tokenCounter, this.sessionId, this.maxOutputTokens);
@@ -3856,6 +3869,18 @@ Follow these instructions carefully:
                   toolResultContent = toolResultContent.split(wsPrefix).join('');
                 }
 
+                // Extract raw output blocks and pass them through to output buffer (before truncation)
+                // This prevents LLM from processing/hallucinating large structured output from execute_plan
+                if (this._outputBuffer) {
+                  const { cleanedContent, extractedBlocks } = extractRawOutputBlocks(toolResultContent, this._outputBuffer);
+                  if (extractedBlocks.length > 0) {
+                    toolResultContent = cleanedContent;
+                    if (this.debug) {
+                      console.log(`[DEBUG] Extracted ${extractedBlocks.length} raw output blocks (${extractedBlocks.reduce((sum, b) => sum + b.length, 0)} chars) to output buffer`);
+                    }
+                  }
+                }
+
                 // Truncate if output exceeds token limit
                 try {
                   const truncateResult = await truncateIfNeeded(toolResultContent, this.tokenCounter, this.sessionId, this.maxOutputTokens);
@@ -4740,12 +4765,19 @@ Convert your previous response content into actual JSON data that follows this s
       // Append DSL output buffer directly to response (bypasses LLM rewriting)
       if (this._outputBuffer && this._outputBuffer.items.length > 0 && !options._schemaFormatted) {
         const outputContent = this._outputBuffer.items.join('\n\n');
-        finalResult = (finalResult || '') + '\n\n' + outputContent;
+        if (options.schema) {
+          // Schema response — the finalResult is JSON. Wrap output in RAW_OUTPUT
+          // delimiters so clients (visor, etc.) can extract and propagate the
+          // content separately from the JSON.
+          finalResult = (finalResult || '') + '\n<<<RAW_OUTPUT>>>\n' + outputContent + '\n<<<END_RAW_OUTPUT>>>';
+        } else {
+          finalResult = (finalResult || '') + '\n\n' + outputContent;
+        }
         if (options.onStream) {
           options.onStream('\n\n' + outputContent);
         }
         if (this.debug) {
-          console.log(`[DEBUG] Appended ${this._outputBuffer.items.length} output buffer items (${outputContent.length} chars) to final result`);
+          console.log(`[DEBUG] Appended ${this._outputBuffer.items.length} output buffer items (${outputContent.length} chars) to final result${options.schema ? ' (with RAW_OUTPUT delimiters)' : ''}`);
         }
         this._outputBuffer.items = [];
       }

package/src/agent/dsl/environment.js

@@ -227,9 +227,21 @@ export function generateSandboxGlobals(options) {
   }
 
   // LLM() built-in — delegate already has its own OTEL, but we add a DSL-level span
+  // When schema is provided, auto-parse the JSON result for easier downstream processing
   if (llmCall) {
     const rawLLM = async (instruction, data, opts = {}) => {
-      return llmCall(instruction, data, opts);
+      const result = await llmCall(instruction, data, opts);
+      // Auto-parse JSON when schema is provided and result is a string
+      if (opts.schema && typeof result === 'string') {
+        try {
+          return JSON.parse(result);
+        } catch (e) {
+          // If parsing fails, return the raw string (may have formatting issues)
+          logFn?.('[LLM] Warning: schema provided but result is not valid JSON');
+          return result;
+        }
+      }
+      return result;
     };
     globals.LLM = traceToolCall('LLM', rawLLM, tracer, logFn);
   }
@@ -306,6 +318,81 @@ export function generateSandboxGlobals(options) {
     return chunks;
   };
 
+  // chunkByKey() - chunk data ensuring same-key items stay together
+  // - Chunks CAN have multiple keys (customers)
+  // - But same key NEVER splits across chunks
+  globals.chunkByKey = (data, keyFn, maxTokens = 20000) => {
+    const CHARS_PER_TOKEN = 4;
+    const maxChars = maxTokens * CHARS_PER_TOKEN;
+    const text = typeof data === 'string' ? data : JSON.stringify(data);
+
+    // Find all File: markers
+    const blockRegex = /^File: ([^\n]+)/gm;
+    const markers = [];
+    let match;
+    while ((match = blockRegex.exec(text)) !== null) {
+      markers.push({ index: match.index, file: match[1].trim() });
+    }
+
+    // No File: headers - fallback to regular chunk
+    if (markers.length === 0) {
+      return globals.chunk(data, maxTokens);
+    }
+
+    const chunks = [];
+    let currentChunk = '';
+    let currentSize = 0;
+    let keysInChunk = new Set();  // Track which keys are in current chunk
+
+    // Process each block
+    for (let i = 0; i < markers.length; i++) {
+      const start = markers[i].index;
+      const end = i + 1 < markers.length ? markers[i + 1].index : text.length;
+      const block = text.slice(start, end).trim();
+      const file = markers[i].file;
+      const key = typeof keyFn === 'function' ? keyFn(file) : file;
+
+      const blockSize = block.length + 2; // +2 for \n\n separator
+      const wouldOverflow = currentSize + blockSize > maxChars;
+      const keyAlreadyInChunk = keysInChunk.has(key);
+
+      // Decision logic:
+      // - If key already in chunk: MUST add (never split a key)
+      // - If new key and would overflow: flush first, then add
+      // - If new key and fits: add to current chunk
+
+      if (!keyAlreadyInChunk && wouldOverflow && currentChunk) {
+        // New key would overflow - flush current chunk first
+        chunks.push(currentChunk.trim());
+        currentChunk = '';
+        currentSize = 0;
+        keysInChunk = new Set();
+      }
+
+      // Add block to current chunk
+      if (currentChunk) currentChunk += '\n\n';
+      currentChunk += block;
+      currentSize += blockSize;
+      keysInChunk.add(key);
+    }
+
+    // Flush final chunk
+    if (currentChunk.trim()) {
+      chunks.push(currentChunk.trim());
+    }
+
+    return chunks.length > 0 ? chunks : [''];
+  };
+
+  // extractPaths() - extract unique file paths from search results
+  // Parses File: headers and returns deduplicated array of paths
+  globals.extractPaths = (searchResults) => {
+    const text = typeof searchResults === 'string' ? searchResults : JSON.stringify(searchResults);
+    const matches = text.match(/^File: ([^\n]+)/gm) || [];
+    const paths = matches.map(m => m.replace('File: ', '').trim());
+    return [...new Set(paths)];
+  };
+
   // Utility functions (pure, no async)
   globals.log = (message) => {
     // Collected by the runtime for the execution log

package/src/tools/executePlan.js

@@ -420,6 +420,51 @@ RULES REMINDER:
   });
 }
 
+// Delimiters for raw output passthrough - prevents LLM from processing/hallucinating large structured output
+export const RAW_OUTPUT_START = '<<<RAW_OUTPUT>>>';
+export const RAW_OUTPUT_END = '<<<END_RAW_OUTPUT>>>';
+
+/**
+ * Extract raw output blocks from tool result content and pass them through to the output buffer.
+ * This prevents parent LLMs from processing/hallucinating large structured output.
+ *
+ * @param {string} content - The tool result content
+ * @param {Object} [outputBuffer] - The output buffer to append extracted content to
+ * @returns {{ cleanedContent: string, extractedBlocks: string[] }} - Content with blocks removed and extracted blocks
+ */
+export function extractRawOutputBlocks(content, outputBuffer = null) {
+  if (typeof content !== 'string') {
+    return { cleanedContent: content, extractedBlocks: [] };
+  }
+
+  const extractedBlocks = [];
+  const regex = new RegExp(`${RAW_OUTPUT_START}\\n([\\s\\S]*?)\\n${RAW_OUTPUT_END}`, 'g');
+
+  let cleanedContent = content;
+  let match;
+
+  // Extract all blocks
+  while ((match = regex.exec(content)) !== null) {
+    extractedBlocks.push(match[1]);
+  }
+
+  // Remove the blocks and any following instruction line from content
+  cleanedContent = content
+    .replace(new RegExp(`${RAW_OUTPUT_START}\\n[\\s\\S]*?\\n${RAW_OUTPUT_END}`, 'g'), '')
+    .replace(/\n\n\[The above raw output \(\d+ chars\) will be passed directly to the final response\. Do NOT repeat, summarize, or modify it\.\]/g, '')
+    .trim();
+
+  // If output buffer provided, append extracted content
+  if (outputBuffer && extractedBlocks.length > 0) {
+    for (const block of extractedBlocks) {
+      outputBuffer.items = outputBuffer.items || [];
+      outputBuffer.items.push(block);
+    }
+  }
+
+  return { cleanedContent, extractedBlocks };
+}
+
 function formatSuccess(result, description, attempt, outputBuffer) {
   let output = '';
 
@@ -452,10 +497,12 @@ function formatSuccess(result, description, attempt, outputBuffer) {
     }
   }
 
-  // If output buffer has content, tell the LLM the data was written to direct output
+  // If output buffer has content, wrap it in delimiters for passthrough
+  // This prevents parent LLMs from processing/hallucinating the raw data
   if (outputBuffer && outputBuffer.items && outputBuffer.items.length > 0) {
-    const totalChars = outputBuffer.items.reduce((sum, item) => sum + item.length, 0);
-    output += `\n\n[Output buffer: ${totalChars} chars written via output(). This content will be appended directly to your response. Do NOT repeat or summarize it.]`;
+    const rawContent = outputBuffer.items.join('\n');
+    output += `\n\n${RAW_OUTPUT_START}\n${rawContent}\n${RAW_OUTPUT_END}`;
+    output += `\n\n[The above raw output (${rawContent.length} chars) will be passed directly to the final response. Do NOT repeat, summarize, or modify it.]`;
   }
 
   return output;