src/lib/server/agent/PiRpcClient.ts

import { randomUUID } from "node:crypto";
import { WebSocket, type RawData } from "ws";

export type PiRpcEvent = Record<string, unknown> & { type?: string };

export interface PiRpcReplayGap {
	expected: number;
	received: number;
	earliestAvailable?: number;
	latestAvailable?: number;
}

export const PI_RPC_REPLAY_GAP_MESSAGE =
	"Agent output was truncated because part of the sandbox event replay was no longer available. This generation was marked as failed to avoid saving an incomplete transcript.";

export class PiRpcReplayGapError extends Error {
	readonly name = "PiRpcReplayGapError";

	constructor(readonly gap: PiRpcReplayGap) {
		super(PI_RPC_REPLAY_GAP_MESSAGE);
	}
}

interface PendingCommand {
	resolve: (response: PiRpcResponse) => void;
	reject: (error: Error) => void;
	timer: ReturnType<typeof setTimeout>;
}

export interface PiRpcResponse extends Record<string, unknown> {
	id?: string;
	type: "response";
	command?: string;
	success: boolean;
	data?: unknown;
	error?: string;
}

/** A correlated, durable command rejection (as opposed to an ambiguous transport failure). */
export class PiRpcCommandRejectedError extends Error {
	readonly name = "PiRpcCommandRejectedError";

	constructor(readonly response: PiRpcResponse) {
		super(response.error ?? `Pi RPC command '${response.command ?? "unknown"}' failed`);
	}
}

export interface PiRpcClientOptions {
	headers: Record<string, string>;
	lastSequence?: number;
	connectTimeoutMs?: number;
	commandTimeoutMs?: number;
	minReconnectDelayMs?: number;
	maxReconnectDelayMs?: number;
	onGap?: (gap: PiRpcReplayGap) => void;
	onStatus?: (status: "connecting" | "connected" | "reconnecting" | "closed") => void;
}

type EventListener = (event: PiRpcEvent, sequence: number) => void;
type ReplayGapListener = (error: PiRpcReplayGapError) => void;

/**
 * Reconnecting RPC client for the sandbox bridge. Commands use pi's native
 * correlation ids while bridge sequence numbers make reconnect replay exact.
 */
export class PiRpcClient {
	private socket: WebSocket | undefined;
	private readonly listeners = new Set<EventListener>();
	private readonly replayGapListeners = new Set<ReplayGapListener>();
	private readonly pending = new Map<string, PendingCommand>();
	private reconnectTimer: ReturnType<typeof setTimeout> | undefined;
	private socketReadyTimer: { socket: WebSocket; timer: ReturnType<typeof setTimeout> } | undefined;
	private connectionDeadlineTimer: ReturnType<typeof setTimeout> | undefined;
	private connectionPromise: Promise<void> | undefined;
	private connectionResolve: (() => void) | undefined;
	private connectionReject: ((error: Error) => void) | undefined;
	private readySocket: WebSocket | undefined;
	private replayTarget: { socket: WebSocket; sequence: number } | undefined;
	private replayGapFailure: PiRpcReplayGapError | undefined;
	private lastConnectionError: Error | undefined;
	private hasConnected = false;
	private disposed = false;
	private reconnectAttempts = 0;
	private lastSequenceValue: number;

	constructor(
		private readonly bridgeUrl: string,
		private readonly options: PiRpcClientOptions
	) {
		this.lastSequenceValue = options.lastSequence ?? 0;
	}

	get lastSequence(): number {
		return this.lastSequenceValue;
	}

	get connected(): boolean {
		return (
			this.socket !== undefined &&
			this.readySocket === this.socket &&
			this.socket.readyState === WebSocket.OPEN
		);
	}

	get replayGapError(): PiRpcReplayGapError | undefined {
		return this.replayGapFailure;
	}

