surgent 0.7.0-alpha.0

package/src/optimizers/languages/types.ts

@@ -0,0 +1,134 @@
+import type { Language, SyntaxNode } from "tree-sitter";
+
+export type SymbolKind = "func" | "class" | "method" | "decl" | "deps" | "public";
+
+export interface SymbolKindRule {
+  kind: SymbolKind;
+  parent?: string;
+  container?: string;
+  topLevelOnly?: boolean;
+}
+
+export interface LanguageSymbol {
+  name: string;
+  path: string;
+  kind: SymbolKind;
+  node: SyntaxNode;
+  range?: [number, number];
+  public?: boolean;
+}
+
+export interface LanguageProfile {
+  extensions: Set<string>;
+  loadLanguage(extension: string): Promise<Language>;
+  readNodeName(node: SyntaxNode): string | undefined;
+  readContainerName(node: SyntaxNode): string | undefined;
+  isPublicSymbol(node: SyntaxNode): boolean;
+  resolveSymbolKind(node: SyntaxNode): SymbolKind | undefined;
+  shouldSkipSymbol(node: SyntaxNode): boolean;
+}
+
+export interface SupportedLanguageBucket {
+  bucketName: string;
+  packageName: string;
+  version: string;
+  extensions: Set<string>;
+}
+
+export interface GrammarInstallSettings {
+  allowedRepos?: string[];
+  deniedRepos?: string[];
+}
+
+export abstract class RuleBasedLanguageProfile implements LanguageProfile {
+  constructor(
+    readonly extensions: Set<string>,
+    protected readonly nameFieldByType: { __default__: string } & Record<string, string>,
+    protected readonly topLevelRoots: Set<string>,
+    protected readonly topLevelParents: Set<string>,
+    protected readonly typeRule: Map<string, SymbolKindRule[]>,
+  ) {}
+
+  abstract loadLanguage(extension: string): Promise<Language>;
+
+  abstract readNodeName(node: SyntaxNode): string | undefined;
+
+  protected findContainerNode(node: SyntaxNode) {
+    let current = node.parent;
+    while (current) {
+      if (this.typeRule.has(current.type)) {
+        return current;
+      }
+      current = current.parent;
+    }
+  }
+
+  protected matchesTopLevelRule(node: SyntaxNode) {
+    let current = node.parent;
+
+    while (current) {
+      if (this.topLevelRoots.has(current.type)) {
+        return true;
+      }
+      if (this.topLevelParents.size > 0 && !this.topLevelParents.has(current.type)) {
+        return false;
+      }
+      current = current.parent;
+    }
+
+    return false;
+  }
+
+  protected readFieldText(node: SyntaxNode, fieldName: string) {
+    return this.readNodeText(node.childForFieldName(fieldName));
+  }
+
+  protected readNameField(node: SyntaxNode) {
+    const fieldName = this.nameFieldByType[node.type] ?? this.nameFieldByType.__default__;
+    return this.readFieldText(node, fieldName);
+  }
+
+  protected readNodeText(node: SyntaxNode | null | undefined) {
+    if (!node) return;
+
+    const nodeText = node.text.trim().replaceAll("\n", " ");
+    return nodeText.length > 0 ? nodeText : undefined;
+  }
+
+  readContainerName(node: SyntaxNode) {
+    const containerNode = this.findContainerNode(node);
+    return containerNode ? this.readNodeName(containerNode) : undefined;
+  }
+
+  isPublicSymbol(node: SyntaxNode) {
+    return this.findContainerNode(node)?.type === "export_statement";
+  }
+
+  resolveSymbolKind(node: SyntaxNode) {
+    const symbolKindRules = this.typeRule.get(node.type);
+    if (!symbolKindRules) return;
+
+    for (const symbolKindRule of symbolKindRules) {
+      if (symbolKindRule.parent && node.parent?.type !== symbolKindRule.parent) {
+        continue;
+      }
+
+      if (symbolKindRule.container) {
+        const containerNode = this.findContainerNode(node);
+        if (!containerNode || containerNode.type !== symbolKindRule.container) {
+          continue;
+        }
+      }
+
+      if (symbolKindRule.topLevelOnly && !this.matchesTopLevelRule(node)) {
+        continue;
+      }
+
+      return symbolKindRule.kind;
+    }
+  }
+
+  shouldSkipSymbol(_node: SyntaxNode) {
+    return false;
+  }
+}

