pi-ast-sgrep 1.3.2 → 2.0.0

package/dist/codemode/dispatch.js

@@ -0,0 +1,320 @@
+/**
+ * Same-tick call coalescing + typed batch / sticky-serve dispatch.
+ *
+ * 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 { mkdtemp, writeFile, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+const MAX_WAVE = 32;
+const MUTATING_TOOLS = new Set(["index_repo"]);
+const abortError = () => Object.assign(new Error("codemode aborted"), { name: "AbortError" });
+function rejectWave(wave, cause) {
+    for (const item of wave)
+        item.reject(cause);
+}
+function sharedBatchOptions(wave) {
+    const signal = wave[0]?.options?.signal;
+    return signal && wave.every((item) => item.options?.signal === signal) ? { signal } : undefined;
+}
+function isSharedAbort(cause, options) {
+    return options?.signal !== undefined
+        && (options.signal.aborted || (cause instanceof Error && cause.name === "AbortError"));
+}
+/**
+ * Wraps a host so Promise.all([asgrep.search, asgrep.defs, …]) collapses into
+ * one microtask wave. Prefers sticky serve → one-shot batch → overlapped spawn.
+ */
+export function createCodemodeDispatcher(host) {
+    let pending = [];
+    let scheduled = false;
+    let stats = emptyStats();
+    const flush = async () => {
+        const wave = pending.filter((item) => !item.settled);
+        pending = [];
+        scheduled = false;
+        if (wave.length === 0)
+            return;
+        stats.waves += 1;
+        stats.calls += wave.length;
+        const waveStarted = Date.now();
+        try {
+            if (wave.length === 1) {
+                await settleOne(host, wave[0], stats);
+                return;
+            }
+            // 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);
+            }
+        }
+        finally {
+            stats.wallMs += Date.now() - waveStarted;
+        }
+    };
+    const enqueue = (item) => new Promise((resolve, reject) => {
+        const signal = item.options?.signal;
+        const cleanup = () => signal?.removeEventListener("abort", onAbort);
+        item.resolve = (value) => {
+            if (item.settled)
+                return;
+            item.settled = true;
+            cleanup();
+            resolve(value);
+        };
+        item.reject = (reason) => {
+            if (item.settled)
+                return;
+            item.settled = true;
+            cleanup();
+            reject(reason);
+        };
+        const onAbort = () => item.reject(abortError());
+        if (signal?.aborted) {
+            item.reject(abortError());
+            return;
+        }
+        signal?.addEventListener("abort", onAbort, { once: true });
+        pending.push(item);
+        if (!scheduled) {
+            scheduled = true;
+            queueMicrotask(() => {
+                void flush();
+            });
+        }
+    });
+    const dispatchHost = {
+        call(tool, args, context, options) {
+            const item = {
+                tool,
+                args,
+                context,
+                settled: false,
+                resolve: () => undefined,
+                reject: () => undefined,
+            };
+            if (options)
+                item.options = options;
+            return enqueue(item);
+        },
+    };
+    return {
+        host: dispatchHost,
+        stats: () => ({ ...stats }),
+        resetStats: () => {
+            stats = emptyStats();
+        },
+    };
+}
+async function settleOne(host, item, stats) {
+    try {
+        if (host.sticky) {
+            stats.stickyCalls += 1;
+            item.resolve(await host.sticky.call(item.tool, item.args, item.options));
+            return;
+        }
+        // N=1 without sticky: direct CLI (batch-of-1 is tempfile + protocol for no gain).
+        const args = argvFor(item.tool, item.args);
+        item.resolve(await host.run(args, item.context, item.options));
+    }
+    catch (err) {
+        item.reject(err);
+    }
+}
+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);
+            stats.stickyCalls += wave.length;
+            settleFromBatch(wave, batch);
+            return;
+        }
+        catch (cause) {
+            if (isSharedAbort(cause, transportOptions)) {
+                rejectWave(wave, cause);
+                return;
+            }
+            if (wave.some((item) => MUTATING_TOOLS.has(item.tool))) {
+                // The worker may have committed a mutation before its transport died.
+                // Replaying the wave through another transport would be ambiguous.
+                rejectWave(wave, cause);
+                return;
+            }
+            // Sticky died mid-program — fall through to one-shot / spawn.
+        }
+    }
+    const batchWave = wave.filter((item) => !item.settled);
+    if (batchWave.length === 0)
+        return;
+    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);
+            stats.batchedCalls += batchWave.length;
+            settleFromBatch(batchWave, batch);
+            return;
+        }
+        catch (cause) {
+            if (isSharedAbort(cause, transportOptions)) {
+                rejectWave(batchWave, cause);
+                return;
+            }
+            if (batchWave.some((item) => MUTATING_TOOLS.has(item.tool))) {
+                rejectWave(batchWave, cause);
+                return;
+            }
+            // Transport failure only — do NOT re-run when per-call ok:false.
+        }
+    }
+    const spawnWave = batchWave.filter((item) => !item.settled);
+    stats.parallelSpawnCalls += spawnWave.length;
+    await Promise.all(spawnWave.map(async (item) => {
+        try {
+            const args = argvFor(item.tool, item.args);
+            item.resolve(await host.run(args, item.context, item.options));
+        }
+        catch (err) {
+            item.reject(err);
+        }
+    }));
+}
+function settleFromBatch(wave, batch) {
+    const byId = new Map(batch.results.map((r) => [r.id, r]));
+    for (let i = 0; i < wave.length; i++) {
+        const item = wave[i];
+        const result = byId.get(String(i));
+        if (!result) {
+            item.reject(new Error(`codemode-batch missing result id=${i}`));
+            continue;
+        }
+        if (!result.ok) {
+            item.reject(new Error(result.error ?? `codemode call ${i} failed`));
+            continue;
+        }
+        item.resolve(asEnvelope(result.value, item.tool));
+    }
+}
+function emptyStats() {
+    return {
+        waves: 0,
+        calls: 0,
+        batchedCalls: 0,
+        parallelSpawnCalls: 0,
+        stickyCalls: 0,
+        wallMs: 0,
+    };
+}
+const ARGV_SPEC = {
+    search: { form: "capsule", key: "query" },
+    semantic: { form: "semantic" },
+    chain: { form: "chain" },
+    defs: { form: "capsule", key: "symbol", prefix: "defs" },
+    callers: { form: "capsule", key: "symbol", prefix: "callers" },
+    imports: { form: "capsule", key: "module", prefix: "imports" },
+    index_status: { form: "status" },
+    index_repo: { form: "index_repo" },
+};
+function argStr(args, key) {
+    return String(args[key] ?? "");
+}
+export function argvFor(tool, args) {
+    const spec = ARGV_SPEC[tool];
+    if (!spec)
+        throw new Error(`codemode tool has no direct CLI fallback: ${tool}`);
+    if (spec.form === "status")
+        return ["status", ".", "--json"];
+    if (spec.form === "index_repo") {
+        const command = args.force === true ? "reindex" : "index";
+        const paths = Array.isArray(args.paths)
+            ? args.paths.filter((path) => typeof path === "string")
+            : [];
+        return [command, ".", "--json", ...paths.flatMap((path) => ["--path", path])];
+    }
+    const limit = num(args.limit, 8);
+    if (spec.form === "chain") {
+        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)];
+    if (spec.form === "semantic") {
+        return ["semantic", argStr(args, "query"), ".", ...capsule];
+    }
+    // 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 num(value, fallback) {
+    if (typeof value === "number" && Number.isFinite(value))
+        return Math.trunc(value);
+    return fallback;
+}
+export function asEnvelope(value, command) {
+    if (value &&
+        typeof value === "object" &&
+        value.tool === "asgrep" &&
+        typeof value.ok === "boolean") {
+        return value;
+    }
+    const record = value && typeof value === "object" ? value : { value };
+    // Overrides AFTER spread so tool/ok/schema cannot be clobbered by payload fields.
+    return {
+        ...record,
+        tool: "asgrep",
+        ...(command ? { command } : {}),
+        schema_version: "1.0.0",
+        ok: true,
+    };
+}
+/** One-shot batch via stdin (no tempfile) when spawn-with-stdin is available. */
+export async function runNativeBatch(run, calls, context, options, writeBatch) {
+    const body = JSON.stringify({
+        root: context.cwd,
+        // Auto/serial warm by default — N parallel SQLite opens are usually slower.
+        parallel_mode: "auto",
+        calls,
+    });
+    if (writeBatch) {
+        const envelope = await writeBatch(body, context, options);
+        return envelopeToBatch(envelope);
+    }
+    // Fallback: tempfile + pi.exec (no stdin).
+    const dir = await mkdtemp(join(tmpdir(), "asgrep-codemode-"));
+    const requestsPath = join(dir, "requests.json");
+    try {
+        await writeFile(requestsPath, body, "utf8");
+        const envelope = await run(["codemode-batch", "--requests", requestsPath, "--json"], context, options);
+        return envelopeToBatch(envelope);
+    }
+    finally {
+        await rm(dir, { recursive: true, force: true }).catch(() => undefined);
+    }
+}
+function envelopeToBatch(envelope) {
+    const results = Array.isArray(envelope.results)
+        ? envelope.results
+        : [];
+    const out = { results };
+    if (typeof envelope.mode === "string")
+        out.mode = envelope.mode;
+    if (typeof envelope.wall_ms === "number")
+        out.wall_ms = envelope.wall_ms;
+    if (typeof envelope.all_ok === "boolean")
+        out.all_ok = envelope.all_ok;
+    return out;
+}

