surgent 0.7.0-alpha.0

package/src/optimizers/languages/grammar.ts

@@ -0,0 +1,200 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { existsSync } from "node:fs";
+import { readFile } from "node:fs/promises";
+import { resolve } from "node:path";
+import { LANGUAGE_REGISTRY } from "./index.js";
+import type { GrammarInstallSettings } from "./types.js";
+import { getPiPath, isMissingFileError, readJson, runCommand, writeJson } from "../../utils.js";
+import { createRequire } from "node:module";
+import { pathToFileURL } from "node:url";
+
+const GRAMMAR_COVERAGE_TARGET = 0.9;
+const GRAMMAR_INSTALL_KEY = "grammarInstall";
+
+async function ensureGrammarCachePackage() {
+  const grammarCacheDir = getPiPath("grammars");
+  const jsonPath = resolve(grammarCacheDir, "package.json");
+
+  try {
+    await readFile(jsonPath, "utf8");
+  } catch (error) {
+    if (!isMissingFileError(error)) {
+      throw error;
+    }
+    await writeJson(jsonPath, {
+      name: "tree-sitter-grammars",
+      private: true,
+      dependencies: {},
+    });
+  }
+}
+
+async function getRepositoryRoot(cwd: string) {
+  try {
+    const { stdout } = await runCommand(cwd, "git", ["rev-parse", "--show-toplevel"]);
+    const repoRoot = stdout.trim();
+    return repoRoot.length > 0 ? repoRoot : undefined;
+  } catch {}
+}
+
+async function selectSupportedBuckets(cwd: string) {
+  const bucketCounts = new Map<string, number>();
+  let totalFileCount = 0;
+
+  const { stdout } = await runCommand(cwd, "bash", [
+    "-lc",
+    "git ls-files | sed 's/.*\\.//' | sort | uniq -c | sort -rn",
+  ]);
+
+  for (const line of stdout.split(/\r?\n/)) {
+    const lineMatch = line.trim().match(/^(\d+)\s+(.+)$/);
+    if (!lineMatch) continue;
+
+    const rawCount = lineMatch[1];
+    const rawExtension = lineMatch[2];
+    if (!rawCount || !rawExtension) continue;
+
+    const count = Number.parseInt(rawCount, 10);
+    const extension = `.${rawExtension.toLowerCase()}`;
+    if (!Number.isFinite(count) || count <= 0) continue;
+
+    totalFileCount += count;
+    for (const languageEntry of LANGUAGE_REGISTRY) {
+      if (!languageEntry.profile.extensions.has(extension)) continue;
+      const nextCount = bucketCounts.get(languageEntry.bucketName) ?? 0;
+      bucketCounts.set(languageEntry.bucketName, nextCount + count);
+      break;
+    }
+  }
+
+  if (totalFileCount === 0) return [];
+  const rankedBuckets = LANGUAGE_REGISTRY.map((languageEntry) => ({
+    count: bucketCounts.get(languageEntry.bucketName) ?? 0,
+    packageName: languageEntry.packageName,
+    version: languageEntry.version,
+  }))
+    .filter((bucket) => bucket.count > 0)
+    .sort((leftBucket, rightBucket) => rightBucket.count - leftBucket.count);
+
+  let coveredCount = 0;
+  const pickedBuckets: Array<{ packageName: string; version: string }> = [];
+  for (const bucket of rankedBuckets) {
+    if (coveredCount / totalFileCount >= GRAMMAR_COVERAGE_TARGET) break;
+    coveredCount += bucket.count;
+    pickedBuckets.push({
+      packageName: bucket.packageName,
+      version: bucket.version,
+    });
+  }
+
+  return pickedBuckets;
+}
+
+function getMissingGrammarPackages(pickedBuckets: Array<{ packageName: string; version: string }>) {
+  const grammarCacheDir = getPiPath("grammars");
+
+  return pickedBuckets
+    .filter((bucket) => {
+      const packagePath = resolve(
+        grammarCacheDir,
+        "node_modules",
+        bucket.packageName,
+        "package.json",
+      );
+      return !existsSync(packagePath);
+    })
+    .map((bucket) => `${bucket.packageName}@${bucket.version}`);
+}
+
+function readGrammarInstallSettings(settings: Record<string, unknown>) {
+  const rawSettings = settings[GRAMMAR_INSTALL_KEY];
+  if (!rawSettings || typeof rawSettings !== "object") return {};
+
+  const settingsMap = rawSettings as Record<string, unknown>;
+  const allowed = Array.isArray(settingsMap.allowedRepos)
+    ? settingsMap.allowedRepos.filter((repo) => typeof repo === "string")
+    : undefined;
+  const denied = Array.isArray(settingsMap.deniedRepos)
+    ? settingsMap.deniedRepos.filter((repo) => typeof repo === "string")
+    : undefined;
+
+  return { allowedRepos: allowed, deniedRepos: denied } satisfies GrammarInstallSettings;
+}
+
+async function updateGrammarInstallSettings(
+  settings: Record<string, unknown>,
+  installSettings: GrammarInstallSettings,
+) {
+  settings[GRAMMAR_INSTALL_KEY] = installSettings;
+  await writeJson(getPiPath("settings"), settings);
+}
+
+export async function ensureGrammarCache(ctx: ExtensionContext) {
+  await ensureGrammarCachePackage();
+
+  const repoRoot = await getRepositoryRoot(ctx.cwd);
+  if (!repoRoot) return;
+
+  const settings = await readJson<Record<string, unknown>>(getPiPath("settings"), {});
+  const installSettings = readGrammarInstallSettings(settings);
+  const allowed = new Set(installSettings.allowedRepos ?? []);
+  const denied = new Set(installSettings.deniedRepos ?? []);
+  if (denied.has(repoRoot)) return;
+
+  const pickedBuckets = await selectSupportedBuckets(ctx.cwd);
+  if (pickedBuckets.length === 0) return;
+
+  const missingPkgs = getMissingGrammarPackages(pickedBuckets);
+  if (missingPkgs.length === 0) return;
+
+  let canInstall = allowed.has(repoRoot);
+  if (!canInstall) {
+    if (!ctx.hasUI) return;
+
+    canInstall = await ctx.ui.confirm(
+      "Install tree-sitter grammars for this repo?",
+      "Surgent installs missing parser grammars for optimized code read, input tokens and context optimizations.",
+    );
+
+    if (!canInstall) {
+      denied.add(repoRoot);
+      await updateGrammarInstallSettings(settings, {
+        allowedRepos: [...allowed],
+        deniedRepos: [...denied],
+      });
+      return;
+    }
+
+    allowed.add(repoRoot);
+    await updateGrammarInstallSettings(settings, {
+      allowedRepos: [...allowed],
+      deniedRepos: [...denied],
+    });
+  }
+
+  await runCommand(getPiPath("grammars"), "npm", ["install", "--save-exact", ...missingPkgs]);
+}
+
+export async function loadGrammarModule(pkgName: string) {
+  const grammarCacheDir = getPiPath("grammars");
+  const cacheRequire = createRequire(resolve(grammarCacheDir, "package.json"));
+
+  let entryPath = "";
+  try {
+    entryPath = cacheRequire.resolve(pkgName);
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    throw new Error(
+      `grammar package not installed in cache (${grammarCacheDir}): ${pkgName}. ${message}`,
+    );
+  }
+
+  try {
+    return await import(pathToFileURL(entryPath).href);
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    throw new Error(
+      `failed loading grammar module from cache (${grammarCacheDir}): ${pkgName}. ${message}`,
+    );
+  }
+}

