sproutclaw/.pi/agent/extensions/webui/backend/rpc/pi-client.ts

import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
import type { ServerResponse } from "node:http";
import { resolvePiRpcLaunch, type WebUiPaths } from "../config/paths.ts";
import type { RunSnapshot, SendCmd, SubmitPromptOptions } from "../types/context.ts";

export interface PiClient {
	sendCmd: SendCmd;
	submitPrompt: (options: SubmitPromptOptions) => void;
	sendExtensionUiResponse: (id: string, response: Record<string, unknown>) => void;
	getRunSnapshot: () => RunSnapshot;
	connectSseClient: (res: ServerResponse) => void;
	removeSseClient: (res: ServerResponse) => void;
	child: ChildProcessWithoutNullStreams;
}

const BUFFERED_EVENT_TYPES = new Set([
	"agent_start",
	"agent_end",
	"message_start",
	"message_update",
	"message_end",
	"tool_execution_start",
	"tool_execution_update",
	"tool_execution_end",
	"compaction_start",
	"compaction_end",
	"queue_update",
	"auto_retry_start",
	"auto_retry_end",
	"bash_update",
]);

const DEFAULT_CMD_TIMEOUT_MS = 60_000;
const PROMPT_PREFLIGHT_TIMEOUT_MS = 30_000;
const PI_STDERR_BUFFER_MAX = 8192;

function formatSpawnArgs(args: string[]): string {
	return args.map((arg) => JSON.stringify(arg)).join(" ");
}

function summarizeStderr(buffer: string): string {
	const trimmed = buffer.trim();
	if (!trimmed) return "(无 stderr 输出)";
	const lines = trimmed.split(/\r?\n/).filter(Boolean);
	return lines.slice(-8).join("\n");
}

export function createPiClient(paths: WebUiPaths, onExit: (code: number | null) => void): PiClient {
	const piLaunch = resolvePiRpcLaunch(paths);
	console.log(
		`[webui] 启动 pi RPC (${piLaunch.mode}): ${piLaunch.command} ${formatSpawnArgs(piLaunch.args)}`,
	);

	const pi = spawn(piLaunch.command, piLaunch.args, {
		cwd: paths.repoRoot,
		stdio: ["pipe", "pipe", "pipe"],
		env: {
			...process.env,
			PI_CODING_AGENT_DIR: paths.agentDir,
		},
	});

	// Keep stdin pipe open; closing it on Windows can trigger RPC shutdown via stdin "end".
	pi.stdin.write("");

	let stderrBuffer = "";
	pi.stderr.on("data", (data: Buffer) => {
		const chunk = data.toString();
		stderrBuffer = (stderrBuffer + chunk).slice(-PI_STDERR_BUFFER_MAX);
		process.stderr.write(`[pi] ${data}`);
	});
	pi.on("exit", (code) => {
		console.log(`[webui] pi 退出, code=${code}`);
		if (code !== 0 && code !== null) {
			console.error(`[webui] pi RPC 启动失败 (code=${code}):\n${summarizeStderr(stderrBuffer)}`);
		}
		onExit(code);
	});

	let buffer = "";
	const pending = new Map<string, { resolve: (value: any) => void; reject: (error: Error) => void }>();
	const sseClients = new Set<ServerResponse>();
	let reqId = 0;

	let isStreaming = false;
	let sessionFile: string | undefined;
	let turnEventBuffer: Record<string, unknown>[] = [];

	function getRunSnapshot(): RunSnapshot {
		return {
			isStreaming,
			sessionFile,
			replay: [...turnEventBuffer],
		};
	}

	function broadcastSse(msg: Record<string, unknown>): void {
		const data = `data: ${JSON.stringify(msg)}\n\n`;
		for (const res of sseClients) res.write(data);
	}

	function trackAgentEvent(msg: Record<string, unknown>): void {
		const type = msg.type;
		if (typeof type !== "string") return;

		if (type === "agent_start") {
			isStreaming = true;
			turnEventBuffer = [msg];
			return;
		}

		if (type === "agent_end") {
			turnEventBuffer.push(msg);
			isStreaming = false;
			turnEventBuffer = [];
			return;
		}

		if (isStreaming && BUFFERED_EVENT_TYPES.has(type)) {
			turnEventBuffer.push(msg);
		}
	}

	function registerPending(
		id: string,
		handlers: { resolve: (value: any) => void; reject: (error: Error) => void },
		timeoutMs: number,
		timeoutLabel: string,
	): void {
		pending.set(id, handlers);
		setTimeout(() => {
			if (!pending.has(id)) return;
			pending.delete(id);
			handlers.reject(new Error(`命令超时: ${timeoutLabel}`));
		}, timeoutMs);
	}

	function onLine(line: string): void {
		if (!line.trim()) return;
		try {
			const msg = JSON.parse(line) as Record<string, unknown>;
			if (msg.type === "response" && msg.id && pending.has(String(msg.id))) {
				const id = String(msg.id);
				const p = pending.get(id)!;
				pending.delete(id);
				if (msg.command === "prompt" && msg.success === false) {
					broadcastSse({
						type: "prompt_rejected",
						error: String(msg.error || "prompt rejected"),
					});
				}
				if (msg.command === "get_state" && msg.success === true) {
					const data = msg.data as { sessionFile?: string } | undefined;
					if (data?.sessionFile) {
						sessionFile = data.sessionFile;
					}
				}
				p.resolve(msg);
				return;
			}
			trackAgentEvent(msg);
			broadcastSse(msg);
		} catch {
			/* ignore non-JSON lines */
		}
	}

	pi.stdout.on("data", (chunk: Buffer) => {
		buffer += chunk.toString();
		const lines = buffer.split("\n");
		buffer = lines.pop() || "";
		for (const line of lines) onLine(line);
	});

	const sendCmd: SendCmd = (command) =>
		new Promise((resolvePromise, reject) => {
			const id = `req_${++reqId}`;
			const line = JSON.stringify({ ...command, id }) + "\n";
			registerPending(
				id,
				{ resolve: resolvePromise, reject },
				DEFAULT_CMD_TIMEOUT_MS,
				String(command.type),
			);
			pi.stdin.write(line);
		});

	function submitPrompt(options: SubmitPromptOptions): void {
		const id = `req_${++reqId}`;
		const payload: Record<string, unknown> = {
			type: "prompt",
			message: options.message,
			images: options.images,
			id,
		};
		if (options.streamingBehavior) {
			payload.streamingBehavior = options.streamingBehavior;
		}
		const line = JSON.stringify(payload) + "\n";
		registerPending(
			id,
			{
				resolve: () => {
					/* preflight success: agent events arrive via SSE */
				},
				reject: (err) => {
					broadcastSse({ type: "prompt_rejected", error: err.message });
				},
			},
			PROMPT_PREFLIGHT_TIMEOUT_MS,
			"prompt",
		);
		pi.stdin.write(line);
	}

	function sendExtensionUiResponse(id: string, response: Record<string, unknown>): void {
		const line = JSON.stringify({ type: "extension_ui_response", id, ...response }) + "\n";
		pi.stdin.write(line);
	}

	function connectSseClient(res: ServerResponse): void {
		const snapshot = getRunSnapshot();
		res.write(
			`data: ${JSON.stringify({
				type: "connected",
				isStreaming: snapshot.isStreaming,
				replay: snapshot.replay,
			})}\n\n`,
		);
		sseClients.add(res);
	}

	return {
		sendCmd,
		submitPrompt,
		sendExtensionUiResponse,
		getRunSnapshot,
		connectSseClient,
		removeSseClient: (res) => sseClients.delete(res),
		child: pi,
	};
}