package/dist/codemode/index.d.ts

@@ -0,0 +1,18 @@
+/**
+ * Code Mode for ast-sgrep in Pi.
+ *
+ * Pattern (Cloudflare / Anthropic PTC / OpenAI PTC / OpenCode):
+ * the model writes JavaScript that calls typed `asgrep.*` methods. Intermediate
+ * results stay in the program; only the shaped return value re-enters the model
+ * context. Parallel calls use `Promise.all` against one warm in-process session.
+ *
+ * Code Mode and MCP are sibling front ends on the same core — pick one per
+ * client. They never import each other. Do not install both for the same agent.
+ */
+export { createAsgrepConnector, type AsgrepConnector, type ConnectorHost, type DispatchSurface, type ConnectorBundle, } from "./connector.js";
+export { runCodemode, normalizeCode, type CodemodeRunResult, type CodemodeRunSuccess, type CodemodeRunFailure } from "./runner.js";
+export { CODEMODE_TYPES_FOR_MODEL, type SearchArgs, type ChainArgs } from "./types.js";
+export { createCodemodeDispatcher, runNativeBatch, argvFor, asEnvelope, type DispatchStats, type BatchCapableHost, type StickyWorker, type BatchResult, } from "./dispatch.js";
+export { startStickyWorker, runBatchViaStdin } from "./worker.js";
+export { NativeSessionPool, sharedNativePool } from "./session-pool.js";
+export { loadCodemodeNative, nativeAvailable, resetNativeCache, type CodemodeNativeBinding, type NativeSession, } from "./native.js";

