@aiready/ast-mcp-server 0.1.2 → 0.1.3

package/dist/index.js

@@ -329,22 +329,101 @@ var SymbolIndex = class {
 };
 var symbolIndex = new SymbolIndex();
 
+// src/worker/pool.ts
+import { Worker } from "worker_threads";
+import path3 from "path";
+import { fileURLToPath } from "url";
+var WorkerPool = class {
+  constructor(poolSize) {
+    this.poolSize = poolSize;
+  }
+  workers = [];
+  available = [];
+  queue = [];
+  activeJobs = /* @__PURE__ */ new Map();
+  taskId = 0;
+  async init() {
+    const __dirname2 = path3.dirname(fileURLToPath(import.meta.url));
+    const workerPath = path3.join(__dirname2, "ast-worker.js");
+    for (let i = 0; i < this.poolSize; i++) {
+      const worker = new Worker(workerPath);
+      worker.on("message", (msg) => this.handleResult(msg));
+      worker.on("error", (err) => this.handleWorkerError(worker, err));
+      this.workers.push(worker);
+      this.available.push(worker);
+    }
+  }
+  async execute(type, payload) {
+    return new Promise((resolve, reject) => {
+      const id = String(++this.taskId);
+      const task = { id, type, payload, resolve, reject };
+      const worker = this.available.pop();
+      if (worker) {
+        this.dispatch(worker, task);
+      } else {
+        this.queue.push(task);
+      }
+    });
+  }
+  dispatch(worker, task) {
+    this.activeJobs.set(task.id, task);
+    worker.postMessage({ id: task.id, type: task.type, payload: task.payload });
+  }
+  handleResult(msg) {
+    const task = this.activeJobs.get(msg.id);
+    if (!task) return;
+    this.activeJobs.delete(msg.id);
+    if (msg.error) {
+      task.reject(new Error(msg.error));
+    } else {
+      task.resolve(msg.result);
+    }
+    const worker = this.workers.find(
+      (w) => !this.available.includes(w) && ![...this.activeJobs.values()].some((t) => t.id === msg.id)
+      // simplistic
+    );
+  }
+  handleWorkerError(worker, err) {
+    const idx = this.workers.indexOf(worker);
+    if (idx !== -1) {
+      worker.terminate();
+      const __dirname2 = path3.dirname(fileURLToPath(import.meta.url));
+      const workerPath = path3.join(__dirname2, "ast-worker.js");
+      const newWorker = new Worker(workerPath);
+      newWorker.on("message", (msg) => this.handleResult(msg));
+      newWorker.on("error", (e) => this.handleWorkerError(newWorker, e));
+      this.workers[idx] = newWorker;
+      this.available.push(newWorker);
+    }
+  }
+  async terminate() {
+    await Promise.all(this.workers.map((w) => w.terminate()));
+    this.workers = [];
+    this.available = [];
+  }
+};
+
 // src/adapters/typescript-adapter.ts
 var TypeScriptAdapter = class {
-  async resolveDefinition(symbolName, path3) {
-    validateWorkspacePath(path3);
+  pool;
+  constructor() {
+    const poolSize = parseInt(process.env.AST_WORKER_POOL_SIZE || "2");
+    this.pool = new WorkerPool(poolSize);
+  }
+  async resolveDefinition(symbolName, path4) {
+    validateWorkspacePath(path4);
     const indexHits = symbolIndex.lookup(symbolName);
     if (indexHits.length > 0) {
       const results = [];
       for (const hit of indexHits) {
         const tsconfig2 = await projectManager.findNearestTsConfig(hit.file);
         if (tsconfig2) {
-          const project2 = projectManager.ensureProject(tsconfig2);
-          const sourceFile2 = project2.addSourceFileAtPathIfExists(hit.file);
-          if (sourceFile2) {
-            const exported2 = sourceFile2.getExportedDeclarations().get(symbolName);
-            if (exported2 && exported2.length > 0) {
-              results.push(this.mapToDefinitionLocation(exported2[0]));
+          const project = projectManager.ensureProject(tsconfig2);
+          const sourceFile = project.addSourceFileAtPathIfExists(hit.file);
+          if (sourceFile) {
+            const exported = sourceFile.getExportedDeclarations().get(symbolName);
+            if (exported && exported.length > 0) {
+              results.push(this.mapToDefinitionLocation(exported[0]));
               continue;
             }
           }
@@ -360,20 +439,32 @@ var TypeScriptAdapter = class {
       }
       return results;
     }
-    if (fs3.statSync(path3).isDirectory()) {
+    if (fs3.statSync(path4).isDirectory()) {
       return [];
     }
-    const tsconfig = await projectManager.findNearestTsConfig(path3);
+    const tsconfig = await projectManager.findNearestTsConfig(path4);
     if (!tsconfig) return [];
-    const project = projectManager.ensureProject(tsconfig);
-    const sourceFile = project.addSourceFileAtPathIfExists(path3);
-    if (!sourceFile) return [];
-    const exported = sourceFile.getExportedDeclarations().get(symbolName);
-    if (!exported) return [];
-    return exported.map((decl) => this.mapToDefinitionLocation(decl));
+    try {
+      const result = await this.pool.execute(
+        "resolve_definition",
+        {
+          tsconfig,
+          file: path4,
+          symbol: symbolName
+        }
+      );
+      return result;
+    } catch {
+      const project = projectManager.ensureProject(tsconfig);
+      const sourceFile = project.addSourceFileAtPathIfExists(path4);
+      if (!sourceFile) return [];
+      const exported = sourceFile.getExportedDeclarations().get(symbolName);
+      if (!exported) return [];
+      return exported.map((decl) => this.mapToDefinitionLocation(decl));
+    }
   }
-  async findReferences(symbolName, path3, limit = 50, offset = 0) {
-    validateWorkspacePath(path3);
+  async findReferences(symbolName, path4, limit = 50, offset = 0) {
+    validateWorkspacePath(path4);
     const hits = symbolIndex.lookup(symbolName);
     if (hits.length === 0) return { references: [], total_count: 0 };
     const hit = hits[0];
@@ -390,13 +481,12 @@ var TypeScriptAdapter = class {
       const { searchCode: searchCode2 } = await import("./search-code-V3LACKQ6.js");
       const searchResults = await searchCode2(
         symbolName,
-        path3,
+        path4,
         "*.{ts,tsx,js,jsx}",
         1e3,
         false
       );
       const filesToLoad = [...new Set(searchResults.map((r) => r.file))];
-      console.log("Search code files to load:", filesToLoad);
       for (const file of filesToLoad) {
         project.addSourceFileAtPathIfExists(file);
       }
@@ -423,57 +513,69 @@ var TypeScriptAdapter = class {
       total_count: unique.length
     };
   }
-  async findImplementations(symbolName, path3, limit = 50, offset = 0) {
-    validateWorkspacePath(path3);
+  async findImplementations(symbolName, path4, limit = 50, offset = 0) {
+    validateWorkspacePath(path4);
     const hits = symbolIndex.lookup(symbolName);
     if (hits.length === 0) return { implementations: [], total_count: 0 };
     const hit = hits[0];
     const tsconfig = await projectManager.findNearestTsConfig(hit.file);
     if (!tsconfig) return { implementations: [], total_count: 0 };
-    const project = projectManager.ensureProject(tsconfig);
-    const sourceFile = project.addSourceFileAtPathIfExists(hit.file);
-    if (!sourceFile) return { implementations: [], total_count: 0 };
-    const exported = sourceFile.getExportedDeclarations().get(symbolName);
-    if (!exported || exported.length === 0)
-      return { implementations: [], total_count: 0 };
-    const targetNode = exported[0];
-    if (!Node2.isClassDeclaration(targetNode) && !Node2.isInterfaceDeclaration(targetNode)) {
-      return { implementations: [], total_count: 0 };
-    }
     try {
-      const { searchCode: searchCode2 } = await import("./search-code-V3LACKQ6.js");
-      const searchResults = await searchCode2(
-        symbolName,
-        path3,
-        "*.{ts,tsx,js,jsx}",
-        1e3,
-        false
-      );
-      const filesToLoad = [...new Set(searchResults.map((r) => r.file))];
-      for (const file of filesToLoad) {
-        project.addSourceFileAtPathIfExists(file);
+      const result = await this.pool.execute("find_implementations", {
+        tsconfig,
+        file: hit.file,
+        symbol: symbolName
+      });
+      return {
+        implementations: result.implementations.slice(offset, offset + limit),
+        total_count: result.total_count
+      };
+    } catch {
+      const project = projectManager.ensureProject(tsconfig);
+      const sourceFile = project.addSourceFileAtPathIfExists(hit.file);
+      if (!sourceFile) return { implementations: [], total_count: 0 };
+      const exported = sourceFile.getExportedDeclarations().get(symbolName);
+      if (!exported || exported.length === 0)
+        return { implementations: [], total_count: 0 };
+      const targetNode = exported[0];
+      if (!Node2.isClassDeclaration(targetNode) && !Node2.isInterfaceDeclaration(targetNode)) {
+        return { implementations: [], total_count: 0 };
       }
-    } catch (_e) {
-    }
-    const results = [];
-    const implementations = targetNode.getImplementations?.();
-    if (implementations) {
-      for (const impl of implementations) {
-        const sf = impl.getSourceFile();
-        const lc = sf.getLineAndColumnAtPos(impl.getTextSpan().getStart());
-        results.push({
-          file: sf.getFilePath(),
-          line: lc.line,
-          column: lc.column,
-          text: impl.getNode().getParent()?.getText() || impl.getNode().getText()
-        });
+      try {
+        const { searchCode: searchCode2 } = await import("./search-code-V3LACKQ6.js");
+        const searchResults = await searchCode2(
+          symbolName,
+          path4,
+          "*.{ts,tsx,js,jsx}",
+          1e3,
+          false
+        );
+        const filesToLoad = [...new Set(searchResults.map((r) => r.file))];
+        for (const file of filesToLoad) {
+          project.addSourceFileAtPathIfExists(file);
+        }
+      } catch (_e) {
+      }
+      const results = [];
+      const implementations = targetNode.getImplementations?.();
+      if (implementations) {
+        for (const impl of implementations) {
+          const sf = impl.getSourceFile();
+          const lc = sf.getLineAndColumnAtPos(impl.getTextSpan().getStart());
+          results.push({
+            file: sf.getFilePath(),
+            line: lc.line,
+            column: lc.column,
+            text: impl.getNode().getParent()?.getText() || impl.getNode().getText()
+          });
+        }
       }
+      const unique = this.deduplicateLocations(results);
+      return {
+        implementations: unique.slice(offset, offset + limit),
+        total_count: unique.length
+      };
     }
-    const unique = this.deduplicateLocations(results);
-    return {
-      implementations: unique.slice(offset, offset + limit),
-      total_count: unique.length
-    };
   }
   async getFileStructure(filePath) {
     const safePath = validateWorkspacePath(filePath);
@@ -500,6 +602,9 @@ var TypeScriptAdapter = class {
     };
     return structure;
   }
+  async shutdown() {
+    await this.pool.terminate();
+  }
   mapToDefinitionLocation(node) {
     const sourceFile = node.getSourceFile();
     const lineAndColumn = sourceFile.getLineAndColumnAtPos(node.getStart());
@@ -615,20 +720,20 @@ var TypeScriptAdapter = class {
 var typescriptAdapter = new TypeScriptAdapter();
 
 // src/tools/resolve-definition.ts
-async function resolveDefinition(symbol, path3) {
-  return await typescriptAdapter.resolveDefinition(symbol, path3);
+async function resolveDefinition(symbol, path4) {
+  return await typescriptAdapter.resolveDefinition(symbol, path4);
 }
 
 // src/tools/find-references.ts
-async function findReferences(symbol, path3, limit = 50, offset = 0) {
-  return await typescriptAdapter.findReferences(symbol, path3, limit, offset);
+async function findReferences(symbol, path4, limit = 50, offset = 0) {
+  return await typescriptAdapter.findReferences(symbol, path4, limit, offset);
 }
 
 // src/tools/find-implementations.ts
-async function findImplementations(symbol, path3, limit = 50, offset = 0) {
+async function findImplementations(symbol, path4, limit = 50, offset = 0) {
   return await typescriptAdapter.findImplementations(
     symbol,
-    path3,
+    path4,
     limit,
     offset
   );
@@ -668,8 +773,8 @@ async function getSymbolDocs(symbol, filePath) {
 }
 
 // src/tools/build-symbol-index.ts
-async function buildSymbolIndex(path3) {
-  return await symbolIndex.buildIndex(path3);
+async function buildSymbolIndex(path4) {
+  return await symbolIndex.buildIndex(path4);
 }
 
 // src/index.ts
@@ -800,8 +905,8 @@ var ASTExplorerServer = class {
       try {
         switch (name) {
           case "resolve_definition": {
-            const { symbol, path: path3 } = ResolveDefinitionSchema.parse(args);
-            const results = await resolveDefinition(symbol, path3);
+            const { symbol, path: path4 } = ResolveDefinitionSchema.parse(args);
+            const results = await resolveDefinition(symbol, path4);
             return {
               content: [
                 { type: "text", text: JSON.stringify(results, null, 2) }
@@ -809,8 +914,8 @@ var ASTExplorerServer = class {
             };
           }
           case "find_references": {
-            const { symbol, path: path3, limit, offset } = FindReferencesSchema.parse(args);
-            const results = await findReferences(symbol, path3, limit, offset);
+            const { symbol, path: path4, limit, offset } = FindReferencesSchema.parse(args);
+            const results = await findReferences(symbol, path4, limit, offset);
             return {
               content: [
                 { type: "text", text: JSON.stringify(results, null, 2) }
@@ -818,10 +923,10 @@ var ASTExplorerServer = class {
             };
           }
           case "find_implementations": {
-            const { symbol, path: path3, limit, offset } = FindImplementationsSchema.parse(args);
+            const { symbol, path: path4, limit, offset } = FindImplementationsSchema.parse(args);
             const results = await findImplementations(
               symbol,
-              path3,
+              path4,
               limit,
               offset
             );
@@ -841,10 +946,10 @@ var ASTExplorerServer = class {
             };
           }
           case "search_code": {
-            const { pattern, path: path3, filePattern, limit, regex } = SearchCodeSchema.parse(args);
+            const { pattern, path: path4, filePattern, limit, regex } = SearchCodeSchema.parse(args);
             const results = await searchCode(
               pattern,
-              path3,
+              path4,
               filePattern,
               limit,
               regex
@@ -856,15 +961,15 @@ var ASTExplorerServer = class {
             };
           }
           case "get_symbol_docs": {
-            const { symbol, path: path3 } = GetSymbolDocsSchema.parse(args);
-            const docs = await getSymbolDocs(symbol, path3);
+            const { symbol, path: path4 } = GetSymbolDocsSchema.parse(args);
+            const docs = await getSymbolDocs(symbol, path4);
             return {
               content: [{ type: "text", text: JSON.stringify(docs, null, 2) }]
             };
           }
           case "build_symbol_index": {
-            const { path: path3 } = BuildSymbolIndexSchema.parse(args);
-            const stats = await buildSymbolIndex(path3);
+            const { path: path4 } = BuildSymbolIndexSchema.parse(args);
+            const stats = await buildSymbolIndex(path4);
             return {
               content: [{ type: "text", text: JSON.stringify(stats, null, 2) }]
             };