pi-ast-sgrep 2.0.2 → 2.1.1

package/README.md

@@ -57,7 +57,7 @@ Pi can make one `asgrep` call like this:
 
 ```json
 {
-  "code": "async () => {\n  const seed = await asgrep.search({ query: 'where are access tokens refreshed?', limit: 5 });\n  const symbol = seed.hits?.[0]?.symbol;\n  if (!symbol) return { seed };\n  const [defs, callers] = await Promise.all([\n    asgrep.defs({ symbol, limit: 5 }),\n    asgrep.callers({ symbol, limit: 10 }),\n  ]);\n  return { symbol, defs: defs.hits, callers: callers.hits };\n}"
+  "code": "async () => {\n  const seed = await asgrep.search('where are access tokens refreshed?', { limit: 5 });\n  const hit = seed.hits?.[0];\n  if (!hit) return { seed, next: seed.suggested_next };\n  const [defs, window] = await Promise.all([\n    asgrep.defs(hit.symbol, { limit: 5 }),\n    asgrep.read({ refs: [hit.ref] }),\n  ]);\n  return { symbol: hit.symbol, defs: defs.hits, window };\n}"
 }
 ```
 
@@ -65,32 +65,28 @@ This workflow narrows the first result, runs independent follow-up searches toge
 
 ### Code Mode API
 
-The Code Mode program receives these asynchronous methods on `asgrep`:
+The Code Mode program receives these asynchronous methods on `asgrep`. Positional arguments work; object forms remain valid.
 
 | Method | Use |
 |---|---|
-| `asgrep.search({ query, limit?, excerptLines? })` | Search by intent, symbol, or a prefixed structural query. |
-| `asgrep.semantic({ query, limit?, excerptLines? })` | Search local semantic embeddings directly. |
-| `asgrep.defs({ symbol, limit? })` | Find definitions for one symbol. |
-| `asgrep.callers({ symbol, limit? })` | Find call sites for one symbol. |
-| `asgrep.imports({ module, limit? })` | Find imports of one module. |
-| `asgrep.chain({ query, limit? })` | Trace related symbols and graph edges. |
-| `asgrep.indexStatus()` | Read index and backend state. |
-| `asgrep.indexRepo({ force? })` | Create, refresh, or rebuild the index. |
-| `asgrep.catalogSearch({ query })` | Discover less common ast-sgrep operations. |
-| `asgrep.catalogDescribe({ name })` | Read the schema for a discovered operation. |
+| `asgrep.search("query", { limit?, in?, lang? })` | Hybrid search: intent, symbol, or prefixed `defs:` / `callers:` / `pattern:` query. Bound a tree with `in`. |
+| `asgrep.find("token")` | Lexical / identifier lookup (`word:`). Prefixed queries pass through. |
+| `asgrep.defs("Symbol")` | Jump to definitions. |
+| `asgrep.callers("Symbol")` | Reverse-walk callers. |
+| `asgrep.read({ path, start, end }` or `{ refs }`) | Batched line windows from the index. Prefer one call with `refs`. |
+| `asgrep.edit(path, oldText, newText)` | Unique string replace jailed to the project root, then targeted reindex. |
+| `asgrep.indexStatus()` / `asgrep.doctor()` | Index health without a fifth outer Pi tool. |
 
-Use `Promise.all` for independent calls. Filter, map, sort, and slice intermediate values in JavaScript. Return only the evidence needed for the next reasoning step.
+Use `Promise.all` for independent calls. Filter, map, sort, and slice intermediate values in JavaScript. Return only the evidence needed for the next reasoning step. A search with 0 hits includes `suggested_next`.
 
-Code Mode runs in a disposable worker with a restricted `node:vm` context that exposes only a serialized `asgrep.*` bridge and console. String and WebAssembly code generation are disabled, ambient Node globals such as `process` and `require` are not exposed, and terminating the worker contains synchronous and microtask CPU loops. Node does not consider `vm` an adversarial-code security boundary, however, and the installed Pi package has full OS-user access; do not treat Code Mode as an OS jail. Prefer Code Mode **or** MCP for a client, never both.
+Code Mode runs **in-process** in a restricted `node:vm` context (no Worker sandbox, no OS jail). `asgrep` and `console` are built inside the context; the host only exposes a JSON bridge and a log sink so host `Function` cannot leak. Return shapes are declared on `asgrep.*` (muscle memory). `find({ query: "blast:Symbol" })` reverse-walks callers; `blast:path/to/file.ts` uses imports. Same trust boundary as Pi `bash`. Prefer Code Mode **or** MCP for a client, never both.
 
 The bridge rejects oversized call arguments and serialized results, allows at
 most 256 host calls per program, and caps collected console output before it
