@probelabs/probe 0.6.0-rc205 → 0.6.0-rc206

package/build/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

package/build/agent/index.js

@@ -9718,10 +9718,10 @@ var init_vercel = __esm({
             let extractOptions = { cwd: effectiveCwd };
             if (input_content) {
               const { writeFileSync: writeFileSync2, unlinkSync } = await import("fs");
-              const { join: join4 } = await import("path");
-              const { tmpdir } = await import("os");
-              const { randomUUID: randomUUID8 } = await import("crypto");
-              tempFilePath = join4(tmpdir(), `probe-extract-${randomUUID8()}.txt`);
+              const { join: join5 } = await import("path");
+              const { tmpdir: tmpdir2 } = await import("os");
+              const { randomUUID: randomUUID9 } = await import("crypto");
+              tempFilePath = join5(tmpdir2(), `probe-extract-${randomUUID9()}.txt`);
               writeFileSync2(tempFilePath, input_content);
               if (debug) {
                 console.error(`Created temporary file for input content: ${tempFilePath}`);
@@ -67117,10 +67117,88 @@ var init_contextCompactor = __esm({
   }
 });
 
+// src/agent/outputTruncator.js
+import { writeFile, mkdir } from "fs/promises";
+import { tmpdir } from "os";
+import { join as join4 } from "path";
+import { randomUUID as randomUUID4 } from "crypto";
+function validateTokenLimit(value) {
+  const num = Number(value);
+  if (isNaN(num) || num <= 0) {
+    return DEFAULT_MAX_OUTPUT_TOKENS;
+  }
+  return num;
+}
+function getMaxOutputTokens(constructorValue) {
+  if (constructorValue !== void 0 && constructorValue !== null) {
+    const validated = validateTokenLimit(constructorValue);
+    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;
+}
+async function truncateIfNeeded(content, tokenCounter, sessionId, maxTokens) {
+  const limit = validateTokenLimit(maxTokens);
+  const tokenCount = tokenCounter.countTokens(content);
+  if (tokenCount <= limit) {
+    return { truncated: false, content };
+  }
+  const maxChars = limit * CHARS_PER_TOKEN;
+  const truncatedContent = content.substring(0, maxChars);
+  let tempFilePath = null;
+  let fileError = null;
+  try {
+    const tempDir = join4(tmpdir(), "probe-output");
+    await mkdir(tempDir, { recursive: true });
+    tempFilePath = join4(tempDir, `tool-output-${sessionId || "unknown"}-${randomUUID4()}.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 || void 0,
+    originalTokens: tokenCount,
+    error: fileError || void 0
+  };
+}
+var DEFAULT_MAX_OUTPUT_TOKENS, CHARS_PER_TOKEN;
+var init_outputTruncator = __esm({
+  "src/agent/outputTruncator.js"() {
+    "use strict";
+    DEFAULT_MAX_OUTPUT_TOKENS = 2e4;
+    CHARS_PER_TOKEN = 4;
+  }
+});
+
 // src/agent/mcp/built-in-server.js
 import { createServer } from "http";
 import { EventEmitter as EventEmitter3 } from "events";
-import { randomUUID as randomUUID4 } from "crypto";
+import { randomUUID as randomUUID5 } from "crypto";
 import { Server as MCPServer } from "@modelcontextprotocol/sdk/server/index.js";
 import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
@@ -67372,7 +67450,7 @@ var init_built_in_server = __esm({
             }
             const eventStore = new InMemoryEventStore();
             transport = new StreamableHTTPServerTransport({
-              sessionIdGenerator: () => randomUUID4(),
+              sessionIdGenerator: () => randomUUID5(),
               eventStore,
               // Enable resumability
               onsessioninitialized: (newSessionId) => {
@@ -68593,7 +68671,7 @@ import { createOpenAI as createOpenAI2 } from "@ai-sdk/openai";
 import { createGoogleGenerativeAI as createGoogleGenerativeAI2 } from "@ai-sdk/google";
 import { createAmazonBedrock as createAmazonBedrock2 } from "@ai-sdk/amazon-bedrock";
 import { streamText as streamText2 } from "ai";
-import { randomUUID as randomUUID5 } from "crypto";
+import { randomUUID as randomUUID6 } from "crypto";
 import { EventEmitter as EventEmitter5 } from "events";
 import { existsSync as existsSync6 } from "fs";
 import { readFile as readFile3, stat, readdir as readdir3 } from "fs/promises";
@@ -68622,6 +68700,7 @@ var init_ProbeAgent = __esm({
     init_FallbackManager();
     init_contextCompactor();
     init_error_types();
+    init_outputTruncator();
     init_tasks();
     dotenv2.config();
     MAX_TOOL_ITERATIONS = (() => {
@@ -68681,9 +68760,10 @@ var init_ProbeAgent = __esm({
        * @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 = {}) {
-        this.sessionId = options.sessionId || randomUUID5();
+        this.sessionId = options.sessionId || randomUUID6();
         this.customPrompt = options.systemPrompt || options.customPrompt || null;
         this.promptType = options.promptType || "code-explorer";
         this.allowEdit = !!options.allowEdit;
@@ -68742,6 +68822,7 @@ var init_ProbeAgent = __esm({
         this.clientApiKey = null;
         this.clientApiUrl = null;
         this.tokenCounter = new TokenCounter();
+        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}`);
@@ -70833,7 +70914,21 @@ You are working with a repository located at: ${searchDirectory}
 `);
                     }
                     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);
+                    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) {
+                      console.error(`[WARN] Tool output truncation failed: ${truncateError.message}`);
+                    }
                     if (this.debug) {
                       const preview = toolResultContent.length > 500 ? toolResultContent.substring(0, 500) + "..." : toolResultContent;
                       console.error(`[DEBUG] ========================================`);
@@ -70982,7 +71077,21 @@ ${errorXml}
                       throw toolError;
                     }
                     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);
+                    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) {
+                      console.error(`[WARN] Tool output truncation failed: ${truncateError.message}`);
+                    }
                     const toolResultMessage = `<tool_result>
 ${toolResultContent}
 </tool_result>`;
@@ -71591,7 +71700,7 @@ Convert your previous response content into actual JSON data that follows this s
        */
       clone(options = {}) {
         const {
-          sessionId = randomUUID5(),
+          sessionId = randomUUID6(),
           stripInternalMessages = true,
           keepSystemMessage = true,
           deepCopy = true,
@@ -71818,7 +71927,7 @@ import { readFileSync as readFileSync2, existsSync as existsSync7 } from "fs";
 import { resolve as resolve7 } from "path";
 
 // src/agent/acp/server.js
-import { randomUUID as randomUUID6 } from "crypto";
+import { randomUUID as randomUUID7 } from "crypto";
 
 // src/agent/acp/connection.js
 import { EventEmitter as EventEmitter6 } from "events";
@@ -72310,7 +72419,7 @@ var ACPServer = class {
    * Handle new session request
    */
   async handleNewSession(params) {
-    const sessionId = params?.sessionId || randomUUID6();
+    const sessionId = params?.sessionId || randomUUID7();
     const mode = params?.mode || SessionMode.NORMAL;
     const session = new ACPSession(sessionId, mode);
     this.sessions.set(sessionId, session);
@@ -72482,7 +72591,7 @@ var ACPServer = class {
 };
 
 // src/agent/acp/tools.js
-import { randomUUID as randomUUID7 } from "crypto";
+import { randomUUID as randomUUID8 } from "crypto";
 
 // src/agent/index.js
 dotenv3.config();

package/build/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
+  };
+}