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;
	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",
]);

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

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

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

	pi.stderr.on("data", (data) => process.stderr.write(`[pi] ${data}`));
	pi.on("exit", (code) => {
		console.log(`[webui] pi 退出, code=${code}`);
		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 line =
			JSON.stringify({
				type: "prompt",
				message: options.message,
				images: options.images,
				id,
			}) + "\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 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,
		getRunSnapshot,
		connectSseClient,
		removeSseClient: (res) => sseClients.delete(res),
		child: pi,
	};
}