smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rpc): preserve v2 snapshot semantics

Browse Source
This commit is contained in:
Wolfgang Schoenberger
2026-07-22 19:00:48 -07:00
parent ce8c9dc75f
commit a0cb946057
10 changed files with 173 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/rpc.md
+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; an individually oversized message 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. 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. Hosts that render incrementally can call `getMessagesPage()` or `get_messages_page()` directly.
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.
### Login
packages/coding-agent/src/modes/rpc/rpc-client.ts
+26 -22
View File
@@ -13,7 +13,7 @@ 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 type { RpcMessagesPage, RpcMessagesPageOptions } from "./rpc-messages";
import { RPC_MESSAGES_PAGE_BUSY_ERROR, type RpcMessagesPage, type RpcMessagesPageOptions } from "./rpc-messages";
import type {
RpcAvailableCommandsUpdateFrame,
RpcAvailableSlashCommand,
@@ -787,27 +787,31 @@ export class RpcClient {
/** Get all messages, draining stable pages when protocol v2 is available. */
async getMessages(): Promise<AgentMessage[]> {
if (this.#protocolVersion === 2) {
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;
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
+12 -5
View File
@@ -265,11 +265,18 @@ export class RpcFrameEncoder {
encode(frame: object): string {
if (isRecord(frame) && frame.type === "agent_start") this.#streamedMessages = [];
const json = JSON.stringify(frame);
const encoded =
this.#protocolVersion === 2 && serializedFrameBytes(json) > MAX_RPC_FRAME_BYTES
? encodeChunkedRpcFrame(frame, `rpc-${++this.#chunkCounter}`)
: 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 =
packages/coding-agent/src/modes/rpc/rpc-messages.ts
+2
View File
@@ -6,6 +6,8 @@ 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;
packages/coding-agent/src/modes/rpc/rpc-mode.ts
+2 -2
View File
@@ -36,7 +36,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 } from "./rpc-messages";
import { pageRpcMessages, RPC_MESSAGES_PAGE_BUSY_ERROR } from "./rpc-messages";
import { RpcSubagentRegistry, readRpcSubagentTranscript } from "./rpc-subagents";
import type {
RpcCommand,
@@ -1296,7 +1296,7 @@ export async function runRpcMode(
case "get_messages_page": {
if (session.isStreaming || session.isCompacting)
return error(id, "get_messages_page", "Cannot page messages while the session is changing");
return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR);
const messages = session.messages;
try {
return success(
packages/coding-agent/test/fixtures/mock-rpc-agent.ts
Vendored
+24
View File
@@ -80,6 +80,16 @@ for await (const raw of console) {
continue;
}
if (frame.type === "get_messages_page") {
if (Bun.env.MOCK_RPC_PAGE_BUSY === "1") {
writeFrame({
id,
type: "response",
command: frame.type,
success: false,
error: "Cannot page messages while the session is changing",
});
continue;
}
const first = frame.cursor === undefined;
writeFrame({
id,
@@ -99,6 +109,20 @@ for await (const raw of console) {
});
continue;
}
if (frame.type === "get_messages" && Bun.env.MOCK_RPC_PAGE_BUSY === "1") {
writeFrame({
id,
type: "response",
command: frame.type,
success: true,
data: {
messages: [
{ role: "assistant", content: [{ type: "text", text: "streaming snapshot" }], timestamp: 3 },
],
},
});
continue;
}
writeFrame({
id,
type: "response",
packages/coding-agent/test/rpc-client.restart.test.ts
+13
View File
@@ -30,6 +30,19 @@ describe("RpcClient lifecycle (issue #4079 B)", () => {
]);
}, 20_000);
test("preserves getMessages snapshot behavior while a v2 page walk is unavailable", async () => {
using client = new RpcClient({
cliPath: MOCK_AGENT,
env: { MOCK_RPC_V2: "1", MOCK_RPC_PAGE_BUSY: "1" },
});
await client.start();
await expect(client.getMessagesPage()).rejects.toThrow("Cannot page messages while the session is changing");
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,
packages/coding-agent/test/rpc-frame.test.ts
+15
View File
@@ -208,6 +208,21 @@ describe("RPC frame encoding", () => {
expect(decoded).toEqual(frame);
});
it("preserves terminal message counts above the protocol v2 ceiling", () => {
const encoder = new RpcFrameEncoder();
encoder.setProtocolVersion(2);
const encoded = encoder.encode({
type: "agent_end",
messages: [{ role: "assistant", content: "x".repeat(MAX_RPC_REASSEMBLED_BYTES) }],
});
expect(decode(encoded)).toEqual({
type: "agent_end",
messages: [],
messageCount: 1,
});
});
it("rejects protocol v2 logical frames above the advertised reassembly ceiling", () => {
const encoder = new RpcFrameEncoder();
encoder.setProtocolVersion(2);
python/omp-rpc/src/omp_rpc/client.py
+32 -21
View File
@@ -117,6 +117,7 @@ _TODO_STATUS_VALUES = frozenset({"pending", "in_progress", "completed", "abandon
_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"
@dataclass(slots=True)
@@ -1007,29 +1008,39 @@ class RpcClient:
def get_messages(self) -> tuple[AgentMessage, ...]:
if self._protocol_version == 2:
messages: list[AgentMessage] = []
seen_cursors: set[str] = set()
total_messages: int | None = None
cursor: str | None = None
while True:
page = self.get_messages_page(cursor=cursor, limit=256)
if total_messages is not None and page.total_messages != total_messages:
try:
messages: list[AgentMessage] = []
seen_cursors: set[str] = set()
total_messages: int | None = None
cursor: str | None = None
while True:
page = self.get_messages_page(cursor=cursor, limit=256)
if (
total_messages is not None
and page.total_messages != total_messages
):
raise RpcError(
"RPC message pagination returned an inconsistent total"
)
total_messages = page.total_messages
messages.extend(page.messages)
cursor = page.next_cursor
if cursor is None:
break
if cursor in seen_cursors:
raise RpcError("RPC message pagination repeated a cursor")
seen_cursors.add(cursor)
if len(messages) != total_messages:
raise RpcError(
"RPC message pagination returned an inconsistent total"
"RPC message pagination ended before the advertised total"
)
total_messages = page.total_messages
messages.extend(page.messages)
cursor = page.next_cursor
if cursor is None:
break
if cursor in seen_cursors:
raise RpcError("RPC message pagination repeated a cursor")
seen_cursors.add(cursor)
if len(messages) != total_messages:
raise RpcError(
"RPC message pagination ended before the advertised total"
)
return tuple(messages)
return tuple(messages)
except RpcCommandError as error:
if (
error.command != "get_messages_page"
or error.error != _RPC_MESSAGES_PAGE_BUSY_ERROR
):
raise
payload = self._request("get_messages")
return parse_agent_messages(cast(JsonValue | None, payload.get("messages")))
python/omp-rpc/tests/test_client.py
+45
View File
@@ -447,6 +447,7 @@ V2_MESSAGES_SERVER = textwrap.dedent(
"""
import base64
import json
import os
import sys
message = {
@@ -507,6 +508,17 @@ V2_MESSAGES_SERVER = textwrap.dedent(
}
)
elif command_type == "get_messages_page":
if os.environ.get("V2_MESSAGES_BUSY") == "1":
emit(
{
"id": request_id,
"type": "response",
"command": command_type,
"success": False,
"error": "Cannot page messages while the session is changing",
}
)
continue
emit(
{
"id": request_id,
@@ -520,6 +532,26 @@ V2_MESSAGES_SERVER = textwrap.dedent(
},
}
)
elif command_type == "get_messages":
emit(
{
"id": request_id,
"type": "response",
"command": command_type,
"success": True,
"data": {
"messages": [
{
"role": "assistant",
"content": [
{"type": "text", "text": "streaming snapshot"}
],
"timestamp": 3,
}
]
},
}
)
else:
emit(
{
@@ -956,6 +988,19 @@ class RpcClientTests(unittest.TestCase):
self.assertEqual(len(messages), 1)
self.assertEqual(len(messages[0]["content"][0]["text"]), 1024 * 1024)
def test_protocol_v2_get_messages_falls_back_to_streaming_snapshot(self) -> None:
with self.make_client(
server=V2_MESSAGES_SERVER, env={"V2_MESSAGES_BUSY": "1"}
) as client:
with self.assertRaisesRegex(
RpcCommandError, "Cannot page messages while the session is changing"
):
client.get_messages_page()
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")