@probelabs/probe 0.6.0-rc205 → 0.6.0-rc207

package/package.json

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

package/src/agent/ProbeAgent.js

@@ -70,6 +70,7 @@ import { RetryManager, createRetryManagerFromEnv } from './RetryManager.js';
 import { FallbackManager, createFallbackManagerFromEnv, buildFallbackProvidersFromEnv } from './FallbackManager.js';
 import { handleContextLimitError } from './contextCompactor.js';
 import { formatErrorForAI, ParameterError } from '../utils/error-types.js';
+import { truncateIfNeeded, getMaxOutputTokens } from './outputTruncator.js';
 import {
   TaskManager,
   createTaskTool,
@@ -90,6 +91,45 @@ const MAX_TOOL_ITERATIONS = (() => {
 })();
 const MAX_HISTORY_MESSAGES = 100;
 
+/**
+ * Extract tool name from wrapped_tool:toolName format
+ * @param {string} wrappedToolError - Error string in format 'wrapped_tool:toolName'
+ * @returns {string} The extracted tool name or 'unknown' if format is invalid
+ */
+function extractWrappedToolName(wrappedToolError) {
+  if (!wrappedToolError || typeof wrappedToolError !== 'string') {
+    return 'unknown';
+  }
+  const colonIndex = wrappedToolError.indexOf(':');
+  return colonIndex !== -1 ? wrappedToolError.slice(colonIndex + 1) : 'unknown';
+}
+
+/**
+ * Check if an error indicates a wrapped tool format error
+ * @param {string|null} error - Error from detectUnrecognizedToolCall
+ * @returns {boolean} True if it's a wrapped tool error
+ */
+function isWrappedToolError(error) {
+  return error && typeof error === 'string' && error.startsWith('wrapped_tool:');
+}
+
+/**
+ * Create error message for wrapped tool format issues
+ * @param {string} wrappedToolName - The tool name that was incorrectly wrapped
+ * @returns {string} User-friendly error message with correct format instructions
+ */
+function createWrappedToolErrorMessage(wrappedToolName) {
+  return `Your response contained an incorrectly formatted tool call (${wrappedToolName} wrapped in XML tags). This cannot be used.
+
+Please use the CORRECT format:
+
+<${wrappedToolName}>
+Your content here
+</${wrappedToolName}>
+
+Do NOT wrap in other tags like <api_call>, <tool_name>, <function>, etc.`;
+}
+
 // Supported image file extensions (imported from shared config)
 
 // Maximum image file size (20MB) to prevent OOM attacks
@@ -145,6 +185,7 @@ export class ProbeAgent {
    * @param {boolean} [options.fallback.stopOnSuccess=true] - Stop on first success
    * @param {number} [options.fallback.maxTotalAttempts=10] - Maximum total attempts across all providers
    * @param {string} [options.completionPrompt] - Custom prompt to run after attempt_completion for validation/review (runs before mermaid/JSON validation)
+   * @param {number} [options.maxOutputTokens] - Maximum tokens for tool output before truncation (default: 20000, can also be set via PROBE_MAX_OUTPUT_TOKENS env var)
    */
   constructor(options = {}) {
     // Basic configuration
@@ -237,6 +278,9 @@ export class ProbeAgent {
     // Initialize token counter
     this.tokenCounter = new TokenCounter();
 
+    // Maximum output tokens for tool results (truncate if exceeded)
+    this.maxOutputTokens = getMaxOutputTokens(options.maxOutputTokens);
+
     if (this.debug) {
       console.log(`[DEBUG] Generated session ID for agent: ${this.sessionId}`);
       console.log(`[DEBUG] Maximum tool iterations configured: ${MAX_TOOL_ITERATIONS}`);
@@ -2537,6 +2581,11 @@ Follow these instructions carefully:
         }
       }
 
+      // Circuit breaker for repeated format errors
+      let lastFormatErrorType = null;
+      let sameFormatErrorCount = 0;
+      const MAX_REPEATED_FORMAT_ERRORS = 3;
+
       // Tool iteration loop (only for non-CLI engines like Vercel/Anthropic/OpenAI)
       while (currentIteration < maxIterations && !completionAttempted) {
         currentIteration++;
@@ -2830,7 +2879,28 @@ Follow these instructions carefully:
               );
 
               if (lastAssistantMessage) {
-                finalResult = lastAssistantMessage.content;
+                const prevContent = lastAssistantMessage.content;
+
+                // Check for patterns indicating a failed/wrapped tool call attempt
+                // Use detectUnrecognizedToolCall for consistent detection logic
+                const wrappedToolError = detectUnrecognizedToolCall(prevContent, validTools);
+
+                if (isWrappedToolError(wrappedToolError)) {
+                  // Previous response was a broken tool call attempt - don't reuse it
+                  const wrappedToolName = extractWrappedToolName(wrappedToolError);
+                  if (this.debug) {
+                    console.log(`[DEBUG] Previous response contains wrapped tool '${wrappedToolName}' - rejecting for __PREVIOUS_RESPONSE__`);
+                  }
+                  currentMessages.push({ role: 'assistant', content: assistantResponseContent });
+                  currentMessages.push({
+                    role: 'user',
+                    content: createWrappedToolErrorMessage(wrappedToolName)
+                  });
+                  completionAttempted = false;
+                  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)}...`);
               } else {
                 finalResult = 'Error: No previous response found to use as completion.';
@@ -2882,7 +2952,24 @@ Follow these instructions carefully:
                 // Execute MCP tool through the bridge
                 const executionResult = await this.mcpBridge.mcpTools[toolName].execute(params);
 
-                const toolResultContent = typeof executionResult === 'string' ? executionResult : JSON.stringify(executionResult, null, 2);
+                let toolResultContent = typeof executionResult === 'string' ? executionResult : JSON.stringify(executionResult, null, 2);
+
+                // Truncate if output exceeds token limit
+                try {
+                  const truncateResult = await truncateIfNeeded(toolResultContent, this.tokenCounter, this.sessionId, this.maxOutputTokens);
+                  if (truncateResult.truncated) {
+                    toolResultContent = truncateResult.content;
+                    if (this.debug) {
+                      console.log(`[DEBUG] Tool output truncated: ${truncateResult.originalTokens} tokens -> saved to ${truncateResult.tempFilePath || 'N/A'}`);
+                      if (truncateResult.error) {
+                        console.log(`[DEBUG] Truncation file error: ${truncateResult.error}`);
+                      }
+                    }
+                  }
+                } catch (truncateError) {
+                  // If truncation fails entirely, log and continue with original content
+                  console.error(`[WARN] Tool output truncation failed: ${truncateError.message}`);
+                }
 
                 // Log MCP tool result in debug mode
                 if (this.debug) {
@@ -3059,10 +3146,28 @@ Follow these instructions carefully:
                 
                 // Add assistant response and tool result to conversation
                 currentMessages.push({ role: 'assistant', content: assistantResponseContent });
-                
-                const toolResultContent = typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult, null, 2);
+
+                let toolResultContent = typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult, null, 2);
+
+                // Truncate if output exceeds token limit
+                try {
+                  const truncateResult = await truncateIfNeeded(toolResultContent, this.tokenCounter, this.sessionId, this.maxOutputTokens);
+                  if (truncateResult.truncated) {
+                    toolResultContent = truncateResult.content;
+                    if (this.debug) {
+                      console.log(`[DEBUG] Tool output truncated: ${truncateResult.originalTokens} tokens -> saved to ${truncateResult.tempFilePath || 'N/A'}`);
+                      if (truncateResult.error) {
+                        console.log(`[DEBUG] Truncation file error: ${truncateResult.error}`);
+                      }
+                    }
+                  }
+                } catch (truncateError) {
+                  // If truncation fails entirely, log and continue with original content
+                  console.error(`[WARN] Tool output truncation failed: ${truncateError.message}`);
+                }
+
                 const toolResultMessage = `<tool_result>\n${toolResultContent}\n</tool_result>`;
-                
+
                 currentMessages.push({
                   role: 'user',
                   content: toolResultMessage
@@ -3125,7 +3230,32 @@ Follow these instructions carefully:
           const unrecognizedTool = detectUnrecognizedToolCall(assistantResponseContent, validTools);
 
           let reminderContent;
-          if (unrecognizedTool) {
+          if (isWrappedToolError(unrecognizedTool)) {
+            // AI wrapped a valid tool name in arbitrary XML tags - provide clear format error
+            const wrappedToolName = extractWrappedToolName(unrecognizedTool);
+            if (this.debug) {
+              console.log(`[DEBUG] Detected wrapped tool '${wrappedToolName}' in assistant response - wrong XML format.`);
+            }
+            const toolError = new ParameterError(
+              `Tool '${wrappedToolName}' found but in WRONG FORMAT - do not wrap tools in other XML tags.`,
+              {
+                suggestion: `Use the tool tag DIRECTLY without any wrapper:
+
+CORRECT FORMAT:
+<${wrappedToolName}>
+<param>value</param>
+</${wrappedToolName}>
+
+WRONG (what you did - do not wrap in other tags):
+<api_call><tool_name>${wrappedToolName}</tool_name>...</api_call>
+<function>${wrappedToolName}</function>
+<call name="${wrappedToolName}">...</call>
+
+Remove ALL wrapper tags and use <${wrappedToolName}> directly as the outermost tag.`
+              }
+            );
+            reminderContent = `<tool_result>\n${formatErrorForAI(toolError)}\n</tool_result>`;
+          } else if (unrecognizedTool) {
             // AI tried to use a tool that's not available - provide clear error
             if (this.debug) {
               console.log(`[DEBUG] Detected unrecognized tool '${unrecognizedTool}' in assistant response.`);
@@ -3135,6 +3265,33 @@ Follow these instructions carefully:
             });
             reminderContent = `<tool_result>\n${formatErrorForAI(toolError)}\n</tool_result>`;
           } else {
+            // No tool call detected at all - check if this is the last iteration
+            // On the last iteration, if the AI gave a substantive response without using
+            // attempt_completion, accept it as the final answer rather than losing the content
+            if (currentIteration >= maxIterations) {
+              // Clean up the response - remove thinking tags
+              let cleanedResponse = assistantResponseContent;
+              // Remove <thinking>...</thinking> blocks
+              cleanedResponse = cleanedResponse.replace(/<thinking>[\s\S]*?<\/thinking>/gi, '').trim();
+              // Also remove unclosed thinking tags
+              cleanedResponse = cleanedResponse.replace(/<thinking>[\s\S]*$/gi, '').trim();
+
+              // Only use if there's substantial content (not just a failed tool call attempt)
+              const hasSubstantialContent = cleanedResponse.length > 50 &&
+                !cleanedResponse.includes('<api_call>') &&
+                !cleanedResponse.includes('<tool_name>') &&
+                !cleanedResponse.includes('<function>');
+
+              if (hasSubstantialContent) {
+                if (this.debug) {
+                  console.log(`[DEBUG] Max iterations reached - accepting AI response as final answer (${cleanedResponse.length} chars)`);
+                }
+                finalResult = cleanedResponse;
+                completionAttempted = true;
+                break;
+              }
+            }
+
             // Standard reminder - no tool call detected at all
             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.
 
@@ -3166,6 +3323,31 @@ Note: <attempt_complete></attempt_complete> reuses your PREVIOUS assistant messa
               console.log(`[DEBUG] No tool call detected in assistant response. Prompting for tool use.`);
             }
           }
+
+          // Circuit breaker: track repeated format errors and break early
+          // For wrapped_tool errors, track them as a category (any wrapped_tool counts)
+          // For other errors, track the exact error type
+          if (unrecognizedTool) {
+            const isWrapped = isWrappedToolError(unrecognizedTool);
+            const errorCategory = isWrapped ? 'wrapped_tool' : unrecognizedTool;
+
+            if (errorCategory === lastFormatErrorType) {
+              sameFormatErrorCount++;
+              if (sameFormatErrorCount >= MAX_REPEATED_FORMAT_ERRORS) {
+                const errorDesc = isWrapped ? 'wrapped tool format' : unrecognizedTool;
+                console.error(`[ERROR] Format error category '${errorCategory}' repeated ${sameFormatErrorCount} times. Breaking loop early to prevent infinite iteration.`);
+                finalResult = `Error: Unable to complete request. The AI model repeatedly used incorrect tool call format (${errorDesc}). Please try rephrasing your question or using a different model.`;
+                break;
+              }
+            } else {
+              lastFormatErrorType = errorCategory;
+              sameFormatErrorCount = 1;
+            }
+          } else {
+            // Reset counter if it's a different kind of "no tool call" situation
+            lastFormatErrorType = null;
+            sameFormatErrorCount = 0;
+          }
         }
 
         // Keep message history manageable

package/src/agent/outputTruncator.js

@@ -0,0 +1,108 @@
+import { writeFile, mkdir } from 'fs/promises';
+import { tmpdir } from 'os';
+import { join } from 'path';
+import { randomUUID } from 'crypto';
+
+const DEFAULT_MAX_OUTPUT_TOKENS = 20000;
+const CHARS_PER_TOKEN = 4; // Conservative approximation
+
+/**
+ * Validate and normalize a token limit value.
+ * Returns the default if the value is invalid (NaN, negative, zero).
+ * @param {any} value - The value to validate
+ * @returns {number} A valid positive token limit
+ */
+function validateTokenLimit(value) {
+  const num = Number(value);
+  if (isNaN(num) || num <= 0) {
+    return DEFAULT_MAX_OUTPUT_TOKENS;
+  }
+  return num;
+}
+
+/**
+ * Get the maximum output tokens limit based on priority:
+ * 1. Constructor value (if provided and valid)
+ * 2. Environment variable PROBE_MAX_OUTPUT_TOKENS (if valid)
+ * 3. Default (20000)
+ * @param {number|undefined} constructorValue - Value passed to ProbeAgent constructor
+ * @returns {number} The maximum output tokens limit (always a valid positive number)
+ */
+export function getMaxOutputTokens(constructorValue) {
+  if (constructorValue !== undefined && constructorValue !== null) {
+    const validated = validateTokenLimit(constructorValue);
+    // Only use constructor value if it was valid; otherwise fall through to env/default
+    if (validated !== DEFAULT_MAX_OUTPUT_TOKENS || Number(constructorValue) === DEFAULT_MAX_OUTPUT_TOKENS) {
+      return validated;
+    }
+  }
+  if (process.env.PROBE_MAX_OUTPUT_TOKENS) {
+    return validateTokenLimit(process.env.PROBE_MAX_OUTPUT_TOKENS);
+  }
+  return DEFAULT_MAX_OUTPUT_TOKENS;
+}
+
+/**
+ * Truncate tool output if it exceeds the token limit.
+ * When truncated, saves full output to a temp file and returns a message with the file path.
+ * If file system operations fail, returns truncated content without file reference.
+ *
+ * @param {string} content - The tool output content to potentially truncate
+ * @param {Object} tokenCounter - TokenCounter instance with countTokens method
+ * @param {string} sessionId - Session ID for naming temp files
+ * @param {number} maxTokens - Maximum tokens allowed (defaults to 20000)
+ * @returns {Promise<{truncated: boolean, content: string, tempFilePath?: string, originalTokens?: number, error?: string}>}
+ */
+export async function truncateIfNeeded(content, tokenCounter, sessionId, maxTokens) {
+  const limit = validateTokenLimit(maxTokens);
+  const tokenCount = tokenCounter.countTokens(content);
+
+  if (tokenCount <= limit) {
+    return { truncated: false, content };
+  }
+
+  // Truncate to approximately maxTokens worth of characters
+  const maxChars = limit * CHARS_PER_TOKEN;
+  const truncatedContent = content.substring(0, maxChars);
+
+  // Try to write full output to temp file
+  let tempFilePath = null;
+  let fileError = null;
+
+  try {
+    const tempDir = join(tmpdir(), 'probe-output');
+    await mkdir(tempDir, { recursive: true });
+    tempFilePath = join(tempDir, `tool-output-${sessionId || 'unknown'}-${randomUUID()}.txt`);
+    await writeFile(tempFilePath, content, 'utf8');
+  } catch (err) {
+    fileError = err.message || 'Unknown file system error';
+    tempFilePath = null;
+  }
+
+  let message;
+  if (tempFilePath) {
+    message = `Output exceeded maximum size (${tokenCount} tokens, limit: ${limit}).
+Full output saved to: ${tempFilePath}
+
+--- Truncated Output (first ${limit} tokens approx) ---
+${truncatedContent}
+...
+--- End of Truncated Output ---`;
+  } else {
+    message = `Output exceeded maximum size (${tokenCount} tokens, limit: ${limit}).
+Warning: Could not save full output to file (${fileError}).
+
+--- Truncated Output (first ${limit} tokens approx) ---
+${truncatedContent}
+...
+--- End of Truncated Output ---`;
+  }
+
+  return {
+    truncated: true,
+    content: message,
+    tempFilePath: tempFilePath || undefined,
+    originalTokens: tokenCount,
+    error: fileError || undefined
+  };
+}

package/src/tools/common.js

@@ -617,6 +617,37 @@ export function detectUnrecognizedToolCall(xmlString, validTools) {
 		}
 	}
 
+	// Check if any valid tool name appears inside specific wrapper patterns
+	// This catches cases where AI wraps tools in arbitrary tags like:
+	// <api_call><tool_name>attempt_completion</tool_name>...</api_call>
+	// <function>search</function>
+	// <call name="extract">...</call>
+	// Only match specific wrapper patterns to avoid false positives with normal text
+	const allToolNames = [...new Set([...knownToolNames, ...validTools])];
+	for (const toolName of allToolNames) {
+		// Escape regex metacharacters in tool name to prevent regex errors
+		const escapedToolName = toolName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+		// Match specific wrapper patterns that indicate a tool call attempt:
+		// 1. <tool_name>toolName</tool_name> - common Claude API-style wrapper
+		// 2. <function>toolName</function> - function call style
+		// 3. <name>toolName</name> - generic name wrapper
+		// 4. <call><name>toolName - partial wrapper patterns
+		const wrapperPatterns = [
+			new RegExp(`<tool_name>\\s*${escapedToolName}\\s*</tool_name>`, 'i'),
+			new RegExp(`<function>\\s*${escapedToolName}\\s*</function>`, 'i'),
+			new RegExp(`<name>\\s*${escapedToolName}\\s*</name>`, 'i'),
+			// Also check for tool name immediately after api_call or call opening tag
+			new RegExp(`<(?:api_call|call)[^>]*>[\\s\\S]*?<tool_name>\\s*${escapedToolName}`, 'i')
+		];
+
+		for (const pattern of wrapperPatterns) {
+			if (pattern.test(xmlString)) {
+				return `wrapped_tool:${toolName}`;
+			}
+		}
+	}
+
 	return null;
 }