	async connect(): Promise<void> {
		if (this.replayGapFailure) throw this.replayGapFailure;
		if (this.disposed) throw new Error("Pi RPC client is closed");
		if (this.connected) return;
		if (this.connectionPromise) return this.connectionPromise;
		this.connectionPromise = new Promise<void>((resolve, reject) => {
			this.connectionResolve = resolve;
			this.connectionReject = reject;
		});
		const connectionPromise = this.connectionPromise;
		const timeoutMs = this.options.connectTimeoutMs ?? 15_000;
		this.connectionDeadlineTimer = setTimeout(() => {
			const detail = this.lastConnectionError?.message;
			const stopInitialRetries = !this.hasConnected;
			this.rejectConnection(
				new Error(
					`Timed out connecting to pi RPC bridge after ${timeoutMs}ms${detail ? `: ${detail}` : ""}`
				)
			);
			if (stopInitialRetries && this.reconnectTimer) clearTimeout(this.reconnectTimer);
			if (stopInitialRetries) this.reconnectTimer = undefined;
			if (this.socket && this.readySocket !== this.socket) this.socket.terminate();
		}, timeoutMs);
		if (!this.socket && !this.reconnectTimer) this.openSocket(false);
		return connectionPromise;
	}

	subscribe(listener: EventListener): () => void {
		this.listeners.add(listener);
		return () => this.listeners.delete(listener);
	}

	subscribeReplayGap(listener: ReplayGapListener): () => void {
		if (this.replayGapFailure) {
			listener(this.replayGapFailure);
			return () => undefined;
		}
		this.replayGapListeners.add(listener);
		return () => this.replayGapListeners.delete(listener);
	}

	async command<TData = unknown>(
		command: Record<string, unknown> & { type: string },
		options: { timeoutMs?: number } = {}
	): Promise<PiRpcResponse & { data?: TData }> {
		await this.connect();
		if (this.replayGapFailure) throw this.replayGapFailure;
		const socket = this.socket;
		if (!socket || socket.readyState !== WebSocket.OPEN) {
			throw new Error("Pi RPC bridge is not connected");
		}

		const id = typeof command.id === "string" ? command.id : randomUUID();
		const timeoutMs = options.timeoutMs ?? this.options.commandTimeoutMs ?? 30_000;
		const response = new Promise<PiRpcResponse>((resolve, reject) => {
			const timer = setTimeout(() => {
				this.pending.delete(id);
				reject(new Error(`Pi RPC command '${command.type}' timed out after ${timeoutMs}ms`));
			}, timeoutMs);
			this.pending.set(id, { resolve, reject, timer });
		});

		try {
			socket.send(JSON.stringify({ ...command, id }));
		} catch (error) {
			const pending = this.pending.get(id);
			if (pending) {
				clearTimeout(pending.timer);
				this.pending.delete(id);
			}
			throw error;
		}

		const result = (await response) as PiRpcResponse & { data?: TData };
		if (!result.success) {
			throw new PiRpcCommandRejectedError(result);
		}
		return result;
	}

	/**
	 * Send a protocol message that intentionally has no correlated response.
	 * Legacy callers may use this for request-id-only `extension_ui_response`;
	 * durable Agent controls use the patched correlated form through `command`.
	 */
	async send(message: Record<string, unknown> & { type: string }): Promise<void> {
		await this.connect();
		if (this.replayGapFailure) throw this.replayGapFailure;
		const socket = this.socket;
		if (!socket || socket.readyState !== WebSocket.OPEN) {
			throw new Error("Pi RPC bridge is not connected");
		}
		socket.send(JSON.stringify(message));
	}

	waitFor(
		predicate: (event: PiRpcEvent) => boolean,
		options: { signal?: AbortSignal; timeoutMs?: number } = {}
	): Promise<{ event: PiRpcEvent; sequence: number }> {
		if (this.replayGapFailure) return Promise.reject(this.replayGapFailure);
		return new Promise((resolve, reject) => {
			let timer: ReturnType<typeof setTimeout> | undefined;
			const cleanup = () => {
				this.listeners.delete(listener);
				this.replayGapListeners.delete(onReplayGap);
				if (timer) clearTimeout(timer);
				options.signal?.removeEventListener("abort", onAbort);
			};
			const listener: EventListener = (event, sequence) => {
				if (!predicate(event)) return;
				cleanup();
				resolve({ event, sequence });
			};
			const onAbort = () => {
				cleanup();
				reject(new DOMException("The operation was aborted", "AbortError"));
			};
			const onReplayGap: ReplayGapListener = (error) => {
				cleanup();
				reject(error);
			};
			this.listeners.add(listener);
			this.replayGapListeners.add(onReplayGap);
			if (this.replayGapFailure) return onReplayGap(this.replayGapFailure);
			if (options.signal?.aborted) return onAbort();
			options.signal?.addEventListener("abort", onAbort, { once: true });
			if (options.timeoutMs) {
				timer = setTimeout(() => {
					cleanup();
					reject(new Error(`Timed out waiting for pi event after ${options.timeoutMs}ms`));
				}, options.timeoutMs);
			}
		});
	}

