@probelabs/probe 0.6.0-rc247 → 0.6.0-rc249

package/package.json

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

package/src/agent/ProbeAgent.js

@@ -85,7 +85,7 @@ import {
   validateAndFixMermaidResponse,
   tryAutoWrapForSimpleSchema
 } from './schemaUtils.js';
-import { removeThinkingTags } from './xmlParsingUtils.js';
+import { removeThinkingTags, extractThinkingContent } from './xmlParsingUtils.js';
 import { predefinedPrompts } from './shared/prompts.js';
 import {
   MCPXmlBridge,
@@ -2904,10 +2904,11 @@ Follow these instructions carefully:
       // Track initial history length for storage
       const oldHistoryLength = this.history.length;
 
-      // Reset output buffer for this answer() call — but NOT during schema correction recursion
-      // When _schemaFormatted is true, this is a recursive call to fix JSON formatting,
-      // and we must preserve the output buffer so the parent call can append it
-      if (this._outputBuffer && !options?._schemaFormatted) {
+      // Reset output buffer for this answer() call — but NOT during recursive calls.
+      // _schemaFormatted: recursive call to fix JSON formatting
+      // _completionPromptProcessed: recursive call for completionPrompt follow-up
+      // Both must preserve the output buffer so the parent call can append it.
+      if (this._outputBuffer && !options?._schemaFormatted && !options?._completionPromptProcessed) {
         this._outputBuffer.items = [];
       }
 
@@ -3554,8 +3555,25 @@ Follow these instructions carefully:
                   continue; // Don't use broken response, continue the loop
                 }
 
-                finalResult = prevContent;
-                if (this.debug) console.log(`[DEBUG] Using previous response as completion: ${finalResult.substring(0, 100)}...`);
+                // Pre-strip thinking tags to avoid losing content at final cleanup stage
+                const strippedContent = removeThinkingTags(prevContent);
+                if (strippedContent.length > 50) {
+                  // Enough content outside thinking tags — use stripped version directly
+                  finalResult = strippedContent;
+                  if (this.debug) console.log(`[DEBUG] Using previous response (thinking-stripped) as completion: ${finalResult.substring(0, 100)}...`);
+                } else {
+                  // Content was mostly/entirely inside thinking tags.
+                  // Extract thinking content and use it as the actual answer.
+                  const thinkingContent = extractThinkingContent(prevContent);
+                  if (thinkingContent && thinkingContent.length > 50) {
+                    finalResult = thinkingContent;
+                    if (this.debug) console.log(`[DEBUG] Previous response was mostly in thinking tags — using thinking content as completion: ${finalResult.substring(0, 100)}...`);
+                  } else {
+                    // Neither stripped nor thinking content is substantive — use raw as fallback
+                    finalResult = prevContent;
+                    if (this.debug) console.log(`[DEBUG] Using previous response as completion (raw): ${finalResult.substring(0, 100)}...`);
+                  }
+                }
               } else {
                 finalResult = 'Error: No previous response found to use as completion.';
                 if (this.debug) console.log(`[DEBUG] No suitable previous response found for attempt_complete shorthand`);
@@ -4296,10 +4314,16 @@ After reviewing, provide your final answer using attempt_completion.`;
 
           // Make a follow-up call with the completion prompt
           // Pass _completionPromptProcessed to prevent infinite loops
+          // Save output buffer — the recursive answer() must not destroy DSL output() content
+          const savedOutputItems = this._outputBuffer ? [...this._outputBuffer.items] : [];
           const completionResult = await this.answer(completionPromptMessage, [], {
             ...options,
             _completionPromptProcessed: true
           });
+          // Restore output buffer so the parent call can append it to the final result
+          if (this._outputBuffer) {
+            this._outputBuffer.items = savedOutputItems;
+          }
 
           // Update finalResult with the result from the completion prompt
           finalResult = completionResult;

package/src/agent/dsl/environment.js

@@ -8,6 +8,7 @@
 
 import {
   searchSchema,
+  searchAllSchema,
   querySchema,
   extractSchema,
   bashSchema,
@@ -16,6 +17,7 @@ import {
 // Map of native tool names to their Zod schemas
 const NATIVE_TOOL_SCHEMAS = {
   search: searchSchema,
+  searchAll: searchAllSchema,
   query: querySchema,
   extract: extractSchema,
   bash: bashSchema,
@@ -23,7 +25,7 @@ const NATIVE_TOOL_SCHEMAS = {
 
 // Tools that are inherently async (make network/LLM calls)
 const ALWAYS_ASYNC = new Set([
-  'search', 'query', 'extract', 'listFiles', 'searchFiles', 'bash',
+  'search', 'searchAll', 'query', 'extract', 'listFiles', 'searchFiles', 'bash',
   'LLM', 'map',
 ]);
 

package/src/tools/common.js

@@ -13,10 +13,19 @@ export const searchSchema = z.object({
 	query: z.string().describe('Search query with Elasticsearch syntax. Use quotes for exact matches, AND/OR for boolean logic, - for negation.'),
 	path: z.string().optional().default('.').describe('Path to search in. For dependencies use "go:github.com/owner/repo", "js:package_name", or "rust:cargo_name" etc.'),
 	exact: z.boolean().optional().default(false).describe('Default (false) enables stemming and keyword splitting for exploratory search - "getUserData" matches "get", "user", "data", etc. Set true for precise symbol lookup where "getUserData" matches only "getUserData". Use true when you know the exact symbol name.'),
+	maxTokens: z.number().nullable().optional().describe('Maximum tokens to return. Default is 20000. Set to null for unlimited results.'),
 	session: z.string().optional().describe('Session ID for result caching and pagination. Pass the session ID from a previous search to get additional results (next page). Results already shown in a session are automatically excluded. Omit for a fresh search.'),
 	nextPage: z.boolean().optional().default(false).describe('Set to true when requesting the next page of results. Requires passing the same session ID from the previous search output.')
 });
 
+export const searchAllSchema = z.object({
+	query: z.string().describe('Search query with Elasticsearch syntax. Use quotes for exact matches, AND/OR for boolean logic, - for negation.'),
+	path: z.string().optional().default('.').describe('Path to search in.'),
+	exact: z.boolean().optional().default(false).describe('Use exact matching instead of stemming.'),
+	maxTokensPerPage: z.number().optional().default(20000).describe('Tokens per page when paginating. Default 20000.'),
+	maxPages: z.number().optional().default(50).describe('Maximum pages to retrieve. Default 50 (safety limit).')
+});
+
 export const querySchema = z.object({
 	pattern: z.string().describe('AST pattern to search for. Use $NAME for variable names, $$$PARAMS for parameter lists, etc.'),
 	path: z.string().optional().default('.').describe('Path to search in'),

package/src/tools/executePlan.js

@@ -85,6 +85,8 @@ function buildToolImplementations(configOptions) {
         if (!searchPaths || searchPaths.length === 0) {
           searchPaths = [cwd || '.'];
         }
+        // Allow maxTokens to be passed through (null = unlimited, undefined = default 20000)
+        const maxTokens = params.maxTokens !== undefined ? params.maxTokens : 20000;
         return await search({
           query: params.query,
           path: searchPaths.join(' '),
@@ -92,7 +94,7 @@ function buildToolImplementations(configOptions) {
           allowTests: true,
           exact: params.exact || false,
           json: false,
-          maxTokens: 20000,
+          maxTokens,
           session: sessionId,
           timeout: 60,
         });
@@ -102,6 +104,70 @@ function buildToolImplementations(configOptions) {
     },
   };
 
+  // searchAll: auto-paginating search that retrieves ALL results
+  // Calls search() repeatedly with same sessionId until no more results
+  tools.searchAll = {
+    execute: async (params) => {
+      try {
+        let searchPaths;
+        if (params.path) {
+          searchPaths = parseAndResolvePaths(params.path, cwd);
+        }
+        if (!searchPaths || searchPaths.length === 0) {
+          searchPaths = [cwd || '.'];
+        }
+        const pathStr = searchPaths.join(' ');
+        const maxTokensPerPage = params.maxTokensPerPage || 20000;
+        const maxPages = params.maxPages || 50; // Safety limit
+
+        let allResults = '';
+        let pageCount = 0;
+
+        while (pageCount < maxPages) {
+          const pageResult = await search({
+            query: params.query,
+            path: pathStr,
+            cwd,
+            allowTests: true,
+            exact: params.exact || false,
+            json: false,
+            maxTokens: maxTokensPerPage,
+            session: sessionId,
+            timeout: 60,
+          });
+
+          pageCount++;
+
+          // Check if we got results
+          if (!pageResult || pageResult.trim().length === 0) {
+            break;
+          }
+
+          // Check for "All results retrieved" or "No results found" signals
+          if (pageResult.includes('All results retrieved') ||
+              pageResult.includes('No results found') ||
+              pageResult.includes('No matching code blocks found')) {
+            // Include this final page if it has content beyond the message
+            if (pageResult.trim().length > 50) {
+              allResults += (allResults ? '\n\n' : '') + pageResult;
+            }
+            break;
+          }
+
+          allResults += (allResults ? '\n\n' : '') + pageResult;
+        }
+
+        if (pageCount >= maxPages) {
+          allResults += `\n\n[Warning: Reached maximum page limit (${maxPages}). Some results may be omitted.]`;
+        }
+
+        return allResults || 'No results found.';
+      } catch (e) {
+        return `SearchAll error: ${e.message}`;
+      }
+    },
+  };
+
   tools.query = {
     execute: async (params) => {
       try {
@@ -345,7 +411,8 @@ ${lastError}
 
 RULES REMINDER:
 - search(query) is KEYWORD SEARCH — pass a search query, NOT a filename. Use extract(filepath) to read file contents.
-- search(), query(), extract(), listFiles(), bash() all return STRINGS, not arrays.
+- search() returns up to 20K tokens by default. Use search(query, path, {maxTokens: null}) for unlimited, or searchAll(query) to auto-paginate ALL results.
+- search(), searchAll(), query(), extract(), listFiles(), bash() all return STRINGS, not arrays.
 - Use chunk(stringData) to split a string into an array of chunks.
 - Use map(array, fn) only with arrays. Do NOT pass strings to map().
 - Do NOT use .map(), .forEach(), .filter(), .join() — use for..of loops instead.
@@ -684,7 +751,8 @@ return table;
 ${funcList}
 
 **Return types — IMPORTANT:**
-- \`search(query)\` → **keyword search** — pass a search query (e.g. "error handling"), NOT a filename. Returns a **string** (matching code snippets). To process parts, use \`chunk()\` to split it.
+- \`search(query)\` → **keyword search** — pass a search query (e.g. "error handling"), NOT a filename. Returns a **string** (matching code snippets, up to 20K tokens by default). Use \`{maxTokens: null}\` for unlimited.
+- \`searchAll(query)\` → **exhaustive keyword search** — auto-paginates to retrieve ALL matching results. Returns a **string** (all matching code snippets concatenated). Use for bulk analysis.
 - \`query(pattern)\` → **AST search** — pass a tree-sitter pattern. Returns a **string** (matching code elements).
 - \`extract(targets)\` → **read file contents** — pass a file path like "src/main.js" or "src/main.js:42". Use this to read specific files found by listFiles(). Returns a **string**.
 - \`listFiles(pattern)\` → **list files** — pass a glob pattern like "**/*.md". Returns an **array** of file path strings. Use directly with \`for (const f of listFiles("**/*.md"))\`.