Merge PR #6330: feat(coding-agent): add lossless paged RPC transport (@wolfiesch)

packages/coding-agent/src/modes/rpc/rpc-client.ts

@@ -12,6 +12,8 @@ import { isRecord, ptree, readJsonl } from "@oh-my-pi/pi-utils";
 import type { FileSink } from "bun";
 import type { BashResult } from "../../exec/bash-executor";
 import type { AgentSessionEvent, SessionStats } from "../../session/agent-session";
+import { MAX_RPC_FRAME_BYTES, MAX_RPC_REASSEMBLED_BYTES, RpcFrameDecoder, type RpcProtocolVersion } from "./rpc-frame";
+import { RPC_MESSAGES_PAGE_BUSY_ERROR, type RpcMessagesPage, type RpcMessagesPageOptions } from "./rpc-messages";
 import type {
 	RpcAvailableCommandsUpdateFrame,
 	RpcAvailableSlashCommand,
@@ -136,6 +138,16 @@ function isRpcResponse(value: unknown): value is RpcResponse {
 	return true;
 }
 
+function supportsRpcProtocolV2(value: Record<string, unknown>): boolean {
+	return (
+		value.type === "ready" &&
+		Array.isArray(value.supportedProtocolVersions) &&
+		value.supportedProtocolVersions.includes(2) &&
+		value.maxFrameBytes === MAX_RPC_FRAME_BYTES &&
+		value.maxReassembledFrameBytes === MAX_RPC_REASSEMBLED_BYTES
+	);
+}
+
 function isAgentEvent(value: unknown): value is AgentEvent {
 	if (!isRecord(value)) return false;
 	const type = value.type;
@@ -218,6 +230,7 @@ export class RpcClient {
 	#customTools: RpcClientCustomTool[] = [];
 	#pendingHostToolCalls = new Map<string, { controller: AbortController }>();
 	#requestId = 0;
+	#protocolVersion: RpcProtocolVersion = 1;
 	#extensionUiListeners: Set<(req: RpcExtensionUIRequest) => void> = new Set();
 	#abortController = new AbortController();
 
@@ -242,6 +255,7 @@ export class RpcClient {
 		// Mint a fresh controller so a previous stop()'s abort does not
 		// short-circuit the new stdout reader (issue #4079).
 		this.#abortController = new AbortController();
+		this.#protocolVersion = 1;
 
 		const cliPath = this.options.cliPath ?? "dist/cli.js";
 		const args = ["--mode", "rpc"];
@@ -269,6 +283,9 @@ export class RpcClient {
 		// Wait for the "ready" signal or process exit
 		const { promise: readyPromise, resolve: readyResolve, reject: readyReject } = Promise.withResolvers<void>();
 		let readySettled = false;
+		let protocolV2Supported = false;
+		let protocolV2Enabled = false;
+		const frameDecoder = new RpcFrameDecoder();
 
 		const reapAfterOutputFailure = async (error: Error) => {
 			if (this.#process !== child) return;
@@ -294,11 +311,15 @@ export class RpcClient {
 		void (async () => {
 			for await (const line of lines) {
 				if (!readySettled && isRecord(line) && line.type === "ready") {
+					protocolV2Supported = supportsRpcProtocolV2(line);
 					readySettled = true;
 					readyResolve();
 					continue;
 				}
-				this.#handleLine(line);
+				if (isRecord(line) && line.type === "rpc_chunk" && !protocolV2Enabled)
+					throw new Error("RPC chunk received before protocol negotiation");
+				const decoded = frameDecoder.push(line);
+				if (decoded) this.#handleLine(decoded);
 			}
 			// A closed stdout is terminal even if the child remains alive. Startup
 			// failures are reaped by the readyPromise catch below; established
@@ -359,6 +380,18 @@ export class RpcClient {
 
 		try {
 			await readyPromise;
+			if (protocolV2Supported) {
+				protocolV2Enabled = true;
+				const response = await this.#send({ type: "negotiate_protocol", protocolVersion: 2 });
+				if (
+					!response.success ||
+					response.command !== "negotiate_protocol" ||
+					!isRecord(response.data) ||
+					response.data.protocolVersion !== 2
+				)
+					throw new Error("RPC protocol v2 negotiation failed");
+				this.#protocolVersion = 2;
+			}
 			if (this.#customTools.length > 0) {
 				await this.setCustomTools(this.#customTools);
 			}
@@ -744,9 +777,42 @@ export class RpcClient {
 	}
 
 	/**
-	 * Get all messages in the session.
+	 * Get one stable, byte-bounded message page.
 	 */
+	async getMessagesPage(options: RpcMessagesPageOptions = {}): Promise<RpcMessagesPage> {
+		const response = await this.#send({ type: "get_messages_page", ...options });
+		return this.#getData<RpcMessagesPage>(response);
+	}
+
+	/** Get all messages, draining stable pages when protocol v2 is available. */
 	async getMessages(): Promise<AgentMessage[]> {
+		if (this.#protocolVersion === 2) {
+			try {
+				const messages: AgentMessage[] = [];
+				const seenCursors = new Set<string>();
+				let totalMessages: number | undefined;
+				let cursor: string | undefined;
+				do {
+					const page = await this.getMessagesPage({ cursor, limit: 256 });
+					if (
+						!Number.isSafeInteger(page.totalMessages) ||
+						page.totalMessages < 0 ||
+						(totalMessages !== undefined && page.totalMessages !== totalMessages)
+					)
+						throw new Error("RPC message pagination returned an inconsistent total");
+					totalMessages = page.totalMessages;
+					messages.push(...page.messages);
+					cursor = page.nextCursor;
+					if (cursor && seenCursors.has(cursor)) throw new Error("RPC message pagination repeated a cursor");
+					if (cursor) seenCursors.add(cursor);
+				} while (cursor);
+				if (messages.length !== totalMessages)
+					throw new Error("RPC message pagination ended before the advertised total");
+				return messages;
+			} catch (error) {
+				if (!(error instanceof Error) || error.message !== RPC_MESSAGES_PAGE_BUSY_ERROR) throw error;
+			}
+		}
 		const response = await this.#send({ type: "get_messages" });
 		return this.#getData<{ messages: AgentMessage[] }>(response).messages;
 	}

packages/coding-agent/src/modes/rpc/rpc-frame.ts

@@ -1,8 +1,24 @@
 import { isDeepStrictEqual } from "node:util";
 import { isRecord } from "@oh-my-pi/pi-utils";
+import type { RpcChunkFrame } from "./rpc-types";
 
 /** Maximum UTF-8 size of one newline-delimited RPC frame, including the newline. */
 export const MAX_RPC_FRAME_BYTES = 1024 * 1024;
+/** Maximum UTF-8 size of one logical frame reassembled by protocol v2. */
+export const MAX_RPC_REASSEMBLED_BYTES = 64 * 1024 * 1024;
+
+const RPC_CHUNK_PAYLOAD_BYTES = 256 * 1024;
+
+export type RpcProtocolVersion = 1 | 2;
+
+interface PendingRpcChunks {
+	chunkId: string;
+	count: number;
+	byteLength: number;
+	nextIndex: number;
+	chunks: Buffer[];
+	receivedBytes: number;
+}
 
 interface ShrinkPass {
 	stringCap: number;
@@ -68,6 +84,103 @@ function encodedMessageSnapshot(encoded: string): { message: unknown } | undefin
 		: undefined;
 }
 
+function encodeChunkedRpcFrame(frame: object, chunkId: string): string {
+	const json = JSON.stringify(frame);
+	const bytes = Buffer.from(json, "utf8");
+	if (bytes.byteLength > MAX_RPC_REASSEMBLED_BYTES) return `${JSON.stringify(overflowFrame(frame))}\n`;
+	const count = Math.ceil(bytes.byteLength / RPC_CHUNK_PAYLOAD_BYTES);
+	let encoded = "";
+	for (let index = 0; index < count; index++) {
+		const chunk: RpcChunkFrame = {
+			type: "rpc_chunk",
+			chunkId,
+			index,
+			count,
+			byteLength: bytes.byteLength,
+			data: bytes
+				.subarray(index * RPC_CHUNK_PAYLOAD_BYTES, (index + 1) * RPC_CHUNK_PAYLOAD_BYTES)
+				.toString("base64"),
+		};
+		const line = `${JSON.stringify(chunk)}\n`;
+		if (serializedFrameBytes(line.slice(0, -1)) > MAX_RPC_FRAME_BYTES)
+			throw new Error("RPC chunk exceeded the transport limit");
+		encoded += line;
+	}
+	return encoded;
+}
+
+function isRpcChunkFrame(value: unknown): value is RpcChunkFrame {
+	return isRecord(value) && value.type === "rpc_chunk";
+}
+
+function decodeBase64(data: unknown): Buffer {
+	if (
+		typeof data !== "string" ||
+		data.length === 0 ||
+		!/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(data)
+	)
+		throw new Error("invalid rpc chunk data");
+	const bytes = Buffer.from(data, "base64");
+	if (bytes.toString("base64") !== data) throw new Error("invalid rpc chunk data");
+	return bytes;
+}
+
+/** Reassemble protocol v2 chunk frames after each JSONL line has been parsed. */
+export class RpcFrameDecoder {
+	#pending?: PendingRpcChunks;
+
+	push(value: unknown): object | undefined {
+		if (!isRpcChunkFrame(value)) {
+			if (this.#pending) throw new Error("rpc chunk sequence interrupted");
+			if (!isRecord(value)) throw new Error("rpc frame must be an object");
+			return value;
+		}
+		const { chunkId, index, count, byteLength } = value;
+		if (
+			typeof chunkId !== "string" ||
+			chunkId.length === 0 ||
+			chunkId.length > 128 ||
+			!Number.isSafeInteger(index) ||
+			!Number.isSafeInteger(count) ||
+			!Number.isSafeInteger(byteLength) ||
+			index < 0 ||
+			count < 2 ||
+			count > Math.ceil(MAX_RPC_REASSEMBLED_BYTES / RPC_CHUNK_PAYLOAD_BYTES) ||
+			index >= count ||
+			byteLength < MAX_RPC_FRAME_BYTES ||
+			byteLength > MAX_RPC_REASSEMBLED_BYTES
+		)
+			throw new Error("invalid rpc chunk metadata");
+		const bytes = decodeBase64(value.data);
+		if (bytes.byteLength > RPC_CHUNK_PAYLOAD_BYTES) throw new Error("rpc chunk payload exceeds the transport limit");
+
+		if (!this.#pending) {
+			if (index !== 0) throw new Error("rpc chunk sequence must start at index 0");
+			this.#pending = { chunkId, count, byteLength, nextIndex: 0, chunks: [], receivedBytes: 0 };
+		}
+		const pending = this.#pending;
+		if (
+			pending.chunkId !== chunkId ||
+			pending.count !== count ||
+			pending.byteLength !== byteLength ||
+			pending.nextIndex !== index
+		)
+			throw new Error("rpc chunk sequence mismatch");
+		pending.chunks.push(bytes);
+		pending.receivedBytes += bytes.byteLength;
+		pending.nextIndex++;
+		if (pending.receivedBytes > pending.byteLength) throw new Error("rpc chunk sequence exceeds declared length");
+		if (pending.nextIndex < pending.count) return undefined;
+		if (pending.receivedBytes !== pending.byteLength) throw new Error("rpc chunk sequence length mismatch");
+
+		this.#pending = undefined;
+		const decoded = new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(pending.chunks));
+		const frame: unknown = JSON.parse(decoded);
+		if (!isRecord(frame)) throw new Error("rpc frame must be an object");
+		return frame;
+	}
+}
+
 function compactTerminalFrame(
 	frame: object,
 	streamedMessageCount: number,
@@ -142,13 +255,34 @@ export function encodeRpcFrame(frame: object, streamedMessageCount = 0, streamed
 /** Stateful encoder that tracks which messages a client has already received. */
 export class RpcFrameEncoder {
 	#streamedMessages: unknown[] = [];
+	#protocolVersion: RpcProtocolVersion = 1;
+	#chunkCounter = 0;
+
+	setProtocolVersion(version: number): void {
+		if (version !== 1 && version !== 2) throw new Error(`Unsupported RPC protocol version: ${version}`);
+		this.#protocolVersion = version;
+	}
 
 	encode(frame: object): string {
 		if (isRecord(frame) && frame.type === "agent_start") this.#streamedMessages = [];
-		const encoded = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages);
+		let json = JSON.stringify(frame);
+		let encoded: string;
+		if (this.#protocolVersion === 2 && serializedFrameBytes(json) > MAX_RPC_FRAME_BYTES) {
+			const compacted = compactTerminalFrame(frame, this.#streamedMessages.length, this.#streamedMessages);
+			json = JSON.stringify(compacted);
+			encoded =
+				serializedFrameBytes(json) > MAX_RPC_FRAME_BYTES
+					? encodeChunkedRpcFrame(compacted, `rpc-${++this.#chunkCounter}`)
+					: `${json}\n`;
+		} else {
+			encoded = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages);
+		}
 		if (!isRecord(frame)) return encoded;
 		if (frame.type === "message_end") {
-			const snapshot = encodedMessageSnapshot(encoded);
+			const snapshot =
+				this.#protocolVersion === 2 && Object.hasOwn(frame, "message")
+					? { message: jsonSnapshot(frame.message) }
+					: encodedMessageSnapshot(encoded);
 			if (snapshot) this.#streamedMessages.push(snapshot.message);
 		} else if (frame.type === "agent_end" && frame.willContinue !== true) this.#streamedMessages = [];
 		return encoded;

packages/coding-agent/src/modes/rpc/rpc-messages.ts

@@ -0,0 +1,111 @@
+import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
+import { isRecord } from "@oh-my-pi/pi-utils";
+
+const DEFAULT_RPC_MESSAGE_PAGE_LIMIT = 100;
+const MAX_RPC_MESSAGE_PAGE_LIMIT = 256;
+const MAX_RPC_MESSAGE_PAGE_BYTES = 768 * 1024;
+const MAX_RPC_MESSAGE_CURSOR_CHARS = 2048;
+
+export const RPC_MESSAGES_PAGE_BUSY_ERROR = "Cannot page messages while the session is changing";
+
+export interface RpcMessageSnapshot {
+	sessionId: string;
+	leafId: string | null;
+	messageCount: number;
+}
+
+export interface RpcMessagesPage {
+	messages: AgentMessage[];
+	nextCursor?: string;
+	totalMessages: number;
+}
+
+interface RpcMessageCursorPayload extends RpcMessageSnapshot {
+	version: 1;
+	offset: number;
+}
+
+export interface RpcMessagesPageOptions {
+	cursor?: string;
+	limit?: number;
+}
+
+function encodeCursor(snapshot: RpcMessageSnapshot, offset: number): string {
+	const payload: RpcMessageCursorPayload = { version: 1, ...snapshot, offset };
+	return Buffer.from(JSON.stringify(payload), "utf8").toString("base64url");
+}
+
+function decodeCursor(cursor: string): RpcMessageCursorPayload {
+	if (cursor.length === 0 || cursor.length > MAX_RPC_MESSAGE_CURSOR_CHARS || !/^[A-Za-z0-9_-]+$/.test(cursor))
+		throw new Error("Invalid RPC message cursor");
+	const bytes = Buffer.from(cursor, "base64url");
+	if (bytes.toString("base64url") !== cursor) throw new Error("Invalid RPC message cursor");
+	let value: unknown;
+	try {
+		value = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes));
+	} catch {
+		throw new Error("Invalid RPC message cursor");
+	}
+	if (!isRecord(value)) throw new Error("Invalid RPC message cursor");
+	const { version, sessionId, leafId, messageCount, offset } = value;
+	if (
+		version !== 1 ||
+		typeof sessionId !== "string" ||
+		sessionId.length === 0 ||
+		sessionId.length > 256 ||
+		!((typeof leafId === "string" && leafId.length > 0 && leafId.length <= 256) || leafId === null) ||
+		typeof messageCount !== "number" ||
+		!Number.isSafeInteger(messageCount) ||
+		messageCount < 0 ||
+		typeof offset !== "number" ||
+		!Number.isSafeInteger(offset) ||
+		offset < 0 ||
+		offset > messageCount
+	)
+		throw new Error("Invalid RPC message cursor");
+	return { version, sessionId, leafId, messageCount, offset };
+}
+
+function sameSnapshot(cursor: RpcMessageCursorPayload, snapshot: RpcMessageSnapshot): boolean {
+	return (
+		cursor.sessionId === snapshot.sessionId &&
+		cursor.leafId === snapshot.leafId &&
+		cursor.messageCount === snapshot.messageCount
+	);
+}
+
+/** Page one stable in-memory message snapshot without crossing the v1 frame budget. */
+export function pageRpcMessages(
+	messages: readonly AgentMessage[],
+	snapshot: RpcMessageSnapshot,
+	options: RpcMessagesPageOptions = {},
+): RpcMessagesPage {
+	if (snapshot.messageCount !== messages.length)
+		throw new Error("RPC message snapshot does not match current messages");
+	const limit = options.limit ?? DEFAULT_RPC_MESSAGE_PAGE_LIMIT;
+	if (!Number.isSafeInteger(limit) || limit < 1 || limit > MAX_RPC_MESSAGE_PAGE_LIMIT)
+		throw new Error(`RPC message page limit must be between 1 and ${MAX_RPC_MESSAGE_PAGE_LIMIT}`);
+	let offset = 0;
+	if (options.cursor !== undefined) {
+		const cursor = decodeCursor(options.cursor);
+		if (!sameSnapshot(cursor, snapshot)) throw new Error("RPC message cursor is stale");
+		offset = cursor.offset;
+	}
+
+	const page: AgentMessage[] = [];
+	let pageBytes = 2;
+	while (offset + page.length < messages.length && page.length < limit) {
+		const message = messages[offset + page.length];
+		const messageBytes = Buffer.byteLength(JSON.stringify(message), "utf8") + (page.length === 0 ? 0 : 1);
+		if (page.length > 0 && pageBytes + messageBytes > MAX_RPC_MESSAGE_PAGE_BYTES) break;
+		page.push(message);
+		pageBytes += messageBytes;
+	}
+
+	const nextOffset = offset + page.length;
+	return {
+		messages: page,
+		...(nextOffset < messages.length ? { nextCursor: encodeCursor(snapshot, nextOffset) } : {}),
+		totalMessages: messages.length,
+	};
+}

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -34,8 +34,9 @@ import type { EventBus } from "../../utils/event-bus";
 import { initializeExtensions } from "../runtime-init";
 import { isRpcHostToolResult, isRpcHostToolUpdate, RpcHostToolBridge } from "./host-tools";
 import { isRpcHostUriResult, RpcHostUriBridge } from "./host-uris";
-import { RpcFrameEncoder } from "./rpc-frame";
+import { MAX_RPC_FRAME_BYTES, MAX_RPC_REASSEMBLED_BYTES, RpcFrameEncoder } from "./rpc-frame";
 import { claimRpcInput } from "./rpc-input";
+import { pageRpcMessages, RPC_MESSAGES_PAGE_BUSY_ERROR } from "./rpc-messages";
 import { RpcSubagentRegistry, readRpcSubagentTranscript } from "./rpc-subagents";
 import type {
 	RpcCommand,
@@ -619,9 +620,19 @@ export async function runRpcMode(
 	process.env.PI_NOTIFICATIONS = "off";
 
 	const frameEncoder = new RpcFrameEncoder();
-	process.stdout.write(frameEncoder.encode({ type: "ready" }));
+	process.stdout.write(
+		frameEncoder.encode({
+			type: "ready",
+			protocolVersion: 1,
+			supportedProtocolVersions: [1, 2],
+			maxFrameBytes: MAX_RPC_FRAME_BYTES,
+			maxReassembledFrameBytes: MAX_RPC_REASSEMBLED_BYTES,
+		}),
+	);
 	const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => {
 		process.stdout.write(frameEncoder.encode(obj));
+		if (isRecord(obj) && obj.type === "response" && obj.command === "negotiate_protocol" && obj.success === true)
+			frameEncoder.setProtocolVersion(2);
 	};
 	const emitRpcTitles = shouldEmitRpcTitles();
 
@@ -936,6 +947,12 @@ export async function runRpcMode(
 		const id = command.id;
 
 		switch (command.type) {
+			case "negotiate_protocol": {
+				if (command.protocolVersion !== 2)
+					return error(id, "negotiate_protocol", `Unsupported RPC protocol version: ${command.protocolVersion}`);
+				return success(id, "negotiate_protocol", { protocolVersion: 2 });
+			}
+
 			// =================================================================
 			// Prompting
 			// =================================================================
@@ -1289,6 +1306,33 @@ export async function runRpcMode(
 				return success(id, "get_messages", { messages: session.messages });
 			}
 
+			case "get_messages_page": {
+				if (session.isStreaming || session.isCompacting)
+					return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR);
+				const messages = session.messages;
+				try {
+					return success(
+						id,
+						"get_messages_page",
+						pageRpcMessages(
+							messages,
+							{
+								sessionId: session.sessionId,
+								leafId: session.sessionManager.getLeafId(),
+								messageCount: messages.length,
+							},
+							{ cursor: command.cursor, limit: command.limit },
+						),
+					);
+				} catch (pageError) {
+					return error(
+						id,
+						"get_messages_page",
+						pageError instanceof Error ? pageError.message : String(pageError),
+					);
+				}
+			}
+
 			// =================================================================
 			// Login
 			// =================================================================

packages/coding-agent/src/modes/rpc/rpc-types.ts

@@ -19,12 +19,16 @@ import type {
 	SubagentProgressPayload,
 } from "../../task";
 import type { TodoPhase } from "../../tools/todo";
+import type { RpcMessagesPage } from "./rpc-messages";
 
 // ============================================================================
 // RPC Commands (stdin)
 // ============================================================================
 
 export type RpcCommand =
+	// Protocol
+	| { id?: string; type: "negotiate_protocol"; protocolVersion: number }
+
 	// Prompting
 	| { id?: string; type: "prompt"; message: string; images?: ImageContent[]; streamingBehavior?: "steer" | "followUp" }
 	| { id?: string; type: "steer"; message: string; images?: ImageContent[] }
@@ -81,6 +85,7 @@ export type RpcCommand =
 
 	// Messages
 	| { id?: string; type: "get_messages" }
+	| { id?: string; type: "get_messages_page"; cursor?: string; limit?: number }
 
 	// Login
 	| { id?: string; type: "get_login_providers" }
@@ -132,6 +137,23 @@ export interface RpcPromptResultFrame {
 	agentInvoked: boolean;
 }
 
+export interface RpcReadyFrame {
+	type: "ready";
+	protocolVersion: 1;
+	supportedProtocolVersions: [1, 2];
+	maxFrameBytes: number;
+	maxReassembledFrameBytes: number;
+}
+
+export interface RpcChunkFrame {
+	type: "rpc_chunk";
+	chunkId: string;
+	index: number;
+	count: number;
+	byteLength: number;
+	data: string;
+}
+
 export interface RpcHandoffResult {
 	savedPath?: string;
 }
@@ -168,6 +190,15 @@ export interface RpcSubagentMessagesResult {
 
 // Success responses with data
 export type RpcResponse =
+	// Protocol
+	| {
+			id?: string;
+			type: "response";
+			command: "negotiate_protocol";
+			success: true;
+			data: { protocolVersion: 2 };
+	  }
+
 	// Prompting (async - events follow)
 	| { id?: string; type: "response"; command: "prompt"; success: true; data?: { agentInvoked: boolean } }
 	| { id?: string; type: "response"; command: "steer"; success: true }
@@ -284,6 +315,7 @@ export type RpcResponse =
 
 	// Messages
 	| { id?: string; type: "response"; command: "get_messages"; success: true; data: { messages: AgentMessage[] } }
+	| { id?: string; type: "response"; command: "get_messages_page"; success: true; data: RpcMessagesPage }
 
 	// Login
 	| {