	close(): void {
		if (this.disposed) return;
		this.disposed = true;
		if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
		this.clearSocketReadyTimer();
		if (this.connectionDeadlineTimer) clearTimeout(this.connectionDeadlineTimer);
		this.rejectConnection(new Error("Pi RPC client closed"));
		for (const [id, pending] of this.pending) {
			clearTimeout(pending.timer);
			pending.reject(new Error(`Pi RPC client closed before response '${id}'`));
		}
		this.pending.clear();
		if (this.socket?.readyState === WebSocket.OPEN) this.socket.close(1000, "Client closed");
		else this.socket?.terminate();
		this.socket = undefined;
		this.readySocket = undefined;
		this.replayTarget = undefined;
		this.options.onStatus?.("closed");
	}

	private openSocket(reconnecting: boolean): void {
		if (this.disposed || this.replayGapFailure || this.socket) return;
		if (!reconnecting) this.options.onStatus?.("connecting");
		let socket: WebSocket;
		try {
			const url = new URL(this.bridgeUrl);
			url.searchParams.set("fromSeq", String(this.lastSequenceValue));
			socket = new WebSocket(url, {
				headers: this.options.headers,
				handshakeTimeout: this.options.connectTimeoutMs ?? 15_000,
			});
		} catch (error) {
			this.lastConnectionError = asError(error, "Unable to open pi RPC bridge");
			this.scheduleReconnect();
			return;
		}
		this.socket = socket;
		this.socketReadyTimer = {
			socket,
			timer: setTimeout(() => {
				if (this.socket !== socket || this.readySocket === socket) return;
				this.lastConnectionError = new Error("Pi RPC bridge readiness timed out");
				socket.terminate();
			}, this.options.connectTimeoutMs ?? 15_000),
		};

		socket.on("message", (raw: RawData) => this.handleMessage(socket, raw));
		socket.on("error", (error: Error) => {
			if (this.socket === socket) this.lastConnectionError = error;
		});
		socket.on("close", () => {
			this.clearSocketReadyTimer(socket);
			if (this.socket !== socket) return;
			this.socket = undefined;
			this.readySocket = undefined;
			if (this.replayTarget?.socket === socket) this.replayTarget = undefined;
			if (this.disposed || this.replayGapFailure) return;
			this.scheduleReconnect();
		});
	}

	private handleMessage(socket: WebSocket, raw: RawData): void {
		if (this.socket !== socket) return;
		let envelope: Record<string, unknown>;
		try {
			envelope = JSON.parse(raw.toString()) as Record<string, unknown>;
		} catch {
			return;
		}

		if (envelope.type === "bridge_ready") {
			const earliest = asInteger(envelope.fromSeq);
			const latest = asInteger(envelope.toSeq);
			if (earliest === undefined || latest === undefined) {
				this.lastConnectionError = new Error("Pi RPC bridge sent invalid replay metadata");
				socket.terminate();
				return;
			}
			if (latest < this.lastSequenceValue) {
				this.failReplayGap({
					expected: this.lastSequenceValue + 1,
					received: latest,
					earliestAvailable: earliest,
					latestAvailable: latest,
				});
				return;
			}
			if (earliest > this.lastSequenceValue + 1) {
				this.failReplayGap({
					expected: this.lastSequenceValue + 1,
					received: earliest,
					earliestAvailable: earliest,
					latestAvailable: latest,
				});
				return;
			}
			this.replayTarget = { socket, sequence: latest };
			if (latest <= this.lastSequenceValue) this.markSocketReady(socket);
			return;
		}

		if (envelope.type !== "bridge_event") return;
		const sequence = asInteger(envelope.seq);
		const payload = envelope.payload;
		if (sequence === undefined || !isRecord(payload)) return;
		if (sequence <= this.lastSequenceValue) return;
		if (sequence > this.lastSequenceValue + 1) {
			this.failReplayGap({
				expected: this.lastSequenceValue + 1,
				received: sequence,
				latestAvailable: this.replayTarget?.sequence,
			});
			return;
		}
		this.lastSequenceValue = sequence;

		if (payload.type === "response" && typeof payload.id === "string") {
			const pending = this.pending.get(payload.id);
			if (pending) {
				clearTimeout(pending.timer);
				this.pending.delete(payload.id);
				pending.resolve(payload as PiRpcResponse);
			}
		}
		for (const listener of this.listeners) listener(payload, sequence);
		if (
			this.replayTarget?.socket === socket &&
			this.lastSequenceValue >= this.replayTarget.sequence
		) {
			this.markSocketReady(socket);
		}
	}

