@probelabs/probe 0.6.0-rc137 → 0.6.0-rc139

package/index.d.ts

@@ -350,13 +350,21 @@ export interface QueryParams {
  */
 export interface ExtractParams {
   /** Files and line numbers or symbols to extract */
-  files: string[];
+  files?: string[];
+  /** Path to a file containing unstructured text to extract file paths from */
+  inputFile?: string;
+  /** Content to pipe to stdin (e.g., git diff output). Alternative to inputFile or files. */
+  content?: string | Buffer;
   /** Path to search in */
   path?: string;
   /** Number of context lines */
   contextLines?: number;
   /** Output format */
-  format?: 'markdown' | 'plain' | 'json';
+  format?: 'markdown' | 'plain' | 'json' | 'xml' | 'color' | 'outline-xml' | 'outline-diff';
+  /** Include test files */
+  allowTests?: boolean;
+  /** Return results as parsed JSON instead of string */
+  json?: boolean;
   /** Session ID */
   sessionId?: string;
 }

package/package.json

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

package/src/agent/ProbeAgent.js

@@ -1270,6 +1270,9 @@ When troubleshooting:
               temperature: 0.3,
             });
 
+            // Get the promise reference BEFORE consuming stream (doesn't lock it)
+            const usagePromise = result.usage;
+
             // Collect the streamed response - stream all content for now
             for await (const delta of result.textStream) {
               assistantResponseContent += delta;
@@ -1279,8 +1282,8 @@ When troubleshooting:
               }
             }
 
-            // Record token usage
-            const usage = await result.usage;
+            // Record token usage - await the promise AFTER stream is consumed
+            const usage = await usagePromise;
             if (usage) {
               this.tokenCounter.recordUsage(usage, result.experimental_providerMetadata);
             }

package/src/extract.js

@@ -3,7 +3,7 @@
  * @module extract
  */
 
-import { exec } from 'child_process';
+import { exec, spawn } from 'child_process';
 import { promisify } from 'util';
 import { getBinaryPath, buildCliArgs, escapeString } from './utils.js';
 
