codebase-contextualizer-cli 1.0.0

package/src/walker.js

@@ -0,0 +1,214 @@
+const fs = require("fs/promises");
+const path = require("path");
+const ignore = require("ignore");
+const { throwIfAborted } = require("./abort");
+const { CACHE_DIR } = require("./cache");
+const {
+  MAX_FILE_SIZE,
+  getFileSizeSkipReason,
+  getSourcePathSkipReason,
+  isBinaryExtension,
+} = require("./file-filter");
+const { toRelativePosixPath, toPosixPath } = require("./paths");
+
+const HARD_IGNORED_DIRS = new Set([".git", "node_modules", CACHE_DIR]);
+const SOURCE_EXTENSIONS = new Set([
+  ".c",
+  ".cc",
+  ".cpp",
+  ".cs",
+  ".cxx",
+  ".go",
+  ".java",
+  ".js",
+  ".jsx",
+  ".kt",
+  ".mjs",
+  ".php",
+  ".py",
+  ".rb",
+  ".rs",
+  ".swift",
+  ".ts",
+  ".tsx",
+]);
+
+async function readGitignore(directory, errors) {
+  const gitignorePath = path.join(directory, ".gitignore");
+
+  try {
+    const contents = await fs.readFile(gitignorePath, "utf8");
+    const matcher = ignore().add(contents.split(/\r?\n/));
+
+    return {
+      base: directory,
+      matcher,
+    };
+  } catch (error) {
+    if (error.code !== "ENOENT") {
+      errors.push({
+        path: gitignorePath,
+        message: error.message,
+      });
+    }
+
+    return null;
+  }
+}
+
+function isInside(base, target) {
+  const relative = path.relative(base, target);
+  return (
+    relative &&
+    relative !== ".." &&
+    !relative.startsWith(`..${path.sep}`) &&
+    !path.isAbsolute(relative)
+  );
+}
+
+function isIgnoredByContexts(contexts, target, isDirectory) {
+  for (const context of contexts) {
+    if (!isInside(context.base, target)) {
+      continue;
+    }
+
+    const relativePath = toPosixPath(path.relative(context.base, target));
+    const candidate = isDirectory ? `${relativePath}/` : relativePath;
+
+    if (context.matcher.ignores(candidate)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+function isSourceFile(fileName) {
+  return SOURCE_EXTENSIONS.has(path.extname(fileName).toLowerCase());
+}
+
+async function walkSourceFiles(root, options = {}) {
+  const files = [];
+  const errors = [];
+  const visitedDirectories = new Set();
+
+  function addSkippedFile(absolutePath, reason) {
+    errors.push({
+      path: absolutePath,
+      message: `Skipped: ${reason}`,
+    });
+  }
+
+  async function getDirectoryKey(directory) {
+    const realPath = await fs.realpath(directory);
+    const stat = await fs.stat(realPath);
+    return `${stat.dev}:${stat.ino}:${process.platform === "win32" ? realPath.toLowerCase() : realPath}`;
+  }
+
+  async function walkDirectory(directory, contexts) {
+    throwIfAborted(options.signal);
+
+    let directoryKey;
+
+    try {
+      directoryKey = await getDirectoryKey(directory);
+    } catch (error) {
+      errors.push({
+        path: directory,
+        message: error.message,
+      });
+      return;
+    }
+
+    if (visitedDirectories.has(directoryKey)) {
+      return;
+    }
+
+    visitedDirectories.add(directoryKey);
+
+    const localContext = await readGitignore(directory, errors);
+    const activeContexts = localContext ? [...contexts, localContext] : contexts;
+    let entries;
+
+    try {
+      entries = await fs.readdir(directory, { withFileTypes: true });
+    } catch (error) {
+      errors.push({
+        path: directory,
+        message: error.message,
+      });
+      return;
+    }
+
+    entries.sort((left, right) => left.name.localeCompare(right.name));
+
+    for (const entry of entries) {
+      throwIfAborted(options.signal);
+
+      const absolutePath = path.join(directory, entry.name);
+
+      if (entry.isSymbolicLink()) {
+        continue;
+      }
+
+      if (entry.isDirectory()) {
+        if (HARD_IGNORED_DIRS.has(entry.name) || isIgnoredByContexts(activeContexts, absolutePath, true)) {
+          continue;
+        }
+
+        await walkDirectory(absolutePath, activeContexts);
+        continue;
+      }
+
+      if (!entry.isFile() || isIgnoredByContexts(activeContexts, absolutePath, false)) {
+        continue;
+      }
+
+      if (isBinaryExtension(entry.name) || !isSourceFile(entry.name)) {
+        continue;
+      }
+
+      const pathSkipReason = getSourcePathSkipReason(entry.name);
+
+      if (pathSkipReason) {
+        addSkippedFile(absolutePath, pathSkipReason);
+        continue;
+      }
+
+      try {
+        const stat = await fs.stat(absolutePath);
+        const sizeSkipReason = getFileSizeSkipReason(stat.size);
+
+        if (sizeSkipReason) {
+          addSkippedFile(absolutePath, sizeSkipReason);
+          continue;
+        }
+
+        files.push({
+          absolutePath,
+          relativePath: toRelativePosixPath(root, absolutePath),
+          size: stat.size,
+          mtimeMs: stat.mtimeMs,
+        });
+      } catch (error) {
+        errors.push({
+          path: absolutePath,
+          message: error.message,
+        });
+      }
+    }
+  }
+
+  await walkDirectory(root, []);
+
+  return {
+    files,
+    errors,
+  };
+}
+
+module.exports = {
+  SOURCE_EXTENSIONS,
+  MAX_FILE_SIZE,
+  walkSourceFiles,
+};

package/src/worker-pool.js

@@ -0,0 +1,178 @@
+const os = require("os");
+const path = require("path");
+const { Worker } = require("worker_threads");
+const { registerCleanup } = require("./shutdown");
+
+const MAX_WORKER_COUNT = 4;
+
+function normalizeWorkerCount(count) {
+  return Math.min(MAX_WORKER_COUNT, Math.max(1, count));
+}
+
+function getDefaultWorkerCount() {
+  return normalizeWorkerCount(os.cpus().length - 1);
+}
+
+class WorkerPool {
+  constructor(options = {}) {
+    this.workerPath = options.workerPath || path.join(__dirname, "worker.js");
+    this.size = normalizeWorkerCount(options.size || getDefaultWorkerCount());
+    this.queue = [];
+    this.workers = new Set();
+    this.closed = false;
+    this.closePromise = null;
+    this.unregisterCleanup = registerCleanup(() => this.close());
+
+    for (let index = 0; index < this.size; index += 1) {
+      this.spawnWorker();
+    }
+  }
+
+  run(task) {
+    if (this.closed) {
+      return Promise.reject(new Error("Worker pool is closed"));
+    }
+
+    return new Promise((resolve, reject) => {
+      this.queue.push({
+        task,
+        resolve,
+        reject,
+      });
+      this.dispatch();
+    });
+  }
+
+  async close() {
+    if (this.closePromise) {
+      return this.closePromise;
+    }
+
+    this.closed = true;
+
+    this.closePromise = (async () => {
+      if (this.unregisterCleanup) {
+        this.unregisterCleanup();
+        this.unregisterCleanup = null;
+      }
+
+      for (const queuedTask of this.queue.splice(0)) {
+        queuedTask.reject(new Error("Worker pool closed before task started"));
+      }
+
+      await Promise.allSettled(
+        Array.from(this.workers, (record) => record.worker.terminate()),
+      );
+
+      for (const record of this.workers) {
+        if (record.currentTask) {
+          record.currentTask.reject(new Error("Worker pool closed before task completed"));
+          record.currentTask = null;
+        }
+      }
+
+      this.workers.clear();
+    })();
+
+    return this.closePromise;
+  }
+
+  spawnWorker() {
+    if (this.closed) {
+      return null;
+    }
+
+    const worker = new Worker(this.workerPath);
+    const record = {
+      worker,
+      idle: true,
+      currentTask: null,
+      replaced: false,
+    };
+
+    this.workers.add(record);
+
+    worker.on("message", (message) => {
+      this.finishTask(record, null, message);
+    });
+
+    worker.on("error", (error) => {
+      this.failWorker(record, error);
+    });
+
+    worker.on("exit", (code) => {
+      this.workers.delete(record);
+
+      if (record.currentTask) {
+        record.currentTask.reject(new Error(`Worker exited before completing task with code ${code}`));
+        record.currentTask = null;
+      }
+
+      if (!this.closed && code !== 0 && !record.replaced) {
+        this.spawnWorker();
+        this.dispatch();
+      }
+    });
+
+    return record;
+  }
+
+  failWorker(record, error) {
+    this.workers.delete(record);
+    record.idle = false;
+    record.replaced = true;
+
+    if (record.currentTask) {
+      record.currentTask.reject(error);
+      record.currentTask = null;
+    }
+
+    if (!this.closed) {
+      this.spawnWorker();
+      this.dispatch();
+    }
+  }
+
+  finishTask(record, error, message) {
+    if (!record.currentTask) {
+      return;
+    }
+
+    const completedTask = record.currentTask;
+    record.currentTask = null;
+    record.idle = true;
+
+    if (error) {
+      completedTask.reject(error);
+    } else if (message && message.ok === false) {
+      completedTask.reject(new Error(message.error || "Worker task failed"));
+    } else {
+      completedTask.resolve(message.result);
+    }
+
+    this.dispatch();
+  }
+
+  dispatch() {
+    if (this.closed || this.queue.length === 0) {
+      return;
+    }
+
+    for (const record of this.workers) {
+      if (!record.idle || this.queue.length === 0) {
+        continue;
+      }
+
+      const queuedTask = this.queue.shift();
+      record.idle = false;
+      record.currentTask = queuedTask;
+      record.worker.postMessage(queuedTask.task);
+    }
+  }
+}
+
+module.exports = {
+  MAX_WORKER_COUNT,
+  WorkerPool,
+  getDefaultWorkerCount,
+};

package/src/worker.js

@@ -0,0 +1,124 @@
+const fs = require("fs/promises");
+const { parentPort } = require("worker_threads");
+const { pipeline } = require("@xenova/transformers");
+const { getContentSkipReason } = require("./file-filter");
+const { getParser, parseSourceFile } = require("./parser");
+
+const MODEL_ID = "Xenova/all-MiniLM-L6-v2";
+
+let embedderPromise;
+
+function getEmbedder() {
+  if (!embedderPromise) {
+    embedderPromise = pipeline("feature-extraction", MODEL_ID);
+  }
+
+  return embedderPromise;
+}
+
+function toFloat32Array(value) {
+  if (value instanceof Float32Array) {
+    return new Float32Array(value);
+  }
+
+  if (ArrayBuffer.isView(value)) {
+    return Float32Array.from(value);
+  }
+
+  if (Array.isArray(value)) {
+    return Float32Array.from(value);
+  }
+
+  throw new Error("Embedding output could not be converted to Float32Array");
+}
+
+async function embedChunk(embedder, chunk) {
+  const output = await embedder(chunk.semanticText, {
+    pooling: "mean",
+    normalize: true,
+  });
+
+  return {
+    ...chunk,
+    embedding: toFloat32Array(output.data),
+  };
+}
+
+async function processFile(task) {
+  if (!getParser(task.filePath)) {
+    return {
+      filePath: task.filePath,
+      relativePath: task.relativePath,
+      skipped: true,
+      reason: "Unsupported source extension",
+      chunks: [],
+    };
+  }
+
+  const code = await fs.readFile(task.filePath, "utf8");
+  const contentSkipReason = getContentSkipReason(code);
+
+  if (contentSkipReason) {
+    return {
+      filePath: task.filePath,
+      relativePath: task.relativePath,
+      skipped: true,
+      reason: contentSkipReason,
+      chunks: [],
+    };
+  }
+
+  const parseResult = parseSourceFile({
+    filePath: task.filePath,
+    relativePath: task.relativePath,
+    code,
+  });
+  const chunks = parseResult.chunks;
+
+  if (chunks.length === 0) {
+    return {
+      filePath: task.filePath,
+      relativePath: task.relativePath,
+      skipped: false,
+      chunks: [],
+    };
+  }
+
+  const embedder = await getEmbedder();
+  const embeddedChunks = [];
+
+  for (const chunk of chunks) {
+    embeddedChunks.push(await embedChunk(embedder, chunk));
+  }
+
+  return {
+    filePath: task.filePath,
+    relativePath: task.relativePath,
+    skipped: false,
+    chunks: embeddedChunks,
+  };
+}
+
+parentPort.on("message", async (task) => {
+  try {
+    const result = await processFile(task);
+    const transferList = [];
+
+    for (const chunk of result.chunks) {
+      transferList.push(chunk.embedding.buffer);
+    }
+
+    parentPort.postMessage(
+      {
+        ok: true,
+        result,
+      },
+      transferList,
+    );
+  } catch (error) {
+    parentPort.postMessage({
+      ok: false,
+      error: error.stack || error.message,
+    });
+  }
+});