package/src/optimizers/languages/index.ts

@@ -0,0 +1,75 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { ensureGrammarCache } from "./grammar.js";
+import type { LanguageProfile } from "./types.js";
+import { TypeScriptLanguageProfile } from "./typescript.js";
+import { PythonLanguageProfile } from "./python.js";
+import { GoLanguageProfile } from "./go.js";
+import { JavaLanguageProfile } from "./java.js";
+import { RustLanguageProfile } from "./rust.js";
+
+export const LANGUAGE_REGISTRY: Array<{
+  bucketName: string;
+  packageName: string;
+  version: string;
+  profile: LanguageProfile;
+}> = [
+  {
+    bucketName: "typescript",
+    packageName: "tree-sitter-typescript",
+    version: "0.23.2",
+    profile: new TypeScriptLanguageProfile(),
+  },
+  {
+    bucketName: "python",
+    packageName: "tree-sitter-python",
+    version: "0.25.0",
+    profile: new PythonLanguageProfile(),
+  },
+  {
+    bucketName: "go",
+    packageName: "tree-sitter-go",
+    version: "0.25.0",
+    profile: new GoLanguageProfile(),
+  },
+  {
+    bucketName: "java",
+    packageName: "tree-sitter-java",
+    version: "0.23.5",
+    profile: new JavaLanguageProfile(),
+  },
+  {
+    bucketName: "rust",
+    packageName: "tree-sitter-rust",
+    version: "0.24.0",
+    profile: new RustLanguageProfile(),
+  },
+];
+
+export function getLanguageProfile(extension: string) {
+  return LANGUAGE_REGISTRY.find(({ profile }) => profile.extensions.has(extension))?.profile;
+}
+
+export function getSupportedExtensions() {
+  const supported = new Set<string>();
+  LANGUAGE_REGISTRY.forEach(({ profile }) => {
+    profile.extensions.forEach((extension) => supported.add(extension));
+  });
+  return supported;
+}
+
+export default function (pi: ExtensionAPI) {
+  pi.on("session_start", async (_event, ctx) => {
+    try {
+      await ensureGrammarCache(ctx);
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      if (ctx.hasUI) {
+        ctx.ui.notify(`tree-sitter grammar install failed: ${message}`, "error");
+      }
+    }
+  });
+}
+
+export { collectSymbols } from "./symbols.js";
+export type { LanguageProfile, LanguageSymbol } from "./types.js";
+export const SYMBOL_KINDS = ["func", "class", "method", "decl", "deps", "public"] as const;