-reaches the extension host. Raw-memory and WebAssembly globals are unavailable;
-worker heap/stack limits contain the remaining accidental memory growth. Native
-tool values are capped at 1 MiB each and complete batch responses at 4 MiB before
-Node-API converts them into extension-host objects. These bounds do not turn `node:vm` into an OS
-sandbox.
+reaches the extension host. Raw-memory and WebAssembly globals are unavailable.
+Native tool values are capped at 1 MiB each and complete batch responses at 4 MiB
+before Node-API converts them into extension-host objects. These bounds do not
+turn `node:vm` into an OS sandbox.
 
 ## Direct one-shot search
 

package/dist/code-mode.d.ts

@@ -1,4 +1,4 @@
-import { type AstSgrepRuntime, type MachineEnvelope, type RunOptions, type RuntimeContext } from "./runtime.js";
+import { type AstSgrepRuntime, type MachineEnvelope, type RunOptions, type RuntimeContext } from "./runtime/runtime.js";
 export type SgrepKind = "asgrep" | "def" | "caller" | "graph" | "anchor" | "import" | "pattern" | "embed";
 export type SgrepSignal = "exact" | "structural" | "semantic";
 export type SgrepRef = `${string}#L${number}-L${number}`;

package/dist/code-mode.js

@@ -1,7 +1,7 @@
 import { constants } from "node:fs";
 import { lstat, open, realpath } from "node:fs/promises";
 import { isAbsolute, relative, resolve, sep } from "node:path";
-import { RuntimeError } from "./runtime.js";
+import { RuntimeError } from "./runtime/runtime.js";
 const DEFAULT_LIMIT = 20;
 const MAX_LIMIT = 100;
 const MAX_EXCERPT_LINES = 100;

package/dist/codemode/connector.d.ts