package/src/optimizers/languages/typescript.ts

@@ -0,0 +1,116 @@
+import type { Language, SyntaxNode } from "tree-sitter";
+import { loadGrammarModule } from "./grammar.js";
+import { RuleBasedLanguageProfile, type SymbolKindRule } from "./types.js";
+
+export class TypeScriptLanguageProfile extends RuleBasedLanguageProfile {
+  constructor() {
+    super(
+      new Set([".ts", ".tsx", ".js", ".mjs", ".cjs", ".jsx"]),
+      { __default__: "name" },
+      new Set(["program"]),
+      new Set([
+        "lexical_declaration",
+        "variable_declaration",
+        "variable_declarator",
+        "export_statement",
+      ]),
+      new Map<string, SymbolKindRule[]>([
+        ["function_declaration", [{ kind: "func", topLevelOnly: true }]],
+        ["arrow_function", [{ kind: "func", topLevelOnly: true }]],
+        ["function_expression", [{ kind: "func", topLevelOnly: true }]],
+        ["abstract_class_declaration", [{ kind: "class" }]],
+        ["class_declaration", [{ kind: "class" }]],
+        ["abstract_method_signature", [{ kind: "method", parent: "class_body" }]],
+        [
+          "method_definition",
+          [
+            { kind: "method", parent: "class_body" },
+            { kind: "method", parent: "object" },
+          ],
+        ],
+        ["variable_declarator", [{ kind: "decl", topLevelOnly: true }]],
+        ["import_specifier", [{ kind: "deps" }]],
+        ["namespace_import", [{ kind: "deps" }]],
+        ["identifier", [{ kind: "deps", parent: "import_clause" }]],
+        ["export_statement", [{ kind: "public", topLevelOnly: true }]],
+      ]),
+    );
+  }
+
+  async loadLanguage(extension: string) {
+    const languagePack = await loadGrammarModule("tree-sitter-typescript");
+    const languageExports =
+      typeof languagePack.default === "object" && languagePack.default !== null
+        ? languagePack.default
+        : languagePack;
+
+    const typeScriptLanguage = languageExports.typescript as Language | undefined;
+    if (!typeScriptLanguage) {
+      throw new Error("tree-sitter-typescript missing typescript export");
+    }
+
+    if (extension === ".tsx" || extension === ".jsx") {
+      return (languageExports.tsx as Language | undefined) ?? typeScriptLanguage;
+    }
+    return typeScriptLanguage;
+  }
+
+  readNodeName(node: SyntaxNode) {
+    if (node.type === "export_statement") {
+      return "exports";
+    }
+
+    if (node.type === "import_specifier") {
+      const aliasText = this.readFieldText(node, "alias");
+      if (aliasText) {
+        return aliasText;
+      }
+    }
+
+    const nodeName = this.readNameField(node);
+    if (nodeName) {
+      return nodeName;
+    }
+
+    if (node.type === "namespace_import") {
+      return this.readNodeText(node.namedChild(0));
+    }
+
+    if (node.type === "identifier") {
+      return this.readNodeText(node);
+    }
+
+    if (
+      node.type === "arrow_function" ||
+      node.type === "function_expression" ||
+      node.type === "function_declaration"
+    ) {
+      const declaratorNode = node.parent?.type === "variable_declarator" ? node.parent : undefined;
+      if (declaratorNode) {
+        const declaratorName = this.readFieldText(declaratorNode, "name");
+        if (declaratorName) {
+          return declaratorName;
+        }
+      }
+
+      return "anonymous";
+    }
+  }
+
+  shouldSkipSymbol(node: SyntaxNode) {
+    if (node.type === "export_statement") {
+      const inlineNode = node.namedChild(0);
+      return (
+        node.childForFieldName("declaration") !== null ||
+        inlineNode?.type === "arrow_function" ||
+        inlineNode?.type === "function_expression" ||
+        inlineNode?.type === "function_declaration" ||
+        (inlineNode?.type === "class_declaration" && inlineNode.childForFieldName("name") !== null)
+      );
+    }
+    if (node.type === "variable_declarator") {
+      return node.childForFieldName("value")?.type === "arrow_function";
+    }
+    return node.type === "function_expression" && node.childForFieldName("name") !== null;
+  }
+}