package/dist/codemode/index.js

@@ -0,0 +1,18 @@
+/**
+ * Code Mode for ast-sgrep in Pi.
+ *
+ * Pattern (Cloudflare / Anthropic PTC / OpenAI PTC / OpenCode):
+ * the model writes JavaScript that calls typed `asgrep.*` methods. Intermediate
+ * results stay in the program; only the shaped return value re-enters the model
+ * context. Parallel calls use `Promise.all` against one warm in-process session.
+ *
+ * Code Mode and MCP are sibling front ends on the same core — pick one per
+ * client. They never import each other. Do not install both for the same agent.
+ */
+export { createAsgrepConnector, } from "./connector.js";
+export { runCodemode, normalizeCode } from "./runner.js";
+export { CODEMODE_TYPES_FOR_MODEL } from "./types.js";
+export { createCodemodeDispatcher, runNativeBatch, argvFor, asEnvelope, } from "./dispatch.js";
+export { startStickyWorker, runBatchViaStdin } from "./worker.js";
+export { NativeSessionPool, sharedNativePool } from "./session-pool.js";
+export { loadCodemodeNative, nativeAvailable, resetNativeCache, } from "./native.js";

package/dist/codemode/native.d.ts

@@ -0,0 +1,55 @@
+/**
+ * Load the in-process Code Mode NAPI addon.
+ *
+ * Same model as MCP: Rust `CodeModeSession` runs inside the Node process.
+ * No `asgrep` CLI spawn on the hot path.
+ *
+ * Resolution order:
+ * 1. `ASGREP_CODEMODE_NAPI_PATH` (dev override)
+ * 2. `@ast-sgrep/<platform>/ast-sgrep-codemode.node` via launcher (release install)
+ * 3. Local `extension/native/` / cargo `target/release` (dev builds)
+ */
+export declare const CODEMODE_BINDING_VERSION = "2.0.0";
+export type NativeSessionConfig = {
+    root?: string;
+    indexPath?: string;
+    limit?: number;
+    useEmbed?: boolean;
+};
+export type NativeBatchCall = {
+    id: string;
+    tool: string;
+    args?: Record<string, unknown>;
+};
+export type NativeBatchResult = {
+    id: string;
+    ok: boolean;
+    value?: unknown;
+    error?: string;
+};
+export type NativeBatchResponse = {
+    allOk: boolean;
+    results: NativeBatchResult[];
+    callCount: number;
+    wallMs: number;
+    mode: string;
+};
+export type NativeSession = {
+    call(tool: string, args?: Record<string, unknown>, signal?: AbortSignal): Promise<unknown>;
+    /** Sync bounded metadata/symbol lookup; omitted on older addons. Throws if busy. */
+    callNow?(tool: string, args?: Record<string, unknown>): unknown;
+    batch(calls: NativeBatchCall[], signal?: AbortSignal): Promise<NativeBatchResponse>;
+    readonly callCount: number;
+    readonly root: string;
+};
+export type CodemodeNativeBinding = {
+    Session: new (config?: NativeSessionConfig) => NativeSession;
+    bindingVersion(): string;
+    isNative(): boolean;
+    asyncApiVersion(): number;
+};
+/** Load the NAPI binding once. Returns null if unavailable on this host. */
+export declare function loadCodemodeNative(): CodemodeNativeBinding | null;
+export declare function nativeAvailable(): boolean;
+/** Reset cache (tests). */
+export declare function resetNativeCache(): void;