@@ -1,6 +1,6 @@
-import type { MachineEnvelope } from "../runtime.js";
-import type { ChainArgs, SearchArgs } from "./types.js";
-import { type BatchCapableHost, type DispatchStats } from "./dispatch.js";
+import type { MachineEnvelope } from "../runtime/runtime.js";
+import type { ChainArgs, EditArgs, FindArgs, ReadArgs, SearchArgs } from "./types.js";
+import { type BatchCapableHost, type DispatchCall, type DispatchStats } from "./dispatch.js";
 /**
  * Spawn/CLI transport. Hosts provide argv `run` only — never a typed twin.
  * Typed entry lives solely on {@link DispatchSurface} (dispatcher output).
@@ -27,6 +27,15 @@ export type AsgrepConnector = {
     search(input: SearchArgs, options?: {
         signal?: AbortSignal;
     }): Promise<MachineEnvelope>;
+    find(input: FindArgs, options?: {
+        signal?: AbortSignal;
+    }): Promise<MachineEnvelope>;
+    read(input: ReadArgs, options?: {
+        signal?: AbortSignal;
+    }): Promise<MachineEnvelope>;
+    edit(input: EditArgs, options?: {
+        signal?: AbortSignal;
+    }): Promise<MachineEnvelope>;
     semantic(input: SearchArgs, options?: {
         signal?: AbortSignal;
     }): Promise<MachineEnvelope>;
@@ -73,10 +82,15 @@ export type AsgrepConnector = {
     }, options?: {
         signal?: AbortSignal;
     }): Promise<MachineEnvelope>;
+    doctor(options?: {
+        signal?: AbortSignal;
+    }): Promise<MachineEnvelope>;
 };
 export type ConnectorBundle = {
     asgrep: AsgrepConnector;
     stats: () => DispatchStats;
+    /** Per-call dispatch trace for the current run (lane + ms + ok, capped). */
+    trace: () => DispatchCall[];
     resetStats: () => void;
 };
 /**

package/dist/codemode/connector.js

@@ -1,3 +1,5 @@
+import { coerceHostArgs } from "./guest-api.js";
+import { defined } from "./types.js";
 import { createCodemodeDispatcher, } from "./dispatch.js";
 const DEFAULT_LIMIT = 8;
 function clampLimit(limit) {
@@ -32,48 +34,74 @@ export function createAsgrepConnector(host, context, options = {}) {
         return { signal: combined };
     };
     const call = (tool, args, signal) => dispatcher.host.call(tool, args, context, callOptions(signal));
-    // Bound function properties (not methods) so vm call sites cannot lose `this`.
-    const asgrep = {
-        search: (input, callOptions) => call("search", {
-            query: input.query,
-            limit: clampLimit(input.limit),
-            excerpt_lines: clampExcerpt(input.excerptLines),
-            format: input.format === "agent" ? "agent" : "capsule",
-        }, callOptions?.signal),
-        semantic: (input, callOptions) => call("semantic", {
-            query: input.query,
+    const searchPayload = (method, input) => {
+        const scoped = coerceHostArgs(method, { ...input });
+        const payload = {
+            query: scoped.query,
             limit: clampLimit(input.limit),
             excerpt_lines: clampExcerpt(input.excerptLines),
             format: input.format === "agent" ? "agent" : "capsule",
-        }, callOptions?.signal),
-        chain: (input, callOptions) => call("chain", {
-            query: input.query,
-            limit: clampLimit(input.limit),
-            top_n: 20,
-        }, callOptions?.signal),
-        defs: (input, callOptions) => call("defs", {
-            symbol: input.symbol,
-            limit: clampLimit(input.limit),
-            excerpt_lines: clampExcerpt(input.excerptLines),
-        }, callOptions?.signal),
-        callers: (input, callOptions) => call("callers", {
-            symbol: input.symbol,
-            limit: clampLimit(input.limit),
-            excerpt_lines: clampExcerpt(input.excerptLines),
-        }, callOptions?.signal),
-        imports: (input, callOptions) => call("imports", {
-            module: input.module,
-            limit: clampLimit(input.limit),
-            excerpt_lines: clampExcerpt(input.excerptLines),
-        }, callOptions?.signal),
+        };
+        if (typeof scoped.lang === "string" && scoped.lang.trim())
+            payload.lang = scoped.lang.trim();
+        return payload;
+    };
+    // Bound function properties (not methods) so vm call sites cannot lose `this`.
+    const asgrep = {
+        search: (input, callOptions) => call("search", searchPayload("search", input), callOptions?.signal),
+        find: (input, callOptions) => call("find", searchPayload("find", input), callOptions?.signal),
+        read: (input, callOptions) => call("read", defined({
+            path: input.path,
+            start: input.start,
+            end: input.end,
+            ref: input.ref,
+            refs: input.refs,
+            context_lines: input.contextLines,
+            max_chars: input.maxChars,
+        }), callOptions?.signal),
+        edit: (input, callOptions) => {
+            // Multi-edit wire contract: every entry carries its own path; the
+            // top-level path is the default for entries that omit it.
+            const edits = input.edits?.map((entry) => ({
+                ...(typeof input.path === "string" ? { path: input.path } : {}),
+                ...entry,
+            }));
+            return call("edit", defined({
+                path: input.path,
+                oldText: input.oldText,
+                newText: input.newText,
+                edits,
+            }), callOptions?.signal);
+        },
+        semantic: (input, callOptions) => call("semantic", searchPayload("semantic", input), callOptions?.signal),
+        chain: (input, callOptions) => call("chain", { query: input.query, limit: clampLimit(input.limit), top_n: 20 }, callOptions?.signal),
+        defs: (input, callOptions) => {
+            const scoped = coerceHostArgs("defs", { ...input });
+            return call("defs", {
+                symbol: scoped.symbol,
+                limit: clampLimit(input.limit),
+                excerpt_lines: clampExcerpt(input.excerptLines),
+            }, callOptions?.signal);
+        },
+        callers: (input, callOptions) => {
+            const scoped = coerceHostArgs("callers", { ...input });
+            return call("callers", {
+                symbol: scoped.symbol,
+                limit: clampLimit(input.limit),
+                excerpt_lines: clampExcerpt(input.excerptLines),
+            }, callOptions?.signal);
+        },
+        imports: (input, callOptions) => call("imports", defined({ module: input.module, limit: clampLimit(input.limit), excerpt_lines: clampExcerpt(input.excerptLines) }), callOptions?.signal),
         indexStatus: (callOptions) => call("index_status", {}, callOptions?.signal),
         indexRepo: (input = {}, callOptions) => call("index_repo", { force: input.force === true }, callOptions?.signal),
         catalogSearch: (input, callOptions) => call("catalog_search", { query: input.query }, callOptions?.signal),
         catalogDescribe: (input, callOptions) => call("catalog_describe", { name: input.name }, callOptions?.signal),
+        doctor: (callOptions) => host.run(["doctor", ".", "--json"], context, callOptions),
     };
     return {
         asgrep,
         stats: dispatcher.stats,
+        trace: dispatcher.trace,
         resetStats: dispatcher.resetStats,
     };
 }

package/dist/codemode/dispatch.d.ts

@@ -4,7 +4,7 @@
  * Amdahl: serial cost is process spawn + SQLite open. Sticky serve kills spawn
  * for the whole Code Mode program; batch coalescing kills it per Promise.all wave.
  */