package/src/optimizers/mapper/files.ts

@@ -0,0 +1,94 @@
+import picomatch from "picomatch";
+import { runCommand, unique } from "../../utils.js";
+
+const SKIPPED_DIRECTORIES = new Set([".git", ".pi", "build", "coverage", "dist", "node_modules"]);
+
+function normalizePath(pathValue: string) {
+  const normalizedPath = pathValue.replaceAll("\\", "/");
+  return normalizedPath.startsWith("./") ? normalizedPath.slice(2) : normalizedPath;
+}
+
+async function rgFiles(
+  projectPath: string,
+  globTargets: string[],
+  signal?: AbortSignal,
+) {
+  const args = ["--files", "--hidden"];
+
+  for (const skipped of SKIPPED_DIRECTORIES) {
+    args.push("--glob", `!**/${skipped}/**`);
+  }
+  for (const globTarget of globTargets) {
+    args.push("--glob", globTarget);
+  }
+
+  args.push(".");
+  const commandResult = await runCommand(projectPath, "rg", args, {
+    signal,
+    successExitCodes: [0, 1],
+    abortMessage: "mapper aborted",
+  });
+  return commandResult.stdout
+    .split("\n")
+    .map((line) => normalizePath(line.trim()))
+    .filter((line) => line.length > 0);
+}
+
+async function grepFiles(
+  projectPath: string,
+  globTargets: string[],
+  signal?: AbortSignal,
+) {
+  const args = ["-r", "-I", "-l"];
+
+  for (const skipped of SKIPPED_DIRECTORIES) {
+    args.push("--exclude-dir", skipped);
+  }
+
+  args.push("-e", "", ".");
+  const commandResult = await runCommand(projectPath, "grep", args, {
+    signal,
+    successExitCodes: [0, 1],
+    abortMessage: "mapper aborted",
+  });
+  const paths = commandResult.stdout
+    .split("\n")
+    .map((line) => normalizePath(line.trim()))
+    .filter((line) => line.length > 0);
+  const matchers = globTargets.map((globTarget) => picomatch(globTarget, { dot: true }));
+  return paths.filter((pathValue) => matchers.some((matcher) => matcher(pathValue)));
+}
+
+export async function resolveTargetPaths(
+  projectPath: string,
+  targets: string[],
+  signal?: AbortSignal,
+) {
+  const globTargets = targets.flatMap((target) => {
+    const normalized = normalizePath(target);
+    if (normalized.length === 0) return [];
+    if (normalized === ".") return ["**"];
+    if (/[*?[\]{}]/.test(normalized)) return [normalized];
+    return [normalized, `${normalized}/**`];
+  });
+
+  if (globTargets.length === 0) return [];
+
+  try {
+    const paths = await rgFiles(projectPath, globTargets, signal);
+    return unique(paths).sort();
+  } catch (rgError) {
+    const rgMessage = rgError instanceof Error ? rgError.message : String(rgError);
+    if (rgMessage === "mapper aborted") {
+      throw rgError;
+    }
+
+    try {
+      const paths = await grepFiles(projectPath, globTargets, signal);
+      return unique(paths).sort();
+    } catch (grepError) {
+      const grepMessage = grepError instanceof Error ? grepError.message : String(grepError);
+      throw new Error(`mapper file scan failed: rg=${rgMessage}; grep=${grepMessage}`);
+    }
+  }
+}

package/src/optimizers/mapper/index.ts

