@estebanforge/pi-antigravity-bridge 1.0.0

package/src/runner.ts

@@ -0,0 +1,386 @@
+// Core agy turn runner: spawn `agy -p`, poll its conversation DB for new steps,
+// decode each, and emit structured events via a callback. Shared by the
+// standalone CLI (scripts/run-agy.ts) and the pi provider (src/provider.ts).
+//
+// Streaming contract:
+//   - agy writes steps to SQLite incrementally as it works.
+//   - We poll the DB every POLL_INTERVAL_MS (250), decode new rows, emit.
+//   - After agy exits, TRAILING_POLLS x 100ms catch late flushes.
+//   - agy does NOT print the conversation id; we bind it by snapshot/diff
+//     (see discovery.ts).
+//
+// agy runs its OWN closed tool loop (read/write/edit/exec against --add-dir).
+// We cannot bridge those tools to pi. Tool steps surface as "tool" events so
+// the UI can show "agy: editing foo.ts"  -  the edit itself already landed.
+
+import { spawn, type ChildProcess } from "node:child_process";
+import { bridgeMcpConfigDir, bridgeMcpConfigExists } from "./mcp-server.js";
+import { ConversationPoller, type Step } from "./poller.js";
+import {
+	CONVERSATIONS_DIR,
+	conversationDbPath,
+	newConversationId,
+	snapshotConversations,
+} from "./discovery.js";
+import { extractAgentText, extractToolCall, extractTitle } from "./protobuf.js";
+
+// --- tuning -----------------------------------------------------------------
+
+export const POLL_INTERVAL_MS = 250;
+export const TRAILING_POLLS = 3;
+export const TRAILING_POLL_MS = 100;
+export const DEFAULT_TIMEOUT_MIN = 10;
+const GRACE_AFTER_TIMEOUT_MS = 5000;
+// Cap on poll-tick retries to bind the conversation id when the snapshot is
+// ambiguous (a concurrent agy also created a .db). Each retry can run a full
+// /proc FD scan, so we stop after this many ticks (~15s at 250ms) and let the
+// turn fail safe rather than scan indefinitely. Generous enough to cover a
+// slow agy startup (OAuth refresh + cold init).
+const MAX_BIND_ATTEMPTS = 60;
+
+// agy conversation ids are UUID DB-stems. First char must be alphanumeric so a
+// leading-dash value can't misbind on agy's arg parser as the token after
+// --conversation (flag injection). Hyphens allowed in the body (real UUIDs).
+const CONV_ID_RE = /^[A-Za-z0-9][A-Za-z0-9-]{0,127}$/;
+
+// --- events -----------------------------------------------------------------
+
+export type AgyEvent =
+	| { kind: "text"; text: string }
+	| { kind: "thinking"; text: string }
+	| { kind: "tool"; name: string; inputJson: string }
+	| { kind: "title"; title: string };
+
+// step_type values observed in real DBs. Tool steps share the same payload
+// layout (field 5 -> toolCall), so we decode them uniformly. Unknown types
+// are skipped (return null)  -  forward-compatible with future agy additions.
+const TOOL_STEP_TYPES = new Set([5, 7, 8, 9, 17, 21, 33, 101, 132, 138]);
+
+/** Map a raw step row to a decoded event, or null when nothing to emit.
+ *  Pure: no I/O, no side effects. Exported for testing. */
+export function decodeStep(step: Step): AgyEvent | null {
+	if (step.stepType === 15) {
+		const t = extractAgentText(step.payload);
+		return t ? { kind: "text", text: t.text } : null;
+	}
+	if (step.stepType === 14) {
+		// Thinking steps reuse the agentText layout (field 20.1) in observed DBs.
+		const t = extractAgentText(step.payload);
+		return t ? { kind: "thinking", text: t.text } : null;
+	}
+	if (step.stepType === 23) {
+		const title = extractTitle(step.payload);
+		return title ? { kind: "title", title } : null;
+	}
+	if (TOOL_STEP_TYPES.has(step.stepType)) {
+		const tc = extractToolCall(step.payload);
+		if (tc?.name) return { kind: "tool", name: tc.name, inputJson: tc.inputJson };
+		return null;
+	}
+	return null;
+}
+
+// --- options + result -------------------------------------------------------
+
+export interface AgyRunOptions {
+	/** Workspace root agy operates in (passed as --add-dir). */
+	cwd: string;
+	/** Exact agy model string, e.g. "Gemini 3.6 Flash (Medium)". Caller
+	 *  resolves aliases; the runner passes this through verbatim. */
+	model?: string;
+	/** agy execution mode. accept-edits = agy applies edits; plan = review-only. */
+	mode?: "accept-edits" | "plan";
+	/** Pass --dangerously-skip-permissions so commands don't hang on an
+	 *  unanswerable y/n prompt in non-interactive `-p` mode. Default true.
+	 *  Required for accept-edits to function; harmless under plan mode. */
+	skipPermissions?: boolean;
+	/** The prompt. Required. */
+	prompt: string;
+	/** Existing conversation id to resume. When set, agy reuses it (no snapshot). */
+	conversationId?: string | null;
+	/** Highest step idx already streamed in a prior turn (resume only). The
+	 *  poller starts reading AFTER this idx so resumed turns don't replay
+	 *  history. Default -1 (read everything). */
+	baseStepIdx?: number;
+	/** Hard cap on the run, in minutes. */
+	timeoutMin?: number;
+	/** Optional AbortSignal for cancellation. */
+	signal?: AbortSignal;
+	/** Conversations dir override (testing / isolation). */
+	conversationsDir?: string;
+	/** agy binary path. Defaults to AGY_BIN env or "agy". */
+	binary?: string;
+	/** Extra args appended (split from AGY_EXTRA_ARGS by the caller). */
+	extraArgs?: string[];
+}
+
+export interface AgyRunResult {
+	exitCode: number;
+	conversationId: string | null;
+	lastIdx: number;
+	aborted: boolean;
+	timedOut: boolean;
+	stderr: string;
+	durationMs: number;
+}
+
+// --- spawn helpers ----------------------------------------------------------
+
+function resolveBinary(explicit?: string): string {
+	return explicit || process.env.AGY_BIN || "agy";
+}
+
+function extraArgsFromEnv(): string[] {
+	const raw = process.env.AGY_EXTRA_ARGS;
+	return raw ? raw.split(/\s+/).filter((s) => s.length > 0) : [];
+}
+
+const sleep = (ms: number) => new Promise<void>((r) => setTimeout(r, ms));
+
+// --- the turn ---------------------------------------------------------------
+
+/** Spawn agy and stream decoded events until it exits. Calls `onEvent` for
+ *  every decoded step in order. Returns the run outcome (exit code, discovered
+ *  conversation id, last step idx seen). Never throws on agy failure  - 
+ *  surfaces non-zero exit / timeout / abort in the result. */
+export async function runAgyTurn(
+	opts: AgyRunOptions,
+	onEvent: (event: AgyEvent) => void,
+): Promise<AgyRunResult> {
+	const start = Date.now();
+	const dir = opts.conversationsDir || CONVERSATIONS_DIR;
+	const mode = opts.mode ?? "accept-edits";
+	const timeoutMin = opts.timeoutMin ?? DEFAULT_TIMEOUT_MIN;
+	const binary = resolveBinary(opts.binary);
+	const extra = [...extraArgsFromEnv(), ...(opts.extraArgs ?? [])];
+
+	const rawConvId = opts.conversationId ?? null;
+	const isContinuation =
+		typeof rawConvId === "string" && rawConvId.length > 0 && CONV_ID_RE.test(rawConvId);
+	const snapshot = isContinuation ? null : snapshotConversations(dir);
+	let lastIdx = opts.baseStepIdx ?? -1;
+
+	// Build argv. --add-dir first so agy binds the workspace before anything.
+	const args = ["--add-dir", opts.cwd, ...extra];
+	// When the MCP tool bridge is running, add its config dir so this agy
+	// discovers pi's tools (memory, codegraph, search, ...). agy reads
+	// .agents/mcp_config.json from --add-dir dirs. AskAntigravity does NOT pass
+	// this dir, so its agy stays plain (no recursion).
+	if (bridgeMcpConfigExists()) args.push("--add-dir", bridgeMcpConfigDir());
+	if (opts.model) args.push("--model", opts.model);
+	args.push("--mode", mode);
+	// Without this, any run_command triggers an interactive permission prompt
+	// that hangs forever in non-interactive print mode (no TTY to answer y/n).
+	// accept-edits only auto-approves file writes, NOT commands. See PLAN.md.
+	if (opts.skipPermissions !== false) args.push("--dangerously-skip-permissions");
+	if (isContinuation && rawConvId) args.push("--conversation", rawConvId);
+	args.push("--print-timeout", `${timeoutMin}m`);
+	args.push("-p", opts.prompt);
+
+	let stderr = "";
+	let boundId = isContinuation ? rawConvId : null;
+	// Counts bind attempts so the /proc FD scan can't run on every tick forever
+	// when the snapshot stays ambiguous (see MAX_BIND_ATTEMPTS).
+	let bindAttempts = 0;
+
+	// One poller per turn, lazily opened once we know the conversation id.
+	let poller: ConversationPoller | null = null;
+	// agy extends the step it is currently writing in place (same idx, growing
+	// text). poll() returns only idx > lastIdx, so we re-read the last
+	// text/thinking step each tick and emit the grown suffix as a delta.
+	// flushStreamStep also runs BEFORE the loop so a step boundary landing
+	// inside this tick doesn't drop the outgoing step's final in-place tail
+	// (the loop may switch streamIdx to a new idx before we'd re-read).
+	let streamIdx = -1;
+	let streamKind: "text" | "thinking" | null = null;
+	let streamEmitted = 0;
+	const flushStreamStep = (): void => {
+		if (streamIdx < 0 || !poller || !streamKind) return;
+		const s = poller.readStepAt(streamIdx);
+		if (!s) return;
+		try {
+			const t = extractAgentText(s.payload);
+			if (t && t.text.length > streamEmitted) {
+				onEvent({ kind: streamKind, text: t.text.slice(streamEmitted) });
+				streamEmitted = t.text.length;
+			}
+		} catch {
+			/* torn re-read; next poll retries */
+		}
+	};
+	const pollOnce = (): boolean => {
+		// Bind the conversation id on a fresh run (agy doesn't print it).
+		// pid lets newConversationId disambiguate when a concurrent agy also
+		// drops a new .db in the dir during our turn (see discovery.ts).
+		if (!boundId && snapshot && bindAttempts < MAX_BIND_ATTEMPTS) {
+			boundId = newConversationId(dir, snapshot, {
+				pid: proc?.pid,
+				// Only the genuinely-ambiguous case (>1 new DB, unresolved) counts
+				// against the cap. The ordinary "agy hasn't written its DB yet" wait
+				// must keep polling for the full turn timeout, not burn a budget
+				// meant to bound the expensive /proc FD-scan retries.
+				onAmbiguous: () => { bindAttempts++; },
+			});
+		}
+		if (!boundId) return false;
+
+		if (!poller) {
+			poller = new ConversationPoller(conversationDbPath(boundId, dir), lastIdx);
+		}
+		if (!poller.isOpen && !poller.tryOpen()) return false;
+
+		// Coalesce: one data_version check per tick gates BOTH the in-place
+		// re-read (flushStreamStep -> readStepAt) and the new-row read. While
+		// agy is thinking and hasn't committed, hasChanged() is false and we
+		// skip both SELECTs, so no row read fires on an idle tick.
+		if (!poller.hasChanged()) return false;
+
+		// Catch the currently-tracked step's final in-place growth BEFORE the
+		// loop may switch tracking to a new step.
+		flushStreamStep();
+		const steps = poller.readNewSteps();
+		for (const step of steps) {
+			// A torn read (agy mid-write) can throw RangeError out of the protobuf
+			// walker. Drop that step rather than aborting the whole turn  -  the
+			// row settles on the next poll. (agy-acp database.ts pattern, lifted
+			// to the decode layer where the throw actually originates.)
+			try {
+				const event = decodeStep(step);
+				if (event) {
+					onEvent(event);
+					if (event.kind === "text" || event.kind === "thinking") {
+						streamIdx = step.idx;
+						streamKind = event.kind;
+						streamEmitted = event.text.length;
+					}
+				}
+			} catch {
+				/* drop undecodable step; lastIdx still advances past it */
+			}
+		}
+		flushStreamStep();
+		lastIdx = poller.lastIdx;
+		return steps.length > 0;
+	};
+
+	const result: AgyRunResult = {
+		exitCode: 0,
+		conversationId: null,
+		lastIdx: -1,
+		aborted: false,
+		timedOut: false,
+		stderr: "",
+		durationMs: 0,
+	};
+
+	let proc: ChildProcess | null = null;
+	let settled = false;
+	let timedOut = false;
+	let sigkillTimer: ReturnType<typeof setTimeout> | undefined;
+	let watchdog: ReturnType<typeof setTimeout> | undefined;
+
+	const killTree = () => {
+		try {
+			if (proc?.pid) process.kill(-proc.pid, "SIGTERM");
+		} catch {
+			/* process group already gone */
+		}
+		if (!sigkillTimer) {
+			sigkillTimer = setTimeout(() => {
+				try {
+					if (proc?.pid) process.kill(-proc.pid, "SIGKILL");
+				} catch {
+					/* give up */
+				}
+			}, GRACE_AFTER_TIMEOUT_MS);
+		}
+	};
+
+	const onAbort = () => killTree();
+
+	try {
+		await new Promise<void>((resolveP, rejectP) => {
+			// detached: true so we can signal the whole process group. agy
+			// spawns its own exec subprocesses in -p mode; a direct kill would
+			// orphan those grandchildren.
+			proc = spawn(binary, args, {
+				cwd: opts.cwd,
+				stdio: ["ignore", "ignore", "pipe"],
+				shell: false,
+				detached: true,
+			});
+			proc.stderr?.setEncoding("utf8");
+			proc.stderr?.on("data", (d: string) => (stderr += d));
+
+			// Drive the DB poll concurrently with the running process. THIS is the
+			// streaming: without it, no event reaches the caller until agy exits,
+			// defeating the whole point of the provider. Cleared in cleanup().
+			const pollTimer = setInterval(pollOnce, POLL_INTERVAL_MS);
+			pollTimer.unref?.();
+
+			const cleanup = () => {
+				clearInterval(pollTimer);
+				if (watchdog) clearTimeout(watchdog);
+				if (sigkillTimer) clearTimeout(sigkillTimer);
+				if (opts.signal) opts.signal.removeEventListener("abort", onAbort);
+			};
+
+			// Enforce the timeout ourselves (agy's --print-timeout is advisory).
+			watchdog = setTimeout(() => {
+				timedOut = true;
+				killTree();
+			}, timeoutMin * 60_000);
+
+			if (opts.signal) {
+				if (opts.signal.aborted) killTree();
+				else opts.signal.addEventListener("abort", onAbort, { once: true });
+			}
+
+			const finish = (code: number | null) => {
+				if (settled) return;
+				settled = true;
+				cleanup();
+				resolveP();
+				void code; // exit code read off proc below
+			};
+
+			proc.on("error", (err) => {
+				cleanup();
+				rejectP(err);
+			});
+			proc.on("close", finish);
+			proc.on("exit", finish);
+		});
+	} catch (err) {
+		// spawn ENOENT etc.  -  surface as a non-zero result, don't throw.
+		stderr += err instanceof Error ? err.message : String(err);
+	}
+
+	// Capture abort state before the trailing-poll loop. On cancel we skip the
+	// trailing polls (3 x 100ms) so the stream finalizes promptly with whatever
+	// was already streamed, instead of stalling ~300ms after agy was killed.
+	const wasAborted = !!opts.signal?.aborted;
+
+	// Trailing polls: agy may flush a final step moments after exit. The
+	// agy-acp pattern (3 x 100ms) catches these without adding noticeable latency.
+	if (!wasAborted) {
+		for (let i = 0; i < TRAILING_POLLS; i++) {
+			pollOnce();
+			await sleep(TRAILING_POLL_MS);
+		}
+	}
+
+	// CFA note: poller/proc are assigned inside closures that TS can't track
+	// through the await, so they narrow to `null` here. Casts break the
+	// narrowing without lying about the runtime type.
+	(poller as ConversationPoller | null)?.close();
+
+	result.exitCode = (proc as ChildProcess | null)?.exitCode ?? (stderr ? 1 : 0);
+	result.aborted = wasAborted;
+	result.timedOut = timedOut;
+	result.conversationId = boundId;
+	result.lastIdx = lastIdx;
+	result.stderr = stderr;
+	result.durationMs = Date.now() - start;
+	return result;
+}