package/src/optimizers/languages/java.ts

@@ -0,0 +1,64 @@
+import type { Language, SyntaxNode } from "tree-sitter";
+import { loadGrammarModule } from "./grammar.js";
+import { RuleBasedLanguageProfile, type SymbolKindRule } from "./types.js";
+
+export class JavaLanguageProfile extends RuleBasedLanguageProfile {
+  constructor() {
+    super(
+      new Set([".java"]),
+      { __default__: "name", exports_module_directive: "package" },
+      new Set(["program"]),
+      new Set<string>(),
+      new Map<string, SymbolKindRule[]>([
+        ["annotation_type_declaration", [{ kind: "class" }]],
+        ["class_declaration", [{ kind: "class" }]],
+        ["enum_declaration", [{ kind: "class" }]],
+        ["interface_declaration", [{ kind: "class" }]],
+        ["record_declaration", [{ kind: "class" }]],
+        [
+          "constructor_declaration",
+          [
+            { kind: "method", parent: "class_body" },
+            { kind: "method", parent: "enum_body_declarations" },
+          ],
+        ],
+        [
+          "method_declaration",
+          [
+            { kind: "method", parent: "class_body" },
+            { kind: "method", parent: "enum_body_declarations" },
+            { kind: "method", parent: "interface_body" },
+          ],
+        ],
+        ["import_declaration", [{ kind: "deps" }]],
+        ["exports_module_directive", [{ kind: "public", parent: "module_body" }]],
+      ]),
+    );
+  }
+
+  async loadLanguage(_extension: string) {
+    const languagePack = await loadGrammarModule("tree-sitter-java");
+    const languageExport = (languagePack.default ?? languagePack) as Language;
+    return languageExport;
+  }
+
+  readNodeName(node: SyntaxNode) {
+    const nodeName = this.readNameField(node);
+    if (nodeName) {
+      return nodeName;
+    }
+
+    if (node.type === "import_declaration") {
+      const importText = node.text
+        .replace(/^import\s+/, "")
+        .replace(/\s*;\s*$/, "")
+        .trim()
+        .replaceAll("\n", " ");
+      return importText.length > 0 ? importText : undefined;
+    }
+
+    if (node.type === "identifier" || node.type === "scoped_identifier") {
+      return this.readNodeText(node);
+    }
+  }
+}

package/src/optimizers/languages/python.ts

