@probelabs/probe 0.6.0-rc124 → 0.6.0-rc126

package/index.d.ts

@@ -31,6 +31,10 @@ export interface ProbeAgentOptions {
   mcpConfig?: any;
   /** @deprecated Use mcpConfig instead */
   mcpServers?: any[];
+  /** Pluggable storage adapter for conversation history */
+  storageAdapter?: StorageAdapter;
+  /** Hook callbacks for event-driven integration */
+  hooks?: Record<string, (data: any) => void | Promise<void>>;
 }
 
 /**
@@ -99,6 +103,97 @@ export interface ChatMessage {
   metadata?: any;
 }
 
+/**
+ * Storage adapter base class for pluggable history storage
+ */
+export declare class StorageAdapter {
+  /**
+   * Load conversation history for a session
+   */
+  loadHistory(sessionId: string): Promise<ChatMessage[]>;
+
+  /**
+   * Save a message to storage
+   */
+  saveMessage(sessionId: string, message: ChatMessage): Promise<void>;
+
+  /**
+   * Clear history for a session
+   */
+  clearHistory(sessionId: string): Promise<void>;
+
+  /**
+   * Optional: Get session metadata
+   */
+  getSessionMetadata(sessionId: string): Promise<any | null>;
+
+  /**
+   * Optional: Update session activity timestamp
+   */
+  updateSessionActivity(sessionId: string): Promise<void>;
+}
+
+/**
+ * Default in-memory storage adapter
+ */
+export declare class InMemoryStorageAdapter extends StorageAdapter {
+  constructor();
+}
+
+/**
+ * Hook types for ProbeAgent event system
+ */
+export declare const HOOK_TYPES: {
+  readonly AGENT_INITIALIZED: 'agent:initialized';
+  readonly AGENT_CLEANUP: 'agent:cleanup';
+  readonly MESSAGE_USER: 'message:user';
+  readonly MESSAGE_ASSISTANT: 'message:assistant';
+  readonly MESSAGE_SYSTEM: 'message:system';
+  readonly TOOL_START: 'tool:start';
+  readonly TOOL_END: 'tool:end';
+  readonly TOOL_ERROR: 'tool:error';
+  readonly AI_STREAM_START: 'ai:stream:start';
+  readonly AI_STREAM_DELTA: 'ai:stream:delta';
+  readonly AI_STREAM_END: 'ai:stream:end';
+  readonly STORAGE_LOAD: 'storage:load';
+  readonly STORAGE_SAVE: 'storage:save';
+  readonly STORAGE_CLEAR: 'storage:clear';
+  readonly ITERATION_START: 'iteration:start';
+  readonly ITERATION_END: 'iteration:end';
+};
+
+/**
+ * Hook manager for event-driven integration
+ */
+export declare class HookManager {
+  constructor();
+
+  /**
+   * Register a hook callback
+   */
+  on(hookName: string, callback: (data: any) => void | Promise<void>): () => void;
+
+  /**
+   * Register a one-time hook callback
+   */
+  once(hookName: string, callback: (data: any) => void | Promise<void>): void;
+
+  /**
+   * Unregister a hook callback
+   */
+  off(hookName: string, callback: (data: any) => void | Promise<void>): void;
+
+  /**
+   * Emit a hook event
+   */
+  emit(hookName: string, data: any): Promise<void>;
+
+  /**
+   * Clear all hooks or hooks for a specific event
+   */
+  clear(hookName?: string): void;
+}
+
 /**
  * Answer options
  */

package/package.json

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

package/src/agent/ProbeAgent.js

@@ -10,6 +10,8 @@ import { existsSync } from 'fs';
 import { readFile, stat } from 'fs/promises';
 import { resolve, isAbsolute, dirname } from 'path';
 import { TokenCounter } from './tokenCounter.js';
