@probelabs/probe
Version:
Node.js wrapper for the probe code search tool
1,362 lines (1,355 loc) • 260 kB
JavaScript
#!/usr/bin/env node
var __defProp = Object.defineProperty;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __esm = (fn, res) => function __init() {
return fn && (res = (0, fn[__getOwnPropNames(fn)[0]])(fn = 0)), res;
};
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
// src/agent/tokenCounter.js
import { get_encoding } from "tiktoken";
var TokenCounter;
var init_tokenCounter = __esm({
"src/agent/tokenCounter.js"() {
"use strict";
TokenCounter = class {
constructor() {
try {
this.tokenizer = get_encoding("cl100k_base");
this.contextSize = 0;
this.history = [];
this.requestTokens = 0;
this.responseTokens = 0;
this.currentRequestTokens = 0;
this.currentResponseTokens = 0;
this.cacheCreationTokens = 0;
this.cacheReadTokens = 0;
this.currentCacheCreationTokens = 0;
this.currentCacheReadTokens = 0;
this.cachedPromptTokens = 0;
this.currentCachedPromptTokens = 0;
} catch (error) {
console.error("Error initializing tokenizer:", error);
this.tokenizer = null;
this.contextSize = 0;
this.requestTokens = 0;
this.responseTokens = 0;
this.currentRequestTokens = 0;
this.currentResponseTokens = 0;
this.cacheCreationTokens = 0;
this.cacheReadTokens = 0;
this.currentCacheCreationTokens = 0;
this.currentCacheReadTokens = 0;
this.cachedPromptTokens = 0;
this.currentCachedPromptTokens = 0;
this.history = [];
}
this.debug = process.env.DEBUG === "1";
}
/**
* Count tokens in a string using tiktoken or fallback method
* @param {string} text - The text to count tokens for
* @returns {number} - The number of tokens
*/
countTokens(text) {
if (typeof text !== "string") {
text = String(text);
}
if (this.tokenizer) {
try {
const tokens = this.tokenizer.encode(text);
return tokens.length;
} catch (error) {
return Math.ceil(text.length / 4);
}
} else {
return Math.ceil(text.length / 4);
}
}
/**
* Add to request token count (manual counting, less used now with recordUsage)
* @param {string|number} input - The text to count tokens for or the token count directly
*/
addRequestTokens(input) {
let tokenCount = 0;
if (typeof input === "number") {
tokenCount = input;
} else if (typeof input === "string") {
tokenCount = this.countTokens(input);
} else {
console.warn("[WARN] Invalid input type for addRequestTokens:", typeof input);
return;
}
this.requestTokens += tokenCount;
this.currentRequestTokens = tokenCount;
if (this.debug) {
console.log(`[DEBUG] (Manual) Added ${tokenCount} request tokens. Total: ${this.requestTokens}, Current: ${this.currentRequestTokens}`);
}
}
/**
* Add to response token count (manual counting, less used now with recordUsage)
* @param {string|number} input - The text to count tokens for or the token count directly
*/
addResponseTokens(input) {
let tokenCount = 0;
if (typeof input === "number") {
tokenCount = input;
} else if (typeof input === "string") {
tokenCount = this.countTokens(input);
} else {
console.warn("[WARN] Invalid input type for addResponseTokens:", typeof input);
return;
}
this.responseTokens += tokenCount;
this.currentResponseTokens = tokenCount;
if (this.debug) {
console.log(`[DEBUG] (Manual) Added ${tokenCount} response tokens. Total: ${this.responseTokens}, Current: ${this.currentResponseTokens}`);
}
}
/**
* Record token usage from the AI SDK's result for a single LLM call.
* This resets 'current' counters and updates totals.
* @param {Object} usage - The usage object { promptTokens, completionTokens, totalTokens }
* @param {Object} providerMetadata - Metadata possibly containing cache info
*/
recordUsage(usage, providerMetadata) {
if (!usage) {
console.warn("[WARN] No usage information provided to recordUsage");
return;
}
this.currentRequestTokens = 0;
this.currentResponseTokens = 0;
this.currentCacheCreationTokens = 0;
this.currentCacheReadTokens = 0;
this.currentCachedPromptTokens = 0;
const promptTokens = Number(usage.promptTokens) || 0;
const completionTokens = Number(usage.completionTokens) || 0;
this.currentRequestTokens = promptTokens;
this.currentResponseTokens = completionTokens;
this.requestTokens += promptTokens;
this.responseTokens += completionTokens;
if (providerMetadata?.anthropic) {
const cacheCreation = Number(providerMetadata.anthropic.cacheCreationInputTokens) || 0;
const cacheRead = Number(providerMetadata.anthropic.cacheReadInputTokens) || 0;
this.currentCacheCreationTokens = cacheCreation;
this.currentCacheReadTokens = cacheRead;
this.cacheCreationTokens += cacheCreation;
this.cacheReadTokens += cacheRead;
if (this.debug) {
console.log(`[DEBUG] Anthropic cache tokens (current): creation=${cacheCreation}, read=${cacheRead}`);
}
}
if (providerMetadata?.openai) {
const cachedPrompt = Number(providerMetadata.openai.cachedPromptTokens) || 0;
this.currentCachedPromptTokens = cachedPrompt;
this.cachedPromptTokens += cachedPrompt;
if (this.debug) {
console.log(`[DEBUG] OpenAI cached prompt tokens (current): ${cachedPrompt}`);
}
}
if (this.debug) {
console.log(
`[DEBUG] Recorded usage: current(req=${this.currentRequestTokens}, resp=${this.currentResponseTokens}), total(req=${this.requestTokens}, resp=${this.responseTokens})`
);
console.log(`[DEBUG] Total cache tokens: Anthropic(create=${this.cacheCreationTokens}, read=${this.cacheReadTokens}), OpenAI(prompt=${this.cachedPromptTokens})`);
}
}
/**
* Calculate the current context window size based on provided messages or internal history.
* @param {Array|null} messages - Optional messages array to use for calculation. If null, uses internal this.history.
* @returns {number} - Total tokens estimated in the context window.
*/
calculateContextSize(messages = null) {
const msgsToCount = messages !== null ? messages : this.history;
let totalTokens = 0;
if (this.debug && messages === null) {
console.log(`[DEBUG] Calculating context size from internal history (${this.history.length} messages)`);
}
for (const msg of msgsToCount) {
let messageTokens = 0;
messageTokens += 4;
if (typeof msg.content === "string") {
messageTokens += this.countTokens(msg.content);
} else if (Array.isArray(msg.content)) {
for (const item of msg.content) {
if (item.type === "text" && typeof item.text === "string") {
messageTokens += this.countTokens(item.text);
} else {
messageTokens += this.countTokens(JSON.stringify(item));
}
}
} else if (msg.content) {
messageTokens += this.countTokens(JSON.stringify(msg.content));
}
if (msg.toolCalls) {
messageTokens += this.countTokens(JSON.stringify(msg.toolCalls));
messageTokens += 5;
}
if (msg.role === "tool" && msg.toolCallId) {
messageTokens += this.countTokens(msg.toolCallId);
messageTokens += 5;
}
if (msg.toolCallResults) {
messageTokens += this.countTokens(JSON.stringify(msg.toolCallResults));
messageTokens += 5;
}
totalTokens += messageTokens;
}
if (messages === null) {
this.contextSize = totalTokens;
if (this.debug) {
console.log(`[DEBUG] Updated internal context size: ${this.contextSize} tokens`);
}
}
return totalTokens;
}
/**
* Update internal history and recalculate internal context window size.
* @param {Array} messages - New message history array.
*/
updateHistory(messages) {
if (!Array.isArray(messages)) {
console.warn("[WARN] updateHistory called with non-array:", messages);
this.history = [];
} else {
this.history = [...messages];
}
this.calculateContextSize();
if (this.debug) {
console.log(`[DEBUG] History updated (${this.history.length} messages). Recalculated context size: ${this.contextSize}`);
}
}
/**
* Clear all counters and internal history. Reset context size.
*/
clear() {
this.requestTokens = 0;
this.responseTokens = 0;
this.currentRequestTokens = 0;
this.currentResponseTokens = 0;
this.cacheCreationTokens = 0;
this.cacheReadTokens = 0;
this.currentCacheCreationTokens = 0;
this.currentCacheReadTokens = 0;
this.cachedPromptTokens = 0;
this.currentCachedPromptTokens = 0;
this.history = [];
this.contextSize = 0;
if (this.debug) {
console.log("[DEBUG] TokenCounter cleared: usage, history, and context size reset.");
}
}
/**
* Start a new conversation turn - reset CURRENT token counters.
* Calculates context size based on history *before* the new turn.
*/
startNewTurn() {
this.currentRequestTokens = 0;
this.currentResponseTokens = 0;
this.currentCacheCreationTokens = 0;
this.currentCacheReadTokens = 0;
this.currentCachedPromptTokens = 0;
this.calculateContextSize();
if (this.debug) {
console.log("[DEBUG] TokenCounter: New turn started. Current counters reset.");
console.log(`[DEBUG] Context size at start of turn: ${this.contextSize} tokens`);
}
}
/**
* Get the current token usage state including context size.
* Recalculates context size from internal history before returning.
* @returns {Object} - Object containing current turn, total session, and context window usage.
*/
getTokenUsage() {
const currentContextSize = this.calculateContextSize();
const currentCacheRead = this.currentCacheReadTokens + this.currentCachedPromptTokens;
const currentCacheWrite = this.currentCacheCreationTokens;
const totalCacheRead = this.cacheReadTokens + this.cachedPromptTokens;
const totalCacheWrite = this.cacheCreationTokens;
const usageData = {
contextWindow: currentContextSize,
// Use the freshly calculated value
current: {
// Usage for the *last* LLM call recorded
request: this.currentRequestTokens,
response: this.currentResponseTokens,
total: this.currentRequestTokens + this.currentResponseTokens,
cacheRead: currentCacheRead,
cacheWrite: currentCacheWrite,
cacheTotal: currentCacheRead + currentCacheWrite,
// Keep detailed breakdown if needed
anthropic: {
cacheCreation: this.currentCacheCreationTokens,
cacheRead: this.currentCacheReadTokens
},
openai: {
cachedPrompt: this.currentCachedPromptTokens
}
},
total: {
// Accumulated usage over the session
request: this.requestTokens,
response: this.responseTokens,
total: this.requestTokens + this.responseTokens,
cacheRead: totalCacheRead,
cacheWrite: totalCacheWrite,
cacheTotal: totalCacheRead + totalCacheWrite,
// Keep detailed breakdown if needed
anthropic: {
cacheCreation: this.cacheCreationTokens,
cacheRead: this.cacheReadTokens
},
openai: {
cachedPrompt: this.cachedPromptTokens
}
}
};
if (this.debug) {
}
return usageData;
}
};
}
});
// src/directory-resolver.js
import path from "path";
import os from "os";
import fs from "fs-extra";
import { fileURLToPath } from "url";
async function getPackageBinDir() {
const debug = process.env.DEBUG === "1" || process.env.VERBOSE === "1";
if (debug) {
console.log("DEBUG: Starting probe binary directory resolution");
}
if (process.env.PROBE_BINARY_PATH) {
if (debug) {
console.log(`DEBUG: Checking PROBE_BINARY_PATH: ${process.env.PROBE_BINARY_PATH}`);
}
const binaryPath = process.env.PROBE_BINARY_PATH;
if (await fs.pathExists(binaryPath)) {
const binDir = path.dirname(binaryPath);
if (await canWriteToDirectory(binDir)) {
if (debug) {
console.log(`DEBUG: Using PROBE_BINARY_PATH directory: ${binDir}`);
}
return binDir;
}
} else {
console.warn(`Warning: PROBE_BINARY_PATH ${binaryPath} does not exist`);
}
}
if (process.env.PROBE_CACHE_DIR) {
if (debug) {
console.log(`DEBUG: Checking PROBE_CACHE_DIR: ${process.env.PROBE_CACHE_DIR}`);
}
const cacheDir = path.join(process.env.PROBE_CACHE_DIR, "bin");
if (await ensureDirectory(cacheDir)) {
if (debug) {
console.log(`DEBUG: Using PROBE_CACHE_DIR: ${cacheDir}`);
}
return cacheDir;
}
}
const packageRoot = await findPackageRoot();
if (packageRoot) {
if (debug) {
console.log(`DEBUG: Found package root: ${packageRoot}`);
}
const packageBinDir = path.join(packageRoot, "bin");
if (await ensureDirectory(packageBinDir) && await canWriteToDirectory(packageBinDir)) {
if (debug) {
console.log(`DEBUG: Using package bin directory: ${packageBinDir}`);
}
return packageBinDir;
} else if (debug) {
console.log(`DEBUG: Package bin directory ${packageBinDir} not writable, trying fallbacks`);
}
}
const homeCache = path.join(os.homedir(), ".probe", "bin");
if (debug) {
console.log(`DEBUG: Trying home cache directory: ${homeCache}`);
}
if (await ensureDirectory(homeCache)) {
if (debug) {
console.log(`DEBUG: Using home cache directory: ${homeCache}`);
}
return homeCache;
}
const tempCache = path.join(os.tmpdir(), "probe-cache", "bin");
if (debug) {
console.log(`DEBUG: Trying temp cache directory: ${tempCache}`);
}
if (await ensureDirectory(tempCache)) {
if (debug) {
console.log(`DEBUG: Using temp cache directory: ${tempCache}`);
}
return tempCache;
}
const errorMessage = [
"Could not find a writable directory for probe binary.",
"Tried the following locations:",
packageRoot ? `- Package bin directory: ${path.join(packageRoot, "bin")}` : "- Package root not found",
`- Home cache directory: ${homeCache}`,
`- Temp cache directory: ${tempCache}`,
"",
"You can override the location using environment variables:",
"- PROBE_BINARY_PATH=/path/to/probe (direct path to binary)",
"- PROBE_CACHE_DIR=/path/to/cache (cache directory, binary will be in /bin subdirectory)"
].join("\n");
throw new Error(errorMessage);
}
async function findPackageRoot() {
const debug = process.env.DEBUG === "1" || process.env.VERBOSE === "1";
let currentDir = __dirname;
const rootDir = path.parse(currentDir).root;
if (debug) {
console.log(`DEBUG: Starting package root search from: ${currentDir}`);
}
while (currentDir !== rootDir) {
const packageJsonPath = path.join(currentDir, "package.json");
try {
if (await fs.pathExists(packageJsonPath)) {
const packageJson = await fs.readJson(packageJsonPath);
if (debug) {
console.log(`DEBUG: Found package.json at ${packageJsonPath}, name: ${packageJson.name}`);
}
if (packageJson.name === "@probelabs/probe") {
if (debug) {
console.log(`DEBUG: Found probe package root: ${currentDir}`);
}
return currentDir;
}
}
} catch (err) {
if (debug) {
console.log(`DEBUG: Error reading package.json at ${packageJsonPath}: ${err.message}`);
}
}
currentDir = path.dirname(currentDir);
}
if (debug) {
console.log("DEBUG: Package root not found, reached filesystem root");
}
return null;
}
async function ensureDirectory(dirPath) {
const debug = process.env.DEBUG === "1" || process.env.VERBOSE === "1";
try {
await fs.ensureDir(dirPath);
const testFile = path.join(dirPath, ".probe-write-test");
await fs.writeFile(testFile, "test");
await fs.remove(testFile);
if (debug) {
console.log(`DEBUG: Directory ${dirPath} is writable`);
}
return true;
} catch (error) {
if (debug) {
console.log(`DEBUG: Directory ${dirPath} not writable: ${error.message}`);
}
return false;
}
}
async function canWriteToDirectory(dirPath) {
const debug = process.env.DEBUG === "1" || process.env.VERBOSE === "1";
try {
const exists = await fs.pathExists(dirPath);
if (!exists) {
if (debug) {
console.log(`DEBUG: Directory ${dirPath} does not exist`);
}
return false;
}
const testFile = path.join(dirPath, ".probe-write-test");
await fs.writeFile(testFile, "test");
await fs.remove(testFile);
if (debug) {
console.log(`DEBUG: Directory ${dirPath} is writable`);
}
return true;
} catch (error) {
if (debug) {
console.log(`DEBUG: Directory ${dirPath} not writable: ${error.message}`);
}
return false;
}
}
var __filename, __dirname;
var init_directory_resolver = __esm({
"src/directory-resolver.js"() {
"use strict";
__filename = fileURLToPath(import.meta.url);
__dirname = path.dirname(__filename);
}
});
// src/downloader.js
import axios from "axios";
import fs2 from "fs-extra";
import path2 from "path";
import { createHash } from "crypto";
import { promisify } from "util";
import { exec as execCallback } from "child_process";
import tar from "tar";
import os2 from "os";
import { fileURLToPath as fileURLToPath2 } from "url";
function detectOsArch() {
const osType = os2.platform();
const archType = os2.arch();
let osInfo;
let archInfo;
switch (osType) {
case "linux":
osInfo = {
type: "linux",
keywords: ["linux", "Linux", "gnu"]
};
break;
case "darwin":
osInfo = {
type: "darwin",
keywords: ["darwin", "Darwin", "mac", "Mac", "apple", "Apple", "osx", "OSX"]
};
break;
case "win32":
osInfo = {
type: "windows",
keywords: ["windows", "Windows", "msvc", "pc-windows"]
};
break;
default:
throw new Error(`Unsupported operating system: ${osType}`);
}
switch (archType) {
case "x64":
archInfo = {
type: "x86_64",
keywords: ["x86_64", "amd64", "x64", "64bit", "64-bit"]
};
break;
case "arm64":
archInfo = {
type: "aarch64",
keywords: ["arm64", "aarch64", "arm", "ARM"]
};
break;
default:
throw new Error(`Unsupported architecture: ${archType}`);
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Detected OS: ${osInfo.type}, Architecture: ${archInfo.type}`);
}
return { os: osInfo, arch: archInfo };
}
function constructAssetInfo(version, osInfo, archInfo) {
let platform;
let extension;
switch (osInfo.type) {
case "linux":
platform = `${archInfo.type}-unknown-linux-gnu`;
extension = "tar.gz";
break;
case "darwin":
platform = `${archInfo.type}-apple-darwin`;
extension = "tar.gz";
break;
case "windows":
platform = `${archInfo.type}-pc-windows-msvc`;
extension = "zip";
break;
default:
throw new Error(`Unsupported OS type: ${osInfo.type}`);
}
const assetName = `probe-v${version}-${platform}.${extension}`;
const checksumName = `${assetName}.sha256`;
const baseUrl = `https://github.com/${REPO_OWNER}/${REPO_NAME}/releases/download/v${version}`;
const assetUrl = `${baseUrl}/${assetName}`;
const checksumUrl = `${baseUrl}/${checksumName}`;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Constructed asset URL: ${assetUrl}`);
}
return {
name: assetName,
url: assetUrl,
checksumName,
checksumUrl
};
}
async function getLatestRelease(version) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log("Fetching release information...");
}
try {
let releaseUrl;
if (version) {
releaseUrl = `https://api.github.com/repos/${REPO_OWNER}/${REPO_NAME}/releases/tags/v${version}`;
} else {
releaseUrl = `https://api.github.com/repos/${REPO_OWNER}/${REPO_NAME}/releases`;
}
const response = await axios.get(releaseUrl);
if (response.status !== 200) {
throw new Error(`Failed to fetch release information: ${response.statusText}`);
}
let releaseData;
if (version) {
releaseData = response.data;
} else {
if (!Array.isArray(response.data) || response.data.length === 0) {
throw new Error("No releases found");
}
releaseData = response.data[0];
}
const tag = releaseData.tag_name;
const assets = releaseData.assets.map((asset) => ({
name: asset.name,
url: asset.browser_download_url
}));
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found release: ${tag} with ${assets.length} assets`);
}
return { tag, assets };
} catch (error) {
if (axios.isAxiosError(error) && error.response?.status === 404) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Release v${version} not found, trying to fetch all releases...`);
}
const response = await axios.get(`https://api.github.com/repos/${REPO_OWNER}/${REPO_NAME}/releases`);
if (response.data.length === 0) {
throw new Error("No releases found");
}
let bestRelease = response.data[0];
if (version && version !== "0.0.0") {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Looking for releases matching version: ${version}`);
console.log(`Available releases: ${response.data.slice(0, 5).map((r) => r.tag_name).join(", ")}...`);
}
for (const release of response.data) {
const releaseTag = release.tag_name.startsWith("v") ? release.tag_name.substring(1) : release.tag_name;
if (releaseTag === version) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found exact matching release: ${release.tag_name}`);
}
bestRelease = release;
break;
}
}
if (bestRelease === response.data[0]) {
const versionParts = version.split(/[\.-]/);
const majorMinor = versionParts.slice(0, 2).join(".");
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Looking for releases matching major.minor: ${majorMinor}`);
}
for (const release of response.data) {
const releaseTag = release.tag_name.startsWith("v") ? release.tag_name.substring(1) : release.tag_name;
const releaseVersionParts = releaseTag.split(/[\.-]/);
const releaseMajorMinor = releaseVersionParts.slice(0, 2).join(".");
if (releaseMajorMinor === majorMinor) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found matching major.minor release: ${release.tag_name}`);
}
bestRelease = release;
break;
}
}
}
}
const tag = bestRelease.tag_name;
const assets = bestRelease.assets.map((asset) => ({
name: asset.name,
url: asset.browser_download_url
}));
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Using release: ${tag} with ${assets.length} assets`);
}
return { tag, assets };
}
throw error;
}
}
function findBestAsset(assets, osInfo, archInfo) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Finding appropriate binary for ${osInfo.type} ${archInfo.type}...`);
}
let bestAsset = null;
let bestScore = 0;
for (const asset of assets) {
if (asset.name.endsWith(".sha256") || asset.name.endsWith(".md5") || asset.name.endsWith(".asc")) {
continue;
}
if (osInfo.type === "windows" && asset.name.match(/darwin|linux/)) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Skipping non-Windows binary: ${asset.name}`);
}
continue;
} else if (osInfo.type === "darwin" && asset.name.match(/windows|msvc|linux/)) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Skipping non-macOS binary: ${asset.name}`);
}
continue;
} else if (osInfo.type === "linux" && asset.name.match(/darwin|windows|msvc/)) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Skipping non-Linux binary: ${asset.name}`);
}
continue;
}
let score = 0;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Evaluating asset: ${asset.name}`);
}
let osMatched = false;
for (const keyword of osInfo.keywords) {
if (asset.name.includes(keyword)) {
score += 10;
osMatched = true;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` OS match found (${keyword}): +10, score = ${score}`);
}
break;
}
}
for (const keyword of archInfo.keywords) {
if (asset.name.includes(keyword)) {
score += 5;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` Arch match found (${keyword}): +5, score = ${score}`);
}
break;
}
}
if (asset.name.startsWith(`${BINARY_NAME}-`)) {
score += 3;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` Binary name match: +3, score = ${score}`);
}
}
if (osMatched && score >= 15) {
score += 5;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` OS+Arch bonus: +5, score = ${score}`);
}
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` Final score for ${asset.name}: ${score}`);
}
if (score === 23) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found perfect match: ${asset.name}`);
}
return asset;
}
if (score > bestScore) {
bestScore = score;
bestAsset = asset;
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(` New best asset: ${asset.name} (score: ${score})`);
}
}
}
if (!bestAsset) {
throw new Error(`Could not find a suitable binary for ${osInfo.type} ${archInfo.type}`);
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Selected asset: ${bestAsset.name} (score: ${bestScore})`);
}
return bestAsset;
}
async function downloadAsset(asset, outputDir) {
await fs2.ensureDir(outputDir);
const assetPath = path2.join(outputDir, asset.name);
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Downloading ${asset.name}...`);
}
const assetResponse = await axios.get(asset.url, { responseType: "arraybuffer" });
await fs2.writeFile(assetPath, Buffer.from(assetResponse.data));
const checksumUrl = asset.checksumUrl || `${asset.url}.sha256`;
const checksumFileName = asset.checksumName || `${asset.name}.sha256`;
let checksumPath = null;
try {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Downloading checksum...`);
}
const checksumResponse = await axios.get(checksumUrl);
checksumPath = path2.join(outputDir, checksumFileName);
await fs2.writeFile(checksumPath, checksumResponse.data);
} catch (error) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log("No checksum file found, skipping verification");
}
}
return { assetPath, checksumPath };
}
async function verifyChecksum(assetPath, checksumPath) {
if (!checksumPath) {
return true;
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Verifying checksum...`);
}
const checksumContent = await fs2.readFile(checksumPath, "utf-8");
const expectedChecksum = checksumContent.trim().split(" ")[0];
const fileBuffer = await fs2.readFile(assetPath);
const actualChecksum = createHash("sha256").update(fileBuffer).digest("hex");
if (expectedChecksum !== actualChecksum) {
console.error(`Checksum verification failed!`);
console.error(`Expected: ${expectedChecksum}`);
console.error(`Actual: ${actualChecksum}`);
return false;
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Checksum verified successfully`);
}
return true;
}
async function extractBinary(assetPath, outputDir) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Extracting ${path2.basename(assetPath)}...`);
}
const assetName = path2.basename(assetPath);
const isWindows = os2.platform() === "win32";
const binaryName = isWindows ? `${BINARY_NAME}.exe` : `${BINARY_NAME}-binary`;
const binaryPath = path2.join(outputDir, binaryName);
try {
const extractDir = path2.join(outputDir, "temp_extract");
await fs2.ensureDir(extractDir);
if (assetName.endsWith(".tar.gz") || assetName.endsWith(".tgz")) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Extracting tar.gz to ${extractDir}...`);
}
await tar.extract({
file: assetPath,
cwd: extractDir
});
} else if (assetName.endsWith(".zip")) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Extracting zip to ${extractDir}...`);
}
await exec(`unzip -q "${assetPath}" -d "${extractDir}"`);
} else {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Copying binary directly to ${binaryPath}`);
}
await fs2.copyFile(assetPath, binaryPath);
if (!isWindows) {
await fs2.chmod(binaryPath, 493);
}
await fs2.remove(extractDir);
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Binary installed to ${binaryPath}`);
}
return binaryPath;
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Searching for binary in extracted files...`);
}
const findBinary = async (dir) => {
const entries = await fs2.readdir(dir, { withFileTypes: true });
for (const entry of entries) {
const fullPath = path2.join(dir, entry.name);
if (entry.isDirectory()) {
const result = await findBinary(fullPath);
if (result) return result;
} else if (entry.isFile()) {
if (entry.name === binaryName || entry.name === BINARY_NAME || isWindows && entry.name.endsWith(".exe")) {
return fullPath;
}
}
}
return null;
};
const binaryFilePath = await findBinary(extractDir);
if (!binaryFilePath) {
const allFiles = await fs2.readdir(extractDir, { recursive: true });
console.error(`Binary not found in extracted files. Found: ${allFiles.join(", ")}`);
throw new Error(`Binary not found in the archive.`);
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found binary at ${binaryFilePath}`);
console.log(`Copying binary to ${binaryPath}`);
}
await fs2.copyFile(binaryFilePath, binaryPath);
if (!isWindows) {
await fs2.chmod(binaryPath, 493);
}
await fs2.remove(extractDir);
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Binary successfully installed to ${binaryPath}`);
}
return binaryPath;
} catch (error) {
console.error(`Error extracting binary: ${error instanceof Error ? error.message : String(error)}`);
throw error;
}
}
async function getVersionInfo(binDir) {
try {
const versionInfoPath = path2.join(binDir, "version-info.json");
if (await fs2.pathExists(versionInfoPath)) {
const content = await fs2.readFile(versionInfoPath, "utf-8");
return JSON.parse(content);
}
return null;
} catch (error) {
console.warn(`Warning: Could not read version info: ${error}`);
return null;
}
}
async function saveVersionInfo(version, binDir) {
const versionInfo = {
version,
lastUpdated: (/* @__PURE__ */ new Date()).toISOString()
};
const versionInfoPath = path2.join(binDir, "version-info.json");
await fs2.writeFile(versionInfoPath, JSON.stringify(versionInfo, null, 2));
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Version info saved: ${version} at ${versionInfoPath}`);
}
}
async function getPackageVersion() {
try {
const possiblePaths = [
path2.resolve(__dirname2, "..", "package.json"),
// When installed from npm: src/../package.json
path2.resolve(__dirname2, "..", "..", "package.json")
// In development: src/../../package.json
];
for (const packageJsonPath of possiblePaths) {
try {
if (fs2.existsSync(packageJsonPath)) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found package.json at: ${packageJsonPath}`);
}
const packageJson = JSON.parse(fs2.readFileSync(packageJsonPath, "utf-8"));
if (packageJson.version) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Using version from package.json: ${packageJson.version}`);
}
return packageJson.version;
}
}
} catch (err) {
console.error(`Error reading package.json at ${packageJsonPath}:`, err);
}
}
return "0.0.0";
} catch (error) {
console.error("Error getting package version:", error);
return "0.0.0";
}
}
async function downloadProbeBinary(version) {
try {
const localDir = await getPackageBinDir();
if (!version || version === "0.0.0") {
version = await getPackageVersion();
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Downloading probe binary (version: ${version || "latest"})...`);
console.log(`Using binary directory: ${localDir}`);
}
const isWindows = os2.platform() === "win32";
const binaryName = isWindows ? `${BINARY_NAME}.exe` : `${BINARY_NAME}-binary`;
const binaryPath = path2.join(localDir, binaryName);
if (await fs2.pathExists(binaryPath)) {
const versionInfo = await getVersionInfo(localDir);
if (versionInfo && versionInfo.version === version) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Using existing binary at ${binaryPath} (version: ${versionInfo.version})`);
}
return binaryPath;
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Existing binary version (${versionInfo?.version || "unknown"}) doesn't match requested version (${version}). Downloading new version...`);
}
}
const { os: osInfo, arch: archInfo } = detectOsArch();
let versionToUse = version;
let bestAsset;
let tagVersion;
if (!versionToUse || versionToUse === "0.0.0") {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log("No specific version requested, will use the latest release");
}
const { tag, assets } = await getLatestRelease(void 0);
tagVersion = tag.startsWith("v") ? tag.substring(1) : tag;
bestAsset = findBestAsset(assets, osInfo, archInfo);
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Found release version: ${tagVersion}`);
}
} else {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Direct download for version: ${versionToUse}`);
}
tagVersion = versionToUse;
bestAsset = constructAssetInfo(versionToUse, osInfo, archInfo);
}
const { assetPath, checksumPath } = await downloadAsset(bestAsset, localDir);
const checksumValid = await verifyChecksum(assetPath, checksumPath);
if (!checksumValid) {
throw new Error("Checksum verification failed");
}
const extractedBinaryPath = await extractBinary(assetPath, localDir);
await saveVersionInfo(tagVersion, localDir);
try {
await fs2.remove(assetPath);
if (checksumPath) {
await fs2.remove(checksumPath);
}
} catch (err) {
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Warning: Could not clean up temporary files: ${err}`);
}
}
if (process.env.DEBUG === "1" || process.env.VERBOSE === "1") {
console.log(`Binary successfully installed at ${extractedBinaryPath} (version: ${tagVersion})`);
}
return extractedBinaryPath;
} catch (error) {
console.error("Error downloading probe binary:", error);
throw error;
}
}
var exec, REPO_OWNER, REPO_NAME, BINARY_NAME, __filename2, __dirname2;
var init_downloader = __esm({
"src/downloader.js"() {
"use strict";
init_utils();
init_directory_resolver();
exec = promisify(execCallback);
REPO_OWNER = "probelabs";
REPO_NAME = "probe";
BINARY_NAME = "probe";
__filename2 = fileURLToPath2(import.meta.url);
__dirname2 = path2.dirname(__filename2);
}
});
// src/utils.js
import path3 from "path";
import fs3 from "fs-extra";
import { fileURLToPath as fileURLToPath3 } from "url";
async function getBinaryPath(options = {}) {
const { forceDownload = false, version } = options;
if (probeBinaryPath && !forceDownload && fs3.existsSync(probeBinaryPath)) {
return probeBinaryPath;
}
if (process.env.PROBE_PATH && fs3.existsSync(process.env.PROBE_PATH) && !forceDownload) {
probeBinaryPath = process.env.PROBE_PATH;
return probeBinaryPath;
}
const binDir = await getPackageBinDir();
const isWindows = process.platform === "win32";
const binaryName = isWindows ? "probe.exe" : "probe";
const binaryPath = path3.join(binDir, binaryName);
if (fs3.existsSync(binaryPath) && !forceDownload) {
probeBinaryPath = binaryPath;
return probeBinaryPath;
}
console.log(`${forceDownload ? "Force downloading" : "Binary not found. Downloading"} probe binary...`);
probeBinaryPath = await downloadProbeBinary(version);
return probeBinaryPath;
}
function buildCliArgs(options, flagMap) {
const cliArgs = [];
for (const [key, flag] of Object.entries(flagMap)) {
if (key in options) {
const value = options[key];
if (typeof value === "boolean") {
if (value) {
cliArgs.push(flag);
}
} else if (Array.isArray(value)) {
for (const item of value) {
cliArgs.push(flag, item);
}
} else if (value !== void 0 && value !== null) {
cliArgs.push(flag, value.toString());
}
}
}
return cliArgs;
}
function escapeString(str) {
if (process.platform === "win32") {
return `"${str.replace(/"/g, '\\"')}"`;
} else {
return `'${str.replace(/'/g, "'\\''")}'`;
}
}
var __filename3, __dirname3, probeBinaryPath;
var init_utils = __esm({
"src/utils.js"() {
"use strict";
init_downloader();
init_directory_resolver();
__filename3 = fileURLToPath3(import.meta.url);
__dirname3 = path3.dirname(__filename3);
probeBinaryPath = "";
}
});
// src/search.js
import { exec as exec2 } from "child_process";
import { promisify as promisify2 } from "util";
async function search(options) {
if (!options || !options.path) {
throw new Error("Path is required");
}
if (!options.query) {
throw new Error("Query is required");
}
const binaryPath = await getBinaryPath(options.binaryOptions || {});
const cliArgs = buildCliArgs(options, SEARCH_FLAG_MAP);
if (options.json && !options.format) {
cliArgs.push("--format", "json");
} else if (options.format) {
if (options.format === "json") {
options.json = true;
}
}
if (!options.maxTokens) {
options.maxTokens = 1e4;
cliArgs.push("--max-tokens", "10000");
}
if (!options.timeout) {
options.timeout = 30;
cliArgs.push("--timeout", "30");
}
if (options.language) {
if (!cliArgs.includes("--language")) {
cliArgs.push("--language", options.language);
}
}
if (options.exact) {
if (!cliArgs.includes("--exact")) {
cliArgs.push("--exact");
}
}
if (!options.session && process.env.PROBE_SESSION_ID) {
options.session = process.env.PROBE_SESSION_ID;
}
const queries = Array.isArray(options.query) ? options.query : [options.query];
if (process.env.DEBUG === "1") {
let logMessage = `
Search: query="${queries[0]}" path="${options.path}"`;
if (options.maxResults) logMessage += ` maxResults=${options.maxResults}`;
logMessage += ` maxTokens=${options.maxTokens}`;
logMessage += ` timeout=${options.timeout}`;
if (options.allowTests) logMessage += " allowTests=true";
if (options.language) logMessage += ` language=${options.language}`;
if (options.exact) logMessage += " exact=true";
if (options.session) logMessage += ` session=${options.session}`;
console.error(logMessage);
}
const positionalArgs = [];
if (queries.length > 0) {
positionalArgs.push(escapeString(queries[0]));
}
positionalArgs.push(escapeString(options.path));
const command = `${binaryPath} search ${cliArgs.join(" ")} ${positionalArgs.join(" ")}`;
try {
const { stdout, stderr } = await execAsync(command, {
shell: true,
timeout: options.timeout * 1e3
// Convert seconds to milliseconds
});
if (stderr && process.env.DEBUG) {
console.error(`stderr: ${stderr}`);
}
let resultCount = 0;
let tokenCount = 0;
let bytesCount = 0;
const lines = stdout.split("\n");
for (const line of lines) {
if (line.startsWith("```") && !line.includes("```language")) {
resultCount++;
}
}
const totalBytesMatch = stdout.match(/Total bytes returned:\s*(\d+)/i);
if (totalBytesMatch && totalBytesMatch[1]) {
bytesCount = parseInt(totalBytesMatch[1], 10);
}
const totalTokensMatch = stdout.match(/Total tokens returned:\s*(\d+)/i);
if (totalTokensMatch && totalTokensMatch[1]) {
tokenCount = parseInt(totalTokensMatch[1], 10);
} else {
const tokenMatch = stdout.match(/Tokens:?\s*(\d+)/i) || stdout.match(/(\d+)\s*tokens/i) || stdout.match(/token count:?\s*(\d+)/i);
if (tokenMatch && tokenMatch[1]) {
tokenCount = parseInt(tokenMatch[1], 10);
} else {
tokenCount = options.maxTokens;
}
}
if (process.env.DEBUG === "1") {
let resultsMessage = `
Search results: ${resultCount} matches, ${tokenCount} tokens`;
if (bytesCount > 0) {
resultsMessage += `, ${bytesCount} bytes`;
}
console.error(resultsMessage);
}
if (options.json) {
try {
return JSON.parse(stdout);
} catch (error) {
console.error("Error parsing JSON output:", error);
return stdout;
}
}
return stdout;
} catch (error) {
if (error.code === "ETIMEDOUT" || error.killed) {
const timeoutMessage = `Search operation timed out after ${options.timeout} seconds.
Command: ${command}`;
console.error(timeoutMessage);
throw new Error(timeoutMessage);
}
const errorMessage = `Error executing search command: ${error.message}
Command: ${command}`;
throw new Error(errorMessage);
}
}
var execAsync, SEARCH_FLAG_MAP;
var init_search = __esm({
"src/search.js"() {
"use strict";
init_utils();
execAsync = promisify2(exec2);
SEARCH_FLAG_MAP = {
filesOnly: "--files-only",
ignore: "--ignore",
excludeFilenames: "--exclude-filenames",
reranker: "--reranker",
frequencySearch: "--frequency",
exact: "--exact",
maxResults: "--max-results",
maxBytes: "--max-bytes",
maxTokens: "--max-tokens",
allowTests: "--allow-tests",
noMerge: "--no-merge",
mergeThreshold: "--merge-threshold",
session: "--session",
timeout: "--timeout",
language: "--language",
format: "--format"
};
}
});
// src/query.js
import { exec as exec3 } from "child_process";
import { promisify as promisify3 } from "util";
async function query(options) {
if (!options || !options.path) {
throw new Error("Path is required");
}
if (!options.pattern) {
throw new Error("Pattern is required");
}
const binaryPath = await getBinaryPath(options.binaryOptions || {});
const cliArgs = buildCliArgs(options, QUERY_FLAG_MAP);
if (options.json && !options.format) {
cliArgs.push("--format", "json");
}
cliArgs.push(escapeString(options.pattern), escapeString(options.path));
if (process.env.DEBUG === "1") {
let logMessage = `Query: pattern="${options.pattern}" path="${options.path}"`;
if (options.language) logMessage += ` language=${options.language}`;
if (options.maxResults) logMessage += ` maxResults=${options.maxResults}`;
if (options.allowTests) logMessage += " allowTests=true";
console.error(logMessage);
}
const command = `${binaryPath} query ${cliArgs.join(" ")}`;
try {
const { stdout, stderr } = await execAsync2(command);
if (stderr) {
console.error(`stderr: ${stderr}`);
}
let resultCount = 0;
const lines = stdout.split("\n");
for (const line of lines) {
if (line.startsWith("```") && !line.includes("```language")) {
resultCount++;
}
}
if (process.env.DEBUG === "1") {
console.error(`Query results: ${resultCount} matches`);
}
if (options.json || options.format === "json") {
try {
return JSON.parse(stdout);
} catch (error) {
console.error("Error parsing JSON output:", error);
return stdout;
}
}
return stdout;
} catch (error) {
const errorMessage = `Error executing query command: ${error.message}
Command: ${command}`;
throw new Error(errorMessage);
}
}
var execAsync2, QUERY_FLAG_MAP;
var init_query = __esm({
"src/query.js"() {
"use strict";
init_utils();
execAsync2 = promisify3(exec3);
QUERY_FLAG_MAP = {
language: "--language",
ignore: "--ignore",
allowTests: "--allow-tests",
maxResults: "--max-results",
format: "--format"
};
}
});
// src/extract.js
import { exec as exec4 } from "child_process";
import { promisify as promisify4 } from "util";
async function extract(options) {
if (!options) {
throw new Error("Options object is required");
}
if ((!options.files || !Array.isArray(options.files) || options.files.length === 0) && !options.inputFile) {
throw new Error("Either files array or inputFile must be provided");
}
const binaryPath = await getBinaryPath(options.binaryOptions || {});
const cliArgs = buildCliArgs(options, EXTRACT_FLAG_MAP);
if (options.json && !options.format) {
cliArgs.push("--format", "json");
}
if (options.files && Array.isArray(options.files) && options.files.length > 0) {
for (const file of options.files) {
cliArgs.push(escapeString(file));
}
}
if (process.env.DEBUG === "1") {
let logMessage = `
Extract:`;
if (options.files && options.files.length > 0) {
logMessage += ` files="${options.files.join(", ")}"`;
}
if (options.inputFile) logMessage += ` inputFile="${options.inputFile}"`;
if (options.allowTests) logMessage += " allowTests=true";
if (options.contextLines) logMessage += ` contextLines=${options.contextLines}`;
if (options.format) logMessage += ` format=${options.format}`;
if (options.json) logMessage += " json=true";
console.error(logMessage);
}
const command = `${binaryPath} extract ${cliArgs.join(" ")}`;
try {
const { stdout, stderr } = await execAsync3(command);
if (stderr) {
console.error(`stderr: ${stderr}`);
}
let tokenUsage = {
requestTokens: 0,
responseTokens: 0,
totalTokens: 0
};
if (options.files && Array.isArray(options.files)) {
tokenUsage.requestTokens = options.files.join(" ").length / 4;
} else if (options.inputFile) {
tokenUsage.requestTokens = options.inputFile.length / 4;
}
if (stdout.includes("Total tokens returned:")) {
const tokenMatch = stdout.match(/Total tokens returned: (\d+)