@probelabs/probe 0.6.0-rc204 → 0.6.0-rc206

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc204",
+  "version": "0.6.0-rc206",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",
@@ -79,7 +79,7 @@
     "@ai-sdk/openai": "^2.0.10",
     "@anthropic-ai/claude-agent-sdk": "^0.1.46",
     "@modelcontextprotocol/sdk": "^1.0.0",
-    "@probelabs/maid": "^0.0.22",
+    "@probelabs/maid": "^0.0.23",
     "adm-zip": "^0.5.16",
     "ai": "^5.0.0",
     "ajv": "^8.17.1",

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,
@@ -145,6 +146,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 +239,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}`);
@@ -2882,7 +2887,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 +3081,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
@@ -3146,10 +3186,13 @@ Remember: Use proper XML format with BOTH opening and closing tags:
 
 Available tools: ${validTools.join(', ')}
 
-Or for quick completion if your previous response was already correct and complete:
-<attempt_complete>
+To complete with a direct answer:
+<attempt_completion>Your final answer here</attempt_completion>
+
+Or if your previous response already contains a complete, direct answer (not a thinking block or JSON):
+<attempt_complete></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.`;
+Note: <attempt_complete></attempt_complete> reuses your PREVIOUS assistant message as the final answer. Only use this if that message was already a valid, complete response to the user's question.`;
           }
 
           currentMessages.push({
@@ -4062,9 +4105,9 @@ Convert your previous response content into actual JSON data that follows this s
       return true;
     }
 
-    // Empty attempt_complete reminders
-    if (content.includes('When using <attempt_complete>') &&
-        content.includes('this must be the ONLY content in your response')) {
+    // Empty attempt_complete reminders (legacy and new format)
+    if (content.includes('<attempt_complete></attempt_complete>') &&
+        content.includes('reuses your PREVIOUS assistant message')) {
       return true;
     }
 

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
+  };
+}