@@ -0,0 +1,63 @@
+import type { Language, SyntaxNode } from "tree-sitter";
+import { loadGrammarModule } from "./grammar.js";
+import { RuleBasedLanguageProfile, type SymbolKindRule } from "./types.js";
+
+export class PythonLanguageProfile extends RuleBasedLanguageProfile {
+  constructor() {
+    super(
+      new Set([".py"]),
+      {
+        __default__: "name",
+        assignment: "left",
+        future_import_statement: "name",
+        import_from_statement: "name",
+        import_statement: "name",
+      },
+      new Set(["module"]),
+      new Set(["decorated_definition", "expression_statement"]),
+      new Map<string, SymbolKindRule[]>([
+        ["class_definition", [{ kind: "class" }]],
+        [
+          "function_definition",
+          [
+            { kind: "method", container: "class_definition" },
+            { kind: "func", topLevelOnly: true },
+          ],
+        ],
+        ["assignment", [{ kind: "decl", topLevelOnly: true }]],
+        ["future_import_statement", [{ kind: "deps", topLevelOnly: true }]],
+        ["import_from_statement", [{ kind: "deps", topLevelOnly: true }]],
+        ["import_statement", [{ kind: "deps", topLevelOnly: true }]],
+      ]),
+    );
+  }
+
+  async loadLanguage(_extension: string) {
+    const languagePack = await loadGrammarModule("tree-sitter-python");
+    const languageExport = (languagePack.default ?? languagePack) as Language;
+    return languageExport;
+  }
+
+  resolveSymbolKind(node: SyntaxNode) {
+    if (
+      node.type === "assignment" &&
+      this.matchesTopLevelRule(node) &&
+      this.readFieldText(node, "left") === "__all__"
+    ) {
+      return "public";
+    }
+
+    return super.resolveSymbolKind(node);
+  }
+
+  readNodeName(node: SyntaxNode) {
+    const nodeName = this.readNameField(node);
+    if (nodeName) {
+      return nodeName;
+    }
+
+    if (node.type === "dotted_name" || node.type === "identifier") {
+      return this.readNodeText(node);
+    }
+  }
+}

package/src/optimizers/languages/rust.ts

@@ -0,0 +1,71 @@
+import type { Language, SyntaxNode } from "tree-sitter";
+import { loadGrammarModule } from "./grammar.js";
+import { RuleBasedLanguageProfile, type SymbolKindRule } from "./types.js";
+
+export class RustLanguageProfile extends RuleBasedLanguageProfile {
+  constructor() {
+    super(
+      new Set([".rs"]),
+      { __default__: "name", use_declaration: "argument" },
+      new Set(["source_file"]),
+      new Set(["declaration_list", "foreign_mod_item", "mod_item"]),
+      new Map<string, SymbolKindRule[]>([
+        [
+          "function_item",
+          [
+            { kind: "method", container: "impl_item" },
+            { kind: "method", container: "trait_item" },
+            { kind: "func", topLevelOnly: true },
+          ],
+        ],
+        [
+          "function_signature_item",
+          [
+            { kind: "method", container: "trait_item" },
+            { kind: "func", topLevelOnly: true },
+          ],
+        ],
+        ["struct_item", [{ kind: "class", topLevelOnly: true }]],
+        ["enum_item", [{ kind: "class", topLevelOnly: true }]],
+        ["union_item", [{ kind: "class", topLevelOnly: true }]],
+        ["trait_item", [{ kind: "class", topLevelOnly: true }]],
+        ["type_item", [{ kind: "class", topLevelOnly: true }]],
+        ["mod_item", [{ kind: "class", topLevelOnly: true }]],
+        ["impl_item", [{ kind: "class" }]],
+        ["macro_definition", [{ kind: "func", topLevelOnly: true }]],
+        ["const_item", [{ kind: "decl", topLevelOnly: true }]],
+        ["static_item", [{ kind: "decl", topLevelOnly: true }]],
+        ["use_declaration", [{ kind: "deps", topLevelOnly: true }]],
+      ]),
+    );
+  }
+
+  async loadLanguage(_extension: string) {
+    const languagePack = await loadGrammarModule("tree-sitter-rust");
+    const languageExport = (languagePack.default ?? languagePack) as Language;
+    return languageExport;
+  }
+
+  readNodeName(node: SyntaxNode) {
+    if (node.type === "impl_item") {
+      return this.readFieldText(node, "type");
+    }
+
+    const nodeName = this.readNameField(node);
+    if (nodeName) {
+      return nodeName;
+    }
+
+    if (node.type === "identifier" || node.type === "type_identifier") {
+      return this.readNodeText(node);
+    }
+  }
+
+  isPublicSymbol(node: SyntaxNode) {
+    return node.namedChildren.some((childNode) => childNode.type === "visibility_modifier");
+  }
+
+  shouldSkipSymbol(node: SyntaxNode) {
+    return node.type === "impl_item";
+  }
+}

