diff --git a/docs/rpc.md b/docs/rpc.md index 05bf375ed..7515255a1 100644 --- a/docs/rpc.md +++ b/docs/rpc.md @@ -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 diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 3eb9ef3e6..4269d65de 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -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 { if (this.#protocolVersion === 2) { - const messages: AgentMessage[] = []; - const seenCursors = new Set(); - 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(); + 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; diff --git a/packages/coding-agent/src/modes/rpc/rpc-frame.ts b/packages/coding-agent/src/modes/rpc/rpc-frame.ts index dfbe857af..cce90beea 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-frame.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-frame.ts @@ -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 = diff --git a/packages/coding-agent/src/modes/rpc/rpc-messages.ts b/packages/coding-agent/src/modes/rpc/rpc-messages.ts index 383733d7d..5307456be 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-messages.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-messages.ts @@ -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; diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 051b252e6..84ada11aa 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -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( diff --git a/packages/coding-agent/test/fixtures/mock-rpc-agent.ts b/packages/coding-agent/test/fixtures/mock-rpc-agent.ts index 0e93424dd..918df7227 100755 --- a/packages/coding-agent/test/fixtures/mock-rpc-agent.ts +++ b/packages/coding-agent/test/fixtures/mock-rpc-agent.ts @@ -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", diff --git a/packages/coding-agent/test/rpc-client.restart.test.ts b/packages/coding-agent/test/rpc-client.restart.test.ts index 16aedbba7..9b31b117e 100644 --- a/packages/coding-agent/test/rpc-client.restart.test.ts +++ b/packages/coding-agent/test/rpc-client.restart.test.ts @@ -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, diff --git a/packages/coding-agent/test/rpc-frame.test.ts b/packages/coding-agent/test/rpc-frame.test.ts index 6c860c2cf..f558537f0 100644 --- a/packages/coding-agent/test/rpc-frame.test.ts +++ b/packages/coding-agent/test/rpc-frame.test.ts @@ -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); diff --git a/python/omp-rpc/src/omp_rpc/client.py b/python/omp-rpc/src/omp_rpc/client.py index 7069b09ed..688d7dea0 100644 --- a/python/omp-rpc/src/omp_rpc/client.py +++ b/python/omp-rpc/src/omp_rpc/client.py @@ -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"))) diff --git a/python/omp-rpc/tests/test_client.py b/python/omp-rpc/tests/test_client.py index e9d3b0193..7d5c2b88b 100644 --- a/python/omp-rpc/tests/test_client.py +++ b/python/omp-rpc/tests/test_client.py @@ -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")