pi-ast-sgrep 1.3.2 → 2.0.0

package/dist/codemode/session-pool.js

@@ -0,0 +1,238 @@
+/**
+ * Session-scoped native Code Mode sessions.
+ *
+ * Primary path: in-process NAPI (`CodeModeSession` inside Node) — same model as
+ * MCP linking core. Zero CLI spawn.
+ *
+ * Fallback: sticky `codemode-serve` child only when the `.node` addon is missing
+ * (unsupported host / incomplete install). Doctor reports that as degraded.
+ */
+import { asEnvelope } from "./dispatch.js";
+import { loadCodemodeNative } from "./native.js";
+import { startStickyWorker } from "./worker.js";
+const abortError = () => Object.assign(new Error("native call aborted"), { name: "AbortError" });
+/** Bounded metadata and symbol lookups that may run on the JS thread. */
+const FAST_LOOKUP = new Set([
+    "defs",
+    "callers",
+    "imports",
+    "index_status",
+    "catalog_search",
+    "catalog_describe",
+]);
+function isBusyError(cause) {
+    const message = cause instanceof Error ? cause.message : String(cause);
+    return /session is busy/i.test(message);
+}
+function inProcessWorker(session) {
+    let tail = Promise.resolve();
+    let closed = false;
+    let inflight = 0;
+    const enqueue = (operation, signal) => {
+        if (closed)
+            return Promise.reject(new Error("native session is closed"));
+        if (signal?.aborted)
+            return Promise.reject(abortError());
+        inflight += 1;
+        const slot = tail.then(() => {
+            if (signal?.aborted)
+                throw abortError();
+            return operation();
+        });
+        tail = slot.then(() => {
+            inflight -= 1;
+        }, () => {
+            inflight -= 1;
+        });
+        if (!signal)
+            return slot;
+        return new Promise((resolve, reject) => {
+            const abort = () => reject(abortError());
+            signal.addEventListener("abort", abort, { once: true });
+            slot.then(resolve, reject).finally(() => signal.removeEventListener("abort", abort));
+        });
+    };
+    return {
+        call(tool, args, options) {
+            if (options?.signal?.aborted)
+                return Promise.reject(abortError());
+            const sync = session.callNow;
+            if (inflight === 0 && !closed && sync && FAST_LOOKUP.has(tool)) {
+                try {
+                    return Promise.resolve(asEnvelope(sync.call(session, tool, args ?? {}), tool));
+                }
+                catch (cause) {
+                    if (!isBusyError(cause))
+                        return Promise.reject(cause);
+                }
+            }
+            return enqueue(async () => asEnvelope(await session.call(tool, args, options?.signal), tool), options?.signal);
+        },
+        batch(calls, options) {
+            return enqueue(async () => {
+                const response = await session.batch(calls, options?.signal);
+                const result = {
+                    results: response.results,
+                    all_ok: response.allOk,
+                    wall_ms: response.wallMs,
+                    mode: response.mode,
+                };
+                return result;
+            }, options?.signal);
+        },
+        async end() {
+            closed = true;
+            await tail;
+            // The NAPI session is released when this worker closure is dropped.
+        },
+    };
+}
+export class NativeSessionPool {
+    #entries = new Map();
+    #starting = new Map();
+    #options = null;
+    #generations = new Map();
+    #startFn;
+    #backend = "none";
+    #shutdownPromise = null;
+    constructor(startFn = startStickyWorker) {
+        this.#startFn = startFn;
+    }
+    configure(options) {
+        this.#options = options;
+    }
+    configured() {
+        return this.#options !== null || loadCodemodeNative() !== null;
+    }
+    /** Active backend after first successful acquire. */
+    backend() {
+        return this.#backend;
+    }
+    async acquire(root) {
+        if (this.#shutdownPromise)
+            return null;
+        const existing = this.#entries.get(root);
+        if (existing)
+            return existing.worker;
+        const inFlight = this.#starting.get(root);
+        if (inFlight)
+            return inFlight;
+        const start = this.#start(root);
+        this.#starting.set(root, start);
+        try {
+            return await start;
+        }
+        finally {
+            if (this.#starting.get(root) === start)
+                this.#starting.delete(root);
+        }
+    }
+    async call(root, tool, args = {}, options) {
+        if (options?.signal?.aborted)
+            throw abortError();
+        const worker = await this.acquire(root);
+        if (!worker)
+            throw new Error("native Code Mode backend unavailable");
+        return worker.call(tool, args, options);
+    }
+    async invalidate(root) {
+        this.#generations.set(root, this.#generationFor(root) + 1);
+        const starting = this.#starting.get(root);
+        this.#starting.delete(root);
+        const entry = this.#entries.get(root);
+        this.#entries.delete(root);
+        if (entry)
+            await entry.worker.end().catch(() => undefined);
+        if (starting)
+            await starting.catch(() => null);
+        if (this.#entries.size === 0)
+            this.#backend = "none";
+    }
+    async shutdown() {
+        if (this.#shutdownPromise)
+            return this.#shutdownPromise;
+        const shutdown = this.#shutdownAll();
+        this.#shutdownPromise = shutdown;
+        try {
+            await shutdown;
+        }
+        finally {
+            if (this.#shutdownPromise === shutdown)
+                this.#shutdownPromise = null;
+        }
+    }
+    async #shutdownAll() {
+        const roots = new Set([...this.#entries.keys(), ...this.#starting.keys()]);
+        for (const root of roots)
+            this.#generations.set(root, this.#generationFor(root) + 1);
+        const entries = [...this.#entries.values()];
+        const starting = [...this.#starting.values()];
+        this.#entries.clear();
+        this.#starting.clear();
+        this.#backend = "none";
+        await Promise.all([
+            ...entries.map((e) => e.worker.end().catch(() => undefined)),
+            ...starting.map((start) => start.catch(() => null)),
+        ]);
+    }
+    #generationFor(root) {
+        return this.#generations.get(root) ?? 0;
+    }
+    async #start(root) {
+        const gen = this.#generationFor(root);
+        const opts = this.#options ?? {};
+        // 1) In-process NAPI (preferred — zero spawn).
+        const binding = loadCodemodeNative();
+        if (binding) {
+            try {
+                const config = { root };
+                if (opts.indexPath)
+                    config.indexPath = opts.indexPath;
+                if (opts.limit !== undefined)
+                    config.limit = opts.limit;
+                if (opts.useEmbed !== undefined)
+                    config.useEmbed = opts.useEmbed;
+                const session = new binding.Session(config);
+                const worker = inProcessWorker(session);
+                if (gen !== this.#generationFor(root)) {
+                    await worker.end().catch(() => undefined);
+                    return null;
+                }
+                this.#entries.set(root, { root, worker, generation: gen, backend: "napi" });
+                this.#backend = "napi";
+                return worker;
+            }
+            catch {
+                // Fall through to CLI sticky.
+            }
+        }
+        // 2) CLI sticky fallback (degraded).
+        if (!opts.binary)
+            return null;
+        try {
+            const stickyOpts = {
+                binary: opts.binary,
+                cwd: root,
+            };
+            if (opts.env)
+                stickyOpts.env = opts.env;
+            if (opts.timeoutMs !== undefined)
+                stickyOpts.timeoutMs = opts.timeoutMs;
+            if (opts.maxOutputBytes !== undefined)
+                stickyOpts.maxOutputBytes = opts.maxOutputBytes;
+            const worker = await this.#startFn(stickyOpts);
+            if (gen !== this.#generationFor(root)) {
+                await worker.end().catch(() => undefined);
+                return null;
+            }
+            this.#entries.set(root, { root, worker, generation: gen, backend: "cli" });
+            this.#backend = "cli";
+            return worker;
+        }
+        catch {
+            return null;
+        }
+    }
+}
+/** Singleton for advanced hosts; tools registration uses a local pool. */
+export const sharedNativePool = new NativeSessionPool();

package/dist/codemode/types.d.ts

@@ -0,0 +1,18 @@
+/** Typed surface the model sees inside a Code Mode program (`asgrep.*`). */
+export type SearchArgs = {
+    query: string;
+    limit?: number;
+    excerptLines?: number;
+    format?: "capsule" | "agent";
+};
+export type ChainArgs = {
+    query: string;
+    limit?: number;
+    excerptLines?: number;
+};
+/**
+ * Compact TypeScript declarations for the `asgrep` tool description.
+ * Keep short — every token here is paid on every turn (schema landfill lesson
+ * from pi-codex-conversion: compose inside Code Mode, don't dump 17 schemas).
+ */
+export declare const CODEMODE_TYPES_FOR_MODEL: string;

package/dist/codemode/types.js

@@ -0,0 +1,21 @@
+/** Typed surface the model sees inside a Code Mode program (`asgrep.*`). */
+/**
+ * Compact TypeScript declarations for the `asgrep` tool description.
+ * Keep short — every token here is paid on every turn (schema landfill lesson
+ * from pi-codex-conversion: compose inside Code Mode, don't dump 17 schemas).
+ */
+export const CODEMODE_TYPES_FOR_MODEL = `
+declare const asgrep: {
+  search(input: { query: string; limit?: number; excerptLines?: number }): Promise<unknown>;
+  semantic(input: { query: string; limit?: number; excerptLines?: number }): Promise<unknown>;
+  chain(input: { query: string; limit?: number }): Promise<unknown>;
+  defs(input: { symbol: string; limit?: number }): Promise<unknown>;
+  callers(input: { symbol: string; limit?: number }): Promise<unknown>;
+  imports(input: { module: string; limit?: number }): Promise<unknown>;
+  indexStatus(): Promise<unknown>;
+  indexRepo(input?: { force?: boolean }): Promise<unknown>;
+  catalogSearch(input: { query: string }): Promise<unknown>;
+  catalogDescribe(input: { name: string }): Promise<unknown>;
+};
+/** JS: Promise, JSON, Array, Object, Map, Set, Math. No require/process/fetch/fs. */
+`.trim();

package/dist/codemode/worker.d.ts

@@ -0,0 +1,29 @@
+/**
+ * Sticky NDJSON Code Mode worker (`asgrep codemode-serve`).
+ *
+ * One process, one warm Searcher, for the entire Code Mode program — the biggest
+ * Amdahl win over per-wave `codemode-batch` spawns.
+ */
+import type { MachineEnvelope } from "../runtime.js";
+import { type StickyWorker } from "./dispatch.js";
+export type StickyWorkerOptions = {
+    binary: string;
+    cwd: string;
+    env?: NodeJS.ProcessEnv;
+    signal?: AbortSignal;
+    /** Kill worker when one request exceeds this duration (ms). */
+    timeoutMs?: number;
+    /** Maximum bytes accepted in one NDJSON response. */
+    maxOutputBytes?: number;
+};
+export declare function startStickyWorker(options: StickyWorkerOptions): Promise<StickyWorker>;
+/** One-shot batch via stdin (avoids tempfile). */
+export declare function runBatchViaStdin(options: {
+    binary: string;
+    cwd: string;
+    body: string;
+    env?: NodeJS.ProcessEnv;
+    signal?: AbortSignal;
+    timeoutMs?: number;
+    maxOutputBytes?: number;
+}): Promise<MachineEnvelope>;

package/dist/codemode/worker.js

@@ -0,0 +1,307 @@
+/**
+ * Sticky NDJSON Code Mode worker (`asgrep codemode-serve`).
+ *
+ * One process, one warm Searcher, for the entire Code Mode program — the biggest
+ * Amdahl win over per-wave `codemode-batch` spawns.
+ */
+import { spawn } from "node:child_process";
+import { asEnvelope } from "./dispatch.js";
+const DEFAULT_MAX_OUTPUT_BYTES = 4 * 1024 * 1024;
+export async function startStickyWorker(options) {
+    if (options.signal?.aborted) {
+        throw new Error("codemode-serve aborted before start");
+    }
+    const child = spawn(options.binary, ["--root", options.cwd, "codemode-serve"], {
+        cwd: options.cwd,
+        env: { ...process.env, ...options.env, NO_COLOR: "1" },
+        stdio: ["pipe", "pipe", "pipe"],
+    });
+    const pending = new Map();
+    let nextId = 0;
+    let closed = false;
+    let stderr = "";
+    let stdout = Buffer.alloc(0);
+    const decoder = new TextDecoder("utf-8", { fatal: true });
+    const maxOutputBytes = options.maxOutputBytes ?? DEFAULT_MAX_OUTPUT_BYTES;
+    const failAll = (err) => {
+        for (const p of pending.values())
+            p.reject(err);
+        pending.clear();
+    };
+    const terminate = (err) => {
+        if (closed)
+            return;
+        closed = true;
+        options.signal?.removeEventListener("abort", onAbort);
+        killChild(child);
+        failAll(err);
+        child.stdout.destroy();
+    };
+    child.stdin.on("error", terminate);
+    const handleLine = (line) => {
+        const trimmed = line.trim();
+        if (!trimmed)
+            return;
+        let msg;
+        try {
+            msg = JSON.parse(trimmed);
+        }
+        catch (cause) {
+            terminate(new Error(`codemode-serve bad JSON: ${trimmed.slice(0, 200)}`));
+            void cause;
+            return;
+        }
+        const type = msg.type;
+        if (type === "bye")
+            return;
+        const id = typeof msg.id === "string" ? msg.id : undefined;
+        if (!id)
+            return;
+        const waiter = pending.get(id);
+        if (!waiter)
+            return;
+        pending.delete(id);
+        waiter.resolve(msg);
+    };
+    child.stdout.on("data", (chunk) => {
+        if (closed)
+            return;
+        const bytes = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
+        let offset = 0;
+        let newline;
+        while ((newline = bytes.indexOf(0x0a, offset)) >= 0) {
+            const segment = bytes.subarray(offset, newline);
+            const lineBytes = stdout.length + segment.length + 1;
+            if (lineBytes > maxOutputBytes) {
+                terminate(new Error(`codemode-serve output exceeded ${maxOutputBytes} bytes`));
+                return;
+            }
+            const line = stdout.length === 0
+                ? segment
+                : Buffer.concat([stdout, segment], stdout.length + segment.length);
+            stdout = Buffer.alloc(0);
+            try {
+                handleLine(decoder.decode(line));
+            }
+            catch {
+                terminate(new Error("codemode-serve output is not valid UTF-8"));
+                return;
+            }
+            if (closed)
+                return;
+            offset = newline + 1;
+        }
+        const tail = bytes.subarray(offset);
+        if (stdout.length + tail.length > maxOutputBytes) {
+            terminate(new Error(`codemode-serve output exceeded ${maxOutputBytes} bytes`));
+        }
+        else if (tail.length > 0) {
+            stdout = stdout.length === 0
+                ? Buffer.from(tail)
+                : Buffer.concat([stdout, tail], stdout.length + tail.length);
+        }
+    });
+    child.stderr.on("data", (chunk) => {
+        stderr += String(chunk);
+        if (stderr.length > 8_192)
+            stderr = stderr.slice(-8_192);
+    });
+    child.on("error", (err) => {
+        terminate(err);
+    });
+    child.on("close", (code, signal) => {
+        closed = true;
+        options.signal?.removeEventListener("abort", onAbort);
+        if (pending.size > 0) {
+            failAll(new Error(`codemode-serve exited code=${code ?? "null"} signal=${signal ?? "null"} stderr=${stderr.slice(0, 512)}`));
+        }
+    });
+    const onAbort = () => terminate(new Error("codemode-serve aborted"));
+    options.signal?.addEventListener("abort", onAbort, { once: true });
+    const write = (payload) => {
+        if (closed || !child.stdin.writable) {
+            return Promise.reject(new Error("codemode-serve is closed"));
+        }
+        const id = typeof payload.id === "string" ? payload.id : String(nextId++);
+        payload.id = id;
+        return new Promise((resolve, reject) => {
+            pending.set(id, { resolve, reject });
+            try {
+                child.stdin.write(`${JSON.stringify(payload)}\n`, (err) => {
+                    if (err)
+                        terminate(err);
+                });
+            }
+            catch (cause) {
+                terminate(cause instanceof Error ? cause : new Error(String(cause)));
+            }
+        });
+    };
+    const writeWithControls = (payload, label, signal) => {
+        if (signal?.aborted)
+            return Promise.reject(new Error(`${label} aborted`));
+        const response = write(payload);
+        return new Promise((resolve, reject) => {
+            let timer;
+            let settled = false;
+            const finish = (action) => {
+                if (settled)
+                    return;
+                settled = true;
+                if (timer)
+                    clearTimeout(timer);
+                signal?.removeEventListener("abort", onRequestAbort);
+                action();
+            };
+            const fail = (err) => finish(() => {
+                terminate(err);
+                reject(err);
+            });
+            const onRequestAbort = () => fail(new Error(`${label} aborted`));
+            signal?.addEventListener("abort", onRequestAbort, { once: true });
+            if (options.timeoutMs && options.timeoutMs > 0) {
+                timer = setTimeout(() => fail(new Error(`${label} timed out after ${options.timeoutMs}ms`)), options.timeoutMs);
+            }
+            response.then((value) => finish(() => resolve(value)), (cause) => finish(() => reject(cause)));
+        });
+    };
+    // Probe: empty End would close — instead send a tiny catalog call to verify protocol,
+    // or just return and let first real call fail. Prefer lazy: no probe.
+    return {
+        async call(tool, args, callOptions) {
+            const msg = await writeWithControls({ type: "call", tool, args }, "codemode call", callOptions?.signal);
+            if (msg.type === "error") {
+                throw new Error(typeof msg.error === "string" ? msg.error : "codemode-serve error");
+            }
+            if (msg.ok === false) {
+                throw new Error(typeof msg.error === "string" ? msg.error : `codemode ${tool} failed`);
+            }
+            return asEnvelope(msg.value, tool);
+        },
+        async batch(calls, callOptions) {
+            const msg = await writeWithControls({ type: "batch", calls }, "codemode batch", callOptions?.signal);
+            if (msg.type === "error") {
+                throw new Error(typeof msg.error === "string" ? msg.error : "codemode-serve batch error");
+            }
+            const results = Array.isArray(msg.results)
+                ? msg.results
+                : [];
+            const out = { results };
+            if (typeof msg.mode === "string")
+                out.mode = msg.mode;
+            if (typeof msg.wall_ms === "number")
+                out.wall_ms = msg.wall_ms;
+            if (typeof msg.all_ok === "boolean")
+                out.all_ok = msg.all_ok;
+            return out;
+        },
+        async end() {
+            if (closed)
+                return;
+            terminate(new Error("codemode-serve ended"));
+        },
+    };
+}
+/** One-shot batch via stdin (avoids tempfile). */
+export async function runBatchViaStdin(options) {
+    if (options.signal?.aborted)
+        throw new Error("codemode-batch aborted");
+    return new Promise((resolve, reject) => {
+        const child = spawn(options.binary, ["codemode-batch", "--requests", "-", "--json"], {
+            cwd: options.cwd,
+            env: { ...process.env, ...options.env, NO_COLOR: "1" },
+            stdio: ["pipe", "pipe", "pipe"],
+        });
+        let stdout = "";
+        let stderr = "";
+        let outputBytes = 0;
+        let settled = false;
+        const maxOutputBytes = options.maxOutputBytes ?? DEFAULT_MAX_OUTPUT_BYTES;
+        const cleanup = () => {
+            if (timer)
+                clearTimeout(timer);
+            options.signal?.removeEventListener("abort", onAbort);
+        };
+        const finish = (action) => {
+            if (settled)
+                return;
+            settled = true;
+            cleanup();
+            action();
+        };
+        const fail = (error) => finish(() => {
+            killChild(child);
+            reject(error);
+        });
+        child.stdin.on("error", fail);
+        const timer = options.timeoutMs && options.timeoutMs > 0
+            ? setTimeout(() => {
+                fail(new Error(`codemode-batch timed out after ${options.timeoutMs}ms`));
+            }, options.timeoutMs)
+            : undefined;
+        const onAbort = () => {
+            fail(new Error("codemode-batch aborted"));
+        };
+        options.signal?.addEventListener("abort", onAbort, { once: true });
+        child.stdout.on("data", (c) => {
+            outputBytes += Buffer.byteLength(c);
+            if (outputBytes > maxOutputBytes) {
+                fail(new Error(`codemode-batch output exceeded ${maxOutputBytes} bytes`));
+                return;
+            }
+            stdout += String(c);
+        });
+        child.stderr.on("data", (c) => {
+            outputBytes += Buffer.byteLength(c);
+            if (outputBytes > maxOutputBytes) {
+                fail(new Error(`codemode-batch output exceeded ${maxOutputBytes} bytes`));
+                return;
+            }
+            stderr += String(c);
+        });
+        child.on("error", (err) => {
+            fail(err);
+        });
+        child.on("close", (code) => {
+            if (settled)
+                return;
+            if (code !== 0) {
+                fail(new Error(`codemode-batch exited ${code}: ${stderr.slice(0, 512) || stdout.slice(0, 512)}`));
+                return;
+            }
+            try {
+                const value = JSON.parse(stdout);
+                finish(() => resolve(value));
+            }
+            catch (cause) {
+                fail(cause instanceof Error ? cause : new Error(String(cause)));
+            }
+        });
+        child.stdin.write(options.body, (err) => {
+            if (err) {
+                fail(err);
+                return;
+            }
+            child.stdin.end();
+        });
+    });
+}
+function killChild(child) {
+    if (child.exitCode !== null || child.signalCode !== null)
+        return;
+    try {
+        child.kill("SIGTERM");
+    }
+    catch {
+        return;
+    }
+    setTimeout(() => {
+        try {
+            if (child.exitCode === null && child.signalCode === null)
+                child.kill("SIGKILL");
+        }
+        catch {
+            // ignore
+        }
+    }, 2_000).unref?.();
+}

package/dist/index.d.ts

@@ -1,7 +1,35 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { FreshnessCoordinator, type FreshnessRuntime } from "./runtime.js";
-type RuntimeLike = FreshnessRuntime;
-type FreshnessLike = Pick<FreshnessCoordinator, "ensureFresh" | "markAffectedPath">;
+import { FreshnessCoordinator, type MachineEnvelope, type RunOptions } from "./runtime.js";
+type RuntimeLike = {
+    run(args: readonly string[], context: {
+        cwd: string;
+    }, options?: RunOptions): Promise<MachineEnvelope>;
+    resolveRoot?(context: {
+        cwd: string;
+    }): Promise<string>;
+    resolveBinaryPath?(options?: {
+        env?: NodeJS.ProcessEnv;
+    }): string;
+    nativeEnv?(options?: {
+        env?: NodeJS.ProcessEnv;
+    }): NodeJS.ProcessEnv;
+    config?: {
+        timeoutMs?: number;
+        maxOutputBytes?: number;
+        refreshIntervalMs?: number;
+    };
+    inspectIndexCompatibility?(context: {
+        cwd: string;
+    }): Promise<"ready" | "missing" | "incompatible">;
+    rebuildIncompatibleIndex?(context: {
+        cwd: string;
+    }, options?: RunOptions): Promise<MachineEnvelope>;
+    resolveIndexPath?(root: string): string;
+    watchExternalChanges?: boolean;
+};
+type FreshnessLike = Pick<FreshnessCoordinator, "ensureFresh" | "markAffectedPath"> & {
+    shutdown?(): void;
+};
 export declare function registerAstSgrepTools(pi: ExtensionAPI, runtime?: RuntimeLike, freshness?: FreshnessLike): void;
 export declare function registerAstSgrepCommands(pi: ExtensionAPI, runtime?: RuntimeLike): void;
 export default function astSgrepExtension(pi: ExtensionAPI): void;