package/dist/codemode/native.js

@@ -0,0 +1,119 @@
+/**
+ * Load the in-process Code Mode NAPI addon.
+ *
+ * Same model as MCP: Rust `CodeModeSession` runs inside the Node process.
+ * No `asgrep` CLI spawn on the hot path.
+ *
+ * Resolution order:
+ * 1. `ASGREP_CODEMODE_NAPI_PATH` (dev override)
+ * 2. `@ast-sgrep/<platform>/ast-sgrep-codemode.node` via launcher (release install)
+ * 3. Local `extension/native/` / cargo `target/release` (dev builds)
+ */
+import { createRequire } from "node:module";
+import { existsSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { fileURLToPath } from "node:url";
+export const CODEMODE_BINDING_VERSION = "2.0.0";
+let cached;
+function platformTriple() {
+    const { platform, arch } = process;
+    if (platform === "linux" && arch === "x64")
+        return "linux-x64-gnu";
+    if (platform === "linux" && arch === "arm64")
+        return "linux-arm64-gnu";
+    if (platform === "darwin" && arch === "arm64")
+        return "darwin-arm64";
+    if (platform === "darwin" && arch === "x64")
+        return "darwin-x64";
+    if (platform === "win32" && arch === "x64")
+        return "win32-x64-msvc";
+    return null;
+}
+function platformPackageAddon() {
+    const require = createRequire(import.meta.url);
+    try {
+        const launcher = require("ast-sgrep");
+        if (typeof launcher.resolveCodemodeAddon === "function") {
+            return launcher.resolveCodemodeAddon();
+        }
+    }
+    catch {
+        // Launcher may be unavailable in isolated unit tests.
+    }
+    return null;
+}
+function candidatePaths() {
+    const here = dirname(fileURLToPath(import.meta.url));
+    const triple = platformTriple();
+    const names = triple
+        ? [
+            `ast-sgrep-codemode.${triple}.node`,
+            `ast-sgrep-codemode.node`,
+        ]
+        : [`ast-sgrep-codemode.node`];
+    const dirs = [
+        // Built next to extension (dev / packaged)
+        join(here, "..", "..", "native"),
+        join(here, "..", "native"),
+        // Workspace release output
+        join(here, "..", "..", "..", "..", "target", "release"),
+    ];
+    const cargoTarget = process.env.CARGO_TARGET_DIR;
+    if (cargoTarget)
+        dirs.push(join(cargoTarget, "release"));
+    const out = [];
+    const override = process.env.ASGREP_CODEMODE_NAPI_PATH;
+    if (override)
+        out.push(override);
+    const packaged = platformPackageAddon();
+    if (packaged)
+        out.push(packaged);
+    for (const dir of dirs) {
+        for (const name of names)
+            out.push(join(dir, name));
+        // cargo cdylib name
+        out.push(join(dir, "libast_sgrep_codemode_napi.so"));
+        out.push(join(dir, "libast_sgrep_codemode_napi.dylib"));
+        out.push(join(dir, "ast_sgrep_codemode_napi.dll"));
+    }
+    return out;
+}
+/** Load the NAPI binding once. Returns null if unavailable on this host. */
+export function loadCodemodeNative() {
+    if (cached !== undefined)
+        return cached;
+    // Force CLI sticky / argv path (unit tests, degraded installs).
+    if (process.env.ASGREP_CODEMODE_BACKEND === "cli") {
+        cached = null;
+        return null;
+    }
+    const require = createRequire(import.meta.url);
+    for (const path of candidatePaths()) {
+        if (!existsSync(path))
+            continue;
+        try {
+            const binding = require(path);
+            if (typeof binding?.isNative === "function" &&
+                binding.isNative() &&
+                typeof binding.bindingVersion === "function" &&
+                binding.bindingVersion() === CODEMODE_BINDING_VERSION &&
+                typeof binding.asyncApiVersion === "function" &&
+                binding.asyncApiVersion() === 1) {
+                cached = binding;
+                return cached;
+            }
+        }
+        catch {
+            // try next candidate
+        }
+    }
+    cached = null;
+    return null;
+}
+export function nativeAvailable() {
+    return loadCodemodeNative() !== null;
+}
+/** Reset cache (tests). */
+export function resetNativeCache() {
+    cached = undefined;
+}

package/dist/codemode/runner.d.ts

@@ -0,0 +1,37 @@
+import type { AsgrepConnector } from "./connector.js";
+import type { DispatchStats } from "./dispatch.js";
+/** Closed sum: success|failure — `ok:true` with `error` (or `ok:false` without) is unrepresentable. */
+export type CodemodeRunSuccess = {
+    ok: true;
+    result: unknown;
+    logs: string[];
+    code: string;
+    stats?: DispatchStats;
+    wallMs: number;
+};
+export type CodemodeRunFailure = {
+    ok: false;
+    result: null;
+    error: string;
+    logs: string[];
+    code: string;
+    stats?: DispatchStats;
+    wallMs: number;
+};
+export type CodemodeRunResult = CodemodeRunSuccess | CodemodeRunFailure;
+/** Strip markdown fences and normalize to an async IIFE expression. */
+export declare function normalizeCode(raw: string): string;
+/**
+ * Run model-generated JavaScript against the typed `asgrep` connector.
+ *
+ * Model-generated code is not trusted with the extension host's ambient Node
+ * authority. A dedicated worker contains CPU/microtask denial of service; its
+ * VM hides `process`, module loading, and host constructors, with a JSON bridge
+ * as the only exposed capability. This is not an OS sandbox, so deployments
+ * requiring adversarial-code isolation should still restrict the Pi process.
+ */
+export declare function runCodemode(rawCode: string, asgrep: AsgrepConnector, options?: {
+    timeoutMs?: number;
+    signal?: AbortSignal;
+    stats?: () => DispatchStats;
+}): Promise<CodemodeRunResult>;