fix(rpc): stream v2 chunk frames with backpressure and recover stale page cursors

Two fixes on top of the paged transport:

- Near-limit v2 framing no longer materializes the full base64 transport:
  chunk lines are generated lazily from a single serialization, the 64 MiB
  reassembly ceiling is enforced via Buffer.byteLength before any
  full-payload allocation, and RPC stdout writes drain with backpressure
  one physical line at a time. Peak RSS for a 63 MiB response drops
  ~686 MB -> ~521 MB; a rejected 80 MiB response drops ~507 MB -> ~259 MB
  (parity with the v1 path).

- get_messages_page errors now carry a machine-readable code
  (session_busy | stale_cursor). Both bundled clients' high-level
  getMessages() drains discard partial pages and fall back to the legacy
  snapshot on either code — previously a cursor invalidated by a
  background mutation (e.g. an appended bash message) threw instead of
  falling back. Direct page calls remain strict.
This commit is contained in:
can1357
2026-07-23 12:38:13 +02:00
parent 3e09e4b1ea
commit f833810d0e
11 changed files with 210 additions and 40 deletions
+2 -2
View File
@@ -187,9 +187,9 @@ correlate it via `id`. Ordering across concurrent commands is not guaranteed
- `{ id?, type: "get_messages" }`
- `{ id?, type: "get_messages_page", cursor?: string, limit?: number }`
`get_messages_page` returns a stable chronological page with `messages`, `totalMessages`, and an opaque `nextCursor` when more messages remain. Cursors are bound to the session ID, durable leaf, and message count. The server rejects stale cursors if the session changes between requests, and refuses to start a paging walk while the session is streaming or compacting. Pages contain at most 256 messages and normally stay below the v1 physical-frame ceiling. A v1 caller can page ordinary histories, but an individual message whose response exceeds that ceiling produces an overflow error; retrieving it losslessly requires negotiated v2 framing.
`get_messages_page` returns a stable chronological page with `messages`, `totalMessages`, and an opaque `nextCursor` when more messages remain. Cursors are bound to the session ID, durable leaf, and message count. The server rejects stale cursors if the session changes between requests, and refuses to start a paging walk while the session is streaming or compacting. Failed page requests carry a machine-readable `code` on the error response — `session_busy` (session is streaming or compacting) or `stale_cursor` (the snapshot behind the cursor changed, e.g. a background bash appended a message between pages) — so clients can react without matching error-message text. Pages contain at most 256 messages and normally stay below the v1 physical-frame ceiling. A v1 caller can page ordinary histories, but an individual message whose response exceeds that ceiling produces an overflow error; retrieving it losslessly requires negotiated v2 framing.
The bundled TypeScript `RpcClient.getMessages()` and Python `RpcClient.get_messages()` drain this paged endpoint automatically after negotiating v2. They retain the legacy monolithic command when connected to a v1 server, and fall back to its best-effort snapshot if the session begins streaming or compacting during a page walk. Direct `getMessagesPage()` and `get_messages_page()` calls remain strict so incremental hosts never mix snapshots silently.
The bundled TypeScript `RpcClient.getMessages()` and Python `RpcClient.get_messages()` drain this paged endpoint automatically after negotiating v2. They retain the legacy monolithic command when connected to a v1 server, and on either `session_busy` or `stale_cursor` they discard partial pages and fall back to the legacy best-effort snapshot. Direct `getMessagesPage()` and `get_messages_page()` calls remain strict so incremental hosts never mix snapshots silently.
### Login
+2
View File
@@ -28,6 +28,8 @@
- Rewrote the `/guided-goal` interviewer rubric around loop-engineering: deterministic success criteria, verification commands, attempt caps, scope boundaries, and stop conditions. Ready objectives must use the five-section structured markdown form.
- Added `task.isolation.apply` (default `true`) to choose whether successful isolated `task` runs automatically apply their changes to the parent checkout or retain patch/branch artifacts for later integration.
- Added opt-in RPC protocol v2 negotiation with bounded, lossless chunking for stdout objects up to 64 MiB, plus stable cursor-based message pages for histories that should not travel as one response. Legacy JSONL clients remain on protocol v1, while the bundled TypeScript and Python RPC clients negotiate, reassemble, and drain message pages automatically.
- Fixed protocol v2 chunked framing materializing the whole base64 transport in memory: near-limit logical frames (~63 MiB) peaked around 686 MB RSS and over-ceiling frames allocated the full payload buffer before rejection. Chunk lines are now produced lazily from a single serialization, the 64 MiB ceiling is checked before any full-payload allocation, and RPC stdout writes honor backpressure line by line.
- Fixed the bundled TypeScript and Python RPC clients throwing when a `get_messages_page` cursor went stale mid-walk (e.g. a background bash appending a message between pages): the high-level `getMessages()` drains now discard partial pages and fall back to the legacy snapshot on both `session_busy` and `stale_cursor`, driven by a new machine-readable `code` field on RPC error responses. Direct page calls remain strict.
## [17.0.8] - 2026-07-22
@@ -13,7 +13,12 @@ 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 {
RPC_MESSAGES_PAGE_BUSY_ERROR,
RPC_MESSAGES_PAGE_STALE_ERROR,
type RpcMessagesPage,
type RpcMessagesPageOptions,
} from "./rpc-messages";
import type {
RpcAvailableCommandsUpdateFrame,
RpcAvailableSlashCommand,
@@ -212,6 +217,26 @@ function normalizeToolResult<TDetails>(result: RpcClientToolResult<TDetails>): A
return result;
}
/** Failed RPC command; `code` mirrors the server's machine-readable error code when present. */
export class RpcCommandError extends Error {
constructor(
message: string,
readonly command: string,
readonly code?: string,
) {
super(message);
this.name = "RpcCommandError";
}
}
/** True when a high-level `getMessages()` drain should discard partial pages and fall back to `get_messages`. */
function isPageFallbackError(error: unknown): boolean {
if (!(error instanceof Error)) return false;
if (error instanceof RpcCommandError && (error.code === "session_busy" || error.code === "stale_cursor"))
return true;
return error.message === RPC_MESSAGES_PAGE_BUSY_ERROR || error.message === RPC_MESSAGES_PAGE_STALE_ERROR;
}
// ============================================================================
// RPC Client
// ============================================================================
@@ -810,7 +835,7 @@ export class RpcClient {
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;
if (!isPageFallbackError(error)) throw error;
}
}
const response = await this.#send({ type: "get_messages" });
@@ -1160,7 +1185,7 @@ export class RpcClient {
#getData<T>(response: RpcResponse): T {
if (!response.success) {
const errorResponse = response as Extract<RpcResponse, { success: false }>;
throw new Error(errorResponse.error);
throw new RpcCommandError(errorResponse.error, errorResponse.command, errorResponse.code);
}
// Type assertion: we trust response.data matches T based on the command sent.
// This is safe because each public method specifies the correct T for its command.
@@ -84,19 +84,27 @@ function encodedMessageSnapshot(encoded: string): { message: unknown } | undefin
: undefined;
}
function encodeChunkedRpcFrame(frame: object, chunkId: string): string {
const json = JSON.stringify(frame);
/**
* Emit protocol v2 chunk frames for one pre-serialized logical frame, one physical
* JSONL line at a time so callers can write with backpressure instead of holding the
* whole ~4/3-sized base64 transport in memory. The reassembly ceiling is enforced on
* `Buffer.byteLength` BEFORE any full-payload allocation.
*/
function* encodeChunkedRpcFrames(frame: object, json: string, chunkId: string): Generator<string> {
const byteLength = Buffer.byteLength(json, "utf8");
if (byteLength > MAX_RPC_REASSEMBLED_BYTES) {
yield `${JSON.stringify(overflowFrame(frame))}\n`;
return;
}
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 = "";
const count = Math.ceil(byteLength / RPC_CHUNK_PAYLOAD_BYTES);
for (let index = 0; index < count; index++) {
const chunk: RpcChunkFrame = {
type: "rpc_chunk",
chunkId,
index,
count,
byteLength: bytes.byteLength,
byteLength,
data: bytes
.subarray(index * RPC_CHUNK_PAYLOAD_BYTES, (index + 1) * RPC_CHUNK_PAYLOAD_BYTES)
.toString("base64"),
@@ -104,9 +112,8 @@ function encodeChunkedRpcFrame(frame: object, chunkId: string): string {
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;
yield line;
}
return encoded;
}
function isRpcChunkFrame(value: unknown): value is RpcChunkFrame {
@@ -263,28 +270,47 @@ export class RpcFrameEncoder {
this.#protocolVersion = version;
}
encode(frame: object): string {
/**
* Encode one logical frame into physical JSONL lines. Encoder bookkeeping runs
* eagerly; only chunk emission is lazy, so a chunked result can be streamed to
* stdout with backpressure without holding the whole transport in memory. The
* returned iterable MUST be fully consumed exactly once.
*/
encodeFrames(frame: object): Iterable<string> {
if (isRecord(frame) && frame.type === "agent_start") this.#streamedMessages = [];
let json = JSON.stringify(frame);
let encoded: string;
const json = JSON.stringify(frame);
let frames: Iterable<string>;
let singleFrame: string | undefined;
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`;
// Reuse the original serialization when compaction was a no-op.
const compactedJson = compacted === frame ? json : JSON.stringify(compacted);
if (serializedFrameBytes(compactedJson) > MAX_RPC_FRAME_BYTES) {
frames = encodeChunkedRpcFrames(compacted, compactedJson, `rpc-${++this.#chunkCounter}`);
} else {
singleFrame = `${compactedJson}\n`;
frames = [singleFrame];
}
} else {
encoded = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages);
singleFrame = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages);
frames = [singleFrame];
}
if (!isRecord(frame)) return encoded;
if (!isRecord(frame)) return frames;
if (frame.type === "message_end") {
const snapshot =
this.#protocolVersion === 2 && Object.hasOwn(frame, "message")
? { message: jsonSnapshot(frame.message) }
: encodedMessageSnapshot(encoded);
: singleFrame !== undefined
? encodedMessageSnapshot(singleFrame)
: undefined;
if (snapshot) this.#streamedMessages.push(snapshot.message);
} else if (frame.type === "agent_end" && frame.willContinue !== true) this.#streamedMessages = [];
return frames;
}
encode(frame: object): string {
let encoded = "";
for (const line of this.encodeFrames(frame)) encoded += line;
return encoded;
}
}
@@ -7,6 +7,21 @@ 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 const RPC_MESSAGES_PAGE_STALE_ERROR = "RPC message cursor is stale";
/** Machine-readable reasons a `get_messages_page` request can fail; carried as `code` on the error response. */
export type RpcMessagesPageErrorCode = "session_busy" | "stale_cursor";
/** Paging failure that maps to a structured wire `code`, so clients can react without matching message text. */
export class RpcMessagesPageError extends Error {
constructor(
message: string,
readonly code: RpcMessagesPageErrorCode,
) {
super(message);
this.name = "RpcMessagesPageError";
}
}
export interface RpcMessageSnapshot {
sessionId: string;
@@ -88,7 +103,8 @@ export function pageRpcMessages(
let offset = 0;
if (options.cursor !== undefined) {
const cursor = decodeCursor(options.cursor);
if (!sameSnapshot(cursor, snapshot)) throw new Error("RPC message cursor is stale");
if (!sameSnapshot(cursor, snapshot))
throw new RpcMessagesPageError(RPC_MESSAGES_PAGE_STALE_ERROR, "stale_cursor");
offset = cursor.offset;
}
@@ -10,6 +10,7 @@
* - Events: AgentSessionEvent objects streamed as they occur
* - Extension UI: Extension UI requests are emitted, client responds with extension_ui_response
*/
import { once } from "node:events";
import { getOAuthProviders } from "@oh-my-pi/pi-ai/oauth";
import { isZodSchema, zodToWireSchema } from "@oh-my-pi/pi-ai/utils/schema";
import { $env, isRecord, readLines, Snowflake } from "@oh-my-pi/pi-utils";
@@ -36,7 +37,7 @@ import { isRpcHostToolResult, isRpcHostToolUpdate, RpcHostToolBridge } from "./h
import { isRpcHostUriResult, RpcHostUriBridge } from "./host-uris";
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 { pageRpcMessages, RPC_MESSAGES_PAGE_BUSY_ERROR, RpcMessagesPageError } from "./rpc-messages";
import { RpcSubagentRegistry, readRpcSubagentTranscript } from "./rpc-subagents";
import type {
RpcCommand,
@@ -620,8 +621,22 @@ export async function runRpcMode(
process.env.PI_NOTIFICATIONS = "off";
const frameEncoder = new RpcFrameEncoder();
process.stdout.write(
frameEncoder.encode({
// Ordered stdout writer honoring backpressure: chunked v2 frames are produced
// lazily by the encoder and written one physical line at a time, so a near-limit
// logical frame never materializes its full base64 transport in memory.
let stdoutQueue: Promise<void> = Promise.resolve();
const writeFrames = (frames: Iterable<string>) => {
stdoutQueue = stdoutQueue
.then(async () => {
for (const line of frames) {
if (!process.stdout.write(line)) await once(process.stdout, "drain");
}
})
// stdout gone (host exited) — nothing left to deliver; keep the queue alive.
.catch(() => {});
};
writeFrames(
frameEncoder.encodeFrames({
type: "ready",
protocolVersion: 1,
supportedProtocolVersions: [1, 2],
@@ -630,7 +645,7 @@ export async function runRpcMode(
}),
);
const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => {
process.stdout.write(frameEncoder.encode(obj));
writeFrames(frameEncoder.encodeFrames(obj));
if (isRecord(obj) && obj.type === "response" && obj.command === "negotiate_protocol" && obj.success === true)
frameEncoder.setProtocolVersion(2);
};
@@ -647,8 +662,8 @@ export async function runRpcMode(
return { id, type: "response", command, success: true, data } as RpcResponse;
};
const error = (id: string | undefined, command: string, message: string): RpcResponse => {
return { id, type: "response", command, success: false, error: message };
const error = (id: string | undefined, command: string, message: string, code?: string): RpcResponse => {
return { id, type: "response", command, success: false, error: message, ...(code ? { code } : {}) };
};
const extensionUserMessageTracker = new RpcExtensionUserMessageTracker();
@@ -1308,7 +1323,7 @@ export async function runRpcMode(
case "get_messages_page": {
if (session.isStreaming || session.isCompacting)
return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR);
return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR, "session_busy");
const messages = session.messages;
try {
return success(
@@ -1329,6 +1344,7 @@ export async function runRpcMode(
id,
"get_messages_page",
pageError instanceof Error ? pageError.message : String(pageError),
pageError instanceof RpcMessagesPageError ? pageError.code : undefined,
);
}
}
@@ -327,8 +327,8 @@ export type RpcResponse =
}
| { id?: string; type: "response"; command: "login"; success: true; data: { providerId: string } }
// Error response (any command can fail)
| { id?: string; type: "response"; command: string; success: false; error: string };
// Error response (any command can fail); `code` is an optional machine-readable reason.
| { id?: string; type: "response"; command: string; success: false; error: string; code?: string };
// ============================================================================
// Subagent Events (stdout)
+16 -1
View File
@@ -87,6 +87,18 @@ for await (const raw of console) {
command: frame.type,
success: false,
error: "Cannot page messages while the session is changing",
code: "session_busy",
});
continue;
}
if (Bun.env.MOCK_RPC_PAGE_STALE === "1" && frame.cursor !== undefined) {
writeFrame({
id,
type: "response",
command: frame.type,
success: false,
error: "RPC message cursor is stale",
code: "stale_cursor",
});
continue;
}
@@ -109,7 +121,10 @@ for await (const raw of console) {
});
continue;
}
if (frame.type === "get_messages" && Bun.env.MOCK_RPC_PAGE_BUSY === "1") {
if (
frame.type === "get_messages" &&
(Bun.env.MOCK_RPC_PAGE_BUSY === "1" || Bun.env.MOCK_RPC_PAGE_STALE === "1")
) {
writeFrame({
id,
type: "response",
@@ -43,6 +43,25 @@ describe("RpcClient lifecycle (issue #4079 B)", () => {
]);
}, 20_000);
test("discards partial pages and falls back to get_messages when a cursor goes stale mid-walk", async () => {
using client = new RpcClient({
cliPath: MOCK_AGENT,
env: { MOCK_RPC_V2: "1", MOCK_RPC_PAGE_STALE: "1" },
});
await client.start();
// Direct page walks stay strict: the stale cursor is surfaced to the caller.
const firstPage = await client.getMessagesPage();
expect(firstPage.nextCursor).toBe("second-page");
await expect(client.getMessagesPage({ cursor: firstPage.nextCursor })).rejects.toThrow(
"RPC message cursor is stale",
);
// The high-level drain discards the partial first page and takes the legacy snapshot.
expect((await client.getMessages()) as unknown).toEqual([
{ role: "assistant", content: [{ type: "text", text: "streaming snapshot" }], timestamp: 3 },
]);
}, 20_000);
test("start() succeeds a second time after stop() on the same instance", async () => {
using client = new RpcClient({
cliPath: MOCK_AGENT,
+17 -5
View File
@@ -118,6 +118,8 @@ _MAX_RPC_FRAME_BYTES = 1024 * 1024
_MAX_RPC_REASSEMBLED_BYTES = 64 * 1024 * 1024
_RPC_CHUNK_PAYLOAD_BYTES = 256 * 1024
_RPC_MESSAGES_PAGE_BUSY_ERROR = "Cannot page messages while the session is changing"
_RPC_MESSAGES_PAGE_STALE_ERROR = "RPC message cursor is stale"
_RPC_MESSAGES_PAGE_FALLBACK_CODES = frozenset({"session_busy", "stale_cursor"})
@dataclass(slots=True)
@@ -313,12 +315,16 @@ class RpcConcurrencyError(RpcError):
class RpcCommandError(RpcError):
"""Raised when the RPC server returns `success: false`."""
"""Raised when the RPC server returns `success: false`.
def __init__(self, command: str, error: str):
`code` carries the server's machine-readable error code when present.
"""
def __init__(self, command: str, error: str, code: str | None = None):
super().__init__(f"{command}: {error}")
self.command = command
self.error = error
self.code = code
class RpcProtocolError(RpcError):
@@ -1036,9 +1042,13 @@ class RpcClient:
)
return tuple(messages)
except RpcCommandError as error:
if (
error.command != "get_messages_page"
or error.error != _RPC_MESSAGES_PAGE_BUSY_ERROR
if error.command != "get_messages_page" or not (
error.code in _RPC_MESSAGES_PAGE_FALLBACK_CODES
or error.error
in (
_RPC_MESSAGES_PAGE_BUSY_ERROR,
_RPC_MESSAGES_PAGE_STALE_ERROR,
)
):
raise
payload = self._request("get_messages")
@@ -1385,9 +1395,11 @@ class RpcClient:
raise response
if not bool(response.get("success", False)):
raw_code = response.get("code")
raise RpcCommandError(
command=str(response.get("command", command_type)),
error=str(response.get("error", "")),
code=raw_code if isinstance(raw_code, str) else None,
)
data = response.get("data")
+39
View File
@@ -519,6 +519,34 @@ V2_MESSAGES_SERVER = textwrap.dedent(
"command": command_type,
"success": False,
"error": "Cannot page messages while the session is changing",
"code": "session_busy",
}
)
continue
if os.environ.get("V2_MESSAGES_STALE") == "1":
if command.get("cursor") is not None:
emit(
{
"id": request_id,
"type": "response",
"command": command_type,
"success": False,
"error": "RPC message cursor is stale",
"code": "stale_cursor",
}
)
continue
emit(
{
"id": request_id,
"type": "response",
"command": command_type,
"success": True,
"data": {
"messages": [message],
"totalMessages": 2,
"nextCursor": "page-two",
},
}
)
continue
@@ -1036,6 +1064,17 @@ class RpcClientTests(unittest.TestCase):
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]["content"][0]["text"], "streaming snapshot")
def test_protocol_v2_get_messages_discards_stale_page_walk(self) -> None:
with self.make_client(
server=V2_MESSAGES_SERVER, env={"V2_MESSAGES_STALE": "1"}
) as client:
with self.assertRaisesRegex(RpcCommandError, "RPC message cursor is stale"):
client.get_messages_page(cursor="page-two")
messages = client.get_messages()
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]["content"][0]["text"], "streaming snapshot")
def test_collect_events_returns_turn_events(self) -> None:
with self.make_client() as client:
client.prompt("slow")