dsh-code-index 0.1.0 → 0.2.0

package/dist/index.js

@@ -9,6 +9,14 @@ import path from "path";
 import Parser from "web-tree-sitter";
 var require2 = createRequire(import.meta.url);
 var WASM_DIR = path.dirname(require2.resolve("tree-sitter-wasms/out/tree-sitter-typescript.wasm"));
+var GRAMMAR_NAMES = {
+  typescript: "tree-sitter-typescript",
+  javascript: "tree-sitter-javascript",
+  python: "tree-sitter-python",
+  go: "tree-sitter-go",
+  rust: "tree-sitter-rust",
+  java: "tree-sitter-java"
+};
 var EXT_TO_LANG = {
   ".ts": "typescript",
   ".tsx": "typescript",
@@ -17,7 +25,12 @@ var EXT_TO_LANG = {
   ".js": "javascript",
   ".jsx": "javascript",
   ".mjs": "javascript",
-  ".cjs": "javascript"
+  ".cjs": "javascript",
+  ".py": "python",
+  ".pyi": "python",
+  ".go": "go",
+  ".rs": "rust",
+  ".java": "java"
 };
 function languageForFile(filePath) {
   const ext = path.extname(filePath).toLowerCase();
@@ -43,6 +56,29 @@ var QUERIES = {
     (method_definition) @method
     (class_declaration) @class
     (variable_declarator) @variable
+  `,
+  python: `
+    (function_definition) @function
+    (class_definition) @class
+  `,
+  go: `
+    (function_declaration) @function
+    (method_declaration) @method
+    (type_spec) @type
+  `,
+  rust: `
+    (function_item) @function
+    (function_signature_item) @method
+    (struct_item) @class
+    (enum_item) @enum
+    (trait_item) @interface
+  `,
+  java: `
+    (class_declaration) @class
+    (interface_declaration) @interface
+    (enum_declaration) @enum
+    (method_declaration) @method
+    (constructor_declaration) @method
   `
 };
 var CAPTURE_KINDS = {
@@ -55,6 +91,29 @@ var CAPTURE_KINDS = {
   variable: { kind: "variable" },
   field: { kind: "field" }
 };
+var IMPORT_QUERIES = {
+  typescript: `
+    (import_statement) @import
+    (export_statement) @reexport
+  `,
+  javascript: `
+    (import_statement) @import
+    (export_statement) @reexport
+  `,
+  python: `
+    (import_from_statement) @import
+    (import_statement) @import
+  `,
+  go: `
+    (import_spec) @import
+  `,
+  rust: `
+    (use_declaration) @use
+  `,
+  java: `
+    (import_declaration) @import
+  `
+};
 var parserPromise = null;
 async function getParser() {
   if (!parserPromise) {
@@ -71,82 +130,184 @@ var languageCache = /* @__PURE__ */ new Map();
 function getLanguage(id) {
   let entry = languageCache.get(id);
   if (!entry) {
-    entry = getParser().then(async (parser) => {
-      const grammarName = id === "typescript" ? "tree-sitter-typescript" : "tree-sitter-javascript";
-      const grammarPath = path.join(WASM_DIR, `${grammarName}.wasm`);
+    entry = getParser().then(async () => {
+      const grammarPath = path.join(WASM_DIR, `${GRAMMAR_NAMES[id]}.wasm`);
       const bytes = await readFile(grammarPath);
       const lang = await Parser.Language.load(bytes);
       const q = lang.query(QUERIES[id]);
       q.delete();
+      const iq = lang.query(IMPORT_QUERIES[id]);
+      iq.delete();
       return lang;
     });
     languageCache.set(id, entry);
   }
   return entry;
 }
-async function extractSymbols(code, id) {
+async function extractAll(code, id) {
   const lang = await getLanguage(id);
   const parser = await getParser();
   parser.setLanguage(lang);
   const tree = parser.parse(code);
   try {
-    const query = lang.query(QUERIES[id]);
-    const rows = [];
+    const symbols = [];
+    const symbolQuery = lang.query(QUERIES[id]);
     try {
-      const captures = query.captures(tree.rootNode);
+      const captures = symbolQuery.captures(tree.rootNode);
       for (const cap of captures) {
         const def = CAPTURE_KINDS[cap.name];
         if (!def) continue;
         const node = cap.node;
+        if (def.kind === "variable" && !isModuleLevelVariable(node)) continue;
         const name2 = nameOf(node);
         if (!name2) continue;
-        rows.push({
+        symbols.push({
           name: name2,
-          kind: def.kind,
+          kind: kindFor(id, node, def.kind),
           file: "",
           // set by the caller (extractor is file-agnostic)
           line: node.startPosition.row + 1,
           endLine: node.endPosition.row + 1,
-          exported: isExported(node),
+          exported: isExported(id, node),
           signature: signatureFor(node)
         });
       }
     } finally {
-      query.delete();
+      symbolQuery.delete();
     }
-    rows.sort((a, b) => a.line - b.line);
-    return rows;
+    const imports = [];
+    const importQuery = lang.query(IMPORT_QUERIES[id]);
+    try {
+      for (const cap of importQuery.captures(tree.rootNode)) {
+        const spec = specifierOf(id, cap.node);
+        if (spec) imports.push(spec);
+      }
+    } finally {
+      importQuery.delete();
+    }
+    symbols.sort((a, b) => a.line - b.line);
+    return { symbols, imports };
   } finally {
     tree.delete();
   }
 }
+async function extractSymbols(code, id) {
+  return (await extractAll(code, id)).symbols;
+}
 function nameOf(node) {
   const field = node.childForFieldName?.("name");
   if (field) return field.text.trim();
   return "";
 }
+function specifierOf(id, node) {
+  switch (node.type) {
+    case "import_statement":
+      if (id === "python") return dottedToPath(node.childForFieldName("name")?.text);
+      return stripQuotes(node.childForFieldName("source")?.text);
+    case "export_statement":
+      return stripQuotes(node.childForFieldName("source")?.text);
+    case "import_from_statement": {
+      const raw = node.childForFieldName("module_name")?.text;
+      if (raw == null) return null;
+      return raw.startsWith(".") ? pythonRelative(raw) : dottedToPath(raw);
+    }
+    case "import_spec":
+      return stripQuotes(node.childForFieldName("path")?.text);
+    case "use_declaration":
+      return rustUsePath(node.childForFieldName("argument")?.text);
+    case "import_declaration":
+      return dottedToPath(node.namedChildren[0]?.text);
+    default:
+      return null;
+  }
+}
+function rustUsePath(text) {
+  if (!text) return null;
+  const base = text.split("{")[0].trim();
+  if (!base) return null;
+  const parts = base.split("::").filter(Boolean);
+  if (parts[0] === "crate") parts.shift();
+  return parts.map((p) => p === "super" ? ".." : p === "self" ? "." : p).join("/");
+}
+function stripQuotes(text) {
+  if (text == null || text.length < 2) return null;
+  const first = text[0];
+  const last = text[text.length - 1];
+  if (first === "'" && last === "'" || first === '"' && last === '"') {
+    return text.slice(1, -1);
+  }
+  return null;
+}
+function dottedToPath(text) {
+  if (!text) return null;
+  return text.trim().replace(/\./g, "/");
+}
+function pythonRelative(raw) {
+  const dots = raw.match(/^\.+/)?.[0].length ?? 0;
+  const rest = raw.slice(dots).replace(/\./g, "/");
+  const prefix = dots === 1 ? "./" : "../".repeat(dots - 1);
+  return prefix + rest;
+}
 function signatureFor(node) {
   const name2 = nameOf(node);
-  const params = node.namedChildren.find(
-    (c) => c.type === "formal_parameters" || c.type === "method_parameters"
+  const params = node.childForFieldName?.("parameters") ?? node.namedChildren.find(
+    (c) => c.type === "formal_parameters" || c.type === "method_parameters" || c.type === "parameters" || c.type === "parameter_list"
   );
   if (params) {
-    return `${name2}${params.text}`;
+    return `${name2}${collapseSpace(params.text)}`;
   }
   const first = node.namedChildren[0];
-  return first ? first.text.trim() : name2;
+  return first ? collapseSpace(first.text) : name2;
 }
-function isExported(node) {
-  let parent = node.parent;
-  let depth = 0;
-  while (parent && depth < 3) {
+function collapseSpace(text) {
+  return text.replace(/\s+/g, " ").replace(/\( /g, "(").replace(/ \)/g, ")").replace(/ ,/g, ",").replace(/,\)/g, ")").trim();
+}
+function isModuleLevelVariable(node) {
+  const declaration = node.parent;
+  const container = declaration?.parent;
+  return container?.type === "program" || container?.type === "export_statement";
+}
+function isExported(id, node) {
+  if (id === "python") {
+    return node.parent?.type === "module";
+  }
+  if (id === "go") {
+    const name2 = nameOf(node);
+    return !!name2 && /^[A-Z]/.test(name2);
+  }
+  if (id === "rust") {
+    return node.namedChildren.some(
+      (c) => c.type === "visibility_modifier" && c.text === "pub"
+    );
+  }
+  if (id === "java") {
+    if (node.parent?.type === "interface_body") return true;
+    return node.namedChildren.some((c) => c.type === "modifiers" && /\bpublic\b/.test(c.text));
+  }
+  for (let parent = node.parent; parent; parent = parent.parent) {
     if (parent.type === "export_statement") return true;
-    if (parent.type === "statement_block") return false;
-    parent = parent.parent;
-    depth++;
+    if (parent.type === "statement_block" || parent.type === "class_body") return false;
+    if (parent.type === "program") return false;
   }
   return false;
 }
+function kindFor(id, node, kind) {
+  if (id === "python" && kind === "function") {
+    const inClassBody = node.parent?.type === "block" && node.parent?.parent?.type === "class_definition";
+    if (inClassBody) return "method";
+  }
+  if (id === "go" && node.type === "type_spec") {
+    const type = node.childForFieldName("type")?.type;
+    if (type === "struct_type") return "class";
+    if (type === "interface_type") return "interface";
+    return "type";
+  }
+  if (id === "rust" && node.type === "function_item") {
+    const inImpl = node.parent?.type === "declaration_list" && node.parent?.parent?.type === "impl_item";
+    if (inImpl) return "method";
+  }
+  return kind;
+}
 async function parseFileToSymbols(filePath, repoRoot, code) {
   const lang = languageForFile(filePath);
   if (!lang) return [];
@@ -183,7 +344,12 @@ var SUPPORTED_EXTS = /* @__PURE__ */ new Set([
   ".js",
   ".jsx",
   ".mjs",
-  ".cjs"
+  ".cjs",
+  ".py",
+  ".pyi",
+  ".go",
+  ".rs",
+  ".java"
 ]);
 async function scanRepo(root, options = {}) {
   const excluded = /* @__PURE__ */ new Set([...DEFAULT_EXCLUDED_DIRS, ...options.excludeDirs ?? []]);
@@ -240,11 +406,19 @@ async function loadIndex(cachePath) {
     const raw = await readFile2(cachePath, "utf8");
     const parsed = JSON.parse(raw);
     if (!parsed || typeof parsed.root !== "string" || !Array.isArray(parsed.files)) return null;
-    return parsed;
+    return healSymbolFiles(parsed);
   } catch {
     return null;
   }
 }
+function healSymbolFiles(index) {
+  for (const file of index.files) {
+    for (const symbol of file.symbols) {
+      if (!symbol.file) symbol.file = file.path;
+    }
+  }
+  return index;
+}
 async function saveIndex(cachePath, index) {
   await mkdir(path3.dirname(cachePath), { recursive: true });
   const temporaryPath = `${cachePath}.${process.pid}.${randomUUID()}.tmp`;
@@ -283,12 +457,13 @@ async function buildIndex(root, options = {}, previous = null) {
         } catch {
           return null;
         }
-        const symbols = await extractSymbols(code, lang);
+        const { symbols, imports } = await extractAll(code, lang);
         return {
           path: f.rel,
           lang,
           mtimeMs: f.mtimeMs,
-          symbols: symbols.map((s) => ({ ...s, file: f.rel }))
+          symbols: symbols.map((s) => ({ ...s, file: f.rel })),
+          imports
         };
       })
     );
@@ -367,8 +542,17 @@ function matchScore(name2, query) {
   if (n === q) return 1;
   if (n.startsWith(q)) return 0.8;
   if (n.includes(q)) return 0.5;
+  if (q.length >= 3 && isSubsequence(q, n)) return 0.3;
   return 0;
 }
+function isSubsequence(query, name2) {
+  let i = 0;
+  for (const ch of name2) {
+    if (ch === query[i]) i++;
+    if (i === query.length) return true;
+  }
+  return false;
+}
 function searchSymbols(index, filter, limit = 50) {
   const q = (filter.query ?? "").trim();
   const filePat = filter.file?.trim().toLowerCase();
@@ -397,6 +581,7 @@ function renderHit(hit) {
 }
 
 // src/repomap.ts
+import path5 from "path";
 var KIND_WEIGHT = {
   class: 1,
   interface: 1,
@@ -409,18 +594,25 @@ var KIND_WEIGHT = {
   import: 0.05,
   module: 0.3
 };
+var TEST_PATH_RE = /(^|\/)(tests?|__tests__)(\/|$)|\.(?:spec|test)\.[cm]?[jt]sx?$|(^|\/)(?:test_[^/]+\.py|[^/]+_test\.(?:py|go))$/;
 function scoreFile(file) {
   let score = 0;
   for (const sym of file.symbols) {
     score += KIND_WEIGHT[sym.kind] ?? 0.3;
     if (sym.exported) score += 0.3;
   }
-  return score / (1 + file.symbols.length * 0.04);
+  score /= 1 + file.symbols.length * 0.04;
+  if (TEST_PATH_RE.test(file.path)) score *= 0.2;
+  return score;
 }
 function rankRepoMap(index, options = {}) {
   const topFiles = options.topFiles ?? 24;
   const perFile = options.symbolsPerFile ?? 18;
-  const ranked = index.files.filter((f) => f.symbols.length > 0).map((f) => ({ file: f, score: scoreFile(f) })).sort((a, b) => b.score - a.score || a.file.path.localeCompare(b.file.path)).slice(0, topFiles);
+  const refs = countReferences(index.files);
+  const ranked = index.files.filter((f) => f.symbols.length > 0).map((f) => ({
+    file: f,
+    score: scoreFile(f) + REF_WEIGHT * (refs.get(f.path) ?? 0)
+  })).sort((a, b) => b.score - a.score || a.file.path.localeCompare(b.file.path)).slice(0, topFiles);
   return ranked.map(({ file, score }) => ({
     path: file.path,
     score,
@@ -432,6 +624,45 @@ function rankRepoMap(index, options = {}) {
     }))
   }));
 }
+var REF_WEIGHT = 0.5;
+function countReferences(files) {
+  const fileSet = new Set(files.map((f) => f.path));
+  const counts = /* @__PURE__ */ new Map();
+  for (const file of files) {
+    const targets = /* @__PURE__ */ new Set();
+    for (const spec of file.imports ?? []) {
+      const target = resolveImport(spec, file.path, fileSet);
+      if (target && target !== file.path) targets.add(target);
+    }
+    for (const target of targets) {
+      counts.set(target, (counts.get(target) ?? 0) + 1);
+    }
+  }
+  return counts;
+}
+function resolveImport(spec, fromPath, fileSet) {
+  if (!spec) return null;
+  if (spec.startsWith("./") || spec.startsWith("../")) {
+    const base = path5.posix.normalize(path5.posix.join(path5.posix.dirname(fromPath), spec));
+    return candidateFor(base, fileSet);
+  }
+  const segments = spec.split("/").filter(Boolean);
+  for (let skip = 0; skip < segments.length; skip++) {
+    const hit = candidateFor(segments.slice(skip).join("/"), fileSet);
+    if (hit) return hit;
+  }
+  return null;
+}
+function candidateFor(base, fileSet) {
+  for (const ext of SUPPORTED_EXTS) {
+    if (fileSet.has(`${base}${ext}`)) return `${base}${ext}`;
+  }
+  for (const ext of SUPPORTED_EXTS) {
+    if (fileSet.has(`${base}/index${ext}`)) return `${base}/index${ext}`;
+    if (fileSet.has(`${base}/__init__${ext}`)) return `${base}/__init__${ext}`;
+  }
+  return null;
+}
 function renderRepoMap(entries, options = {}) {
   const maxChars = options.maxChars ?? 3200;
   const lines = ["# repo map"];
@@ -454,13 +685,14 @@ function renderRepoMap(entries, options = {}) {
 
 // src/tools.ts
 import { defineTool } from "@deepseek-ai/dsh-tools";
-import path5 from "path";
+import path6 from "path";
 
 // src/config.ts
 var DEFAULTS = {
   excludeDirs: [],
   mapTopFiles: 24,
   mapMaxChars: 3200,
+  mapTtlMs: 6e4,
   autoInject: true
 };
 var state = { current: { ...DEFAULTS } };
@@ -476,6 +708,9 @@ function applyConfig(partial) {
   if (!Number.isFinite(state.current.mapMaxChars) || state.current.mapMaxChars < 200) {
     state.current.mapMaxChars = DEFAULTS.mapMaxChars;
   }
+  if (!Number.isFinite(state.current.mapTtlMs) || state.current.mapTtlMs < 1e3) {
+    state.current.mapTtlMs = DEFAULTS.mapTtlMs;
+  }
 }
 function getConfig() {
   return state.current;
@@ -507,7 +742,7 @@ function createIndexCache(load, ttlMs = 6e4, now = Date.now) {
   }
   return {
     get(root, force = false) {
-      const resolvedRoot = path5.resolve(root);
+      const resolvedRoot = path6.resolve(root);
       const key = cacheKeyForRoot(resolvedRoot);
       const running = inFlight.get(key);
       if (running) {
@@ -544,7 +779,7 @@ async function resolveRoot(arg, exec) {
   const cwd = exec.agent?.session?.header?.cwd;
   const base = arg ?? cwd ?? process.cwd();
   const root = await findRepoRoot(base);
-  if (!root) throw new Error(`no git repository found from ${path5.resolve(base)}`);
+  if (!root) throw new Error(`no git repository found from ${path6.resolve(base)}`);
   return root;
 }
 var tools = [
@@ -698,7 +933,6 @@ var tools = [
 // src/index.ts
 var name = "dsh-code-index";
 var inject = ["tools", "systemPrompt"];
-var MAP_TTL_MS = 6e4;
 function apply(ctx, pluginConfig) {
   applyConfig(pluginConfig);
   invalidateIndexCache();
@@ -744,7 +978,7 @@ function apply(ctx, pluginConfig) {
         // before tool guidance (100–199), after persona (0)
         text: () => {
           const now = Date.now();
-          if (cached && now - cached.at < MAP_TTL_MS) return cached.text;
+          if (cached && now - cached.at < getConfig().mapTtlMs) return cached.text;
           void warmMap();
           return cached?.text ?? "";
         }
@@ -772,6 +1006,7 @@ export {
   buildIndex,
   buildIndexWithCache,
   defaultCachePath,
+  extractAll,
   extractSymbols,
   findRepoRoot,
   inject,