+import { InMemoryStorageAdapter } from './storage/InMemoryStorageAdapter.js';
+import { HookManager, HOOK_TYPES } from './hooks/HookManager.js';
 import { 
   createTools,
   searchToolDefinition,
@@ -80,6 +82,8 @@ export class ProbeAgent {
    * @param {string} [options.mcpConfigPath] - Path to MCP configuration file
    * @param {Object} [options.mcpConfig] - MCP configuration object (overrides mcpConfigPath)
    * @param {Array} [options.mcpServers] - Deprecated, use mcpConfig instead
+   * @param {Object} [options.storageAdapter] - Custom storage adapter for history management
+   * @param {Object} [options.hooks] - Hook callbacks for events (e.g., {'tool:start': callback})
    */
   constructor(options = {}) {
     // Basic configuration
@@ -95,6 +99,19 @@ export class ProbeAgent {
     this.maxIterations = options.maxIterations || null;
     this.disableMermaidValidation = !!options.disableMermaidValidation;
 
+    // Storage adapter (defaults to in-memory)
+    this.storageAdapter = options.storageAdapter || new InMemoryStorageAdapter();
+
+    // Hook manager
+    this.hooks = new HookManager();
+
+    // Register hooks from options
+    if (options.hooks) {
+      for (const [hookName, callback] of Object.entries(options.hooks)) {
+        this.hooks.on(hookName, callback);
+      }
+    }
+
     // Bash configuration
     this.enableBash = !!options.enableBash;
     this.bashConfig = options.bashConfig || {};
@@ -152,9 +169,29 @@ export class ProbeAgent {
 
   /**
    * Initialize the agent asynchronously (must be called after constructor)
-   * This method initializes MCP and merges MCP tools into the tool list
+   * This method initializes MCP and merges MCP tools into the tool list, and loads history from storage
    */
   async initialize() {
+    // Load history from storage adapter
+    try {
+      const history = await this.storageAdapter.loadHistory(this.sessionId);
+      this.history = history;
+
+      if (this.debug && history.length > 0) {
+        console.log(`[DEBUG] Loaded ${history.length} messages from storage for session ${this.sessionId}`);
+      }
+
+      // Emit storage load hook
+      await this.hooks.emit(HOOK_TYPES.STORAGE_LOAD, {
+        sessionId: this.sessionId,
+        messages: history
+      });
+    } catch (error) {
+      console.error(`[ERROR] Failed to load history from storage:`, error);
+      // Continue with empty history if storage fails
+      this.history = [];
+    }
+
     // Initialize MCP if enabled and not already initialized
     if (this.enableMcp && !this._mcpInitialized) {
       this._mcpInitialized = true; // Prevent multiple initialization attempts
@@ -193,6 +230,12 @@ export class ProbeAgent {
         this.mcpBridge = null;
       }
     }
+
+    // Emit agent initialized hook
+    await this.hooks.emit(HOOK_TYPES.AGENT_INITIALIZED, {
+      sessionId: this.sessionId,
+      agent: this
+    });
   }
 
   /**
@@ -261,7 +304,8 @@ export class ProbeAgent {
     // Get API keys from environment variables
     const anthropicApiKey = process.env.ANTHROPIC_API_KEY;
     const openaiApiKey = process.env.OPENAI_API_KEY;
-    const googleApiKey = process.env.GOOGLE_API_KEY;
+    // Support both GOOGLE_GENERATIVE_AI_API_KEY (official) and GOOGLE_API_KEY (legacy)
+    const googleApiKey = process.env.GOOGLE_GENERATIVE_AI_API_KEY || process.env.GOOGLE_API_KEY;
     const awsAccessKeyId = process.env.AWS_ACCESS_KEY_ID;
     const awsSecretAccessKey = process.env.AWS_SECRET_ACCESS_KEY;
     const awsRegion = process.env.AWS_REGION;
@@ -320,7 +364,7 @@ export class ProbeAgent {
     } else if ((awsAccessKeyId && awsSecretAccessKey && awsRegion) || awsApiKey) {
       this.initializeBedrockModel(awsAccessKeyId, awsSecretAccessKey, awsRegion, awsSessionToken, awsApiKey, awsBedrockBaseUrl, modelName);
     } else {
-      throw new Error('No API key provided. Please set ANTHROPIC_API_KEY, OPENAI_API_KEY, GOOGLE_API_KEY, AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION), or AWS_BEDROCK_API_KEY environment variables.');
+      throw new Error('No API key provided. Please set ANTHROPIC_API_KEY, OPENAI_API_KEY, GOOGLE_GENERATIVE_AI_API_KEY (or GOOGLE_API_KEY), AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION), or AWS_BEDROCK_API_KEY environment variables.');
     }
   }
 
@@ -1080,6 +1124,16 @@ When troubleshooting:
     }
 
     try {
+      // Track initial history length for storage
+      const oldHistoryLength = this.history.length;
+
+      // Emit user message hook
+      await this.hooks.emit(HOOK_TYPES.MESSAGE_USER, {
+        sessionId: this.sessionId,
+        message,
+        images
+      });
+
       // Generate system message
       const systemMessage = await this.getSystemMessage();
 
@@ -1196,9 +1250,13 @@ When troubleshooting:
               temperature: 0.3,
             });
 
-            // Collect the streamed response
+            // Collect the streamed response - stream all content for now
             for await (const delta of result.textStream) {
               assistantResponseContent += delta;
+              // For now, stream everything - we'll handle segmentation after tools execute
+              if (options.onStream) {
+                options.onStream(delta);
+              }
             }
 
             // Record token usage
@@ -1283,6 +1341,16 @@ When troubleshooting:
               const validation = attemptCompletionSchema.safeParse(params);
               if (validation.success) {
                 finalResult = validation.data.result;
+
+                // Stream the final result if callback is provided
+                if (options.onStream && finalResult) {
+                  const chunkSize = 50; // Characters per chunk for smoother streaming
+                  for (let i = 0; i < finalResult.length; i += chunkSize) {
+                    const chunk = finalResult.slice(i, Math.min(i + chunkSize, finalResult.length));
+                    options.onStream(chunk);
+                  }
+                }
+
                 if (this.debug) console.log(`[DEBUG] Task completed successfully with result: ${finalResult.substring(0, 100)}...`);
               } else {
                 console.error(`[ERROR] Invalid attempt_completion parameters:`, validation.error);
@@ -1365,12 +1433,13 @@ When troubleshooting:
                   console.error(`[DEBUG] ========================================\n`);
                 }
 
-                // Emit tool start event
+                // Emit tool start event with stream pause signal
                 this.events.emit('toolCall', {
                   timestamp: new Date().toISOString(),
                   name: toolName,
                   args: toolParams,
-                  status: 'started'
+                  status: 'started',
+                  pauseStream: true  // Signal to pause text streaming
                 });
                 
                 // Execute tool with tracing if available
@@ -1597,6 +1666,21 @@ IMPORTANT: When using <attempt_complete>, this must be the ONLY content in your
       // Update token counter with final history
       this.tokenCounter.updateHistory(this.history);
 
+      // Save new messages to storage (save only the new ones added in this turn)
+      try {
+        const messagesToSave = currentMessages.slice(oldHistoryLength);
+        for (const message of messagesToSave) {
+          await this.storageAdapter.saveMessage(this.sessionId, message);
+          await this.hooks.emit(HOOK_TYPES.STORAGE_SAVE, {
+            sessionId: this.sessionId,
+            message
+          });
+        }
+      } catch (error) {
+        console.error(`[ERROR] Failed to save messages to storage:`, error);
+        // Continue even if storage fails
+      }
+
       // Schema handling - format response according to provided schema
       // Skip schema processing if result came from attempt_completion tool
       // Don't apply schema formatting if we failed due to max iterations
@@ -2040,11 +2124,24 @@ Convert your previous response content into actual JSON data that follows this s
   /**
    * Clear conversation history and reset counters
    */
-  clearHistory() {
+  async clearHistory() {
+    // Clear in storage
+    try {
+      await this.storageAdapter.clearHistory(this.sessionId);
+    } catch (error) {
+      console.error(`[ERROR] Failed to clear history in storage:`, error);
+    }
+
+    // Clear in-memory
     this.history = [];
     this.tokenCounter.clear();
     clearToolExecutionData(this.sessionId);
-    
+
+    // Emit hook
+    await this.hooks.emit(HOOK_TYPES.STORAGE_CLEAR, {
+      sessionId: this.sessionId
+    });
+
     if (this.debug) {
       console.log(`[DEBUG] Cleared conversation history and reset counters for session ${this.sessionId}`);
     }

package/src/agent/hooks/HookManager.js

@@ -0,0 +1,146 @@
+/**
+ * Hook manager for ProbeAgent
+ * Enables event-driven integration with external systems
+ */
+export class HookManager {
+  constructor() {
+    this.hooks = new Map(); // hookName -> Set<callback>
+  }
+
+  /**
+   * Register a hook callback
+   * @param {string} hookName - Name of the hook
+   * @param {Function} callback - Callback function
+   * @returns {Function} Unregister function
+   */
+  on(hookName, callback) {
+    if (!this.hooks.has(hookName)) {
+      this.hooks.set(hookName, new Set());
+    }
+    this.hooks.get(hookName).add(callback);
+
+    // Return unregister function
+    return () => this.off(hookName, callback);
+  }
+
+  /**
+   * Register a one-time hook callback
+   * @param {string} hookName - Name of the hook
+   * @param {Function} callback - Callback function
+   * @returns {Function} Unregister function
+   */
+  once(hookName, callback) {
+    const wrappedCallback = async (data) => {
+      this.off(hookName, wrappedCallback);
+      await callback(data);
+    };
+    return this.on(hookName, wrappedCallback);
+  }
+
+  /**
+   * Unregister a hook callback
+   * @param {string} hookName - Name of the hook
+   * @param {Function} callback - Callback function
+   */
+  off(hookName, callback) {
+    const callbacks = this.hooks.get(hookName);
+    if (callbacks) {
+      callbacks.delete(callback);
+    }
+  }
+
+  /**
+   * Emit a hook event
+   * @param {string} hookName - Name of the hook
+   * @param {any} data - Data to pass to callbacks
+   * @returns {Promise<void>}
+   */
+  async emit(hookName, data) {
+    const callbacks = this.hooks.get(hookName);
+    if (!callbacks || callbacks.size === 0) return;
+
+    // Execute all callbacks in parallel using Promise.allSettled
+    // This ensures one failing hook doesn't break others
+    const promises = Array.from(callbacks).map(callback => {
+      try {
+        return Promise.resolve(callback(data));
+      } catch (error) {
+        // Catch synchronous errors
+        return Promise.reject(error);
+      }
+    });
+
+    const results = await Promise.allSettled(promises);
+
+    // Log any rejected promises
+    results.forEach((result, index) => {
+      if (result.status === 'rejected') {
+        console.error(`[HookManager] Error in hook "${hookName}" (callback ${index + 1}):`, result.reason);
+      }
+    });
+  }
+
+  /**
+   * Clear all hooks or hooks for a specific event
+   * @param {string} [hookName] - Optional hook name to clear
+   */
+  clear(hookName) {
+    if (hookName) {
+      this.hooks.delete(hookName);
+    } else {
+      this.hooks.clear();
+    }
+  }
+
+  /**
+   * Get list of registered hook names
+   * @returns {string[]} Array of hook names
+   */
+  getHookNames() {
+    return Array.from(this.hooks.keys());
+  }
+
+  /**
+   * Get number of callbacks for a hook
+   * @param {string} hookName - Name of the hook
+   * @returns {number} Number of callbacks
+   */
+  getCallbackCount(hookName) {
+    const callbacks = this.hooks.get(hookName);
+    return callbacks ? callbacks.size : 0;
+  }
+}
+
+/**
+ * Available hook types
+ * @type {Object<string, string>}
+ */
+export const HOOK_TYPES = {
+  // Lifecycle hooks
+  AGENT_INITIALIZED: 'agent:initialized',
+  AGENT_CLEANUP: 'agent:cleanup',
+
+  // Message hooks
+  MESSAGE_USER: 'message:user',
+  MESSAGE_ASSISTANT: 'message:assistant',
+  MESSAGE_SYSTEM: 'message:system',
+
+  // Tool execution hooks
+  TOOL_START: 'tool:start',
+  TOOL_END: 'tool:end',
+  TOOL_ERROR: 'tool:error',
+
+  // AI streaming hooks
+  AI_STREAM_START: 'ai:stream:start',
+  AI_STREAM_DELTA: 'ai:stream:delta',
+  AI_STREAM_END: 'ai:stream:end',
+
+  // Storage hooks
+  STORAGE_LOAD: 'storage:load',
+  STORAGE_SAVE: 'storage:save',
+  STORAGE_CLEAR: 'storage:clear',
+
+  // Iteration hooks
+  ITERATION_START: 'iteration:start',
+  ITERATION_END: 'iteration:end',
+};

package/src/agent/hooks/index.js

@@ -0,0 +1 @@
+export { HookManager, HOOK_TYPES } from './HookManager.js';

package/src/agent/probeTool.js

@@ -224,19 +224,27 @@ export const listFilesTool = {
     // Security: Validate path to prevent traversal attacks
     const secureBaseDir = path.resolve(baseCwd);
 
+    // Check if this is a dependency path that should bypass workspace restrictions
+    const isDependencyPath = directory.startsWith('/dep/') ||
+                            directory.startsWith('go:') ||
+                            directory.startsWith('js:') ||
+                            directory.startsWith('rust:');
+
     // If directory is absolute, check if it's within the secure base directory
     // If it's relative, resolve it against the secure base directory
     let targetDir;
     if (path.isAbsolute(directory)) {
       targetDir = path.resolve(directory);
       // Check if the absolute path is within the secure base directory
-      if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      // Allow dependency paths to bypass this restriction
+      if (!isDependencyPath && !targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
         throw new Error(`Path traversal attempt detected. Cannot access directory outside workspace: ${directory}`);
       }
     } else {
       targetDir = path.resolve(secureBaseDir, directory);
       // Double-check the resolved path is still within the secure base directory
-      if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      // Allow dependency paths to bypass this restriction
+      if (!isDependencyPath && !targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
         throw new Error(`Path traversal attempt detected. Access denied: ${directory}`);
       }
     }
@@ -323,19 +331,27 @@ export const searchFilesTool = {
     const baseCwd = workingDirectory || process.cwd();
     const secureBaseDir = path.resolve(baseCwd);
 
+    // Check if this is a dependency path that should bypass workspace restrictions
+    const isDependencyPath = directory.startsWith('/dep/') ||
+                            directory.startsWith('go:') ||
+                            directory.startsWith('js:') ||
+                            directory.startsWith('rust:');
+
     // If directory is absolute, check if it's within the secure base directory
     // If it's relative, resolve it against the secure base directory
     let targetDir;
     if (path.isAbsolute(directory)) {
       targetDir = path.resolve(directory);
       // Check if the absolute path is within the secure base directory
-      if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      // Allow dependency paths to bypass this restriction
+      if (!isDependencyPath && !targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
         throw new Error(`Path traversal attempt detected. Cannot access directory outside workspace: ${directory}`);
       }
     } else {
       targetDir = path.resolve(secureBaseDir, directory);
       // Double-check the resolved path is still within the secure base directory
-      if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      // Allow dependency paths to bypass this restriction
+      if (!isDependencyPath && !targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
         throw new Error(`Path traversal attempt detected. Access denied: ${directory}`);
       }
     }

package/src/agent/storage/InMemoryStorageAdapter.js

@@ -0,0 +1,49 @@
+import { StorageAdapter } from './StorageAdapter.js';
+
+/**
+ * Default in-memory storage adapter
+ * This is the default behavior - stores history in a Map in memory
+ */
+export class InMemoryStorageAdapter extends StorageAdapter {
+  constructor() {
+    super();
+    this.sessions = new Map(); // sessionId -> {messages: [], metadata: {}}
+  }
+
+  async loadHistory(sessionId) {
+    const session = this.sessions.get(sessionId);
+    return session ? session.messages : [];
+  }
+
+  async saveMessage(sessionId, message) {
+    if (!this.sessions.has(sessionId)) {
+      this.sessions.set(sessionId, {
+        messages: [],
+        metadata: {
+          createdAt: new Date().toISOString(),
+          lastActivity: new Date().toISOString()
+        }
+      });
+    }
+
+    const session = this.sessions.get(sessionId);
+    session.messages.push(message);
+    session.metadata.lastActivity = new Date().toISOString();
+  }
+
+  async clearHistory(sessionId) {
+    this.sessions.delete(sessionId);
+  }
+
+  async getSessionMetadata(sessionId) {
+    const session = this.sessions.get(sessionId);
+    return session ? session.metadata : null;
+  }
+
+  async updateSessionActivity(sessionId) {
+    const session = this.sessions.get(sessionId);
+    if (session) {
+      session.metadata.lastActivity = new Date().toISOString();
+    }
+  }
+}

package/src/agent/storage/StorageAdapter.js

@@ -0,0 +1,51 @@
+/**
+ * Base class for storage adapters
+ * Implement this interface to provide custom storage backends for ProbeAgent history
+ */
+export class StorageAdapter {
+  /**
+   * Load conversation history for a session
+   * @param {string} sessionId - Session identifier
+   * @returns {Promise<Array<Object>>} Array of message objects with {role, content, ...}
+   */
+  async loadHistory(sessionId) {
+    throw new Error('StorageAdapter.loadHistory() must be implemented by subclass');
+  }
+
+  /**
+   * Save a message to storage
+   * @param {string} sessionId - Session identifier
+   * @param {Object} message - Message object { role, content, ... }
+   * @returns {Promise<void>}
+   */
+  async saveMessage(sessionId, message) {
+    throw new Error('StorageAdapter.saveMessage() must be implemented by subclass');
+  }
+
+  /**
+   * Clear history for a session
+   * @param {string} sessionId - Session identifier
+   * @returns {Promise<void>}
+   */
+  async clearHistory(sessionId) {
+    throw new Error('StorageAdapter.clearHistory() must be implemented by subclass');
+  }
+
+  /**
+   * Get session metadata (optional)
+   * @param {string} sessionId - Session identifier
+   * @returns {Promise<Object|null>} Session metadata or null
+   */
+  async getSessionMetadata(sessionId) {
+    return null;
+  }
+
+  /**
+   * Update session activity timestamp (optional)
+   * @param {string} sessionId - Session identifier
+   * @returns {Promise<void>}
+   */
+  async updateSessionActivity(sessionId) {
+    // Optional - implement if you want to track session activity
+  }
+}

package/src/agent/storage/index.js

@@ -0,0 +1,2 @@
+export { StorageAdapter } from './StorageAdapter.js';
+export { InMemoryStorageAdapter } from './InMemoryStorageAdapter.js';

package/src/index.js

@@ -36,6 +36,8 @@ import { bashTool } from './tools/bash.js';
 import { ProbeAgent } from './agent/ProbeAgent.js';
 import { SimpleTelemetry, SimpleAppTracer, initializeSimpleTelemetryFromOptions } from './agent/simpleTelemetry.js';
 import { listFilesToolInstance, searchFilesToolInstance } from './agent/probeTool.js';
+import { StorageAdapter, InMemoryStorageAdapter } from './agent/storage/index.js';
+import { HookManager, HOOK_TYPES } from './agent/hooks/index.js';
 
 export {
 	search,
@@ -50,6 +52,12 @@ export {
 	DEFAULT_SYSTEM_MESSAGE,
 	// Export AI Agent (NEW!)
 	ProbeAgent,
+	// Export storage adapters
+	StorageAdapter,
+	InMemoryStorageAdapter,
+	// Export hooks
+	HookManager,
+	HOOK_TYPES,
 	// Export telemetry classes
 	SimpleTelemetry,
 	SimpleAppTracer,