dotfiles/pi/.pi/agent/shared/pi-ask-bridge.ts

/**
 * pi-ask-bridge — Unix-domain socket server that pipes ask requests from
 * the pi-ask-mcp subprocess (running inside Claude CLI inside chat-claude)
 * into pi's native ask UI (askSingleQuestionWithInlineNote / askQuestionsWithTabs).
 *
 * Architecture:
 *
 *     pi process
 *       └── chat-claude extension
 *           ├── AskBridge (here) — listens on $PI_ASK_SOCKET
 *           └── claude -p ... --mcp-config <generated>
 *               └── pi-ask-mcp/server.js
 *                   ↳ on tools/call ask → connects to $PI_ASK_SOCKET,
 *                     sends question, awaits answer
 *
 * Lifecycle:  start one bridge per chat-claude session; close on exit.
 * Concurrency: pi's ui.custom overlay is modal, so asks are serialised
 * across all open connections via an internal promise chain.
 *
 * Wire format (NDJSON, one message per line):
 *
 *   request  → { id, type: "ask", questions: [{id, question, options[], multi?, recommended?}, ...] }
 *   response ← { id, type: "result", results: [{id, selectedOptions[], customInput?}, ...] }
 *   error    ← { id, type: "error", message: "…" }
 */

import { createServer, type Server as NetServer, type Socket } from "node:net";
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { join, dirname } from "node:path";
import { tmpdir } from "node:os";
import { fileURLToPath } from "node:url";
import type { ExtensionUIContext } from "@mariozechner/pi-coding-agent";
import { askSingleQuestionWithInlineNote } from "../extensions/pi-ask-tool/ask-inline-ui.js";
import { askQuestionsWithTabs } from "../extensions/pi-ask-tool/ask-tabs-ui.js";
import type { AskQuestion, AskSelection } from "../extensions/pi-ask-tool/ask-logic.js";

// =============================================================================
// Public API
// =============================================================================

export interface AskBridge {
	/** Path to the generated --mcp-config JSON, suitable for `claude --mcp-config`. */
	mcpConfigPath: string;
	/** Absolute path to the underlying Unix socket (informational). */
	socketPath: string;
	/** How many ask requests this bridge has served so far. */
	count(): number;
	/** Stop accepting connections, remove socket + temp dir. Idempotent. */
	close(): void;
}

export interface StartAskBridgeOptions {
	/** pi UI context (must come from an interactive session). */
	ui: ExtensionUIContext;
	/**
	 * Absolute path to extensions/pi-ask-mcp/server.js. Auto-derived from
	 * import.meta.url when omitted (assumes the conventional layout).
	 */
	mcpServerEntry?: string;
	/**
	 * MCP server name surfaced in the tool prefix. Defaults to "pi", which
	 * makes the tool name `mcp__pi__ask` in Claude's tool stream.
	 */
	serverName?: string;
	/** Optional notification fired whenever a new ask is served. */
	onAsk?: (totalSoFar: number) => void;
}

export function startAskBridge(opts: StartAskBridgeOptions): AskBridge {
	const dir  = mkdtempSync(join(tmpdir(), "pi-ask-"));   // 0700 perms
	const sock = join(dir, "ask.sock");
	let askCount = 0;
	let closed = false;

	const server: NetServer = createServer((conn) =>
		handleConnection(conn, opts, () => {
			askCount += 1;
			opts.onAsk?.(askCount);
			return askCount;
		}),
	);
	server.on("error", () => { /* socket disappeared, etc. — bridge is single-tenant, ignore */ });
	server.listen(sock);

	const mcpEntry   = opts.mcpServerEntry ?? defaultMcpEntry();
	const serverName = opts.serverName     ?? "pi";
	const cfgPath    = join(dir, "mcp.json");
	writeFileSync(cfgPath, JSON.stringify({
		mcpServers: {
			[serverName]: {
				command: "node",
				args:    [mcpEntry],
				env:     { PI_ASK_SOCKET: sock },
			},
		},
	}));

	return {
		socketPath:    sock,
		mcpConfigPath: cfgPath,
		count: () => askCount,
		close: () => {
			if (closed) return;
			closed = true;
			try { server.close(); } catch { /* noop */ }
			try { rmSync(dir, { recursive: true, force: true }); } catch { /* noop */ }
		},
	};
}

// =============================================================================
// Internals
// =============================================================================

function defaultMcpEntry(): string {
	// shared/pi-ask-bridge.ts → ../extensions/pi-ask-mcp/server.js
	const here = dirname(fileURLToPath(import.meta.url));
	return join(here, "..", "extensions", "pi-ask-mcp", "server.js");
}

// pi.ui.custom is modal — only one overlay can be on screen at a time.
// Serialise asks across ALL connections via this single promise chain.
let askChain: Promise<unknown> = Promise.resolve();

function handleConnection(
	conn: Socket,
	opts: StartAskBridgeOptions,
	bumpCount: () => number,
) {
	let buf = "";
	conn.on("data", (data) => {
		buf += data.toString();
		let nl = buf.indexOf("\n");
		while (nl >= 0) {
			const line = buf.slice(0, nl);
			buf = buf.slice(nl + 1);
			handleLine(line, conn, opts, bumpCount);
			nl = buf.indexOf("\n");
		}
	});
	conn.on("error", () => { /* peer might disappear if Claude is killed mid-flight */ });
}

function handleLine(
	line: string,
	conn: Socket,
	opts: StartAskBridgeOptions,
	bumpCount: () => number,
) {
	if (!line.trim()) return;
	let msg: any;
	try { msg = JSON.parse(line); } catch { return; }
	if (msg.type !== "ask") return;
	const id        = String(msg.id ?? "");
	const questions = Array.isArray(msg.questions) ? (msg.questions as AskQuestion[]) : [];

	askChain = askChain.then(async () => {
		bumpCount();
		try {
			const results = await askViaPiUI(opts.ui, questions);
			writeReply(conn, { id, type: "result", results });
		} catch (err) {
			writeReply(conn, {
				id,
				type:    "error",
				message: err instanceof Error ? err.message : String(err),
			});
		}
	});
}

function writeReply(conn: Socket, msg: unknown) {
	try { conn.write(JSON.stringify(msg) + "\n"); conn.end(); } catch { /* gone */ }
}

interface QuestionResult {
	id: string;
	selectedOptions: string[];
	customInput?: string;
}

async function askViaPiUI(
	ui: ExtensionUIContext,
	questions: AskQuestion[],
): Promise<QuestionResult[]> {
	if (questions.length === 0) return [];

	if (questions.length === 1 && !questions[0].multi) {
		const sel: AskSelection = await askSingleQuestionWithInlineNote(ui, questions[0]);
		return [toResult(questions[0], sel)];
	}

	const tab = await askQuestionsWithTabs(ui, questions);
	return questions.map((q, i) => toResult(q, tab.selections[i] ?? { selectedOptions: [] }));
}

function toResult(q: AskQuestion, sel: AskSelection): QuestionResult {
	const out: QuestionResult = { id: q.id, selectedOptions: [...sel.selectedOptions] };
	if (sel.customInput) out.customInput = sel.customInput;
	return out;
}