package/src/sessions.ts

@@ -0,0 +1,159 @@
+// Persisted map: pi session key -> agy conversation id + last streamed step.
+//
+// agy holds its own conversation history in the DB keyed by conversation id.
+// To resume a multi-turn pi conversation, we thread the same agy id across
+// streamSimple calls and skip steps we already streamed last turn.
+//
+// The key is options.sessionId when pi provides it, else the cwd (single-
+// conversation-per-process fallback). Stored at
+// ~/.pi/agent/antigravity-bridge/sessions.json with atomic rename.
+
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+const STORE_PATH = path.join(
+	os.homedir(),
+	".pi",
+	"agent",
+	"antigravity-bridge",
+	"sessions.json",
+);
+
+export interface AgySession {
+	conversationId: string;
+	lastStepIdx: number;
+	/** pi context.messages length captured at the start of the last agy turn.
+	 *  Used by the provider to compute a delta digest of pi-side context agy
+	 *  was not spawned for (compaction summaries, other-provider turns). 0 on a
+	 *  fresh or pre-watermark session. Optional for backward compat with older
+	 *  persisted sessions that predate the watermark. */
+	lastMessageCount?: number;
+}
+
+type StoreMap = Record<string, AgySession>;
+
+/** Narrow untrusted parsed JSON into a clean StoreMap, validating each entry's
+ *  shape instead of trusting the whole structure (the cast previously let a
+ *  hand-edited or corrupt file plant wrong-typed fields into the cache).
+ *  Drops any entry whose `conversationId` isn't a string; falls back to -1 when
+ *  `lastStepIdx` is missing or not a finite number. Returns {} for non-object
+ *  input. */
+function narrowStoreMap(parsed: unknown): StoreMap {
+	if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return {};
+	const clean: StoreMap = {};
+	for (const [key, value] of Object.entries(parsed as Record<string, unknown>)) {
+		if (!value || typeof value !== "object") continue;
+		const conversationId = (value as { conversationId?: unknown }).conversationId;
+		if (typeof conversationId !== "string") continue;
+		const idx = (value as { lastStepIdx?: unknown }).lastStepIdx;
+		const mc = (value as { lastMessageCount?: unknown }).lastMessageCount;
+		clean[key] = {
+			conversationId,
+			lastStepIdx: typeof idx === "number" && Number.isFinite(idx) ? idx : -1,
+			lastMessageCount: typeof mc === "number" && Number.isFinite(mc) ? mc : 0,
+		};
+	}
+	return clean;
+}
+
+/** Atomic-write JSON session store. Writes are serialized through a promise
+ *  chain so two concurrent turns can't interleave renames.
+ *
+ *  Multi-process safety: persist() re-reads the file fresh and overlays only
+ *  the keys this instance has dirtied since load, so two pi processes editing
+ *  different sessions never clobber each other. (Whole-file overwrite from a
+ *  stale in-memory snapshot was the previous behavior  -  last writer won.) */
+export class SessionStore {
+	private readonly path: string;
+	private cache: StoreMap = {};
+	private readonly dirty = new Set<string>();
+	private writeChain: Promise<void> = Promise.resolve();
+
+	constructor(storePath: string = STORE_PATH) {
+		this.path = storePath;
+		this.load();
+	}
+
+	private load(): void {
+		try {
+			const raw = fs.readFileSync(this.path, "utf8");
+			this.cache = narrowStoreMap(JSON.parse(raw));
+		} catch {
+			// missing or corrupt  -  start empty. The first set() will create it.
+			this.cache = {};
+		}
+	}
+
+	get(key: string): AgySession | null {
+		const s = this.cache[key];
+		if (!s || typeof s.conversationId !== "string") return null;
+		return {
+			conversationId: s.conversationId,
+			lastStepIdx: s.lastStepIdx ?? -1,
+			lastMessageCount: s.lastMessageCount ?? 0,
+		};
+	}
+
+	/** Update one session and persist atomically (temp file + rename). The
+	 *  write is queued; callers don't need to await it. */
+	set(key: string, session: AgySession): void {
+		this.cache[key] = session;
+		this.dirty.add(key);
+		// Serialize through a chain so concurrent sets don't interleave renames.
+		this.writeChain = this.writeChain
+			.then(() => this.persist())
+			.catch(() => {
+				/* swallow; next set() retries */
+			});
+	}
+
+	private async persist(): Promise<void> {
+		const dir = path.dirname(this.path);
+		await fs.promises.mkdir(dir, { recursive: true });
+
+		// Re-read the current file so we merge against the latest on-disk state,
+		// not our possibly-stale in-memory snapshot. Overlay only OUR dirty keys
+		// so a concurrent process's writes to other keys survive.
+		let disk: StoreMap = {};
+		try {
+			const raw = await fs.promises.readFile(this.path, "utf8");
+			disk = narrowStoreMap(JSON.parse(raw));
+		} catch {
+			/* missing or corrupt  -  start from empty */
+		}
+		const merged: StoreMap = { ...disk };
+		for (const key of this.dirty) {
+			merged[key] = this.cache[key];
+		}
+
+		const tmp = `${this.path}.${process.pid}.tmp`;
+		await fs.promises.writeFile(tmp, JSON.stringify(merged, null, 2) + "\n", {
+			mode: 0o600,
+		});
+		await fs.promises.rename(tmp, this.path);
+		this.dirty.clear();
+	}
+
+	/** Wipe all session bindings (forces fresh agy conversations on every
+	 *  active session). Used by the /agy clear command. */
+	clear(): void {
+		this.cache = {};
+		this.dirty.clear();
+		this.writeChain = this.writeChain
+			.then(async () => {
+				const tmp = `${this.path}.${process.pid}.tmp`;
+				await fs.promises.mkdir(path.dirname(this.path), { recursive: true });
+				await fs.promises.writeFile(tmp, "{}\n", { mode: 0o600 });
+				await fs.promises.rename(tmp, this.path);
+			})
+			.catch(() => {
+				/* swallow; user can retry */
+			});
+	}
+
+	/** Count persisted session bindings (for status display). */
+	get size(): number {
+		return Object.keys(this.cache).length;
+	}
+}