@probelabs/probe 0.6.0-rc161 → 0.6.0-rc162

package/src/agent/contextCompactor.js

@@ -0,0 +1,271 @@
+/**
+ * Context Window Compactor
+ *
+ * Handles context window overflow by intelligently removing intermediate agentic
+ * monologue sections while preserving user messages, final answers, and the most
+ * recent monologue.
+ */
+
+/**
+ * Context limit error patterns to detect from various AI providers
+ * Simple substring matching similar to RetryManager's approach
+ */
+const CONTEXT_LIMIT_ERROR_PATTERNS = [
+  // Anthropic
+  'context_length_exceeded',
+  'prompt is too long',
+
+  // OpenAI
+  'maximum context length',
+  'context length is',
+
+  // Google/Gemini
+  'input token count exceeds',
+  'token limit exceeded',
+
+  // Generic patterns
+  'context window',
+  'too many tokens',
+  'token limit',
+  'context limit',
+  'exceed',  // Catches "exceeds", "exceed maximum", etc.
+  'over the limit',
+  'maximum tokens'
+];
+
+/**
+ * Check if an error message indicates a context window limit was exceeded
+ * Uses simple substring matching for reliability and maintainability
+ *
+ * @param {Error|string} error - The error object or error message
+ * @returns {boolean} - True if the error indicates context limit exceeded
+ */
+export function isContextLimitError(error) {
+  if (!error) return false;
+
+  // Get error message in various forms
+  const errorMessage = (typeof error === 'string' ? error : (error?.message || '')).toLowerCase();
+  const errorString = error.toString().toLowerCase();
+
+  // Check if any pattern matches
+  for (const pattern of CONTEXT_LIMIT_ERROR_PATTERNS) {
+    const lowerPattern = pattern.toLowerCase();
+    if (errorMessage.includes(lowerPattern) || errorString.includes(lowerPattern)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+/**
+ * Identify message boundaries in conversation history
+ * Structure: <user> -> <internal agentic monologue> -> <final-agent-answer>
+ *
+ * A "segment" is:
+ * - user message (role: 'user')
+ * - followed by 0+ assistant messages (internal monologue)
+ * - ending with tool_result or attempt_completion (final answer)
+ *
+ * @param {Array} messages - Array of message objects with {role, content}
+ * @returns {Array} - Array of segments, each containing {userIndex, monologueIndices, finalIndex}
+ */
+export function identifyMessageSegments(messages) {
+  const segments = [];
+  let currentSegment = null;
+
+  for (let i = 0; i < messages.length; i++) {
+    const msg = messages[i];
+
+    // Skip system messages
+    if (msg.role === 'system') {
+      continue;
+    }
+
+    // User message starts a new segment
+    if (msg.role === 'user') {
+      // Check if this is a tool_result (final answer from previous segment)
+      const content = typeof msg.content === 'string' ? msg.content : '';
+      const isToolResult = content.includes('<tool_result>');
+
+      if (isToolResult && currentSegment) {
+        // This is the final answer for the current segment
+        currentSegment.finalIndex = i;
+        segments.push(currentSegment);
+        currentSegment = null;
+      } else {
+        // Save previous segment if it exists
+        if (currentSegment) {
+          segments.push(currentSegment);
+        }
+
+        // Start new segment
+        currentSegment = {
+          userIndex: i,
+          monologueIndices: [],
+          finalIndex: null
+        };
+      }
+    }
+
+    // Assistant message is part of monologue
+    if (msg.role === 'assistant' && currentSegment) {
+      const content = typeof msg.content === 'string' ? msg.content : '';
+
+      // Check if this contains attempt_completion (marks end of segment)
+      if (content.includes('<attempt_completion>') || content.includes('attempt_completion')) {
+        currentSegment.monologueIndices.push(i);
+        currentSegment.finalIndex = i;
+        segments.push(currentSegment);
+        currentSegment = null;
+      } else {
+        // Regular monologue message
+        currentSegment.monologueIndices.push(i);
+      }
+    }
+  }
+
+  // Save any remaining segment
+  if (currentSegment) {
+    segments.push(currentSegment);
+  }
+
+  return segments;
+}
+
+/**
+ * Compact messages by removing intermediate monologues
+ *
+ * Strategy:
+ * 1. Keep all user messages
+ * 2. Keep all final answers (tool_results, attempt_completion)
+ * 3. Remove intermediate monologue messages from completed segments
+ * 4. Keep the most recent (active) segment intact
+ *
+ * @param {Array} messages - Array of message objects
+ * @param {Object} options - Compaction options
+ * @param {boolean} [options.keepLastSegment=true] - Keep the most recent segment intact
+ * @param {number} [options.minSegmentsToKeep=1] - Minimum number of recent segments to preserve fully
+ * @returns {Array} - Compacted message array
+ */
+export function compactMessages(messages, options = {}) {
+  const {
+    keepLastSegment = true,
+    minSegmentsToKeep = 1
+  } = options;
+
+  if (!messages || messages.length === 0) {
+    return messages;
+  }
+
+  // Identify segments
+  const segments = identifyMessageSegments(messages);
+
+  if (segments.length === 0) {
+    return messages;
+  }
+
+  // Determine which segments to keep fully vs compact
+  const segmentsToPreserve = keepLastSegment
+    ? Math.max(minSegmentsToKeep, 1)
+    : minSegmentsToKeep;
+
+  const compactableSegments = segments.slice(0, -segmentsToPreserve);
+  const preservedSegments = segments.slice(-segmentsToPreserve);
+
+  // Build set of indices to keep
+  const indicesToKeep = new Set();
+
+  // Keep system messages
+  messages.forEach((msg, idx) => {
+    if (msg.role === 'system') {
+      indicesToKeep.add(idx);
+    }
+  });
+
+  // For compactable segments: keep user message and final answer only
+  compactableSegments.forEach(segment => {
+    indicesToKeep.add(segment.userIndex);
+    if (segment.finalIndex !== null) {
+      indicesToKeep.add(segment.finalIndex);
+    }
+  });
+
+  // For preserved segments: keep everything
+  preservedSegments.forEach(segment => {
+    indicesToKeep.add(segment.userIndex);
+    segment.monologueIndices.forEach(idx => indicesToKeep.add(idx));
+    if (segment.finalIndex !== null) {
+      indicesToKeep.add(segment.finalIndex);
+    }
+  });
+
+  // Filter messages
+  const compactedMessages = messages.filter((_, idx) => indicesToKeep.has(idx));
+
+  return compactedMessages;
+}
+
+/**
+ * Calculate reduction statistics
+ * @param {Array} originalMessages - Original message array
+ * @param {Array} compactedMessages - Compacted message array
+ * @returns {Object} - Statistics about the compaction
+ */
+export function calculateCompactionStats(originalMessages, compactedMessages) {
+  const originalCount = originalMessages.length;
+  const compactedCount = compactedMessages.length;
+  const removed = originalCount - compactedCount;
+  const reductionPercent = originalCount > 0
+    ? ((removed / originalCount) * 100).toFixed(1)
+    : 0;
+
+  // Estimate token savings (rough approximation)
+  const estimateTokens = (msgs) => {
+    return msgs.reduce((sum, msg) => {
+      const content = typeof msg.content === 'string'
+        ? msg.content
+        : JSON.stringify(msg.content);
+      // Rough estimate: 1 token ≈ 4 characters
+      return sum + Math.ceil(content.length / 4);
+    }, 0);
+  };
+
+  const originalTokens = estimateTokens(originalMessages);
+  const compactedTokens = estimateTokens(compactedMessages);
+  const tokensSaved = originalTokens - compactedTokens;
+
+  return {
+    originalCount,
+    compactedCount,
+    removed,
+    reductionPercent: parseFloat(reductionPercent),
+    originalTokens,
+    compactedTokens,
+    tokensSaved
+  };
+}
+
+/**
+ * Main compaction handler for ProbeAgent
+ * Detects context limit errors and performs intelligent compaction
+ *
+ * @param {Error} error - The error from the AI provider
+ * @param {Array} messages - Current message array
+ * @param {Object} options - Compaction options
+ * @returns {Object|null} - { compacted: true, messages, stats } or null if not applicable
+ */
+export function handleContextLimitError(error, messages, options = {}) {
+  if (!isContextLimitError(error)) {
+    return null;
+  }
+
+  const compactedMessages = compactMessages(messages, options);
+  const stats = calculateCompactionStats(messages, compactedMessages);
+
+  return {
+    compacted: true,
+    messages: compactedMessages,
+    stats
+  };
+}

package/src/agent/index.js

@@ -134,6 +134,8 @@ function parseArgs() {
     useStdin: false, // New flag to indicate stdin should be used
     outline: false, // New flag to enable outline format
     noMermaidValidation: false, // New flag to disable mermaid validation
+    allowedTools: null, // Tool filtering: ['*'] = all, [] = none, ['tool1', 'tool2'] = specific
+    disableTools: false, // Convenience flag to disable all tools
     // Bash tool configuration
     enableBash: false,
     bashAllow: null,
@@ -189,6 +191,19 @@ function parseArgs() {
       config.outline = true;
     } else if (arg === '--no-mermaid-validation') {
       config.noMermaidValidation = true;
+    } else if (arg === '--allowed-tools' && i + 1 < args.length) {
+      // Parse allowed tools: comma-separated list or special values
+      const toolsArg = args[++i];
+      if (toolsArg === '*' || toolsArg === 'all') {
+        config.allowedTools = ['*'];
+      } else if (toolsArg === 'none' || toolsArg === '') {
+        config.allowedTools = [];
+      } else {
+        config.allowedTools = toolsArg.split(',').map(t => t.trim()).filter(t => t.length > 0);
+      }
+    } else if (arg === '--disable-tools') {
+      // Convenience flag to disable all tools (raw AI mode)
+      config.disableTools = true;
     } else if (arg === '--enable-bash') {
       config.enableBash = true;
     } else if (arg === '--bash-allow' && i + 1 < args.length) {
@@ -243,6 +258,13 @@ Options:
   --model <name>                   Override model name
   --allow-edit                     Enable code modification capabilities
   --enable-delegate                Enable delegate tool for task distribution to subagents
+  --allowed-tools <tools>          Filter available tools (comma-separated list)
+                                   Use '*' or 'all' for all tools (default)
+                                   Use 'none' or '' for no tools (raw AI mode)
+                                   Specific tools: search,query,extract,listFiles,searchFiles
+                                   Supports exclusion: '*,!bash' (all except bash)
+  --disable-tools                  Disable all tools (raw AI mode, no code analysis)
+                                   Convenience flag equivalent to --allowed-tools none
   --verbose                        Enable verbose output
   --outline                        Use outline-xml format for code search results
   --mcp                           Run as MCP server
@@ -284,6 +306,9 @@ Examples:
   probe agent "Analyze codebase" --schema schema.json  # Schema from file
   probe agent "Debug issue" --trace-file ./debug.jsonl --verbose
   probe agent "Analyze code" --trace-remote http://localhost:4318/v1/traces
+  probe agent "Explain this code" --allowed-tools search,extract  # Only search and extract
+  probe agent "What is this project about?" --allowed-tools none  # Raw AI mode (no tools)
+  probe agent "Tell me about this project" --disable-tools        # Raw AI mode (convenience flag)
   probe agent --mcp               # Start MCP server mode
   probe agent --acp               # Start ACP server mode
 
@@ -434,7 +459,9 @@ class ProbeAgentMcpServer {
             allowEdit: !!args.allow_edit,
             debug: process.env.DEBUG === '1',
             maxResponseTokens: args.max_response_tokens,
-            disableMermaidValidation: !!args.no_mermaid_validation
+            disableMermaidValidation: !!args.no_mermaid_validation,
+            allowedTools: args.allowed_tools,
+            disableTools: args.disable_tools
           };
 
           this.agent = new ProbeAgent(agentConfig);
@@ -726,6 +753,8 @@ async function main() {
       outline: config.outline,
       maxResponseTokens: config.maxResponseTokens,
       disableMermaidValidation: config.noMermaidValidation,
+      allowedTools: config.allowedTools,
+      disableTools: config.disableTools,
       enableBash: config.enableBash,
       bashConfig: bashConfig
     };

package/src/agent/schemaUtils.js

@@ -100,6 +100,13 @@ export function decodeHtmlEntities(text) {
 /**
  * Clean AI response by extracting JSON content when response contains JSON
  * Only processes responses that contain JSON structures { or [
+ *
+ * NOTE: This function handles both JSON extraction AND content stripping.
+ * Future improvement: Consider splitting into separate functions for better separation of concerns:
+ * - extractJsonContent() - Find and extract JSON from response
+ * - stripNonJsonContent() - Remove explanatory text while preserving validation-relevant content
+ * This would allow validation to run on full responses before content is discarded.
+ *
  * @param {string} response - Raw AI response
  * @returns {string} - Cleaned response with JSON boundaries extracted if applicable
  */

package/src/agent/xmlParsingUtils.js

@@ -59,10 +59,30 @@ export function extractThinkingContent(xmlString) {
 export function checkAttemptCompleteRecovery(cleanedXmlString, validTools = []) {
   // Check for <attempt_completion> with content (with or without closing tag)
   // This handles: "<attempt_completion>content" or "<attempt_completion>content</attempt_completion>"
-  const attemptCompletionMatch = cleanedXmlString.match(/<attempt_completion>([\s\S]*?)(?:<\/attempt_completion>|$)/);
-  if (attemptCompletionMatch) {
-    const content = attemptCompletionMatch[1].trim();
-    const hasClosingTag = cleanedXmlString.includes('</attempt_completion>');
+
+  // IMPORTANT: Use greedy match ([\s\S]*) instead of non-greedy ([\s\S]*?) to handle cases
+  // where the content contains the string "</attempt_completion>" (e.g., in regex patterns or code examples).
+  // We want to find the LAST occurrence of </attempt_completion>, not the first one.
+  const openTagIndex = cleanedXmlString.indexOf('<attempt_completion>');
+  if (openTagIndex !== -1) {
+    const afterOpenTag = cleanedXmlString.substring(openTagIndex + '<attempt_completion>'.length);
+    const closeTagIndex = cleanedXmlString.lastIndexOf('</attempt_completion>');
+
+    let content;
+    let hasClosingTag = false;
+
+    if (closeTagIndex !== -1 && closeTagIndex >= openTagIndex + '<attempt_completion>'.length) {
+      // Found a closing tag at or after the opening tag - extract content between them
+      content = cleanedXmlString.substring(
+        openTagIndex + '<attempt_completion>'.length,
+        closeTagIndex
+      ).trim();
+      hasClosingTag = true;
+    } else {
+      // No closing tag - use content from opening tag to end of string
+      content = afterOpenTag.trim();
+      hasClosingTag = false;
+    }
 
     if (content) {
       // If there's content after the tag, use it as the result

package/src/tools/common.js

@@ -384,7 +384,22 @@ export function parseXmlToolCall(xmlString, validTools = DEFAULT_VALID_TOOLS) {
 			continue; // Tool not found, try next tool
 		}
 
-		let closeIndex = xmlString.indexOf(closeTag, openIndex + openTag.length);
+		// For attempt_completion, use lastIndexOf to find the LAST occurrence of closing tag
+		// This prevents issues where the content contains the closing tag string (e.g., in regex patterns)
+		// For other tools, use indexOf from the opening tag position
+		let closeIndex;
+		if (toolName === 'attempt_completion') {
+			// Find the last occurrence of the closing tag in the entire string
+			// This assumes attempt_completion doesn't have nested tags of the same name
+			closeIndex = xmlString.lastIndexOf(closeTag);
+			// Make sure the closing tag is after the opening tag
+			if (closeIndex !== -1 && closeIndex <= openIndex + openTag.length) {
+				closeIndex = -1; // Invalid, treat as no closing tag
+			}
+		} else {
+			closeIndex = xmlString.indexOf(closeTag, openIndex + openTag.length);
+		}
+
 		let hasClosingTag = closeIndex !== -1;
 
 		// If no closing tag found, use content until end of string