-import type { MachineEnvelope } from "../runtime.js";
+import type { MachineEnvelope } from "../runtime/runtime.js";
 import type { ConnectorHost, DispatchSurface } from "./connector.js";
 export type CodemodeToolCall = {
     tool: string;
@@ -18,6 +18,15 @@ export type DispatchStats = {
     stickyCalls: number;
     wallMs: number;
 };
+/** One settled host call: which lane carried it and how long it took. */
+export type DispatchCall = {
+    tool: string;
+    /** Short display target: query/symbol/module/path, whichever the call used. */
+    target: string;
+    lane: "sticky" | "batch" | "spawn" | "serial";
+    ok: boolean;
+    ms: number;
+};
 export type BatchResult = {
     results: Array<{
         id: string;
@@ -41,6 +50,8 @@ export type StickyWorker = {
         signal?: AbortSignal;
     }): Promise<BatchResult>;
     end(): Promise<void>;
+    /** True after timeout/crash/end. Missing means "assume live". */
+    closed?: () => boolean;
 };
 export type BatchCapableHost = ConnectorHost & {
     /** One-shot warm batch (codemode-batch). */
@@ -63,6 +74,7 @@ export type BatchCapableHost = ConnectorHost & {
 export declare function createCodemodeDispatcher(host: BatchCapableHost): {
     host: DispatchSurface;
     stats: () => DispatchStats;
+    trace: () => DispatchCall[];
     resetStats: () => void;
 };
 export declare function argvFor(tool: string, args: Record<string, unknown>): string[];

package/dist/codemode/dispatch.js

@@ -7,8 +7,16 @@
 import { mkdtemp, writeFile, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
+function callTarget(args) {
+    const refs = args.refs ?? args.ref;
+    const value = args.query ?? args.symbol ?? args.module ?? args.path ?? args.name
+        ?? (Array.isArray(refs) ? refs.map((r) => typeof r === "string" ? r : r?.path ?? "").filter(Boolean).join(", ")
+            : typeof refs === "string" ? refs : "");
+    const text = String(value).replace(/\s+/g, " ").trim();
+    return text.length > 60 ? text.slice(0, 59) + "…" : text;
+}
 const MAX_WAVE = 32;
-const MUTATING_TOOLS = new Set(["index_repo"]);
+const MUTATING_TOOLS = new Set(["index_repo", "edit"]);
 const abortError = () => Object.assign(new Error("codemode aborted"), { name: "AbortError" });
 function rejectWave(wave, cause) {
     for (const item of wave)
@@ -30,6 +38,21 @@ export function createCodemodeDispatcher(host) {
     let pending = [];
     let scheduled = false;
     let stats = emptyStats();
+    // Per-run call log: which lane each call took and how long it took. Bounded;
+    // surfaced in codemode result details so callers can see the dispatch shape.
+    const calls = [];
+    const TRACE_CAP = 128;
+    const recordCall = (item, ok) => {
+        if (calls.length >= TRACE_CAP)
+            return;
+        calls.push({ tool: item.tool, target: callTarget(item.args), lane: item.lane, ok, ms: Date.now() - item.startedAt });
+    };
+    // Mutations never batch and never overlap each other: edit/index_repo run on
+    // a serial tail so Promise.all([edit, edit]) cannot interleave writes, and a
+    // later wave's mutation queues behind an earlier one. Reads keep wave
+    // batching; ordering vs a concurrent write is intentionally unspecified
+    // (same semantics as today's shared batch — the store stays transactional).
+    let mutationTail = Promise.resolve();
     const flush = async () => {
         const wave = pending.filter((item) => !item.settled);
         pending = [];
@@ -39,18 +62,33 @@ export function createCodemodeDispatcher(host) {
         stats.waves += 1;
         stats.calls += wave.length;
         const waveStarted = Date.now();
+        const mutations = wave.filter((item) => MUTATING_TOOLS.has(item.tool) && !item.settled);
+        const reads = wave.filter((item) => !MUTATING_TOOLS.has(item.tool) && !item.settled);
+        let writesDone = mutationTail;
+        if (mutations.length > 0) {
+            writesDone = mutationTail.then(async () => {
+                for (const item of mutations) {
+                    if (!item.settled)
+                        await settleOne(host, item, stats);
+                }
+            });
+            mutationTail = writesDone;
+        }
         try {
-            if (wave.length === 1) {
-                await settleOne(host, wave[0], stats);
-                return;
+            if (reads.length === 1) {
+                await settleOne(host, reads[0], stats);
             }
-            // Chunk oversized waves (batch max = 32).
-            for (let offset = 0; offset < wave.length; offset += MAX_WAVE) {
-                const chunk = wave.slice(offset, offset + MAX_WAVE).filter((item) => !item.settled);
-                if (chunk.length === 0)
-                    continue;
-                await settleWave(host, chunk, stats);
+            else if (reads.length > 1) {
+                // Chunk oversized waves (batch max = 32).
+                for (let offset = 0; offset < reads.length; offset += MAX_WAVE) {
+                    const chunk = reads.slice(offset, offset + MAX_WAVE).filter((item) => !item.settled);
+                    if (chunk.length === 0)
+                        continue;
+                    await settleWave(host, chunk, stats);
+                }
             }
+            // Reads resolve independently; the wave is not done until queued writes land.
+            await writesDone;
         }
         finally {
             stats.wallMs += Date.now() - waveStarted;
@@ -64,6 +102,7 @@ export function createCodemodeDispatcher(host) {
                 return;
             item.settled = true;
             cleanup();
+            recordCall(item, true);
             resolve(value);
         };
         item.reject = (reason) => {
@@ -71,6 +110,7 @@ export function createCodemodeDispatcher(host) {
                 return;
             item.settled = true;
             cleanup();
+            recordCall(item, false);
             reject(reason);
         };
         const onAbort = () => item.reject(abortError());
@@ -94,6 +134,8 @@ export function createCodemodeDispatcher(host) {
                 args,
                 context,
                 settled: false,
+                startedAt: Date.now(),
+                lane: MUTATING_TOOLS.has(tool) ? "serial" : "spawn",
                 resolve: () => undefined,
                 reject: () => undefined,
             };
@@ -105,14 +147,17 @@ export function createCodemodeDispatcher(host) {
     return {
         host: dispatchHost,
         stats: () => ({ ...stats }),
+        trace: () => calls.slice(),
         resetStats: () => {
             stats = emptyStats();
+            calls.length = 0;
         },
     };
 }
 async function settleOne(host, item, stats) {
     try {
         if (host.sticky) {
+            item.lane = "sticky";
             stats.stickyCalls += 1;
             item.resolve(await host.sticky.call(item.tool, item.args, item.options));
             return;
@@ -125,17 +170,18 @@ async function settleOne(host, item, stats) {
         item.reject(err);
     }
 }
+/** Pending[] -> wire calls ({id, tool, args}) shared by both batch transports. */
+function packCalls(wave) {
+    return wave.map((item, index) => ({ id: String(index), tool: item.tool, args: item.args }));
+}
 async function settleWave(host, wave, stats) {
     if (host.sticky) {
         const transportOptions = sharedBatchOptions(wave);
         try {
-            const calls = wave.map((item, index) => ({
-                id: String(index),
-                tool: item.tool,
-                args: item.args,
-            }));
-            const batch = await host.sticky.batch(calls, transportOptions);
+            const batch = await host.sticky.batch(packCalls(wave), transportOptions);
             stats.stickyCalls += wave.length;
+            for (const item of wave)
+                item.lane = "sticky";
             settleFromBatch(wave, batch);
             return;
         }
@@ -159,13 +205,10 @@ async function settleWave(host, wave, stats) {
     if (host.runBatch) {
         const transportOptions = sharedBatchOptions(batchWave);
         try {
-            const calls = batchWave.map((item, index) => ({
-                id: String(index),
-                tool: item.tool,
-                args: item.args,
-            }));
-            const batch = await host.runBatch(calls, batchWave[0].context, transportOptions);
+            const batch = await host.runBatch(packCalls(batchWave), batchWave[0].context, transportOptions);
             stats.batchedCalls += batchWave.length;
+            for (const item of batchWave)
+                item.lane = "batch";
             settleFromBatch(batchWave, batch);
             return;
         }
@@ -221,6 +264,7 @@ function emptyStats() {
 }
 const ARGV_SPEC = {
     search: { form: "capsule", key: "query" },
+    find: { form: "find" },
     semantic: { form: "semantic" },
     chain: { form: "chain" },
     defs: { form: "capsule", key: "symbol", prefix: "defs" },
@@ -250,15 +294,31 @@ export function argvFor(tool, args) {
         return ["chain", argStr(args, "query"), ".", "--json", "--limit", String(limit)];
     }
     const excerpt = num(args.excerpt_lines ?? args.excerptLines, 0);
-    const capsule = ["--json", "--format", "agent-capsule", "--limit", String(limit), "--excerpt-lines", String(excerpt)];
+    const capsule = withLang(["--json", "--format", "agent-capsule", "--limit", String(limit), "--excerpt-lines", String(excerpt)], args);
     if (spec.form === "semantic") {
         return ["semantic", argStr(args, "query"), ".", ...capsule];
     }
+    if (spec.form === "find") {
+        const raw = argStr(args, "query").trim();
+        let token = raw;
+        if (/^blast:/i.test(raw)) {
+            const target = raw.slice(raw.indexOf(":") + 1).trim();
+            token = /[\\/.]/.test(target) ? `imports:${target}` : `callers:${target}`;
+        }
+        else if (!/^(defs|callers|imports|literal|regex|word|pattern):/i.test(raw)) {
+            token = `word:${raw}`;
+        }
+        return [...capsule, token, "."];
+    }
     // capsule (+ optional prefix for defs/callers/imports)
     const raw = argStr(args, spec.key);
     const token = spec.prefix ? `${spec.prefix}:${raw}` : raw;
     return [...capsule, token, "."];
 }
+function withLang(argv, args) {
+    const lang = typeof args.lang === "string" ? args.lang.trim() : "";
+    return lang ? ["--lang", lang, ...argv] : argv;
+}
 function num(value, fallback) {
     if (typeof value === "number" && Number.isFinite(value))
         return Math.trunc(value);

package/dist/codemode/guest-api.d.ts

@@ -0,0 +1,16 @@
+/** Guest-call packing so the first shape a model tries actually works. */
+import { type CodemodeHostMethod } from "./types.js";
+export declare function resolveHostMethod(method: string): CodemodeHostMethod | undefined;
+/** Turn positional guest calls into the host object shape. */
+export declare function packGuestCall(method: string, args: unknown[]): Record<string, unknown>;
+/** Accept query as a symbol alias and fold in:/fileFilter into the query string. */
+export declare function coerceHostArgs(method: string, input: Record<string, unknown>): Record<string, unknown>;
+export declare function applyQueryScope(query: string, args: Record<string, unknown>): string | undefined;
+export declare function unknownMethodError(method: string): string;
+export declare function timeoutHint(message: string): string;
+/**
+ * Strip fences and invoke a function expression, including non-async arrows.
+ * A single expression with no `return` is returned automatically.
+ * Bare statements still wrap in an async IIFE.
+ */
+export declare function normalizeCode(raw: string): string;