@@ -0,0 +1,133 @@
+import { defineTool, truncateHead } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import { resolveTargetPaths } from "./files.js";
+import { collectSymbols, SYMBOL_KINDS } from "../languages/index.js";
+import type { LanguageSymbol } from "../languages/index.js";
+import type { MapperResult } from "./types.js";
+import { renderCallText } from "../../utils.js";
+
+function collapseGroupedSymbols(symbols: LanguageSymbol[]) {
+  let groupedSymbolKind: "deps" | "public" | undefined;
+  let importsGroupIndex = 0;
+  let exportsGroupIndex = 0;
+
+  return symbols.flatMap((symbol) => {
+    if (symbol.kind !== "deps" && symbol.kind !== "public") {
+      groupedSymbolKind = undefined;
+      return [symbol];
+    }
+    if (symbol.kind === groupedSymbolKind) return [];
+
+    groupedSymbolKind = symbol.kind;
+    if (symbol.kind === "deps") {
+      importsGroupIndex += 1;
+      return [{ ...symbol, name: `imports~${importsGroupIndex}`, range: undefined }];
+    }
+
+    exportsGroupIndex += 1;
+    return [{ ...symbol, name: `exports~${exportsGroupIndex}`, range: undefined }];
+  });
+}
+
+const codeMap = defineTool({
+  name: "code_map",
+  label: "Code map",
+  description:
+    "Fast symbol and code blocks offset/limit indexing. Best for: narrowing targets to inspect/read or code discovery.",
+  parameters: Type.Object({
+    targets: Type.Array(Type.String(), {
+      description: "Paths or globs to scan (relative to cwd). Keep scope narrow.",
+    }),
+    kinds: Type.Optional(
+      Type.Array(Type.Union(SYMBOL_KINDS.map((kind) => Type.Literal(kind))), {
+        description: "Abstraction kinds to include. Omit to include all supported.",
+      }),
+    ),
+  }),
+  async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+    const kinds = new Set(params.kinds ?? SYMBOL_KINDS);
+    const result: MapperResult = { symbols: [], failed: [] };
+    let paths: string[] = [];
+    try {
+      paths = await resolveTargetPaths(ctx.cwd, params.targets, signal);
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      return {
+        isError: true,
+        details: `code_map target scan failed: ${message}`,
+        content: [{ type: "text", text: `code_map target scan failed: ${message}` }],
+      };
+    }
+
+    if (paths.length === 0) {
+      return {
+        isError: false,
+        details: "(no symbols found) targets matched no supported files or symbols",
+        content: [
+          {
+            type: "text",
+            text: "(no symbols found) targets matched no supported files or symbols",
+          },
+        ],
+      };
+    }
+
+    for (const path of paths) {
+      if (signal?.aborted) {
+        return {
+          isError: true,
+          details: "code_map aborted",
+          content: [{ type: "text", text: "code_map aborted" }],
+        };
+      }
+
+      try {
+        const symbols = collapseGroupedSymbols(await collectSymbols(ctx.cwd, path, kinds));
+        result.symbols.push(...symbols);
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        result.failed.push(`${path}: ${message}`);
+      }
+    }
+
+    const outputLines: string[] = [];
+    let outputPath = "";
+
+    for (const symbol of result.symbols) {
+      if (symbol.path !== outputPath) {
+        outputPath = symbol.path;
+        outputLines.push(symbol.path);
+      }
+
+      let line = `  [${symbol.public ? "public " : ""}${symbol.kind}] ${symbol.name}`;
+      if (symbol.range) {
+        line += ` L${symbol.range[0]}-L${symbol.range[1]}`;
+      }
+      outputLines.push(line);
+    }
+
+    if (result.failed.length > 0) {
+      outputLines.push(...result.failed.map((failure) => `failed ${failure}`));
+    }
+    if (outputLines.length === 0) {
+      outputLines.push("(no symbols found) targets matched no supported files or symbols");
+    }
+
+    const truncation = truncateHead(outputLines.join("\n"));
+    const output = truncation.truncated
+      ? `${truncation.content}\n\n[Output truncated at 2000 lines or 50KB. Narrow targets or kinds.]`
+      : truncation.content;
+    return { isError: false, details: output, content: [{ type: "text", text: output }] };
+  },
+  renderCall(args, theme, { isPartial }) {
+    const targets = Array.isArray(args.targets) ? args.targets.join(", ") : "";
+    const kinds = Array.isArray(args.kinds) ? args.kinds.join(", ") : "default";
+
+    return renderCallText(
+      `${theme.fg("toolTitle", "code_map")} ${theme.underline(theme.fg("accent", targets))} ${theme.fg("dim", `[${kinds}]`)}`,
+      isPartial,
+    );
+  },
+});
+
+export default codeMap;