@@ -22,13 +22,14 @@ const EXTRACT_FLAG_MAP = {
 
 /**
  * Extract code blocks from files
- * 
+ *
  * @param {Object} options - Extract options
  * @param {string[]} [options.files] - Files to extract from (can include line numbers with colon, e.g., "/path/to/file.rs:10")
  * @param {string} [options.inputFile] - Path to a file containing unstructured text to extract file paths from
+ * @param {string|Buffer} [options.content] - Content to pipe to stdin (e.g., git diff output). Alternative to inputFile.
  * @param {boolean} [options.allowTests] - Include test files
  * @param {number} [options.contextLines] - Number of context lines to include
- * @param {string} [options.format] - Output format ('markdown', 'plain', 'json')
+ * @param {string} [options.format] - Output format ('markdown', 'plain', 'json', 'xml', 'color', 'outline-xml', 'outline-diff')
  * @param {Object} [options.binaryOptions] - Options for getting the binary
  * @param {boolean} [options.binaryOptions.forceDownload] - Force download even if binary exists
  * @param {string} [options.binaryOptions.version] - Specific version to download
@@ -41,16 +42,22 @@ export async function extract(options) {
 		throw new Error('Options object is required');
 	}
 
-	// Either files or inputFile must be provided
-	if ((!options.files || !Array.isArray(options.files) || options.files.length === 0) && !options.inputFile) {
-		throw new Error('Either files array or inputFile must be provided');
+	// Either files, inputFile, or content must be provided
+	const hasFiles = options.files && Array.isArray(options.files) && options.files.length > 0;
+	const hasInputFile = !!options.inputFile;
+	const hasContent = options.content !== undefined && options.content !== null;
+
+	if (!hasFiles && !hasInputFile && !hasContent) {
+		throw new Error('Either files array, inputFile, or content must be provided');
 	}
 
 	// Get the binary path
 	const binaryPath = await getBinaryPath(options.binaryOptions || {});
 
-	// Build CLI arguments from options
-	const cliArgs = buildCliArgs(options, EXTRACT_FLAG_MAP);
+	// Build CLI arguments from options (excluding content which goes via stdin)
+	const filteredOptions = { ...options };
+	delete filteredOptions.content;
+	const cliArgs = buildCliArgs(filteredOptions, EXTRACT_FLAG_MAP);
 
 	// If json option is true, override format to json
 	if (options.json && !options.format) {
@@ -58,7 +65,7 @@ export async function extract(options) {
 	}
 
 	// Add files as positional arguments if provided
-	if (options.files && Array.isArray(options.files) && options.files.length > 0) {
+	if (hasFiles) {
 		for (const file of options.files) {
 			cliArgs.push(escapeString(file));
 		}
@@ -71,6 +78,7 @@ export async function extract(options) {
 			logMessage += ` files="${options.files.join(', ')}"`;
 		}
 		if (options.inputFile) logMessage += ` inputFile="${options.inputFile}"`;
+		if (options.content) logMessage += ` content=(${typeof options.content === 'string' ? options.content.length : options.content.byteLength} bytes)`;
 		if (options.allowTests) logMessage += " allowTests=true";
 		if (options.contextLines) logMessage += ` contextLines=${options.contextLines}`;
 		if (options.format) logMessage += ` format=${options.format}`;
@@ -78,7 +86,12 @@ export async function extract(options) {
 		console.error(logMessage);
 	}
 
-	// Execute command
+	// If content is provided, use spawn with stdin piping
+	if (hasContent) {
+		return extractWithStdin(binaryPath, cliArgs, options.content, options);
+	}
+
+	// Otherwise use exec for simple command execution
 	const command = `${binaryPath} extract ${cliArgs.join(' ')}`;
 
 	try {
@@ -88,62 +101,132 @@ export async function extract(options) {
 			console.error(`stderr: ${stderr}`);
 		}
 
-		// Parse the output to extract token usage information
-		let tokenUsage = {
-			requestTokens: 0,
-			responseTokens: 0,
-			totalTokens: 0
-		};
-
-		// Calculate approximate request tokens
-		if (options.files && Array.isArray(options.files)) {
-			tokenUsage.requestTokens = options.files.join(' ').length / 4;
-		} else if (options.inputFile) {
-			tokenUsage.requestTokens = options.inputFile.length / 4;
-		}
+		return processExtractOutput(stdout, options);
+	} catch (error) {
+		// Enhance error message with command details
+		const errorMessage = `Error executing extract command: ${error.message}\nCommand: ${command}`;
+		throw new Error(errorMessage);
+	}
+}
 
-		// Try to extract token information from the output
-		if (stdout.includes('Total tokens returned:')) {
-			const tokenMatch = stdout.match(/Total tokens returned: (\d+)/);
-			if (tokenMatch && tokenMatch[1]) {
-				tokenUsage.responseTokens = parseInt(tokenMatch[1], 10);
-				tokenUsage.totalTokens = tokenUsage.requestTokens + tokenUsage.responseTokens;
+/**
+ * Extract with content piped to stdin
+ * @private
+ */
+function extractWithStdin(binaryPath, cliArgs, content, options) {
+	return new Promise((resolve, reject) => {
+		const process = spawn(binaryPath, ['extract', ...cliArgs], {
+			stdio: ['pipe', 'pipe', 'pipe']
+		});
+
+		let stdout = '';
+		let stderr = '';
+
+		// Collect stdout
+		process.stdout.on('data', (data) => {
+			stdout += data.toString();
+		});
+
+		// Collect stderr
+		process.stderr.on('data', (data) => {
+			stderr += data.toString();
+		});
+
+		// Handle process exit
+		process.on('close', (code) => {
+			if (stderr && process.env.DEBUG === '1') {
+				console.error(`stderr: ${stderr}`);
 			}
-		}
-
-		// Add token usage information to the output
-		let output = stdout;
 
-		// Add token usage information at the end if not already present
-		if (!output.includes('Token Usage:')) {
-			output += `\nToken Usage:\n  Request tokens: ${tokenUsage.requestTokens}\n  Response tokens: ${tokenUsage.responseTokens}\n  Total tokens: ${tokenUsage.totalTokens}\n`;
-		}
+			if (code !== 0) {
+				reject(new Error(`Extract command failed with exit code ${code}: ${stderr}`));
+				return;
+			}
 
-		// Parse JSON if requested or if format is json
-		if (options.json || options.format === 'json') {
 			try {
-				const jsonOutput = JSON.parse(stdout);
-
-				// Add token usage to JSON output
-				if (!jsonOutput.token_usage) {
-					jsonOutput.token_usage = {
-						request_tokens: tokenUsage.requestTokens,
-						response_tokens: tokenUsage.responseTokens,
-						total_tokens: tokenUsage.totalTokens
-					};
-				}
-
-				return jsonOutput;
+				const result = processExtractOutput(stdout, options);
+				resolve(result);
 			} catch (error) {
-				console.error('Error parsing JSON output:', error);
-				return output; // Fall back to string output with token usage
+				reject(error);
 			}
+		});
+
+		// Handle errors
+		process.on('error', (error) => {
+			reject(new Error(`Failed to spawn extract process: ${error.message}`));
+		});
+
+		// Write content to stdin and close
+		if (typeof content === 'string') {
+			process.stdin.write(content);
+		} else {
+			process.stdin.write(content);
 		}
+		process.stdin.end();
+	});
+}
 
-		return output;
-	} catch (error) {
-		// Enhance error message with command details
-		const errorMessage = `Error executing extract command: ${error.message}\nCommand: ${command}`;
-		throw new Error(errorMessage);
+/**
+ * Process extract output and add token usage information
+ * @private
+ */
+function processExtractOutput(stdout, options) {
+	// Parse the output to extract token usage information
+	let tokenUsage = {
+		requestTokens: 0,
+		responseTokens: 0,
+		totalTokens: 0
+	};
+
+	// Calculate approximate request tokens
+	if (options.files && Array.isArray(options.files)) {
+		tokenUsage.requestTokens = options.files.join(' ').length / 4;
+	} else if (options.inputFile) {
+		tokenUsage.requestTokens = options.inputFile.length / 4;
+	} else if (options.content) {
+		const contentLength = typeof options.content === 'string'
+			? options.content.length
+			: options.content.byteLength;
+		tokenUsage.requestTokens = contentLength / 4;
+	}
+
+	// Try to extract token information from the output
+	if (stdout.includes('Total tokens returned:')) {
+		const tokenMatch = stdout.match(/Total tokens returned: (\d+)/);
+		if (tokenMatch && tokenMatch[1]) {
+			tokenUsage.responseTokens = parseInt(tokenMatch[1], 10);
+			tokenUsage.totalTokens = tokenUsage.requestTokens + tokenUsage.responseTokens;
+		}
+	}
+
+	// Add token usage information to the output
+	let output = stdout;
+
+	// Add token usage information at the end if not already present
+	if (!output.includes('Token Usage:')) {
+		output += `\nToken Usage:\n  Request tokens: ${tokenUsage.requestTokens}\n  Response tokens: ${tokenUsage.responseTokens}\n  Total tokens: ${tokenUsage.totalTokens}\n`;
+	}
+
+	// Parse JSON if requested or if format is json
+	if (options.json || options.format === 'json') {
+		try {
+			const jsonOutput = JSON.parse(stdout);
+
+			// Add token usage to JSON output
+			if (!jsonOutput.token_usage) {
+				jsonOutput.token_usage = {
+					request_tokens: tokenUsage.requestTokens,
+					response_tokens: tokenUsage.responseTokens,
+					total_tokens: tokenUsage.totalTokens
+				};
+			}
+
+			return jsonOutput;
+		} catch (error) {
+			console.error('Error parsing JSON output:', error);
+			return output; // Fall back to string output with token usage
+		}
 	}
-}
+
+	return output;
+}