	private scheduleReconnect(): void {
		if (
			this.disposed ||
			this.replayGapFailure ||
			this.reconnectTimer ||
			(!this.hasConnected && !this.connectionPromise)
		)
			return;
		const minimum = this.options.minReconnectDelayMs ?? 250;
		const maximum = this.options.maxReconnectDelayMs ?? 5_000;
		const delay = Math.min(maximum, minimum * 2 ** this.reconnectAttempts++);
		this.options.onStatus?.("reconnecting");
		this.reconnectTimer = setTimeout(() => {
			this.reconnectTimer = undefined;
			this.openSocket(true);
		}, delay);
	}

	private rejectConnection(error: Error): void {
		const reject = this.connectionReject;
		this.clearConnectionPromise();
		reject?.(error);
	}

	private resolveConnection(): void {
		const resolve = this.connectionResolve;
		this.clearConnectionPromise();
		resolve?.();
	}

	private markSocketReady(socket: WebSocket): void {
		if (this.socket !== socket || this.replayGapFailure) return;
		this.clearSocketReadyTimer(socket);
		if (this.replayTarget?.socket === socket) this.replayTarget = undefined;
		this.readySocket = socket;
		this.hasConnected = true;
		this.lastConnectionError = undefined;
		this.reconnectAttempts = 0;
		this.options.onStatus?.("connected");
		this.resolveConnection();
	}

	private failReplayGap(gap: PiRpcReplayGap): void {
		if (this.replayGapFailure) return;
		const error = new PiRpcReplayGapError(gap);
		this.replayGapFailure = error;
		this.lastConnectionError = error;
		if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
		this.reconnectTimer = undefined;
		this.clearSocketReadyTimer();
		this.replayTarget = undefined;

		try {
			this.options.onGap?.(gap);
		} catch {
			// Observability callbacks must not weaken the fail-closed boundary.
		}
		for (const listener of this.replayGapListeners) {
			try {
				listener(error);
			} catch {
				// One consumer must not prevent the others from observing the gap.
			}
		}
		this.replayGapListeners.clear();
		this.rejectConnection(error);
		for (const pending of this.pending.values()) {
			clearTimeout(pending.timer);
			pending.reject(error);
		}
		this.pending.clear();

		const socket = this.socket;
		this.socket = undefined;
		this.readySocket = undefined;
		if (socket?.readyState === WebSocket.OPEN) {
			try {
				socket.send(JSON.stringify({ type: "abort" }));
			} catch {
				// The replay gap remains terminal even if the best-effort abort cannot be sent.
			}
			try {
				socket.close(1011, "Replay gap");
			} catch {
				socket.terminate();
			}
		} else {
			socket?.terminate();
		}
		this.options.onStatus?.("closed");
	}

	private clearSocketReadyTimer(socket?: WebSocket): void {
		const active = this.socketReadyTimer;
		if (!active || (socket && active.socket !== socket)) return;
		clearTimeout(active.timer);
		this.socketReadyTimer = undefined;
	}

	private clearConnectionPromise(): void {
		if (this.connectionDeadlineTimer) clearTimeout(this.connectionDeadlineTimer);
		this.connectionDeadlineTimer = undefined;
		this.connectionPromise = undefined;
		this.connectionResolve = undefined;
		this.connectionReject = undefined;
	}
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function asError(value: unknown, fallback: string): Error {
	return value instanceof Error ? value : new Error(fallback);
}

function asInteger(value: unknown): number | undefined {
	return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined;
}