package/src/optimizers/mapper/types.ts

@@ -0,0 +1,6 @@
+import type { LanguageSymbol } from "../languages/types.js";
+
+export interface MapperResult {
+  symbols: LanguageSymbol[];
+  failed: string[];
+}

package/src/optimizers/pruner/cleanup.ts

@@ -0,0 +1,121 @@
+import { getEntryId, getMessage, getParentId, getToolResultMessage } from "../entries.js";
+import type { RemovedEntryRemoval } from "./types.js";
+import { isRecord } from "../../utils.js";
+
+function resolveRemovedId(
+  entryId: string | null,
+  replacementParents: Map<string, string | null>,
+): string | null {
+  const visitedIds = new Set<string>();
+  let resolvedId = entryId;
+  while (resolvedId && replacementParents.has(resolvedId) && !visitedIds.has(resolvedId)) {
+    visitedIds.add(resolvedId);
+    resolvedId = replacementParents.get(resolvedId) ?? null;
+  }
+  return resolvedId;
+}
+
+function hasEmptyResult(message: {
+  content?: unknown;
+  toolName?: unknown;
+}): boolean {
+  if (message.toolName !== "ls" && message.toolName !== "find") {
+    return false;
+  }
+
+  if (!Array.isArray(message.content) || message.content.length !== 1) return false;
+  const content = message.content[0];
+  if (!isRecord(content) || content.type !== "text" || typeof content.text !== "string") return false;
+
+  return (
+    content.text === "(empty directory)" ||
+    content.text === "No files found matching pattern" ||
+    content.text === "No changes found." ||
+    content.text === "No changes in selected files."
+  );
+}
+
+export function getRemovedToolCallId(message: {
+  content?: unknown;
+  isError?: unknown;
+  toolCallId?: unknown;
+  toolName?: unknown;
+}): string | undefined {
+  if (typeof message.toolCallId !== "string") return;
+  if (message.isError === true) return message.toolName === "bash" ? undefined : message.toolCallId;
+  return hasEmptyResult(message) ? message.toolCallId : undefined;
+}
+
+export function removeEntries(entries: Record<string, unknown>[]): RemovedEntryRemoval {
+  const removedToolCallIds = new Set<string>();
+  const replacementParents = new Map<string, string | null>();
+
+  for (const entry of entries) {
+    const message = getToolResultMessage(entry);
+    const entryId = getEntryId(entry);
+    const toolCallId = message ? getRemovedToolCallId(message) : undefined;
+    if (!entryId || !toolCallId) continue;
+
+    removedToolCallIds.add(toolCallId);
+    replacementParents.set(entryId, getParentId(entry));
+  }
+
+  if (removedToolCallIds.size === 0) {
+    return { changed: false, entries, replacementParents };
+  }
+
+  const retainedEntries: Record<string, unknown>[] = [];
+  for (const entry of entries) {
+    const entryId = getEntryId(entry);
+    if (entryId && replacementParents.has(entryId)) continue;
+
+    const message = getMessage(entry);
+    if (message?.role !== "assistant" || !Array.isArray(message.content)) {
+      retainedEntries.push(entry);
+      continue;
+    }
+
+    const retainedContent = message.content.filter(
+      (block) =>
+        !isRecord(block) ||
+        block.type !== "toolCall" ||
+        typeof block.id !== "string" ||
+        !removedToolCallIds.has(block.id),
+    );
+    if (retainedContent.length === message.content.length) {
+      retainedEntries.push(entry);
+      continue;
+    }
+    if (entryId && retainedContent.every((block) => isRecord(block) && block.type === "thinking")) {
+      replacementParents.set(entryId, getParentId(entry));
+      continue;
+    }
+    retainedEntries.push({ ...entry, message: { ...message, content: retainedContent } });
+  }
+
+  const repairedEntries: Record<string, unknown>[] = [];
+  for (const entry of retainedEntries) {
+    const updatedEntry = { ...entry };
+    let changed = false;
+
+    const parentId = getParentId(entry);
+    const repairedParentId = resolveRemovedId(parentId, replacementParents);
+    if (parentId !== repairedParentId) {
+      updatedEntry.parentId = repairedParentId;
+      changed = true;
+    }
+
+    for (const referenceField of ["fromId", "targetId", "firstKeptEntryId"]) {
+      const referenceId = entry[referenceField];
+      if (typeof referenceId !== "string") continue;
+      const repairedReferenceId = resolveRemovedId(referenceId, replacementParents);
+      if (repairedReferenceId && repairedReferenceId !== referenceId) {
+        updatedEntry[referenceField] = repairedReferenceId;
+        changed = true;
+      }
+    }
+    repairedEntries.push(changed ? updatedEntry : entry);
+  }
+
+  return { changed: true, entries: repairedEntries, replacementParents };
+}