package/src/optimizers/languages/symbols.ts

@@ -0,0 +1,95 @@
+import { readFile } from "node:fs/promises";
+import type { SyntaxNode } from "tree-sitter";
+import { getLanguageProfile } from "./index.js";
+import type { LanguageSymbol, SymbolKind } from "./types.js";
+import { extname, resolve } from "node:path";
+import Parser from "tree-sitter";
+
+async function createCodeParser(extension: string) {
+  const languageProfile = getLanguageProfile(extension);
+  if (!languageProfile) {
+    throw new Error(`missing grammar for extension: ${extension}`);
+  }
+  const parser = new Parser();
+  parser.setLanguage(await languageProfile.loadLanguage(extension));
+  return parser;
+}
+
+async function getParserForPath(path: string, parsers: Map<string, Parser>) {
+  const extension = extname(path).toLowerCase();
+  const parser = parsers.get(extension);
+  if (parser) return parser;
+
+  const createdParser = await createCodeParser(extension);
+  parsers.set(extension, createdParser);
+  return createdParser;
+}
+
+async function getRootNode(cwd: string, path: string) {
+  const parsers = new Map<string, Parser>();
+  const absolutePath = resolve(cwd, path);
+
+  let code = "";
+  try {
+    code = await readFile(absolutePath, "utf8");
+  } catch {
+    return;
+  }
+
+  const parser = await getParserForPath(path, parsers);
+  const tree = parser.parse(code);
+  return tree.rootNode;
+}
+
+export async function collectSymbols(cwd: string, path: string, kinds: Set<SymbolKind>) {
+  const extension = extname(path).toLowerCase();
+  const profile = getLanguageProfile(extension);
+  if (!profile) return [];
+
+  const root = await getRootNode(cwd, path);
+  if (!root) return [];
+
+  const symbols: LanguageSymbol[] = [];
+  const pendingNodes: SyntaxNode[] = [root];
+  const symbolIdCounts = new Map<string, number>();
+
+  while (pendingNodes.length > 0) {
+    const currentNode = pendingNodes.pop();
+    if (!currentNode) continue;
+
+    for (let childIndex = currentNode.namedChildCount - 1; childIndex >= 0; childIndex -= 1) {
+      const namedChild = currentNode.namedChild(childIndex);
+      if (namedChild) {
+        pendingNodes.push(namedChild);
+      }
+    }
+
+    const symbolKind = profile.resolveSymbolKind(currentNode);
+    if (!symbolKind || !kinds.has(symbolKind) || profile.shouldSkipSymbol(currentNode)) {
+      continue;
+    }
+
+    const baseName = profile.readNodeName(currentNode);
+    if (!baseName) continue;
+
+    const containerName = profile.readContainerName(currentNode);
+    const symbolName =
+      symbolKind === "method" && containerName ? `${containerName}.${baseName}` : baseName;
+
+    const symbolIdCount = (symbolIdCounts.get(symbolName) ?? 0) + 1;
+    symbolIdCounts.set(symbolName, symbolIdCount);
+    const isAnonymous = symbolName === "anonymous" && !currentNode.childForFieldName("name");
+
+    symbols.push({
+      name:
+        isAnonymous || symbolIdCount > 1 ? `${symbolName}~${symbolIdCount}` : symbolName,
+      path,
+      kind: symbolKind,
+      node: currentNode,
+      range: [currentNode.startPosition.row + 1, currentNode.endPosition.row + 1],
+      public: profile.isPublicSymbol(currentNode),
+    });
+  }
+
+  return symbols;
+}

package/src/optimizers/languages/tree-sitter-languages.d.ts

@@ -0,0 +1,23 @@
+declare module "tree-sitter-go" {
+  import type { Language } from "tree-sitter";
+  const language: Language;
+  export default language;
+}
+
+declare module "tree-sitter-java" {
+  import type { Language } from "tree-sitter";
+  const language: Language;
+  export default language;
+}
+
+declare module "tree-sitter-python" {
+  import type { Language } from "tree-sitter";
+  const language: Language;
+  export default language;
+}
+
+declare module "tree-sitter-rust" {
+  import type { Language } from "tree-sitter";
+  const language: Language;
+  export default language;
+}