package/src/optimizers/pruner/context.ts

@@ -0,0 +1,46 @@
+import { getToolResultMessage } from "../entries.js";
+import { getRemovedToolCallId } from "./cleanup.js";
+import type { ContextPruneResult } from "./types.js";
+
+export function buildPrunerState(entries: Record<string, unknown>[]): Set<string> {
+  const removedToolCallIds = new Set<string>();
+  for (const entry of entries) {
+    const message = getToolResultMessage(entry);
+    const toolCallId = message ? getRemovedToolCallId(message) : undefined;
+    if (toolCallId) removedToolCallIds.add(toolCallId);
+  }
+  return removedToolCallIds;
+}
+
+export function filterContextMessages(
+  messages: ContextPruneResult["messages"],
+  state: Set<string>,
+): ContextPruneResult {
+  if (state.size === 0) return { changed: false, messages };
+
+  const retainedMessages: ContextPruneResult["messages"] = [];
+  let changed = false;
+  for (const message of messages) {
+    if (message.role === "toolResult" && state.has(message.toolCallId)) {
+      changed = true;
+      continue;
+    }
+    if (message.role !== "assistant") {
+      retainedMessages.push(message);
+      continue;
+    }
+
+    const retainedContent = message.content.filter(
+      (block) => block.type !== "toolCall" || !state.has(block.id),
+    );
+    if (retainedContent.length === message.content.length) {
+      retainedMessages.push(message);
+      continue;
+    }
+    changed = true;
+    if (retainedContent.some((block) => block.type !== "thinking")) {
+      retainedMessages.push({ ...message, content: retainedContent });
+    }
+  }
+  return { changed, messages: retainedMessages };
+}

package/src/optimizers/pruner/index.ts

@@ -0,0 +1,45 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { getRemovedToolCallId } from "./cleanup.js";
+import { buildPrunerState, filterContextMessages } from "./context.js";
+import { readSessionEntries } from "../entries.js";
+import { rewritePrunedSessionFile } from "./session.js";
+
+function loadPrunerState(sessionFile: string | undefined): Set<string> {
+  const entries = readSessionEntries(sessionFile);
+  return entries ? buildPrunerState(entries) : new Set<string>();
+}
+
+export default function (pi: ExtensionAPI) {
+  let state = new Set<string>();
+
+  pi.on("session_start", (_event, ctx) => {
+    state = loadPrunerState(ctx.sessionManager.getSessionFile());
+  });
+
+  pi.on("session_tree", (_event, ctx) => {
+    state = loadPrunerState(ctx.sessionManager.getSessionFile());
+  });
+
+  pi.on("agent_end", (event) => {
+    for (const message of event.messages) {
+      if (message.role !== "toolResult") continue;
+      const toolCallId = getRemovedToolCallId(message);
+      if (toolCallId) state.add(toolCallId);
+    }
+  });
+
+  pi.on("session_shutdown", (event, ctx) => {
+    const sessionFile = ctx.sessionManager.getSessionFile();
+    if (sessionFile) {
+      rewritePrunedSessionFile(sessionFile, ctx.sessionManager.getLeafId(), false);
+    }
+    if (event.targetSessionFile && event.targetSessionFile !== sessionFile) {
+      rewritePrunedSessionFile(event.targetSessionFile, null, true);
+    }
+  });
+
+  pi.on("context", (event) => {
+    const pruned = filterContextMessages(event.messages, state);
+    if (pruned.changed) return { messages: pruned.messages };
+  });
+}